from __future__ import annotations import uuid from http import HTTPStatus from typing import TYPE_CHECKING from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from sqlmodel import select from shared.definitions.api_keys import PROVIDER_GROUP_LABELS from shared.definitions.mode_features import provider_allowed from shared.enums.api_key import APIProvider from shared.enums.instance import InstanceMode from shared.models.api_key import ( API_PROVIDER_META, APIKey, APIKeyCreate, APIKeyRead, APIKeyUpdate, ProviderInfo, ) from shared.models.instance_settings import SINGLETON_KEY, InstanceSettings from shared.utils.crypto import encrypt_secret, try_decrypt from shared.utils.datetime import utc_now if TYPE_CHECKING: from fastapi import HTTPException def _http_error(status: HTTPStatus, detail: str) -> HTTPException: from fastapi import HTTPException # noqa: PLC0415 return HTTPException(status_code=status, detail=detail) class APIKeyService: def __init__(self, session: AsyncSession): self.session = session async def _instance_mode(self) -> str: result = await self.session.execute( select(InstanceSettings.mode).where( InstanceSettings.singleton_key == SINGLETON_KEY ) ) return result.scalar_one_or_none() or InstanceMode.BUG_BOUNTY.value def _mask_key(self, key_value: str) -> str: if len(key_value) <= 4: # noqa: PLR2004 return "****" return f"****{key_value[-4:]}" def _to_read(self, api_key: APIKey) -> APIKeyRead: return APIKeyRead( id=api_key.id, provider=api_key.provider, key_value_masked=self._mask_key( try_decrypt(api_key.key_value) or api_key.key_value ), key_meta=api_key.key_meta, is_enabled=api_key.is_enabled, usage_counter=api_key.usage_counter, last_used_at=api_key.last_used_at, last_test_at=api_key.last_test_at, last_test_ok=api_key.last_test_ok, last_test_message=api_key.last_test_message, created_at=api_key.created_at, updated_at=api_key.updated_at, meta=API_PROVIDER_META.get(api_key.provider, {}), ) async def list_keys(self) -> list[APIKeyRead]: result = await self.session.execute(select(APIKey).order_by(APIKey.provider)) return [self._to_read(k) for k in result.scalars().all()] async def list_providers(self) -> list[ProviderInfo]: result = await self.session.execute(select(APIKey)) configured = {k.provider: k for k in result.scalars().all()} providers = [] for provider, meta in API_PROVIDER_META.items(): key = configured.get(provider) providers.append( ProviderInfo( provider=provider, name=meta["name"], description=meta["description"], docs_url=meta["docs_url"], icon=meta.get("icon", "package"), requires_username=meta.get("requires_username", False), group=meta["group"], group_label=PROVIDER_GROUP_LABELS[meta["group"]], configured=key is not None, is_enabled=key.is_enabled if key else False, ) ) mode = await self._instance_mode() return [p for p in providers if provider_allowed(mode, p.provider.value)] async def create_key(self, data: APIKeyCreate) -> APIKeyRead: mode = await self._instance_mode() if not provider_allowed(mode, data.provider.value): raise _http_error( HTTPStatus.FORBIDDEN, f"{data.provider.value} requires Bug bounty mode.", ) if API_PROVIDER_META.get(data.provider, {}).get("requires_username"): username = (data.key_meta or {}).get("username") if not username or not str(username).strip(): raise _http_error( HTTPStatus.BAD_REQUEST, f"{data.provider.value} requires a username.", ) existing = await self.session.execute( select(APIKey).where( APIKey.provider == data.provider, ) ) if existing.scalar_one_or_none(): raise _http_error( HTTPStatus.CONFLICT, f"An API key for {data.provider.value} exists.", ) api_key = APIKey( provider=data.provider, key_value=encrypt_secret(data.key_value), key_meta=data.key_meta, ) self.session.add(api_key) try: await self.session.commit() except IntegrityError as e: await self.session.rollback() raise _http_error( HTTPStatus.CONFLICT, f"An API key for {data.provider.value} exists.", ) from e await self.session.refresh(api_key) return self._to_read(api_key) async def update_key(self, key_id: str, data: APIKeyUpdate) -> APIKeyRead: api_key = await self._get_key_or_404(key_id) if data.key_value is not None: api_key.key_value = encrypt_secret(data.key_value) api_key.last_test_at = None api_key.last_test_ok = None api_key.last_test_message = None if data.is_enabled is not None: api_key.is_enabled = data.is_enabled if data.key_meta is not None: api_key.key_meta = data.key_meta api_key.updated_at = utc_now() self.session.add(api_key) await self.session.commit() await self.session.refresh(api_key) return self._to_read(api_key) async def record_test(self, api_key: APIKey, ok: bool, message: str) -> None: api_key.last_test_at = utc_now() api_key.last_test_ok = ok api_key.last_test_message = message[:500] self.session.add(api_key) await self.session.commit() async def delete_key(self, key_id: str) -> None: api_key = await self._get_key_or_404(key_id) await self.session.delete(api_key) await self.session.commit() async def get_key_for_provider(self, provider: APIProvider) -> str | None: result = await self.session.execute( select(APIKey).where( APIKey.provider == provider, APIKey.is_enabled.is_(True), ) ) api_key = result.scalar_one_or_none() if not api_key: return None return try_decrypt(api_key.key_value) or api_key.key_value async def increment_usage(self, provider: APIProvider) -> None: result = await self.session.execute( select(APIKey).where( APIKey.provider == provider, ) ) api_key = result.scalar_one_or_none() if api_key: api_key.usage_counter += 1 api_key.last_used_at = utc_now() self.session.add(api_key) await self.session.commit() async def disable_key(self, provider: APIProvider) -> None: result = await self.session.execute( select(APIKey).where( APIKey.provider == provider, ) ) api_key = result.scalar_one_or_none() if api_key: api_key.is_enabled = False api_key.updated_at = utc_now() self.session.add(api_key) await self.session.commit() async def _get_key_or_404(self, key_id: str) -> APIKey: try: uuid_id = uuid.UUID(key_id) except ValueError as e: raise _http_error(HTTPStatus.BAD_REQUEST, "Invalid key ID") from e result = await self.session.execute(select(APIKey).where(APIKey.id == uuid_id)) api_key = result.scalar_one_or_none() if not api_key: raise _http_error(HTTPStatus.NOT_FOUND, "API key not found") return api_key