aur-shield/aur_shield/server.py

"""FastAPI server for AUR-Shield."""
from __future__ import annotations

import hashlib
from pathlib import Path
from typing import Any

from fastapi import FastAPI, HTTPException
from fastapi.responses import FileResponse, JSONResponse
from fastapi.staticfiles import StaticFiles

from .config import Config
from .aur_client import aur_sources, aur_info
from .scanner import scan_pkgbuild, ScanVerdict
from .builder import build_package, add_to_repo, BuildResult
from .cache import ScanCache

app = FastAPI(
    title="AUR-Shield",
    description="AI-powered AUR firewall",
    version="0.1.0",
)

# Global config — set by main()
_config: Config | None = None
_scan_cache: ScanCache | None = None


def get_config() -> Config:
    if _config is None:
        raise RuntimeError("Server not initialized")
    return _config


def get_cache() -> ScanCache:
    if _scan_cache is None:
        raise RuntimeError("Cache not initialized")
    return _scan_cache


def _hash_pkgbuild(pkgbuild: str) -> str:
    return hashlib.sha256(pkgbuild.encode()).hexdigest()[:16]


@app.get("/api/status")
async def status() -> dict[str, Any]:
    """Server status and cache info."""
    cfg = get_config()
    cache = get_cache()
    return {
        "status": "running",
        "model": cfg.ollama.model,
        "ollama_url": cfg.ollama.url,
        "repo_dir": cfg.server.repo_dir,
        "cached_scans": len(cache.list_all()),
        "scan_only": cfg.server.scan_only,
    }


@app.get("/api/scan/{package}")
async def scan_package(package: str) -> dict[str, Any]:
    """Scan a package without building."""
    cfg = get_config()
    cache = get_cache()

    # Fetch from AUR
    source = await aur_sources(package)
    if source is None:
        raise HTTPException(404, f"Package '{package}' not found in AUR")

    pkgbuild_hash = _hash_pkgbuild(source.pkgbuild)

    # Check cache
    cached = cache.get(package, pkgbuild_hash)
    if cached:
        return {
            "package": package,
            "cached": True,
            "verdict": cached.verdict,
            "confidence": cached.confidence,
            "findings": cached.findings,
            "reasoning": cached.reasoning,
            "version": source.package.version,
            "scan_only": cfg.server.scan_only,
        }

    # Scan
    result = await scan_pkgbuild(
        package, source.pkgbuild, source.srcinfo,
        cfg.ollama, cfg.security,
    )

    # Cache the result
    cache.put(result, pkgbuild_hash)

    return {
        "package": package,
        "cached": False,
        "verdict": result.verdict.value,
        "confidence": result.confidence,
        "findings": result.findings,
        "reasoning": result.reasoning,
        "sources": result.sources_checked,
        "version": source.package.version,
        "ioc_matches": result.ioc_matches,
        "typosquat_matches": result.typosquat_matches,
        "scan_only": cfg.server.scan_only,
    }


@app.get("/api/build/{package}")
async def build_endpoint(package: str) -> dict[str, Any]:
    """Scan + build + add to repo."""
    cfg = get_config()
    cache = get_cache()

    # 1. Fetch from AUR
    source = await aur_sources(package)
    if source is None:
        raise HTTPException(404, f"Package '{package}' not found in AUR")

    pkgbuild_hash = _hash_pkgbuild(source.pkgbuild)

    # 2. Check cache for scan
    cached = cache.get(package, pkgbuild_hash)
    if cached:
        if cached.verdict == ScanVerdict.MALICIOUS.value:
            raise HTTPException(403, f"Package '{package}' is cached as MALICIOUS")
        verdict = cached.verdict
        findings = cached.findings
        confidence = cached.confidence
    else:
        # 3. Scan
        result = await scan_pkgbuild(
            package, source.pkgbuild, source.srcinfo,
            cfg.ollama, cfg.security,
        )
        cache.put(result, pkgbuild_hash)
        verdict = result.verdict.value
        findings = result.findings
        confidence = result.confidence

    # 4. Block if malicious
    if verdict == ScanVerdict.MALICIOUS.value:
        raise HTTPException(
            403,
            f"Package '{package}' was flagged as MALICIOUS. Findings: {findings}",
        )

    if verdict == ScanVerdict.SUSPICIOUS.value:
        # Allow suspicious but warn
        pass

    # 5. Build (skip if scan_only mode)
    if cfg.server.scan_only:
        return {
            "package": package,
            "verdict": verdict,
            "confidence": confidence,
            "findings": findings,
            "scan_only": True,
            "build_success": False,
            "message": "Scan-only mode — build locally on client",
        }

    # 6. Build
    build_result = build_package(
        source.pkgbuild, package,
        cfg.server, cfg.build,
    )

    if not build_result.success:
        return {
            "package": package,
            "verdict": verdict,
            "confidence": confidence,
            "build_success": False,
            "build_error": build_result.error,
            "build_log": build_result.log[-2000:],
        }

    # 6. Add to repo
    try:
        db_path = add_to_repo(build_result.pkg_files, cfg.server.repo_dir)
    except Exception as e:
        return {
            "package": package,
            "verdict": verdict,
            "confidence": confidence,
            "build_success": True,
            "repo_success": False,
            "repo_error": str(e),
        }

    return {
        "package": package,
        "verdict": verdict,
        "confidence": confidence,
        "findings": findings,
        "build_success": True,
        "repo_success": True,
        "pkg_files": [Path(f).name for f in build_result.pkg_files],
        "repo_db": Path(db_path).name,
    }


@app.get("/api/report/{package}")
async def get_report(package: str) -> dict[str, Any]:
    """Get the last scan report for a package."""
    cache = get_cache()
    entries = cache.list_all()
    for entry in entries:
        if entry["package"] == package:
            return entry
    raise HTTPException(404, f"No scan report for '{package}'")


@app.get("/api/cache")
async def list_cache() -> dict[str, Any]:
    """List all cached scan results."""
    cache = get_cache()
    return {"entries": cache.list_all()}


@app.delete("/api/cache/{package}")
async def clear_cache_entry(package: str) -> dict[str, str]:
    """Clear a cached scan result."""
    cache = get_cache()
    path = cache._cache_path(package)
    if path.exists():
        path.unlink()
        return {"status": "deleted", "package": package}
    raise HTTPException(404, f"No cache for '{package}'")


# Serve the pacman repo files
@app.get("/repo/{filename}")
async def serve_repo_file(filename: str):
    """Serve repo files for pacman."""
    cfg = get_config()
    repo_path = Path(cfg.server.repo_dir) / filename
    if not repo_path.exists():
        raise HTTPException(404, f"File '{filename}' not in repo")
    return FileResponse(repo_path)


def init_server(config: Config) -> None:
    """Initialize the server with config."""
    global _config, _scan_cache
    _config = config
    _scan_cache = ScanCache(
        cache_dir=str(Path(config.server.work_dir) / "cache"),
        ttl_hours=config.cache.ttl_hours,
    )
    # Ensure dirs exist
    Path(config.server.repo_dir).mkdir(parents=True, exist_ok=True)
    Path(config.server.work_dir).mkdir(parents=True, exist_ok=True)