homelab/userbot/panel/backend/app/main.py

from __future__ import annotations

from contextlib import asynccontextmanager
from typing import Annotated

from fastapi import FastAPI, Query, Request
from fastapi.exceptions import RequestValidationError
from fastapi.responses import FileResponse, JSONResponse, Response
from fastapi.staticfiles import StaticFiles

from .auth_service import TelegramAuthService
from .config import settings
from .errors import PanelError
from .kubernetes_service import KubernetesService
from .models import (
    AuthResult,
    CodeSubmit,
    DeleteRequest,
    InstanceSummary,
    PasswordSubmit,
    PhoneStart,
    StringSessionStart,
)


@asynccontextmanager
async def lifespan(app: FastAPI):
    app.state.kubernetes = KubernetesService(settings)
    app.state.telegram = TelegramAuthService(settings.auth_ttl_seconds)
    try:
        app.state.kubernetes.ensure_prerequisites()
    except PanelError as exc:
        print(f'WARNING: userbot prerequisites check failed at startup: {exc.detail}')
    yield
    await app.state.telegram.close()


app = FastAPI(
    title='Userbot Kubernetes Control',
    version='1.0.0',
    lifespan=lifespan,
)


@app.exception_handler(PanelError)
async def panel_error_handler(_request: Request, exc: PanelError) -> JSONResponse:
    return JSONResponse(status_code=exc.status_code, content={'detail': exc.detail})


@app.exception_handler(RequestValidationError)
async def validation_error_handler(
    _request: Request,
    exc: RequestValidationError,
) -> JSONResponse:
    errors = [
        {
            'loc': error.get('loc', ()),
            'msg': error.get('msg', 'Invalid value'),
            'type': error.get('type', 'value_error'),
        }
        for error in exc.errors()
    ]
    return JSONResponse(status_code=422, content={'detail': errors})


def kube(request: Request) -> KubernetesService:
    return request.app.state.kubernetes


def telegram(request: Request) -> TelegramAuthService:
    return request.app.state.telegram


@app.get('/api/health')
def health(request: Request) -> dict[str, object]:
    service = kube(request)
    try:
        service.apps.list_namespaced_deployment(
            settings.namespace,
            limit=1,
        )
    except Exception:
        return {'ok': False, 'kubernetes': False}
    return {'ok': True, 'kubernetes': True}


@app.get('/api/instances', response_model=list[InstanceSummary])
def list_instances(
    request: Request,
    query: str = '',
    status: str = '',
) -> list[InstanceSummary]:
    return kube(request).list_instances(query=query, status=status)


@app.get('/api/instances/{instance_id}', response_model=InstanceSummary)
def get_instance(instance_id: str, request: Request) -> InstanceSummary:
    return kube(request).get_instance(instance_id)


@app.get('/api/instances/{instance_id}/logs')
def get_logs(
    instance_id: str,
    request: Request,
    tail: Annotated[int, Query(ge=1, le=1000)] = 250,
) -> dict[str, str]:
    return {'logs': kube(request).logs(instance_id, tail)}


@app.post('/api/instances/{instance_id}/start', response_model=InstanceSummary)
def start_instance(instance_id: str, request: Request) -> InstanceSummary:
    return kube(request).scale(instance_id, 1)


@app.post('/api/instances/{instance_id}/stop', response_model=InstanceSummary)
def stop_instance(instance_id: str, request: Request) -> InstanceSummary:
    return kube(request).scale(instance_id, 0)


@app.post('/api/instances/{instance_id}/restart', response_model=InstanceSummary)
def restart_instance(instance_id: str, request: Request) -> InstanceSummary:
    return kube(request).restart(instance_id)


@app.post('/api/instances/{instance_id}/delete', status_code=204)
def delete_instance(
    instance_id: str,
    payload: DeleteRequest,
    request: Request,
) -> Response:
    if payload.confirmation != instance_id:
        raise PanelError(422, 'Type the instance id exactly to confirm deletion')
    kube(request).delete(instance_id, delete_data=payload.delete_data)
    return Response(status_code=204)


@app.post('/api/auth/phone/start', response_model=AuthResult)
async def auth_phone_start(
    payload: PhoneStart,
    request: Request,
) -> AuthResult:
    service = kube(request)
    service.assert_available(payload.instance_id)
    flow_id = await telegram(request).start_phone(payload)
    return AuthResult(status='code_required', flow_id=flow_id)


@app.post('/api/auth/phone/{flow_id}/code', response_model=AuthResult)
async def auth_phone_code(
    flow_id: str,
    payload: CodeSubmit,
    request: Request,
) -> AuthResult:
    authorized = await telegram(request).submit_code(flow_id, payload.code)
    if authorized is None:
        return AuthResult(status='password_required', flow_id=flow_id)
    instance = kube(request).provision(authorized.account, authorized.session_string)
    return AuthResult(status='ready', instance=instance)


@app.post('/api/auth/phone/{flow_id}/password', response_model=AuthResult)
async def auth_phone_password(
    flow_id: str,
    payload: PasswordSubmit,
    request: Request,
) -> AuthResult:
    authorized = await telegram(request).submit_password(flow_id, payload.password)
    instance = kube(request).provision(authorized.account, authorized.session_string)
    return AuthResult(status='ready', instance=instance)


@app.post('/api/auth/string-session', response_model=AuthResult)
async def auth_string_session(
    payload: StringSessionStart,
    request: Request,
) -> AuthResult:
    service = kube(request)
    service.assert_available(payload.instance_id)
    authorized = await telegram(request).validate_string_session(payload)
    instance = service.provision(authorized.account, authorized.session_string)
    return AuthResult(status='ready', instance=instance)


static_dir = settings.static_dir
assets_dir = static_dir / 'assets'
if assets_dir.exists():
    app.mount('/assets', StaticFiles(directory=assets_dir), name='assets')


@app.get('/', include_in_schema=False)
def index() -> FileResponse:
    return FileResponse(static_dir / 'index.html')


@app.get('/{path:path}', include_in_schema=False)
def spa_fallback(path: str) -> FileResponse:
    root = static_dir.resolve()
    candidate = (static_dir / path).resolve()
    try:
        candidate.relative_to(root)
    except ValueError:
        return FileResponse(static_dir / 'index.html')
    if candidate.is_file():
        return FileResponse(candidate)
    return FileResponse(static_dir / 'index.html')