xiaozhi-esp32-server/main/manager-api-fastapi/app/repositories/security.py

from __future__ import annotations

from datetime import datetime
from typing import Any

from sqlalchemy import text
from sqlalchemy.ext.asyncio import AsyncSession

from app.core.database import Repository


class SecurityRepository(Repository):
    def __init__(self, session: AsyncSession):
        super().__init__(session)

    async def get_param_value(self, code: str) -> str | None:
        value = await self.scalar(
            "SELECT param_value FROM sys_params WHERE param_code = :code LIMIT 1",
            {"code": code},
        )
        return None if value is None else str(value)

    async def get_mobile_area_items(self) -> list[dict[str, Any]]:
        return await self.fetch_all(
            "SELECT d.dict_label AS name, d.dict_value AS `key` "
            "FROM sys_dict_data d "
            "LEFT JOIN sys_dict_type t ON d.dict_type_id = t.id "
            "WHERE t.dict_type = :dict_type ORDER BY d.sort ASC",
            {"dict_type": "MOBILE_AREA"},
        )

    async def get_user_by_username(self, username: str | None) -> dict[str, Any] | None:
        return await self.fetch_one(
            "SELECT id, username, password, super_admin, status, creator, create_date, updater, update_date "
            "FROM sys_user WHERE username = :username LIMIT 1",
            {"username": username},
        )

    async def get_user_by_id(self, user_id: int) -> dict[str, Any] | None:
        return await self.fetch_one(
            "SELECT id, username, password, super_admin, status, creator, create_date, updater, update_date "
            "FROM sys_user WHERE id = :id LIMIT 1",
            {"id": user_id},
        )

    async def count_users(self) -> int:
        return int(await self.scalar("SELECT COUNT(*) FROM sys_user") or 0)

    async def insert_user(
        self,
        *,
        user_id: int,
        username: str | None,
        password: str,
        super_admin: int,
        now: datetime,
    ) -> None:
        await self.execute(
            "INSERT INTO sys_user "
            "(id, username, password, super_admin, status, creator, create_date, updater, update_date) "
            "VALUES (:id, :username, :password, :super_admin, 1, NULL, :now, NULL, :now)",
            {
                "id": user_id,
                "username": username,
                "password": password,
                "super_admin": super_admin,
                "now": now,
            },
        )

    async def get_token_by_user_id(self, user_id: int, *, for_update: bool = False) -> dict[str, Any] | None:
        sql = (
            "SELECT id, user_id, token, expire_date, update_date, create_date "
            "FROM sys_user_token WHERE user_id = :user_id LIMIT 1 FOR UPDATE"
            if for_update and self._supports_for_update()
            else "SELECT id, user_id, token, expire_date, update_date, create_date "
            "FROM sys_user_token WHERE user_id = :user_id LIMIT 1"
        )
        return await self.fetch_one(
            sql,
            {"user_id": user_id},
        )

    async def insert_token(
        self,
        *,
        token_id: int,
        user_id: int,
        token: str,
        now: datetime,
        expire_date: datetime,
    ) -> None:
        await self.execute(
            "INSERT INTO sys_user_token (id, user_id, token, expire_date, update_date, create_date) "
            "VALUES (:id, :user_id, :token, :expire_date, :now, :now)",
            {
                "id": token_id,
                "user_id": user_id,
                "token": token,
                "expire_date": expire_date,
                "now": now,
            },
        )

    async def update_token(self, *, token_id: int, token: str, now: datetime, expire_date: datetime) -> None:
        await self.execute(
            "UPDATE sys_user_token SET token = :token, expire_date = :expire_date, update_date = :now "
            "WHERE id = :id",
            {"id": token_id, "token": token, "expire_date": expire_date, "now": now},
        )

    async def update_password(
        self,
        user_id: int,
        password_hash: str,
        now: datetime,
        *,
        preserve_audit_fields: bool = False,
    ) -> int:
        return await self.execute(
            "UPDATE sys_user SET password = :password, "
            "update_date = CASE WHEN :preserve_audit = 1 THEN update_date ELSE :now END WHERE id = :id",
            {
                "id": user_id,
                "password": password_hash,
                "now": now,
                "preserve_audit": int(preserve_audit_fields),
            },
        )

    async def expire_user_token(self, user_id: int, expire_date: datetime) -> int:
        return await self.execute(
            "UPDATE sys_user_token SET expire_date = :expire_date WHERE user_id = :user_id",
            {"user_id": user_id, "expire_date": expire_date},
        )

    def _supports_for_update(self) -> bool:
        bind = self.session.get_bind()
        return bind.dialect.name != "sqlite"


async def raw_user_token(session: AsyncSession, token: str) -> dict[str, Any] | None:
    result = await session.execute(
        text(
            "SELECT t.id AS token_id, t.user_id, t.token, t.expire_date, "
            "u.username, u.super_admin, u.status "
            "FROM sys_user_token t JOIN sys_user u ON u.id = t.user_id "
            "WHERE t.token = :token LIMIT 1"
        ),
        {"token": token},
    )
    row = result.mappings().first()
    return dict(row) if row is not None else None