mihir2007/Cyber-Risk

1"""2optimizer.py3------------4Budget-Constrained Security Investment Optimizer.5 6Formulates and solves a 0/1 knapsack-style Mixed-Integer Linear Program7(MILP) with PuLP: given a budget B (₹) and a catalogue of candidate8`SecurityControl` rows, choose the subset of controls that maximizes total9financial risk reduction (ΔEAL) without exceeding the budget.10 11Because the underlying likelihood model in `quant_engine.py` combines12control efficacies *multiplicatively* (stacking survival probabilities),13the true joint ΔEAL of a control set is not a simple linear sum of14per-control effects. To keep the optimization tractable as a linear15program while still respecting this non-linearity, we:16 17  1. Compute each control's *standalone* marginal ΔEAL (its effect if it18     were the only control deployed).19  2. Apply a "marginal efficiency discount factor" to controls that20     overlap in scope with many other candidates targeting the same asset21     tier -- an overlap-based proxy for diminishing returns, since several22     controls hardening the same tier will realize less *combined* benefit23     than the sum of their standalone effects.24  3. Solve the MILP on these discounted, linear coefficients.25  4. Re-run the full Monte Carlo simulation once more with the *actual*26     selected control set turned on together, to report the true27     (non-linear) projected EAL and ROSI -- the MILP's linear objective is28     used purely to *select* a good portfolio, not to report the final29     numbers.30"""31 32from __future__ import annotations33 34from dataclasses import dataclass35 36import pulp37from sqlalchemy.orm import Session38 39from models import Asset, SecurityControl40from quant_engine import MonteCarloRiskEngine, N_ITERATIONS41 42# Discount strength: each additional overlapping control (same target tier,43# or "All"-scoped) reduces this control's linear objective coefficient by44# roughly this fraction, modeling saturating security returns.45OVERLAP_DISCOUNT_RATE = 0.1246 47# In-memory caches to make interactive optimization instantaneous (<5ms)48_marginal_cache: dict[str, tuple[float, dict[str, ControlMarginalValue]]] = {}49_joint_simulation_cache: dict[tuple[str, ...], float] = {}50 51 52def clear_optimizer_cache() -> None:53    """Flushes cached marginal values and joint simulations when catalogue/assets change."""54    _marginal_cache.clear()55    _joint_simulation_cache.clear()56 57 58def _get_catalogue_key(assets: list[Asset], candidate_controls: list[SecurityControl]) -> str:59    asset_sig = tuple(sorted((a.id, getattr(a, 'tier', '')) for a in assets))60    ctrl_sig = tuple(sorted((c.code, c.cost_inr, c.likelihood_reduction) for c in candidate_controls))61    return f"{hash(asset_sig)}_{hash(ctrl_sig)}"62 63 64@dataclass65class ControlMarginalValue:66    control: SecurityControl67    standalone_reduction_inr: float68    overlap_count: int69    discount_factor: float70    adjusted_value_inr: float71 72 73@dataclass74class OptimizationResult:75    selected_controls: list[SecurityControl]76    control_marginal_values: dict[str, ControlMarginalValue]77    total_spent_inr: float78    remaining_budget_inr: float79    baseline_eal_inr: float80    projected_eal_inr: float81    net_risk_reduction_inr: float82    rosi_percent: float83 84 85def _scope_overlaps(a: SecurityControl, b: SecurityControl) -> bool:86    """True if two controls' target tiers meaningfully overlap in scope."""87    if a.code == b.code:88        return False89    if a.target_tier == "All" or b.target_tier == "All":90        return True91    return a.target_tier == b.target_tier92 93 94def _get_or_compute_marginal_values(95    assets: list[Asset],96    candidate_controls: list[SecurityControl],97) -> tuple[float, dict[str, ControlMarginalValue]]:98    """99    Computes (and caches) each candidate control's standalone marginal ΔEAL100    and overlap-based diminishing-returns discount. Because candidate controls101    and assets are invariant with respect to the user's budget, this is cached102    to make subsequent knapsack solves run in single-digit milliseconds.103    """104    cat_key = _get_catalogue_key(assets, candidate_controls)105    if cat_key in _marginal_cache:106        return _marginal_cache[cat_key]107 108    baseline_engine = MonteCarloRiskEngine(assets=assets, active_controls=[], n_iterations=N_ITERATIONS)109    baseline_result = baseline_engine.run()110    baseline_eal = baseline_result.total_eal_inr111 112    marginal_values: dict[str, ControlMarginalValue] = {}113 114    for control in candidate_controls:115        engine = MonteCarloRiskEngine(assets=assets, active_controls=[control], n_iterations=2500)116        result = engine.run()117        standalone_reduction = max(baseline_eal - result.total_eal_inr, 0.0)118 119        overlap_count = sum(120            1 for other in candidate_controls if _scope_overlaps(control, other)121        )122        discount_factor = 1.0 / (1.0 + OVERLAP_DISCOUNT_RATE * overlap_count)123        adjusted_value = standalone_reduction * discount_factor124 125        marginal_values[control.code] = ControlMarginalValue(126            control=control,127            standalone_reduction_inr=standalone_reduction,128            overlap_count=overlap_count,129            discount_factor=discount_factor,130            adjusted_value_inr=adjusted_value,131        )132 133    _marginal_cache[cat_key] = (baseline_eal, marginal_values)134    return baseline_eal, marginal_values135 136 137def optimize_budget_allocation(138    db: Session,139    assets: list[Asset],140    candidate_controls: list[SecurityControl],141    budget_inr: float,142) -> OptimizationResult:143    """144    Solves the 0/1 knapsack MILP: maximize discounted ΔEAL subject to a145    total cost constraint, then re-simulates the chosen portfolio jointly146    to report true (non-linear) projected EAL and ROSI.147    """148    baseline_eal, marginal_values = _get_or_compute_marginal_values(assets, candidate_controls)149 150    # --- Build and solve the 0/1 knapsack MILP --------------------------- #151    problem = pulp.LpProblem("Security_Investment_Optimization", pulp.LpMaximize)152 153    decision_vars: dict[str, pulp.LpVariable] = {154        control.code: pulp.LpVariable(f"select_{control.code}", cat="Binary")155        for control in candidate_controls156    }157 158    # Objective: maximize total discounted risk reduction (ΔEAL, in ₹).159    problem += pulp.lpSum(160        decision_vars[control.code] * marginal_values[control.code].adjusted_value_inr161        for control in candidate_controls162    ), "Total_Discounted_Risk_Reduction"163 164    # Constraint: total deployment cost must not exceed the budget.165    problem += (166        pulp.lpSum(decision_vars[control.code] * control.cost_inr for control in candidate_controls)167        <= budget_inr,168        "Budget_Constraint",169    )170 171    solver = pulp.PULP_CBC_CMD(msg=False)172    problem.solve(solver)173 174    selected_controls = [175        control176        for control in candidate_controls177        if decision_vars[control.code].value() == 1178    ]179    total_spent = sum(c.cost_inr for c in selected_controls)180 181    # --- Re-simulate the actual joint portfolio for true reported numbers --#182    selected_tuple = tuple(sorted(c.code for c in selected_controls))183    if not selected_controls:184        projected_eal = baseline_eal185    elif selected_tuple in _joint_simulation_cache:186        projected_eal = _joint_simulation_cache[selected_tuple]187    else:188        joint_engine = MonteCarloRiskEngine(assets=assets, active_controls=selected_controls, n_iterations=N_ITERATIONS)189        joint_result = joint_engine.run()190        projected_eal = joint_result.total_eal_inr191        _joint_simulation_cache[selected_tuple] = projected_eal192 193    net_risk_reduction = baseline_eal - projected_eal194    rosi_percent = (195        ((net_risk_reduction - total_spent) / total_spent) * 100.0 if total_spent > 0 else 0.0196    )197 198    return OptimizationResult(199        selected_controls=selected_controls,200        control_marginal_values=marginal_values,201        total_spent_inr=total_spent,202        remaining_budget_inr=budget_inr - total_spent,203        baseline_eal_inr=baseline_eal,204        projected_eal_inr=projected_eal,205        net_risk_reduction_inr=net_risk_reduction,206        rosi_percent=rosi_percent,207    )208 209 210def prewarm_optimizer_cache(db: Session, presets: list[float] | None = None) -> None:211    """Pre-computes marginal values and popular budget presets so interactive UI responds in <5ms."""212    assets = db.query(Asset).all()213    candidate_controls = db.query(SecurityControl).all()214    if not assets or not candidate_controls:215        return216    _get_or_compute_marginal_values(assets, candidate_controls)217    default_presets = presets or [2_000_000, 5_000_000, 10_000_000, 20_000_000, 35_000_000]218    for p in default_presets:219        optimize_budget_allocation(db, assets, candidate_controls, float(p))220 221