risk_analyzer.py

"""
risk_analyzer.py
================
Computes contract health metrics and an attack surface score
from stress test results (the __METRICS__ JSON block in stdout).
"""

import json
import re
import logging
from typing import Any, Dict, List, Optional

logger = logging.getLogger(__name__)

# ─────────────────────────────────────────────────────────────────── #
# WEIGHTED ATTACK SURFACE SCORE                                        #
# risk_score = sum(weight_i * normalized_metric_i), capped at 100    #
# ─────────────────────────────────────────────────────────────────── #

_WEIGHTS = {
    "liquidity_risk":    0.25,
    "oracle_risk":       0.20,
    "governance_risk":   0.20,
    "flash_loan_risk":   0.15,
    "code_vuln_risk":    0.20,
}

_SEVERITY_SCORES = {
    "CRITICAL": 1.0,
    "HIGH":     0.75,
    "MEDIUM":   0.5,
    "LOW":      0.25,
}


def compute_risk_score(
    liquidity_drain_pct: float,
    price_impact_pct:    float,
    oracle_lag_detected: bool,
    flash_loan_used:     bool,
    governance_attacked: bool,
    vulnerabilities:     List[Dict[str, str]],
) -> Dict[str, Any]:
    """
    Compute the weighted attack surface score and component breakdown.

    All input metrics are 0–100 (percentages) or booleans.

    Returns:
        {
          "total_score":       float (0–100),
          "risk_level":        "LOW" | "MEDIUM" | "HIGH" | "CRITICAL",
          "components":        {component: score},
          "insolvency_prob":   float (0–100),
          "recommendations":   [str],
        }
    """
    # Normalize each component to [0, 1]
    liquidity_risk  = min(abs(liquidity_drain_pct) / 100, 1.0)
    oracle_risk     = min(abs(price_impact_pct)    / 80,  1.0) if oracle_lag_detected else min(abs(price_impact_pct) / 80, 1.0) * 0.6
    governance_risk = 1.0 if governance_attacked else 0.0
    flash_loan_risk = 1.0 if flash_loan_used else 0.0

    # Code vulnerability score from static analysis
    vuln_score = 0.0
    if vulnerabilities:
        total_weight = sum(_SEVERITY_SCORES.get(v.get("severity", "LOW"), 0.25) for v in vulnerabilities)
        vuln_score   = min(total_weight / (len(vulnerabilities) * 1.5), 1.0)

    components = {
        "liquidity_risk":  round(liquidity_risk  * 100, 1),
        "oracle_risk":     round(oracle_risk      * 100, 1),
        "governance_risk": round(governance_risk  * 100, 1),
        "flash_loan_risk": round(flash_loan_risk  * 100, 1),
        "code_vuln_risk":  round(vuln_score       * 100, 1),
    }

    total = (
        _WEIGHTS["liquidity_risk"]  * liquidity_risk  +
        _WEIGHTS["oracle_risk"]     * oracle_risk      +
        _WEIGHTS["governance_risk"] * governance_risk  +
        _WEIGHTS["flash_loan_risk"] * flash_loan_risk  +
        _WEIGHTS["code_vuln_risk"]  * vuln_score
    ) * 100

    total = round(min(total, 100), 1)

    # Risk level
    if total >= 75:
        risk_level = "CRITICAL"
    elif total >= 50:
        risk_level = "HIGH"
    elif total >= 25:
        risk_level = "MEDIUM"
    else:
        risk_level = "LOW"

    # Insolvency probability (heuristic)
    insolvency_prob = min(
        liquidity_drain_pct * 1.2 +
        (abs(price_impact_pct) * 0.5 if price_impact_pct < 0 else 0) +
        (30 if flash_loan_used else 0),
        100
    )

    # Recommendations
    recs = _generate_recommendations(
        liquidity_risk, oracle_risk, governance_risk, flash_loan_risk, vulnerabilities
    )

    return {
        "total_score":       total,
        "risk_level":        risk_level,
        "components":        components,
        "weights":           _WEIGHTS,
        "insolvency_prob":   round(max(insolvency_prob, 0), 1),
        "recommendations":   recs,
    }


def _generate_recommendations(
    liq: float, oracle: float, gov: float, flash: float,
    vulns: List[Dict]
) -> List[str]:
    recs = []
    if liq > 0.5:
        recs.append("Implement withdrawal rate limiting or circuit breakers to prevent rapid liquidity drain.")
    if oracle > 0.4:
        recs.append("Use a TWAP oracle (e.g., Uniswap V3 TWAP) instead of a spot price oracle to reduce manipulation surface.")
    if gov > 0:
        recs.append("Add a governance timelock (≥48 h) and quorum requirements to prevent flash-loan governance attacks.")
    if flash > 0:
        recs.append("Add reentrancy guards (ReentrancyGuard) and ensure all state updates happen before external calls.")
    for v in vulns:
        sev = v.get("severity", "LOW")
        if sev in ("CRITICAL", "HIGH"):
            recs.append(f"[{sev}] {v.get('description', v.get('type', 'Fix vulnerability'))}")
    if not recs:
        recs.append("No critical issues detected. Continue regular audits and monitoring.")
    return recs[:8]  # cap at 8 recommendations


# ─────────────────────────────────────────────────────────────────── #
# PARSE SCRIPT OUTPUT                                                  #
# ─────────────────────────────────────────────────────────────────── #

def parse_script_output(stdout: str) -> Optional[Dict[str, Any]]:
    """
    Extract the __METRICS__ JSON block from Hardhat script stdout.
    The script emits: __METRICS__:{...json...}
    """
    m = re.search(r'__METRICS__:\s*(\{.*\})', stdout, re.DOTALL)
    if not m:
        return None
    try:
        return json.loads(m.group(1))
    except json.JSONDecodeError as e:
        logger.warning("Failed to parse __METRICS__ JSON: %s", e)
        return None


def full_risk_report(
    script_stdout:     str,
    scenario:          Dict[str, Any],
    contract_analysis: Optional[Dict[str, Any]] = None,
) -> Dict[str, Any]:
    """
    Build complete risk report combining script execution metrics +
    static analysis + weighted score.

    Args:
        script_stdout:     Raw stdout from `npx hardhat run ...`
        scenario:          Structured scenario dict
        contract_analysis: ContractAnalyzer output (optional)

    Returns:
        Complete risk report dict suitable for frontend dashboard.
    """
    raw_metrics = parse_script_output(script_stdout) or {}
    vulns       = (contract_analysis or {}).get("vulnerabilities", [])
    risk_focus  = scenario.get("risk_focus", [])

    liq_drain    = float(raw_metrics.get("liquidityDrainPct",   0))
    price_impact = float(raw_metrics.get("priceImpactPct",      0))
    oracle_lag   = "oracle_manipulation" in risk_focus
    flash_loan   = "flash_loan"          in risk_focus
    governance   = "governance"          in risk_focus

    score_data = compute_risk_score(
        liquidity_drain_pct  = liq_drain,
        price_impact_pct     = price_impact,
        oracle_lag_detected  = oracle_lag,
        flash_loan_used      = flash_loan,
        governance_attacked  = governance,
        vulnerabilities      = vulns,
    )

    # Contract health metrics
    health = _compute_health_metrics(raw_metrics, contract_analysis)

    return {
        "scenario_title":    scenario.get("title", "Stress Test"),
        "scenario_desc":     scenario.get("description", ""),
        "risk_score":        score_data,
        "health_metrics":    health,
        "raw_metrics":       raw_metrics,
        "vulnerabilities":   vulns,
        "risk_focus":        risk_focus,
        "script_output_len": len(script_stdout),
        "execution_success": "__METRICS__:" in script_stdout,
    }


def _compute_health_metrics(
    raw: Dict[str, Any],
    analysis: Optional[Dict[str, Any]]
) -> Dict[str, Any]:
    """Compute protocol health indicators from execution metrics."""
    liq_drain    = float(raw.get("liquidityDrainPct",   0))
    price_impact = float(raw.get("priceImpactPct",      0))

    liquidity_ratio       = max(1 - liq_drain / 100, 0)
    collateral_ratio      = max(1 - abs(price_impact) / 100 * 1.3, 0)
    token_price_impact    = abs(price_impact)
    reentrancy_exposure   = any(
        v.get("type") == "Reentrancy"
        for v in ((analysis or {}).get("vulnerabilities", []))
    )
    oracle_dep_risk       = any(
        "Oracle" in v.get("type", "")
        for v in ((analysis or {}).get("vulnerabilities", []))
    )

    return {
        "liquidity_ratio":          round(liquidity_ratio * 100, 1),
        "collateralization_ratio":  round(collateral_ratio * 100, 1),
        "token_price_impact_pct":   round(token_price_impact, 2),
        "reentrancy_exposure":      reentrancy_exposure,
        "oracle_dependency_risk":   oracle_dep_risk,
        "total_functions":          len((analysis or {}).get("functions", [])),
        "total_vulnerabilities":    len((analysis or {}).get("vulnerabilities", [])),
    }