harness_zbt/backend/app/domains/identity/repository.py

from dataclasses import replace
from typing import Protocol

from app.domains.identity.models import AdminRole, AdminUser, AuthSession, H5User


class IdentityRepository(Protocol):
    def get_h5_user_by_mobile(self, mobile: str) -> H5User | None: ...

    def get_h5_user(self, user_id: str) -> H5User | None: ...

    def list_h5_users(self) -> list[H5User]: ...

    def save_h5_user(self, user: H5User) -> None: ...

    def get_admin_user_by_username(self, username: str) -> AdminUser | None: ...

    def get_admin_user(self, user_id: str) -> AdminUser | None: ...

    def list_admin_users(self) -> list[AdminUser]: ...

    def save_admin_user(self, user: AdminUser) -> None: ...

    def list_roles(self) -> list[AdminRole]: ...

    def save_role(self, role: AdminRole) -> None: ...

    def get_session(self, session_id: str) -> AuthSession | None: ...

    def save_session(self, session: AuthSession) -> None: ...


class InMemoryIdentityRepository:
    def __init__(self) -> None:
        self._h5_users: dict[str, H5User] = {}
        self._h5_user_ids_by_mobile: dict[str, str] = {}
        self._admin_users: dict[str, AdminUser] = {}
        self._admin_user_ids_by_username: dict[str, str] = {}
        self._roles: dict[str, AdminRole] = {}
        self._sessions: dict[str, AuthSession] = {}

    def get_h5_user_by_mobile(self, mobile: str) -> H5User | None:
        user_id = self._h5_user_ids_by_mobile.get(mobile)
        return self._h5_users.get(user_id) if user_id else None

    def get_h5_user(self, user_id: str) -> H5User | None:
        return self._h5_users.get(user_id)

    def list_h5_users(self) -> list[H5User]:
        return sorted(
            self._h5_users.values(),
            key=lambda user: user.created_at,
            reverse=True,
        )

    def save_h5_user(self, user: H5User) -> None:
        self._h5_users[user.id] = user
        self._h5_user_ids_by_mobile[user.mobile] = user.id

    def get_admin_user_by_username(self, username: str) -> AdminUser | None:
        user_id = self._admin_user_ids_by_username.get(username)
        return self._admin_users.get(user_id) if user_id else None

    def get_admin_user(self, user_id: str) -> AdminUser | None:
        return self._admin_users.get(user_id)

    def list_admin_users(self) -> list[AdminUser]:
        return sorted(self._admin_users.values(), key=lambda user: user.username)

    def save_admin_user(self, user: AdminUser) -> None:
        self._admin_users[user.id] = user
        self._admin_user_ids_by_username[user.username] = user.id
        for role_code in user.roles:
            self._roles.setdefault(
                role_code,
                AdminRole(
                    code=role_code,
                    name=role_code,
                    data_scope=user.data_scope,
                    status="ACTIVE",
                    permissions=user.permissions,
                ),
            )

    def list_roles(self) -> list[AdminRole]:
        return sorted(self._roles.values(), key=lambda role: role.code)

    def save_role(self, role: AdminRole) -> None:
        if role.code not in self._roles:
            raise KeyError(role.code)
        self._roles[role.code] = role
        for user_id, user in tuple(self._admin_users.items()):
            if role.code in user.roles:
                assigned_roles = [self._roles[code] for code in user.roles]
                permissions = tuple(
                    sorted(
                        {
                            permission
                            for assigned_role in assigned_roles
                            for permission in assigned_role.permissions
                        }
                    )
                )
                scopes = {assigned_role.data_scope for assigned_role in assigned_roles}
                self._admin_users[user_id] = replace(
                    user,
                    permissions=permissions,
                    data_scope="ALL" if "ALL" in scopes else sorted(scopes)[0],
                )

    def get_session(self, session_id: str) -> AuthSession | None:
        return self._sessions.get(session_id)

    def save_session(self, session: AuthSession) -> None:
        self._sessions[session.id] = session