import secrets import string from sqlalchemy import func, select from sqlalchemy.ext.asyncio import AsyncSession from app.core.roles import ROLE_AGENT, ROLE_OEM from app.models.card_key import CardKey from app.models.user import User from app.schemas.admin_card_key import CardKeyCreate, CardKeyUpdate DEFAULT_PAGE_SIZE = 20 MAX_PAGE_SIZE = 200 _SERIAL_ALPHABET = string.ascii_uppercase + string.digits _SERIAL_LENGTH = 16 class AdminCardKeyError(Exception): def __init__(self, message: str) -> None: self.message = message super().__init__(message) def _generate_serial() -> str: return "".join(secrets.choice(_SERIAL_ALPHABET) for _ in range(_SERIAL_LENGTH)) async def _unique_serial(db: AsyncSession, preferred: str | None = None) -> str: if preferred: existing = await db.scalar( select(CardKey.id).where(CardKey.serial_number == preferred) ) if existing is not None: raise AdminCardKeyError("序列号已存在") return preferred for _ in range(20): candidate = _generate_serial() existing = await db.scalar( select(CardKey.id).where(CardKey.serial_number == candidate) ) if existing is None: return candidate raise AdminCardKeyError("生成序列号失败,请重试") async def _validate_assign_ids( db: AsyncSession, *, oem_id: int | None, agent_id: int | None, ) -> None: if oem_id is not None and agent_id is not None: raise AdminCardKeyError("OEM 与代理只能指定其一") if oem_id is not None: oem = await db.get(User, oem_id) if oem is None or oem.role_id != ROLE_OEM: raise AdminCardKeyError("所选 OEM 不存在") if agent_id is not None: agent = await db.get(User, agent_id) if agent is None or agent.role_id != ROLE_AGENT: raise AdminCardKeyError("所选代理不存在") def _list_card_keys_filters( *, status: str | None = None, username: str | None = None, oem_id: int | None = None, agent_id: int | None = None, ) -> list: filters = [] if status == "unused": filters.append(CardKey.activated_at.is_(None)) elif status == "used": filters.append(CardKey.activated_at.is_not(None)) if username and username.strip(): filters.append(CardKey.username.ilike(f"%{username.strip()}%")) if oem_id is not None: filters.append(CardKey.oem_id == oem_id) if agent_id is not None: filters.append(CardKey.agent_id == agent_id) return filters async def list_card_keys( db: AsyncSession, *, page: int = 1, page_size: int = DEFAULT_PAGE_SIZE, status: str | None = None, username: str | None = None, oem_id: int | None = None, agent_id: int | None = None, ) -> tuple[list[CardKey], int]: page_size = min(max(page_size, 1), MAX_PAGE_SIZE) page = max(page, 1) offset = (page - 1) * page_size filters = _list_card_keys_filters( status=status, username=username, oem_id=oem_id, agent_id=agent_id, ) count_stmt = select(func.count()).select_from(CardKey) list_stmt = select(CardKey) for clause in filters: count_stmt = count_stmt.where(clause) list_stmt = list_stmt.where(clause) total = await db.scalar(count_stmt) or 0 result = await db.scalars( list_stmt.order_by(CardKey.id.desc()).offset(offset).limit(page_size) ) return list(result.all()), total async def create_card_keys(db: AsyncSession, body: CardKeyCreate) -> list[CardKey]: await _validate_assign_ids(db, oem_id=body.oem_id, agent_id=body.agent_id) created: list[CardKey] = [] for index in range(body.count): serial = await _unique_serial( db, body.serial_number if index == 0 else None, ) card = CardKey( serial_number=serial, duration_days=body.duration_days, remark=body.remark, oem_id=body.oem_id, agent_id=body.agent_id, ) db.add(card) created.append(card) await db.flush() for card in created: await db.refresh(card) return created async def update_card_key( db: AsyncSession, card_id: int, body: CardKeyUpdate, ) -> CardKey: card = await db.get(CardKey, card_id) if card is None: raise AdminCardKeyError("卡密不存在") if card.activated_at is not None: if body.duration_days is not None: raise AdminCardKeyError("已激活卡密不能修改时长") elif body.duration_days is not None: card.duration_days = body.duration_days if body.remark is not None: card.remark = body.remark await db.flush() await db.refresh(card) return card async def delete_card_key(db: AsyncSession, card_id: int) -> None: card = await db.get(CardKey, card_id) if card is None: raise AdminCardKeyError("卡密不存在") if card.activated_at is not None: raise AdminCardKeyError("已激活卡密不能删除") await db.delete(card)