uotp/app.py

104 lines
2.9 KiB
Python
Raw Normal View History

2019-10-07 17:13:42 +00:00
import pyotp
import decouple
2019-10-08 18:07:28 +00:00
import os
2019-10-07 17:13:42 +00:00
from os.path import join
from flask import Flask, request
from flask_restful import Resource, Api
from schemas import OTPSchema, CreateOTPSchema, DeleteOTPSchema, ListAccountSchema
2019-10-07 17:13:42 +00:00
from config import etcd_client
2019-10-08 18:07:28 +00:00
from helper import is_valid_otp, create_admin_if_dont_exists
2019-10-07 17:13:42 +00:00
app = Flask(__name__)
api = Api(app)
create_admin_if_dont_exists(etcd_client)
class Verify(Resource):
@staticmethod
def get():
data = request.json
if data:
schema = OTPSchema(data)
if schema.is_valid():
return {"message": "Verified"}, 200
else:
return schema.get_errors(), 400
2019-10-07 17:13:42 +00:00
else:
return {"message": "No Data"}, 400
2019-10-07 17:13:42 +00:00
class Create(Resource):
@staticmethod
def post():
data = request.json
schema = CreateOTPSchema(data)
if schema.is_valid():
_key = join(decouple.config("BASE_PREFIX"), data["name"])
2019-10-07 17:13:42 +00:00
if etcd_client.get(_key) is None:
if not isinstance(data["realm"], list):
realms = [data["realm"]]
else:
realms = data["realm"]
_value = {"seed": pyotp.random_base32(), "realm": realms}
2019-10-07 17:13:42 +00:00
etcd_client.put(_key, _value, value_in_json=True)
2019-10-08 18:07:28 +00:00
return {
"message": "Account Created",
"credentials": {
"name": data["name"],
"realm": _value["realm"],
"seed": _value["seed"]
}
}, 200
2019-10-07 17:13:42 +00:00
else:
2019-10-08 18:07:28 +00:00
return schema.get_errors(), 400
class Delete(Resource):
@staticmethod
def post():
data = request.json
schema = DeleteOTPSchema(data)
if schema.is_valid():
_key = join(decouple.config("BASE_PREFIX"), data["name"])
2019-10-08 18:07:28 +00:00
etcd_client.client.delete(_key)
return {"message": "Account Deleted"}
else:
return schema.get_errors(), 400
class List(Resource):
@staticmethod
def get():
data = request.json
schema = ListAccountSchema(data)
if schema.is_valid():
result = etcd_client.get_prefix(
decouple.config("BASE_PREFIX"), value_in_json=True
)
r = {}
for entry in result:
_name = entry.key.split("/")[-1]
r["{}".format(_name)] = {
"seed": entry.value["seed"],
"realm": entry.value["realm"],
}
2019-10-08 18:07:28 +00:00
return r
else:
return schema.get_errors(), 400
2019-10-07 17:13:42 +00:00
api.add_resource(Verify, "/verify")
api.add_resource(Create, "/create")
2019-10-08 18:07:28 +00:00
api.add_resource(Delete, "/delete")
api.add_resource(List, "/list")
2019-10-07 17:13:42 +00:00
if __name__ == "__main__":
app.run(debug=True, port=decouple.config("PORT", int))