A-Team-Security-Infra-Agent-Engineer/backend/app.py

"""
AI Agent - Infrastructure & Security Engineer - Backend API v1.1
FastAPI hub — routes, middleware, startup/shutdown.
"""
import os, json, uuid, time, asyncio, subprocess
from pathlib import Path
from functools import partial

from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware

from config import VERSION, REPORTS, _chat_executor, _START_TIME, log
from database import db, init_db
from utils import running_reports, running_terraform
from services.genai import _TF_RESOURCE_REF_PATH, _populate_tf_valid_types
from routes.cis_engine import _read_cis_version

# ── FastAPI App ──────────────────────────────────────────────────────────────
app = FastAPI(title="AI Agent - Infrastructure & Security Engineer", version=VERSION)
_CORS_ORIGINS = os.environ.get("CORS_ORIGINS", "").split(",") if os.environ.get("CORS_ORIGINS") else []
app.add_middleware(CORSMiddleware,
                   allow_origins=_CORS_ORIGINS or ["http://localhost:8080", "http://127.0.0.1:8080"],
                   allow_credentials=True,
                   allow_methods=["GET", "POST", "PUT", "DELETE", "OPTIONS"],
                   allow_headers=["Content-Type", "Authorization"])

# ── Routers ──────────────────────────────────────────────────────────────────
from routes.auth import router as auth_router
from routes.users import router as users_router
from routes.oci_config import router as oci_config_router
from routes.oci_explorer import router as oci_explorer_router
from routes.genai import router as genai_router
from routes.mcp import router as mcp_router
from routes.adb import router as adb_router
from routes.embeddings import router as embeddings_router
from routes.reports import router as reports_router
from routes.chat import router as chat_router
from routes.terraform import router as terraform_router
from routes.settings import router as settings_router
from routes.cis_engine import router as cis_engine_router

for r in (auth_router, users_router, oci_config_router, oci_explorer_router,
          genai_router, mcp_router, adb_router, embeddings_router,
          reports_router, chat_router, terraform_router, settings_router,
          cis_engine_router):
    app.include_router(r)

# ── Shutdown ─────────────────────────────────────────────────────────────────
@app.on_event("shutdown")
async def shutdown():
    """Graceful shutdown: terminate running subprocesses and thread pool."""
    for rid, proc in list(running_reports.items()):
        try:
            proc.terminate()
            log.info(f"Terminated report process {rid}")
        except Exception as e:
            log.warning(f"Failed to terminate report process {rid}: {e}")
    running_reports.clear()
    for wid, proc in list(running_terraform.items()):
        try:
            proc.terminate()
            log.info(f"Terminated terraform process {wid}")
        except Exception as e:
            log.warning(f"Failed to terminate terraform process {wid}: {e}")
    running_terraform.clear()
    _chat_executor.shutdown(wait=False)
    log.info("Graceful shutdown completed")

# ── Startup ──────────────────────────────────────────────────────────────────
@app.on_event("startup")
async def startup():
    init_db()
    # Apply timezone from app_settings
    try:
        with db() as c:
            tz_row = c.execute("SELECT value FROM app_settings WHERE key='timezone'").fetchone()
        if tz_row and tz_row["value"]:
            os.environ["TZ"] = tz_row["value"]
            time.tzset()
            log.info(f"Timezone set to {tz_row['value']} (from app_settings)")
    except Exception as e:
        log.warning(f"Failed to set timezone from settings: {e}")
    # Mark orphaned "running" reports as failed
    with db() as c:
        orphaned = c.execute("UPDATE reports SET status='failed', error_msg='Interrompido: container reiniciado' WHERE status='running'").rowcount
    if orphaned:
        log.warning(f"Marked {orphaned} orphaned running report(s) as failed")
    # Clean up orphaned .compliance_generating markers
    for marker in REPORTS.glob("*/.compliance_generating"):
        marker.unlink(missing_ok=True)
        log.info(f"Cleaned orphaned compliance marker: {marker.parent.name}")
    # Auto-register CIS Compliance Scanner MCP server
    with db() as c:
        dupes = c.execute("SELECT id FROM mcp_servers WHERE name='CIS Compliance Scanner' ORDER BY created_at ASC").fetchall()
        if len(dupes) > 1:
            keep = dupes[0]["id"]
            for d in dupes[1:]:
                c.execute("DELETE FROM mcp_servers WHERE id=?", (d["id"],))
            log.info(f"Removed {len(dupes)-1} duplicate CIS Compliance Scanner entries, kept {keep}")
        if not dupes:
            mid = str(uuid.uuid4())
            c.execute(
                "INSERT INTO mcp_servers (id, user_id, name, description, server_type, command, args) VALUES (?,?,?,?,?,?,?)",
                (mid, "system", "CIS Compliance Scanner",
                 "OCI CIS Foundations Benchmark 3.0 - 48 CIS checks + 11 OBP checks via MCP",
                 "stdio", "python3", json.dumps(["/app/mcp_cis_server.py"])))
            log.info(f"Auto-registered CIS Compliance Scanner MCP server (id={mid})")
    # Save CIS engine version
    try:
        vi = _read_cis_version()
        with db() as c:
            c.execute("INSERT INTO app_settings (key,value,updated_at) VALUES ('cis_engine_version',?,datetime('now')) ON CONFLICT(key) DO UPDATE SET value=excluded.value, updated_at=excluded.updated_at",
                      (vi["version"],))
        log.info(f"CIS engine version: {vi['version']} ({vi['date']})")
    except Exception as e:
        log.warning(f"Could not read CIS engine version: {e}")
    # Generate TF resource reference if not present
    ref_path = Path(_TF_RESOURCE_REF_PATH)
    if not ref_path.exists():
        try:
            loop = asyncio.get_event_loop()
            result = await loop.run_in_executor(None, partial(
                subprocess.run, ["python", "/app/gen_tf_reference.py"], capture_output=True, text=True, timeout=300))
            if result.returncode == 0 and ref_path.exists():
                log.info(f"Generated TF resource reference at startup: {ref_path.stat().st_size} bytes")
            else:
                log.warning(f"TF reference generation failed at startup: {result.stderr[:200] if result.stderr else 'unknown'}")
        except Exception as e:
            log.warning(f"Could not generate TF reference at startup: {e}")
    # Populate tf_valid_types table
    if ref_path.exists():
        try:
            count = _populate_tf_valid_types()
            if count:
                log.info(f"TF valid types DB populated: {count} entries")
        except Exception as e:
            log.warning(f"Could not populate tf_valid_types: {e}")

    log.info(f"AI Agent v{VERSION} API started")

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="0.0.0.0", port=8000)