def __exit__(self, *args):
TinyDB.__exit__(self, *args)
FileLock.__exit__(self, *args)
class TinyAuthServer(web.Application):
def __init__(self, data_dir: Path, constant_secret: str,
admin_password: str, user_cache_size: int):
super().__init__()
self.data_dir = data_dir
self.data_dir.mkdir(exist_ok=True, parents=False)
self.db = TinyDB(data_dir / "db.json",
indent=4,
storage=CachingMiddleware(JSONStorage))
self.users = self.db.table('users', cache_size=user_cache_size)
self.setup_root_user(admin_password)
self.encryptor = Fernet(self.load_fernet_key())
self.constant_secret: str = constant_secret
self.add_routes([
web.post('/auth/get_cookie', self.handle_get_cookie),
web.get('/auth', self.handle_is_authenticated),
web.get('/auth/login', self.handle_login_form),
web.post('/auth/logout', self.handle_logout)
])
def __enter__(self):
self.db.__enter__()
def __exit__(self, exc_type, exc_val, exc_tb):
self.db.__exit__(exc_type, exc_val, exc_tb)
@staticmethod
async def handle_logout(request: web.Request):
cookie = request.cookies.get(AUTH_COOKIE_NAME)
if cookie == "":
return web.Response(reason="Not logged in", status=400)
else:
response = web.Response(reason=f"Successfully logged out.",
status=200)
response.set_cookie(AUTH_COOKIE_NAME, "")
return response
def setup_root_user(self, admin_password: str):
user = Query()
element = self.users.get(user.name == 'admin')
if element is None:
self.logger.info("Recreating database.")
self.users.insert({
'name': 'admin',
'password': admin_password,
"permissions": PermissionLevels.ADMIN
})
self.users.update({'password': admin_password}, user.name == 'admin')
async def handle_is_authenticated(self,
request: web.Request) -> web.Response:
cookie = request.cookies.get(AUTH_COOKIE_NAME)
try:
user_name = await self.unpack_encrypted_cookie(cookie)
user = self.users.get(Query().name == user_name)
if user is None:
return web.Response(reason="User not found.", status=401)
else:
self.logger.info(user)
return web.Response(reason="Authenticated", status=200)
except AuthenticationException:
self.logger.info("Could not authenticate.")
return web.Response(reason="Could not authenticate.", status=401)
@staticmethod
async def handle_login_form(request: web.Request):
return web.FileResponse(
Path(__file__).parent.parent / "resources" / 'login.html')
async def handle_get_cookie(self, request: web.Request):
login_info = await request.post()
query = Query()
user = login_info.get("user", None)
if user is None:
self.logger.info(
"get cookie failed, because json posted did not contain user")
return web.Response(reason=f"You must specify the user to login",
status=400)
if self.users.get(query.name == user) is None:
self.logger.info(f"get cookie failed, unknown user {user}")
return web.Response(reason=f"Invalid login credentials",
status=400)
password = login_info.get("password", None)
if password is None:
self.logger.info(
"get cookie failed, because json posted did not contain password"
)
return web.Response(reason=f"You must specify a password",
status=400)
if password != self.users.get(query.name == user)["password"]:
self.logger.info("Get cookie failed, invalid password.")
return web.Response(reason=f"Invalid login credentials",
status=400)
self.logger.info(f"Succesfully authenticated user {user}")
response = web.json_response(data={"authenticated": "true"},
reason="Successfully authenticated.")
cookie: str = await self.create_encrypted_cookie(user)
response.set_cookie(AUTH_COOKIE_NAME,
cookie,
secure=True,
httponly=True)
response.content_type = "application/json"
return response
async def create_encrypted_cookie(self, user: str) -> str:
cookie_salt = secrets.token_hex(AUTH_COOKIE_SALT_LENGTH)
# we use the auth cookie constant secret to quickly check if a cookie is valid after decryption.
valid_until = datetime.datetime.now(
tz=datetime.timezone.utc) + datetime.timedelta(days=60)
valid_until_iso = valid_until.isoformat(timespec="seconds")
raw_cookie = "@@@@".join(
[self.constant_secret, user, valid_until_iso,
cookie_salt]).encode("utf-8")
encrypted_cookie = self.encryptor.encrypt(raw_cookie)
return base64.b64encode(encrypted_cookie).decode("utf-8")
async def unpack_encrypted_cookie(self, cookie: str) -> str:
if cookie is None:
raise AuthenticationException(
"Could not find authentication cookie.")
try:
# returns str to indicate the user, raises AuthenticationException otherwise
encrypted_cookie = base64.b64decode(cookie.encode("utf-8"))
raw_cookie = self.encryptor.decrypt(encrypted_cookie).decode(
"utf-8")
constant_secret, user, valid_until_iso, _ = raw_cookie.split(
"@@@@")
if constant_secret != self.constant_secret:
raise AuthenticationException(
"Constant secret does not match.")
if datetime.datetime.fromisoformat(
valid_until_iso) < datetime.datetime.now(
tz=datetime.timezone.utc):
raise AuthenticationException("Stale authentication cookie.")
else:
return user
except UnicodeDecodeError:
msg = "Could decode auth cookie."
self.logger.info(msg)
raise AuthenticationException(msg)
except ValueError:
msg = "Could not unpack authentication cookie."
self.logger.info(msg)
raise AuthenticationException(msg)
except Exception as e:
msg = "Could not unpack Cookie. Unhandled exception."
self.logger.exception(msg, e)
raise AuthenticationException(msg)
def load_fernet_key(self) -> bytes:
keypath = self.data_dir / "fernet.key"
if not keypath.exists():
new_key = Fernet.generate_key()
with open(keypath, "wb") as fp:
fp.write(new_key)
return new_key
else:
with open(keypath, "rb") as fp:
return fp.read()
class Db:
def __init__(self, path: Path = None):
if not path:
path = Path(_tinydb)
path.parent.mkdir(exist_ok=True, parents=True)
# only a single thread should ready or modify the database at any one time
# so we will use this lock whenever we use self.tinydb
self._lock = threading.Lock()
self.tinydb = TinyDB(str(path))
self.seen = self.tinydb.table("seen")
def __enter__(self):
return self
def __exit__(self, *args):
self.tinydb.__exit__()
def debug_eprint_all(self) -> None:
t = self.seen
with self._lock:
all_ = t.all()
u.eprint("====db.all")
u.eprint_pprint(all_)
u.eprint("db.all====")
def _unique_words_seen_no_lock(self) -> int:
return len(self.seen)
def unique_words_seen(self) -> int:
with self._lock:
return self._unique_words_seen_no_lock()
def add_word(self, word: str) -> Dict:
'''
Usage
=====
>>> is_unusual = add_word(t, 'augend')
Returns
=======
True iff word has not been seen before
'''
_word_sanity_check(word)
t = self.seen
with self._lock:
Q = Query()
q = t.get(Q.word == word)
if not q:
n = self._unique_words_seen_no_lock()
_doc_id = t.insert({'word': word, 'count': 1})
return {"n": n + 1, "is-new": True}
doc_id = q.doc_id
count = t.get(doc_id=doc_id)['count'] + 1
t.update({'count': count}, doc_ids=[doc_id])
return {"is-new": False, 'count': count}
def get_word_count(self, word: str) -> int:
t = self.seen
with self._lock:
Q = Query()
q = t.search(Q.word == word)
lenq = len(q)
if lenq == 0:
return 0
elif lenq == 1:
return q[0]['count']
else:
u.eprint(
f"db.py: corrupt data! word '{word}' occured {lenq} times in db, expected 0 or 1"
)
def total_words_seen(self) -> int:
t = self.seen
with self._lock:
_all = t.all()
total = 0
for x in _all:
total += x["count"]
return total
def words_with_count(self) -> Dict[str, int]:
t = self.seen
d = {}
with self._lock:
for x in t.all():
d[x["word"]] = x["count"]
return d