uotp/schemas.py
2019-10-07 22:13:42 +05:00

154 lines
No EOL
5.2 KiB
Python

import json
import decouple
import pyotp
from os.path import join
from helper import is_valid_otp
from config import etcd_client
class Field:
def __init__(self, _name, _type, _value=None):
self.name = _name
self.value = _value
self.type = _type
self.__errors = []
def validation(self):
return True
def is_valid(self):
if self.value == KeyError:
self.add_error("'{}' field is a required field".format(self.name))
else:
if not isinstance(self.value, self.type):
self.add_error("Incorrect Type for '{}' field".format(self.name))
else:
self.validation()
if self.__errors:
return False
return True
def get_errors(self):
return self.__errors
def add_error(self, error):
self.__errors.append(error)
class BaseSchema:
def __init__(self, data, fields=None):
_ = data # suppress linter warning
self.__errors = []
if fields is None:
self.fields = []
else:
self.fields = fields
def validation(self):
# custom validation is optional
return True
def is_valid(self):
for field in self.fields:
field.is_valid()
self.add_field_errors(field)
for parent in self.__class__.__bases__:
try:
parent.validation(self)
except AttributeError:
pass
if not self.__errors:
self.validation()
if self.__errors:
return False
return True
def get_errors(self):
return {"message": self.__errors}
def add_field_errors(self, field: Field):
self.__errors += field.get_errors()
def add_error(self, error):
self.__errors.append(error)
class OTPSchema(BaseSchema):
def __init__(self, data: dict, fields=None):
self.name = Field("name", str, data.get("name", KeyError))
self.realm = Field("realm", str, data.get("realm", KeyError))
self.token = Field("token", str, data.get("token", KeyError))
self.auth_name = Field("auth-name", str,
data.get("auth-name", KeyError))
self.auth_realm = Field("auth-realm", str,
data.get("auth-realm", KeyError))
self.auth_token = Field("auth-token", str,
data.get("auth-token", KeyError))
self.auth_realm.validation = self.auth_realm_validation
_fields = [self.name, self.realm, self.token,
self.auth_name, self.auth_realm,
self.auth_token]
if fields:
_fields += fields
super().__init__(data=data, fields=_fields)
def auth_realm_validation(self):
if self.auth_realm.value != decouple.config("AUTH_REALM"):
self.add_error(
"Authentication realm must be {}".format(decouple.config("AUTH_REALM"))
)
def validation(self):
if is_valid_otp(etcd_client, self.auth_name.value,
self.auth_realm.value, self.auth_token.value):
if is_valid_otp(etcd_client, self.name.value,
self.realm.value, self.token.value):
_key = join(decouple.config("BASE_PREFIX"),
self.realm.value, self.name.value)
entry = etcd_client.get(_key, value_in_json=True)
if not entry:
self.add_error("No such Account Found")
else:
self.add_error("Invalid OTP Credentials")
else:
self.add_error("Invalid Auth Credentials")
class CreateOTPSchema(BaseSchema):
def __init__(self, data: dict, fields=None):
self.name = Field("name", str, data.get("name", KeyError))
self.realm = Field("realm", str, data.get("realm", KeyError))
self.admin_name = Field("admin-name", str,
data.get("admin-name", KeyError))
self.admin_realm = Field("admin-realm", str,
data.get("admin-realm", KeyError))
self.admin_token = Field("admin-token", str,
data.get("admin-token", KeyError))
self.admin_realm.validation = self.admin_realm_validation
_fields = [self.name, self.realm,
self.admin_name, self.admin_realm,
self.admin_token]
if fields:
_fields += fields
super().__init__(data=data, fields=_fields)
def admin_realm_validation(self):
if self.admin_realm.value != decouple.config("ADMIN_REALM"):
self.add_field_errors(
'Admin must be from {} realm'.format(decouple.config('ADMIN_REALM'))
)
def validation(self):
if is_valid_otp(etcd_client, self.admin_name.value,
self.admin_realm.value, self.admin_token.value):
_key = join(decouple.config("BASE_PREFIX"),
self.realm.value, self.name.value)
if etcd_client.get(_key):
self.add_error("Account already exists")
else:
self.add_error("Invalid Admin OTP Credentials")