import asyncio from zoneinfo import ZoneInfo import httpx from fastapi import HTTPException, status from sqlalchemy import select from sqlalchemy.exc import IntegrityError from sqlalchemy.ext.asyncio import AsyncSession from app.core.crypto import encrypt_secret, try_decrypt from shared.definitions.mode_features import VALID_MODES, capabilities_for from shared.definitions.retention import ( KEEP_FOREVER, SCAN_RETENTION_DAYS, SCREENSHOT_RETENTION_DAYS, ) from shared.enums.instance import AIProvider from shared.http import get_async_client from shared.logging import get_logger from shared.models.instance_settings import ( SINGLETON_KEY, InstanceSettings, InstanceSettingsRead, InstanceSettingsUpdate, ) from shared.services import scan_admission from shared.services.celery_dispatch import dispatch_scan_admission from shared.services.scan_resolve import MASK from shared.utils.datetime import utc_now from shared.utils.net import validate_public_https_url logger = get_logger(__name__) _TEST_TIMEOUT = 10 _OPENAI_BASE = "https://api.openai.com" _ANTHROPIC_BASE = "https://api.anthropic.com" _ANTHROPIC_VERSION = "2023-06-01" _ANTHROPIC_PING_MODEL = "claude-haiku-4-5" _GOOGLE_BASE = "https://generativelanguage.googleapis.com" _HTTP_OK = 200 _HTTP_UNAUTHORIZED = 401 _VALID_AI_PROVIDERS = frozenset(p.value for p in AIProvider) _TAIL = 4 def _mask_ai_key(key: str) -> str: if len(key) <= _TAIL: return MASK return f"{MASK}{key[-_TAIL:]}" def _validate_public_https_url(raw: str) -> str: validate_public_https_url(raw, label="Azure endpoint") return raw def _apply_limit(settings: InstanceSettings, value: int | None) -> bool: if value is None or value == settings.concurrent_scans: return False settings.concurrent_scans = value return True async def _admit_on_change(changed: bool) -> None: if not changed: return try: await asyncio.to_thread(dispatch_scan_admission) except Exception: logger.warning("scan admission dispatch failed", exc_info=True) def _check_window(label: str, days: int, allowed: tuple[int, ...]) -> None: if days not in allowed: choices = ", ".join(str(d) for d in allowed if d != KEEP_FOREVER) raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=f"{label} retention must be {choices} days, or {KEEP_FOREVER} to keep forever.", ) class InstanceSettingsService: def __init__(self, session: AsyncSession): self.session = session async def get_or_create(self) -> InstanceSettings: result = await self.session.execute( select(InstanceSettings).where( InstanceSettings.singleton_key == SINGLETON_KEY ) ) settings = result.scalar_one_or_none() if settings is None: settings = InstanceSettings(singleton_key=SINGLETON_KEY) self.session.add(settings) try: await self.session.commit() await self.session.refresh(settings) except IntegrityError: await self.session.rollback() result = await self.session.execute( select(InstanceSettings).where( InstanceSettings.singleton_key == SINGLETON_KEY ) ) settings = result.scalar_one() return settings def to_read(self, settings: InstanceSettings) -> InstanceSettingsRead: key = try_decrypt(settings.ai_api_key_encrypted) return InstanceSettingsRead( id=settings.id, singleton_key=settings.singleton_key, instance_name=settings.instance_name, timezone=settings.timezone, mode=settings.mode, onboarding_completed=settings.onboarding_completed, onboarding_completed_at=settings.onboarding_completed_at, onboarding_completed_by=settings.onboarding_completed_by, onboarding_step=settings.onboarding_step, onboarding_state=settings.onboarding_state or {}, scan_history_retention_days=settings.scan_history_retention_days, screenshot_retention_days=settings.screenshot_retention_days, cert_recheck_enabled=settings.cert_recheck_enabled, concurrent_scans=settings.concurrent_scans, ai_enabled=settings.ai_enabled, ai_provider=settings.ai_provider, ai_model=settings.ai_model, ai_api_key_masked=_mask_ai_key(key) if key else None, ai_configured=bool(settings.ai_api_key_encrypted), ai_features=settings.ai_features or {}, capabilities=capabilities_for(settings.mode), created_at=settings.created_at, updated_at=settings.updated_at, ) async def update(self, data: InstanceSettingsUpdate) -> InstanceSettingsRead: settings = await self.get_or_create() if data.instance_name is not None: name = data.instance_name.strip() if not name: raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail="Instance name is required.", ) settings.instance_name = name[:120] if data.timezone is not None: try: ZoneInfo(data.timezone) except Exception as exc: msg = f"Invalid timezone '{data.timezone}'." raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=msg ) from exc settings.timezone = data.timezone if data.mode is not None: if data.mode not in VALID_MODES: msg = f"Invalid mode '{data.mode}'. Must be one of {', '.join(sorted(VALID_MODES))}." raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=msg) settings.mode = data.mode if data.scan_history_retention_days is not None: _check_window( "Scan history", data.scan_history_retention_days, SCAN_RETENTION_DAYS ) settings.scan_history_retention_days = data.scan_history_retention_days if data.screenshot_retention_days is not None: _check_window( "Screenshot", data.screenshot_retention_days, SCREENSHOT_RETENTION_DAYS, ) settings.screenshot_retention_days = data.screenshot_retention_days if data.cert_recheck_enabled is not None: settings.cert_recheck_enabled = data.cert_recheck_enabled limit_changed = _apply_limit(settings, data.concurrent_scans) if data.ai_enabled is not None: settings.ai_enabled = data.ai_enabled if data.ai_provider is not None: if data.ai_provider not in _VALID_AI_PROVIDERS: msg = ( f"Invalid ai_provider '{data.ai_provider}'. " f"Must be one of {', '.join(sorted(_VALID_AI_PROVIDERS))}." ) raise HTTPException(status_code=status.HTTP_400_BAD_REQUEST, detail=msg) settings.ai_provider = data.ai_provider if data.ai_model is not None: settings.ai_model = data.ai_model if data.ai_features is not None: settings.ai_features = dict(data.ai_features) if data.ai_api_key is not None and MASK not in data.ai_api_key: settings.ai_api_key_encrypted = ( encrypt_secret(data.ai_api_key) if data.ai_api_key else None ) settings.updated_at = utc_now() await self.session.commit() await self.session.refresh(settings) await _admit_on_change(limit_changed) return await self.with_limit(self.to_read(settings)) async def with_limit(self, read: InstanceSettingsRead) -> InstanceSettingsRead: try: read.concurrent_scans_auto = await asyncio.to_thread( scan_admission.automatic_limit ) except Exception: logger.warning("automatic scan limit unreadable", exc_info=True) return read async def test_ai( self, provider: str, model: str | None, api_key: str | None ) -> dict: if not api_key or MASK in api_key: api_key = try_decrypt((await self.get_or_create()).ai_api_key_encrypted) if not api_key: return {"success": False, "message": "No API key provided."} try: return await self._ping(provider, model, api_key) except (httpx.HTTPError, ValueError, KeyError) as exc: return {"success": False, "message": f"Connection failed: {exc}"} async def _ping(self, provider: str, model: str | None, api_key: str) -> dict: async with get_async_client(timeout=_TEST_TIMEOUT) as client: if provider == AIProvider.OPENAI.value: return await self._ping_openai(client, api_key) if provider == AIProvider.ANTHROPIC.value: return await self._ping_anthropic(client, api_key) if provider == AIProvider.GOOGLE.value: return await self._ping_google(client, api_key) if provider == AIProvider.AZURE_OPENAI.value: return await self._ping_azure(client, model, api_key) return {"success": False, "message": f"Unsupported provider '{provider}'."} async def _ping_openai(self, client: httpx.AsyncClient, api_key: str) -> dict: resp = await client.get( f"{_OPENAI_BASE}/v1/models", headers={"Authorization": f"Bearer {api_key}"}, ) return self._result(resp, "OpenAI") async def _ping_anthropic(self, client: httpx.AsyncClient, api_key: str) -> dict: resp = await client.post( f"{_ANTHROPIC_BASE}/v1/messages", headers={ "x-api-key": api_key, "anthropic-version": _ANTHROPIC_VERSION, }, json={ "model": _ANTHROPIC_PING_MODEL, "max_tokens": 1, "messages": [{"role": "user", "content": "ping"}], }, ) return self._result(resp, "Anthropic") async def _ping_google(self, client: httpx.AsyncClient, api_key: str) -> dict: resp = await client.get( f"{_GOOGLE_BASE}/v1beta/models", params={"key": api_key}, ) return self._result(resp, "Google") async def _ping_azure( self, client: httpx.AsyncClient, model: str | None, api_key: str ) -> dict: if not model: return { "success": False, "message": "Azure OpenAI requires an endpoint URL in the model field.", } url = await asyncio.to_thread(_validate_public_https_url, model.rstrip("/")) resp = await client.get( url, headers={"api-key": api_key}, follow_redirects=False, ) ok = resp.status_code != _HTTP_UNAUTHORIZED if ok: return {"success": True, "message": "Azure OpenAI reachable."} return {"success": False, "message": "Azure OpenAI authentication failed."} def _result(self, resp: httpx.Response, name: str) -> dict: if resp.status_code == _HTTP_OK: return {"success": True, "message": f"{name} connection successful."} if resp.status_code == _HTTP_UNAUTHORIZED: return {"success": False, "message": f"{name} authentication failed."} return { "success": False, "message": f"{name} returned status {resp.status_code}.", }