martynattakit/CodeSentinel-CWE_Classification

1"""2api/main.py3FastAPI application — exposes the vulnerability classification pipeline4as an HTTP API. Single endpoint: POST /classify5 6Run locally:7    uvicorn api.main:app --reload --port 80008 9Run in HF Spaces:10    uvicorn api.main:app --host 0.0.0.0 --port 786011"""12 13from __future__ import annotations14import time15from contextlib import asynccontextmanager16from typing import Optional17 18from fastapi import FastAPI, HTTPException, Request19from fastapi.middleware.cors import CORSMiddleware20from fastapi.responses import JSONResponse21from pydantic import BaseModel, Field22 23from pipeline.router import get_router24 25# ── Lifespan — warm up classifier on startup ──────────────────────────────────26# RoBERTa (125MB) loads fast — warm it up at startup so first request isn't slow27# Qwen (5GB) is lazy-loaded on first code input — too large to preload28 29@asynccontextmanager30async def lifespan(app: FastAPI):31    print("[API] Warming up RoBERTa classifier...")32    router = get_router()33    router._get_classifier()  # preload RoBERTa only34    print("[API] Ready.")35    yield36    print("[API] Shutting down.")37 38 39# ── App ───────────────────────────────────────────────────────────────────────40 41app = FastAPI(42    title="CodeSentinel API",43    description=(44        "Vulnerability classification API. "45        "Classifies code snippets and CVE descriptions into CWE categories, "46        "with optional ATLAS AI/ML attack pattern matching."47    ),48    version="0.1.0",49    lifespan=lifespan,50)51 52# Allow frontend (HF Spaces, localhost) to call the API53app.add_middleware(54    CORSMiddleware,55    allow_origins=["*"],  # tighten this in production56    allow_methods=["POST", "GET"],57    allow_headers=["*"],58)59 60 61# ── Schemas ───────────────────────────────────────────────────────────────────62 63class ClassifyRequest(BaseModel):64    input: str = Field(65        ...,66        min_length=10,67        max_length=8000,68        description="Raw input — code snippet, CVE description, or bug report.",69        examples=[70            "def get_user(name): return db.execute('SELECT * FROM users WHERE name=' + name)"71        ],72    )73 74 75class CWEPrediction(BaseModel):76    cwe_id:      str77    description: str78    severity:    str79    confidence:  float80 81 82class ATLASMatch(BaseModel):83    atlas_id:        str84    technique:       str85    tactic:          str86    confidence:      str87    matched_signals: list[str]88    description:     str89    mitigations:     list[str]90    real_world:      Optional[str]91    reasoning:       str92 93 94class ClassifyResponse(BaseModel):95    # Primary result96    cwe_id:       str97    cwe_name:     str98    severity:     str99    confidence:   float100 101    # Explanation (Qwen's structured description or original input)102    description:  str103 104    # Top-3 alternatives (excluding top-1)105    alternatives: list[CWEPrediction]106 107    # ATLAS match — None if no AI/ML signals detected108    atlas_match:  Optional[ATLASMatch]109 110    # Metadata111    input_type:   str   # "code" or "text"112    warning:      Optional[str]113    elapsed_s:    float114 115 116class HealthResponse(BaseModel):117    status:  str118    version: str119    models:  dict120 121 122class ErrorResponse(BaseModel):123    error:   str124    detail:  Optional[str]125 126 127@app.get("/health", response_model=HealthResponse)128async def health():129    """Health check — returns model load status."""130    router = get_router()131    return {132        "status":  "ok",133        "version": "0.1.0",134        "models": {135            "roberta":       "loaded" if router._classifier is not None else "not loaded",136            "qwen":          "loaded" if router._code_analyzer is not None else "not loaded (lazy)",137            "atlas_matcher": "loaded" if router._atlas_matcher is not None else "not loaded (lazy)",138        }139    }140 141 142@app.post(143    "/classify",144    response_model=ClassifyResponse,145    responses={146        400: {"model": ErrorResponse, "description": "Invalid input"},147        500: {"model": ErrorResponse, "description": "Internal classification error"},148    },149)150async def classify(request: ClassifyRequest):151    """152    Classify a vulnerability input.153 154    - **Code snippets** → Qwen analyzes → RoBERTa classifies → CWE result155    - **CVE/text descriptions** → RoBERTa classifies directly → CWE result156    - **AI/ML-related inputs** → also runs ATLAS pattern matcher157 158    Returns a unified output card with CWE ID, severity, explanation,159    remediation hints, and optional ATLAS technique match.160    """161    try:162        router = get_router()163        result = router.run(request.input)164    except ValueError as e:165        raise HTTPException(status_code=400, detail=str(e))166    except Exception as e:167        raise HTTPException(168            status_code=500,169            detail=f"Classification failed: {str(e)}"170        )171 172    # Build alternatives list173    alternatives = [174        CWEPrediction(175            cwe_id=alt["cwe_id"],176            description=alt["description"],177            severity=alt["severity"],178            confidence=alt["confidence"],179        ) for alt in result.get("alternatives", [])180    ]181 182    # Build ATLAS match if present183    atlas_match = None184    if result.get("atlas_match"):185        atlas_match = ATLASMatch(**result["atlas_match"])186 187    return ClassifyResponse(188        cwe_id=result["cwe_id"],189        cwe_name=result["cwe_name"],190        severity=result["severity"],191        confidence=result["confidence"],192        description=result["description"],193        alternatives=alternatives,194        atlas_match=atlas_match,195        input_type=result["input_type"],196        warning=result.get("warning"),197        elapsed_s=result["elapsed_s"],198    )199 200 201# ── Error handlers ────────────────────────────────────────────────────────────202 203@app.exception_handler(404)204async def not_found_handler(request: Request, exc):205    return JSONResponse(206        status_code=404,207        content={"error": "Not found", "detail": str(request.url)}208    )209 210 211@app.exception_handler(500)212async def server_error_handler(request: Request, exc):213    return JSONResponse(214        status_code=500,215        content={"error": "Internal server error", "detail": "Check server logs."}216    )