roc-fnb-server/roc_fnb/website/database.py

71 lines
2.2 KiB
Python

from os import environ
from typing import Optional, Self
from pdb import set_trace as debugger
from bson.objectid import ObjectId
import pymongo
from pymongo import MongoClient
from roc_fnb.website.models.user import User
KEYLEN = 64
with open('private-key.pem') as pk:
PRIVATE_KEY = pk.read()
with open('public-key.pem') as pk:
PUBLIK_KEY = pk.read()
class Database(MongoClient):
@classmethod
def from_env(cls) -> Self:
return cls(
host=environ.get('DATABASE_HOST', default='localhost'),
port=int(environ.get('DATABASE_PORT', default=27017))
)
@property
def db(self) -> pymongo.database.Database:
if (env := environ.get('ENV_MODE')) and env == 'production':
return self.production
else:
return self.development
def store_user(self, user: User) -> User:
"""Store the given user in the database, set _id on it, and return it"""
result = self.db.users.insert_one(user.document)
user._id = result.inserted_id
return user
def get_user_by_email(self, email: str) -> Optional[User]:
"""
Return the user associated with the given email.
This does not imply authentiation
"""
if result := self.db.users.find_one({'email': email}):
return User(**result)
return None
def get_user_by_name(self, name: str) -> Optional[User]:
if result := self.db.users.find_one({'name': name}):
return User(**result)
return None
def delete_user(self, _id: ObjectId):
self.db.users.delete_one({'_id': _id})
def get_user_by_id(self, id: ObjectId) -> Optional[User]:
if user := self.db.users.find_one({'_id': id}):
return User(**user)
return None
def get_user_from_token(self, token: str) -> Optional[User]:
"""
Verify a user and retreive a their full profile from the database.
This is like User.verify_jwt except it also fetches fields from the
database which are not present in the client-visible token.
"""
if jwt_user := User.verify_jwt(token):
return self.get_user_by_id(jwt_user._id)
return None