forked from AikidoSec/firewall-python
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathinit.py
More file actions
79 lines (64 loc) · 2.03 KB
/
init.py
File metadata and controls
79 lines (64 loc) · 2.03 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
import json
import os
from pymongo import MongoClient
from aikido_zen.aws_lambda import protect
class Dog:
def __init__(self, dog_name, pswd):
self.dog_name = dog_name
self.pswd = pswd
class Dogs:
def __init__(self, client):
self.collection = client['my_database']['dogs'] # Replace with your database name
def find_by(self, dog_name, pswd):
return self.collection.find_one({"dog_name": dog_name, "pswd": pswd})
def persist(self, dog):
self.collection.insert_one({"dog_name": dog.dog_name, "pswd": dog.pswd})
def lambda_handler(event, context):
# Normally you'd use environment variables for this
client = MongoClient(mongo_uri)
try:
return main(client, event)
except Exception as e:
print(f"Error: {e}")
return {
"statusCode": 500,
"body": str(e)
}
finally:
client.close()
handler = protect(lambda_handler)
def main(client, event):
dogs = Dogs(client)
# Ensure a dog exists for testing
dog = dogs.find_by("Doggo 1", "xyz")
if not dog:
dogs.persist(Dog("Doggo 1", "xyz"))
if not event.get("body") or event.get("httpMethod") != "POST":
return {
"statusCode": 405,
"body": "Method Not Allowed",
}
body = json.loads(event["body"])
if not body.get("dog_name") or not body.get("pswd"):
return {
"statusCode": 400,
"body": "Bad Request",
}
# This is just for demo purposes, normally you'd use bcrypt or something
actual_dog = dogs.find_by(body["dog_name"], body["pswd"])
if not actual_dog:
return {
"statusCode": 401,
"body": "Unauthorized",
}
return {
"statusCode": 200,
"headers": {
"Content-Type": "application/json",
},
"body": json.dumps({
"token": "123",
"success": True,
}),
}