sakhteman/auth/repositories/AuthRepository.py

from http.client import HTTPException
from typing import List
from uuid import UUID

from auth.models import UserInDB, UserModel, UIdsInput, UserOptions, SessionModel, PasswordChangeModel, UserUpdateModel, \
    VisualUserUpdate
from auth.models.roles import RoleModel, ModuleModel
from auth.models.units import UnitModel
from auth.utils.mongo import use_auth_collection
from auth.utils.pwd import get_password_hash, verify_password
from core.exceptions import ModifiedCountException, ValidationException
from core.models.base import DELETED_AT_QUERY
from core.utils.app_context import use_user
from core.utils.crudify import mongo_soft_deleter

COLLECTION_UNITS = 'units'
COLLECTION_USERS = 'users'
COLLECTION_ROLES = 'roles'
COLLECTION_SESSIONS = 'sessions'

_projection_user = {
    'mobile': 1,
    'email': 1,
    'username': 1,
    '_id': 1,
    'name': 1
}


class AuthRepository:
    @staticmethod
    async def create(_input: UserInDB) -> dict:
        record = _input.init_id().cu_at().mongo_dict()
        record['password'] = get_password_hash(_input.password)
        insert = await use_auth_collection(COLLECTION_USERS).insert_one(record)
        if not insert.acknowledged:
            raise HTTPException(status_code=400, detail="User already exists")
        return insert.inserted_id

    @staticmethod
    async def user_query(query: str = '') -> List[UserModel]:
        q = {'$or': [{'name': {'$regex': query.strip()}},
                     {'mobile': {'$regex': query.strip()}}
                     ]} if query.strip() else {}
        cursor = use_auth_collection('users').find(q, {'_id': 1, 'name': 1, 'mobile': 1})
        return [UserModel(**user) async for user in cursor]

    @staticmethod
    async def get_units():
        cursor = use_auth_collection(COLLECTION_UNITS).find()
        return [UnitModel(**document) async for document in cursor]

    @staticmethod
    async def get_user(_id: UUID) -> UserModel:
        user = await use_auth_collection(COLLECTION_USERS).find_one({"_id": _id}, _projection_user)
        if user is None:
            raise HTTPException(status_code=404, detail="User not found")
        return UserModel(**user)

    @staticmethod
    async def update_options(options: UserOptions) -> bool:
        user = use_user()
        user = await use_auth_collection(COLLECTION_USERS) \
            .update_one({"_id": user.id}, {"$set": {"options": options.dict()}})
        if user.modified_count == 0:
            raise ModifiedCountException("User not updated")
        return True

    @staticmethod
    async def update_avatar(av_input: VisualUserUpdate) -> bool:
        user = use_user()
        _update_query = {"$set": {"avatar": av_input.storage}} if av_input.storage else {'$unset': {'avatar': 1}}
        user = await use_auth_collection(COLLECTION_USERS) \
            .update_one({"_id": user.id}, _update_query)
        if user.modified_count == 0:
            raise ModifiedCountException("User not updated")
        return True

    @staticmethod
    async def update_password(_input: PasswordChangeModel) -> bool:
        user = use_user()
        user = await use_auth_collection(COLLECTION_USERS).find_one({'_id': user.id})
        if not user:
            raise HTTPException(status_code=404, detail="User not found")
        if not verify_password(_input.old_password, user.get('password', None)):
            raise ValidationException('Old Password not valid')
        user['password'] = get_password_hash(_input.new_password)
        user = await use_auth_collection(COLLECTION_USERS).update_one({'_id': user.get('_id')},
                                                                      {'$set': {
                                                                          'password': user['password']}})
        if user.modified_count == 0:
            raise ModifiedCountException("Password")
        return True

    @staticmethod
    async def update_user(_input: UserUpdateModel) -> bool:
        user = use_user()
        user = await use_auth_collection(COLLECTION_USERS).find_one({'_id': UUID(str(user.id))})
        user = await use_auth_collection(COLLECTION_USERS).update_one({'_id': user.get('_id')},
                                                                      {'$set': _input.u_at().dict()})
        if user.modified_count == 0:
            raise ModifiedCountException("User not update")
        return True

    @staticmethod
    async def get_my_sessions() -> List[SessionModel]:
        user = use_user()
        user = use_auth_collection(COLLECTION_SESSIONS).find({"user_id": user.id, **DELETED_AT_QUERY})
        return [SessionModel(**x) async for x in user]

    @staticmethod
    async def delete_session(session_id: UUID) -> bool:
        return await mongo_soft_deleter(session_id, COLLECTION_SESSIONS, client=use_auth_collection)

    @staticmethod
    async def get_users_by_id(_input: UIdsInput) -> List[UserModel]:
        cursor = use_auth_collection(COLLECTION_USERS) \
            .find({'_id': {'$in': _input.u_ids}})
        return [UserModel(**user) async for user in cursor]

    @staticmethod
    async def create_unit(_input: UnitModel) -> str:
        document = _input.init_id().mongo_dict()
        result = await use_auth_collection(COLLECTION_UNITS).insert_one(document)
        return result.inserted_id

    @staticmethod
    async def get_roles() -> List[dict]:
        roles = await use_auth_collection(COLLECTION_ROLES).find().to_list(None)
        return [role for role in roles]

    @staticmethod
    async def create_role(_input: RoleModel) -> str:
        record = _input.init_id().cu_at().mongo_dict()
        insert = await use_auth_collection(COLLECTION_ROLES).insert_one(record)
        if not insert.acknowledged:
            raise HTTPException("Error creating role")
        return insert.inserted_id

    @staticmethod
    async def get_my_modules() -> List[ModuleModel]:
        user = use_user()
        records = await use_auth_collection(COLLECTION_ROLES).find(
            {'_id': {'$in': user.roles or []}, **DELETED_AT_QUERY}, {'modules': True}).to_list(None)
        if not records:
            return []
        _list = [x.get('modules', []) for x in records]
        if _list:
            return list([item for sublist in _list for item in sublist])
        return []

    @staticmethod
    async def get_f_routes(roles_ids: List[UUID]) -> List[str]:

        roles = await use_auth_collection(COLLECTION_ROLES).find(
            {'_id': {'$in': roles_ids}}) \
            .to_list(None)

        def _union_list(*kwargs):
            final = []
            for kwarg in kwargs:
                final.extend(kwarg)
            return final

        return _union_list(*[role['f_routes'] for role in roles])