CoolFace
Apppublic

mihir2007/Cyber-Risk

sourceHugging Faceupdated 16d agoView on Hugging Face
0likes
main.py299 linesDownload Raw Back to root
1"""2main.py3-------4FastAPI server exposing the CRQ platform's REST API.5 6Run with:7    uvicorn main:app --reload --port 80008 9Interactive API docs are available at /docs once the server is running.10"""11 12from __future__ import annotations13 14import json15 16from fastapi import Depends, FastAPI, HTTPException17from fastapi.middleware.cors import CORSMiddleware18from sqlalchemy.orm import Session19 20import schemas21from compliance import generate_compliance_report22from database import get_db, init_db23from models import Asset, SecurityControl, SimulationRun24from optimizer import (25    optimize_budget_allocation,26    prewarm_optimizer_cache,27    clear_optimizer_cache,28)29from quant_engine import MonteCarloRiskEngine, run_enterprise_simulation30from seeder import seed_database31 32app = FastAPI(33    title="AI-Powered Continuous Cyber Risk Quantification (CRQ) Platform",34    description=(35        "Quantifies cyber risk in INR via Monte Carlo simulation, maps exposure to "36        "RBI CSF / SEBI CSCRF, and optimizes security investment under budget constraints."37    ),38    version="1.0.0",39)40 41app.add_middleware(42    CORSMiddleware,43    allow_origins=["*"],44    allow_credentials=True,45    allow_methods=["*"],46    allow_headers=["*"],47)48 49 50@app.on_event("startup")51def on_startup() -> None:52    """Creates database tables on first launch if they do not already exist."""53    init_db()54    import threading55    def _warm():56        try:57            from database import SessionLocal58            with SessionLocal() as db:59                prewarm_optimizer_cache(db)60        except Exception:61            pass62    threading.Thread(target=_warm, daemon=True).start()63 64 65# --------------------------------------------------------------------------- #66# Seeding67# --------------------------------------------------------------------------- #68@app.post("/api/seed", tags=["Admin"])69def seed(db: Session = Depends(get_db)) -> dict:70    """Seeds the database with synthetic enterprise assets, CVEs, and controls."""71    clear_optimizer_cache()72    summary = seed_database(db)73    return {"status": "ok", "inserted": summary}74 75 76# --------------------------------------------------------------------------- #77# Assets78# --------------------------------------------------------------------------- #79@app.get("/api/assets", response_model=list[schemas.AssetRead], tags=["Assets"])80def list_assets(db: Session = Depends(get_db)) -> list[Asset]:81    """Returns all assets with their vulnerabilities and current estimated EAL."""82    assets = db.query(Asset).all()83    if not assets:84        return []85 86    active_controls = db.query(SecurityControl).filter(SecurityControl.is_active.is_(True)).all()87    engine = MonteCarloRiskEngine(assets=assets, active_controls=active_controls)88    result = engine.run()89    eal_by_asset_id = {r.asset_id: r.eal_inr for r in result.per_asset_results}90 91    output: list[schemas.AssetRead] = []92    for asset in assets:93        asset_read = schemas.AssetRead.model_validate(asset)94        asset_read.estimated_eal_inr = eal_by_asset_id.get(asset.id, 0.0)95        output.append(asset_read)96    return output97 98 99# --------------------------------------------------------------------------- #100# Enterprise Quantification101# --------------------------------------------------------------------------- #102@app.get(103    "/api/quantification/enterprise",104    response_model=schemas.EnterpriseRiskSummary,105    tags=["Quantification"],106)107def enterprise_quantification(db: Session = Depends(get_db)) -> schemas.EnterpriseRiskSummary:108    """109    Runs the full enterprise Monte Carlo simulation, logs the result to110    `SimulationRun`, and returns EAL, VaR 95/99, top-5 riskiest assets,111    regulatory fine exposure, and the loss exceedance curve.112    """113    assets = db.query(Asset).all()114    if not assets:115        raise HTTPException(status_code=400, detail="No assets found. Call POST /api/seed first.")116 117    active_controls = db.query(SecurityControl).filter(SecurityControl.is_active.is_(True)).all()118    result = run_enterprise_simulation(db, assets, active_controls, persist=True)119 120    top_5 = sorted(result.per_asset_results, key=lambda r: r.eal_inr, reverse=True)[:5]121    top_5_contributions = [122        schemas.AssetRiskContribution(123            asset_id=r.asset_id,124            hostname=r.hostname,125            tier=r.tier,126            asset_type=r.asset_type,127            eal_inr=round(r.eal_inr, 2),128            annual_event_frequency=r.annual_event_frequency,129        )130        for r in top_5131    ]132 133    engine = MonteCarloRiskEngine(assets=assets, active_controls=active_controls)134    curve_points = engine.loss_exceedance_curve(result.aggregate_loss_samples)135 136    latest_run = (137        db.query(SimulationRun).order_by(SimulationRun.timestamp.desc()).first()138    )139 140    return schemas.EnterpriseRiskSummary(141        total_eal_inr=round(result.total_eal_inr, 2),142        var_95_inr=round(result.var_95_inr, 2),143        var_99_inr=round(result.var_99_inr, 2),144        top_5_riskiest_assets=top_5_contributions,145        total_regulatory_fine_exposure_inr=round(result.total_regulatory_fine_exposure_inr, 2),146        loss_exceedance_curve=[147            schemas.LossExceedancePoint(return_period_years=t, loss_inr=round(loss, 2))148            for t, loss in curve_points149        ],150        simulation_iterations=result.iterations,151        generated_at=latest_run.timestamp if latest_run else __import__("datetime").datetime.utcnow(),152        unauthorized_subprocessor_count=result.unauthorized_subprocessor_count,153        shadow_leakage_exposure_inr=round(result.shadow_leakage_exposure_inr, 2),154        identified_leak_vectors=result.identified_leak_vectors,155    )156 157 158@app.get(159    "/api/lineage/flows",160    response_model=list[schemas.DataLineageFlowRead],161    tags=["Data Lineage & Provenance"],162)163def get_data_lineage_flows(db: Session = Depends(get_db)) -> list[schemas.DataLineageFlowRead]:164    """Returns all data provenance and transitive sub-processor delegation flows."""165    assets = db.query(Asset).all()166    if not assets:167        raise HTTPException(status_code=400, detail="No assets found. Call POST /api/seed first.")168 169    active_controls = db.query(SecurityControl).filter(SecurityControl.is_active.is_(True)).all()170    from graph_engine import AttackGraphEngine171    engine = AttackGraphEngine(assets=assets, session=db, active_controls=active_controls)172    flows, _ = engine.trace_transitive_data_lineage()173    return flows174 175 176@app.get(177    "/api/quantification/loss-exceedance",178    response_model=list[schemas.LossExceedancePoint],179    tags=["Quantification"],180)181def loss_exceedance_curve(db: Session = Depends(get_db)) -> list[schemas.LossExceedancePoint]:182    """Returns Loss Exceedance Curve data points for frontend plotting."""183    assets = db.query(Asset).all()184    if not assets:185        raise HTTPException(status_code=400, detail="No assets found. Call POST /api/seed first.")186 187    active_controls = db.query(SecurityControl).filter(SecurityControl.is_active.is_(True)).all()188    engine = MonteCarloRiskEngine(assets=assets, active_controls=active_controls)189    result = engine.run()190    curve_points = engine.loss_exceedance_curve(result.aggregate_loss_samples)191 192    return [193        schemas.LossExceedancePoint(return_period_years=t, loss_inr=round(loss, 2))194        for t, loss in curve_points195    ]196 197 198# --------------------------------------------------------------------------- #199# Budget Optimization200# --------------------------------------------------------------------------- #201@app.post(202    "/api/optimize/budget",203    response_model=schemas.OptimizationResponse,204    tags=["Optimization"],205)206def optimize_budget(207    request: schemas.BudgetOptimizationRequest, db: Session = Depends(get_db)208) -> schemas.OptimizationResponse:209    """210    Solves the 0/1 knapsack MILP for the given INR budget, activates the211    selected controls in the database, logs the run to `SimulationRun`,212    and returns the chosen portfolio with ROSI.213    """214    assets = db.query(Asset).all()215    if not assets:216        raise HTTPException(status_code=400, detail="No assets found. Call POST /api/seed first.")217 218    candidate_controls = db.query(SecurityControl).all()219    if not candidate_controls:220        raise HTTPException(status_code=400, detail="No security controls found. Call POST /api/seed first.")221 222    result = optimize_budget_allocation(db, assets, candidate_controls, request.budget_inr)223 224    selected_codes = {c.code for c in result.selected_controls}225    for control in candidate_controls:226        control.is_active = control.code in selected_codes227    db.commit()228 229    run_record = SimulationRun(230        total_eal_inr=result.projected_eal_inr,231        var_95_inr=0.0,232        var_99_inr=0.0,233        allocated_budget_inr=request.budget_inr,234        selected_controls_json=json.dumps(sorted(selected_codes)),235    )236    db.add(run_record)237    db.commit()238 239    recommendations = [240        schemas.ControlRecommendation(241            code=c.code,242            name=c.name,243            cost_inr=c.cost_inr,244            likelihood_reduction=c.likelihood_reduction,245            target_tier=c.target_tier,246            marginal_eal_reduction_inr=round(247                result.control_marginal_values[c.code].adjusted_value_inr, 2248            ),249        )250        for c in result.selected_controls251    ]252 253    return schemas.OptimizationResponse(254        budget_inr=request.budget_inr,255        selected_controls=recommendations,256        total_spent_inr=round(result.total_spent_inr, 2),257        remaining_budget_inr=round(result.remaining_budget_inr, 2),258        baseline_eal_inr=round(result.baseline_eal_inr, 2),259        projected_eal_inr=round(result.projected_eal_inr, 2),260        net_risk_reduction_inr=round(result.net_risk_reduction_inr, 2),261        rosi_percent=round(result.rosi_percent, 2),262    )263 264 265# --------------------------------------------------------------------------- #266# Compliance267# --------------------------------------------------------------------------- #268@app.get(269    "/api/compliance/status",270    response_model=schemas.FrameworkComplianceReport,271    tags=["Compliance"],272)273def compliance_status(db: Session = Depends(get_db)) -> schemas.FrameworkComplianceReport:274    """Returns RBI CSF / SEBI CSCRF / NIST CSF 2.0 compliance scores and gaps."""275    return generate_compliance_report(db)276 277 278# --------------------------------------------------------------------------- #279# Simulation History280# --------------------------------------------------------------------------- #281@app.get(282    "/api/simulations/history",283    response_model=list[schemas.SimulationRunRead],284    tags=["Quantification"],285)286def simulation_history(db: Session = Depends(get_db)) -> list[SimulationRun]:287    """Returns all past simulation runs, most recent first."""288    return db.query(SimulationRun).order_by(SimulationRun.timestamp.desc()).all()289 290 291@app.get("/", tags=["Admin"])292def root() -> dict:293    """Basic health-check / welcome route."""294    return {295        "service": "CRQ Platform API",296        "docs": "/docs",297        "status": "running",298    }299