125 lines
4.3 KiB
Python
125 lines
4.3 KiB
Python
import pyotp
|
|
import decouple
|
|
from os.path import join
|
|
|
|
from flask import Flask, request
|
|
from flask_restful import Resource, Api
|
|
from schemas import OTPSchema, CreateOTPSchema
|
|
|
|
from config import etcd_client
|
|
from helper import is_valid_otp
|
|
|
|
app = Flask(__name__)
|
|
api = Api(app)
|
|
|
|
def create_admin_if_dont_exists(etcd_client):
|
|
_key = join(
|
|
decouple.config('BASE_PREFIX'),
|
|
decouple.config('ADMIN_REALM'),
|
|
'admin',
|
|
)
|
|
if etcd_client.get(_key) is None:
|
|
print('admin does not exists!. So, creating one')
|
|
_value = {
|
|
'seed': pyotp.random_base32(),
|
|
}
|
|
etcd_client.put(_key, _value, value_in_json=True)
|
|
|
|
|
|
|
|
create_admin_if_dont_exists(etcd_client)
|
|
|
|
|
|
class Verify(Resource):
|
|
@staticmethod
|
|
def get():
|
|
data = request.json
|
|
schema = OTPSchema(data)
|
|
if schema.is_valid():
|
|
return {'message': 'Verified'}, 200
|
|
else:
|
|
return schema.get_errors(), 400
|
|
|
|
# try:
|
|
# name = data['name']
|
|
# realm = data['realm']
|
|
# token = data['token']
|
|
|
|
# auth_name = data['auth-name']
|
|
# auth_realm = data['auth-realm']
|
|
# auth_token = data['auth-token']
|
|
# except Exception:
|
|
# return {
|
|
# "message": "Your provided data is not correct."
|
|
# }
|
|
# else:
|
|
# if is_valid_otp(auth_name, auth_realm, auth_token):
|
|
# _key = join(decouple.config("BASE_PREFIX"), realm, name)
|
|
# entry = etcd_client.get(_key, value_in_json=True)
|
|
# if entry:
|
|
# totp = pyotp.TOTP(entry.value['seed'])
|
|
# try:
|
|
# is_token_valid = totp.verify(token)
|
|
# except:
|
|
# return {'message': 'Invalid Data'}, 400
|
|
# else:
|
|
# if is_token_valid:
|
|
# return {'message': 'Verified'}, 200
|
|
# else:
|
|
# return {'message': 'Invalid token'}, 400
|
|
# else:
|
|
# return {"message": "No such Account Found"}, 400
|
|
# else:
|
|
# return {'message': 'Invalid Auth Credentials'}, 400
|
|
|
|
|
|
class Create(Resource):
|
|
@staticmethod
|
|
def post():
|
|
data = request.json
|
|
|
|
schema = CreateOTPSchema(data)
|
|
if schema.is_valid():
|
|
_key = join(decouple.config('BASE_PREFIX'), data['realm'], data['name'])
|
|
|
|
if etcd_client.get(_key) is None:
|
|
_value = {
|
|
'seed': pyotp.random_base32(),
|
|
}
|
|
etcd_client.put(_key, _value, value_in_json=True)
|
|
return {'message': 'Account Created\n'
|
|
'name: {}, realm: {}, seed: {}'.format(data['name'],data['realm'], _value['seed'])}
|
|
else:
|
|
return schema.get_errors()
|
|
# try:
|
|
# name = data['name']
|
|
# realm = data['realm']
|
|
# admin_name = data['admin-name']
|
|
# admin_realm = data['admin-realm']
|
|
# admin_token = data['admin-token']
|
|
# except Exception:
|
|
# return {'message': 'Invalid Data'}, 400
|
|
# else:
|
|
# if admin_realm == decouple.config('ADMIN_REALM'):
|
|
# if is_valid_otp(admin_name, admin_realm, admin_token):
|
|
# _key = join(decouple.config('BASE_PREFIX'), realm, name)
|
|
|
|
# if etcd_client.get(_key) is None:
|
|
# _value = {
|
|
# 'seed': pyotp.random_base32(),
|
|
# }
|
|
# etcd_client.put(_key, _value, value_in_json=True)
|
|
# return {'message': 'Account Created\n'
|
|
# 'name: {}, realm: {}, seed: {}'.format(name, realm, _value['seed'])}
|
|
# else:
|
|
# return {'message': 'Account already exists'}, 400
|
|
# else:
|
|
# return {'message': 'Invalid Admin OTP Credentials'}
|
|
# else:
|
|
# return {'message': 'Admin must be from {} realm'.format(decouple.config('ADMIN_REALM'))}
|
|
|
|
api.add_resource(Verify, "/verify")
|
|
api.add_resource(Create, "/create")
|
|
|
|
if __name__ == "__main__":
|
|
app.run(debug=True)
|