script_generator.py

"""
script_generator.py
====================
Converts a structured scenario dict (output of nlp_engine.py or scenarios.py)
into executable Hardhat JavaScript test scripts that run against the local
Hardhat node on http://127.0.0.1:8545.
"""

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


# ─────────────────────────────────────────────────────────────────── #
# TEMPLATE HELPERS                                                     #
# ─────────────────────────────────────────────────────────────────── #

def _eth(amount: float) -> str:
    """Format ETH amount as ethers.js parseEther string."""
    return f'ethers.parseEther("{amount:.6f}")'


def _usdc(amount: float) -> str:
    """Format USDC as parseUnits with 6 decimals."""
    return f'ethers.parseUnits("{amount:.2f}", 6)'


def _action_to_js(action: Dict[str, Any], idx: int, contract_functions: List[str]) -> str:
    """Convert a single action dict to a Hardhat JS snippet."""
    atype   = action.get("type", "swap")
    token   = action.get("token", "ETH")
    amount  = float(action.get("amount") or 1000)
    speed   = action.get("speed", "instant")
    target  = action.get("target", "pool")

    delay_ms = {"instant": 0, "rapid": 500, "slow": 3000, "gradual": 1000}.get(speed, 0)
    delay_code = f"  await sleep({delay_ms});\n" if delay_ms > 0 else ""

    amount_str = _eth(amount) if token in ("ETH", "TOKEN") else _usdc(amount)

    lines = [f"  // Action {idx + 1}: {atype} {amount} {token} → {target}"]

    if atype == "deposit":
        lines.append(f"  await token.connect(actor).approve(vault.target, {amount_str});")
        lines.append(f"  await vault.connect(actor).deposit({amount_str});")
    elif atype == "withdraw":
        lines.append(f"  await vault.connect(actor).withdraw({amount_str});")
    elif atype == "swap":
        lines.append(f"  await pool.connect(actor).swap(token.target, {amount_str}, 0, actor.address);")
    elif atype == "borrow":
        lines.append(f"  await vault.connect(actor).borrow({amount_str});")
    elif atype == "liquidate":
        lines.append(f"  await vault.connect(actor).liquidate(victim.address);")
    elif atype == "flash_loan":
        lines.append(f"  await flashAttacker.connect(actor).executeAttack(vault.target, {amount_str});")
    elif atype == "oracle_update":
        price = int(amount * 1e18)
        lines.append(f"  await vault.connect(owner).setOraclePrice({price}n);")
    elif atype == "manipulate_price":
        lines.append(f"  // Manipulate price via large swap")
        lines.append(f"  await pool.connect(actor).swap(token.target, {amount_str}, 0, actor.address);")
    elif atype == "drain":
        lines.append(f"  await pool.connect(actor).removeLiquidity(await pool.balanceOf(actor.address), 0, 0);")
    elif atype == "governance_vote":
        lines.append(f"  await governance.connect(actor).castVote(proposalId, 1);")
        lines.append(f"  await governance.connect(actor).execute(proposalId);")
    else:
        lines.append(f"  // Unknown action type: {atype} — skipping")

    if delay_code:
        lines.append(delay_code.rstrip())

    return "\n".join(lines)


def _market_event_to_js(event: Dict[str, Any], step_idx: int) -> str:
    """Convert a market event to Hardhat time/state manipulation code."""
    etype     = event.get("type", "price_drop")
    magnitude = float(event.get("magnitude", 0.5))
    duration  = int(event.get("duration_steps", 10))

    lines = [f"  // Market event {step_idx + 1}: {etype} (magnitude={magnitude})"]

    if etype == "price_drop":
        new_price = int(1000 * (1 - magnitude) * 1e18)
        lines.append(f"  await vault.connect(owner).setOraclePrice({new_price}n);")
        lines.append(f"  console.log('  ⤵ Oracle price dropped {int(magnitude*100)}%');")
    elif etype == "price_spike":
        new_price = int(1000 * magnitude * 1e18)
        lines.append(f"  await vault.connect(owner).setOraclePrice({new_price}n);")
        lines.append(f"  console.log('  ⤴ Oracle price spiked {magnitude}x');")
    elif etype == "oracle_lag":
        lines.append(f"  // Simulate oracle lag: mine {duration} blocks without price update")
        lines.append(f"  await network.provider.send('hardhat_mine', ['{hex(duration)}']);")
        lines.append(f"  console.log('  ⏳ Oracle lag: {duration} blocks mined');")
    elif etype == "liquidity_drain":
        lines.append(f"  console.log('  💧 Simulating liquidity drain {int(magnitude*100)}%...');")
        lines.append(f"  // Whale removes {int(magnitude*100)}% of LP position")
    elif etype == "depeg":
        peg_break = int(1000 * (1 - magnitude) * 1e18)
        lines.append(f"  await vault.connect(owner).setOraclePrice({peg_break}n);")
        lines.append(f"  console.log('  🔴 Stablecoin depeg: {int(magnitude*100)}% below peg');")
    elif etype == "hyperinflation":
        lines.append(f"  // Mine {duration} blocks to simulate hyperinflationary time passage")
        lines.append(f"  await network.provider.send('hardhat_mine', ['{hex(duration)}']);")
        lines.append(f"  console.log('  📈 Hyperinflation phase: {duration} blocks mined');")
    else:
        lines.append(f"  // Event '{etype}' — advance time by {duration} steps")
        lines.append(f"  await network.provider.send('evm_increaseTime', [{duration * 60}]);")
        lines.append(f"  await network.provider.send('evm_mine', []);")

    return "\n".join(lines)


# ─────────────────────────────────────────────────────────────────── #
# MAIN GENERATOR                                                       #
# ─────────────────────────────────────────────────────────────────── #

def generate_hardhat_script(
    scenario: Dict[str, Any],
    contract_analysis: Optional[Dict[str, Any]] = None,
    rpc_url: str = "http://127.0.0.1:8545",
) -> str:
    """
    Generate a runnable Hardhat JavaScript stress-test script from a scenario dict.

    Args:
        scenario:          Structured scenario from nlp_engine.parse_scenario()
                           or scenarios.get_scenario().
        contract_analysis: Optional ContractAnalyzer results (for function names).
        rpc_url:           Hardhat node RPC URL.

    Returns:
        JavaScript string — save to a .js file and run with `npx hardhat run`.
    """
    title       = scenario.get("title", "Custom Stress Test")
    description = scenario.get("description", "")
    actor_type  = scenario.get("actor", "whale")
    actions     = scenario.get("actions", [])
    events      = scenario.get("market_events", [])
    risk_focus  = scenario.get("risk_focus", [])
    steps       = scenario.get("steps", 50)

    fn_names = []
    if contract_analysis:
        fn_names = [f["name"] for f in contract_analysis.get("functions", [])]

    actions_js = "\n\n".join(
        _action_to_js(a, i, fn_names) for i, a in enumerate(actions)
    )
    events_js = "\n\n".join(
        _market_event_to_js(e, i) for i, e in enumerate(events)
    )

    script = f'''/**
 * Stress Test: {title}
 * {description}
 *
 * Actor:       {actor_type}
 * Risk focus:  {", ".join(risk_focus) or "general"}
 * Steps:       {steps}
 * Generated:   DeFi Risk Simulation Lab — Smart Contract Stress Testing Engine
 *
 * Run with:
 *   npx hardhat run scripts/stress_<name>.js --network localhost
 */

const {{ ethers, network }} = require("hardhat");

// ──────────────────────────── HELPERS ────────────────────────────── //

function sleep(ms) {{
  return new Promise(resolve => setTimeout(resolve, ms));
}}

async function getGasReport(tx) {{
  const receipt = await tx.wait();
  return {{
    gasUsed:      receipt.gasUsed.toString(),
    effectiveGasPrice: receipt.gasPrice?.toString() ?? "—",
    blockNumber:  receipt.blockNumber,
  }};
}}

async function poolSnapshot(vault) {{
  try {{
    const [usdc, eth, price] = await vault.getPoolState();
    return {{ usdc: usdc.toString(), eth: eth.toString(), oraclePrice: price.toString() }};
  }} catch {{ return {{}}; }}
}}

// ────────────────────────── MAIN TEST ────────────────────────────── //

async function main() {{
  console.log("\\n🔥 Stress Test: {title}");
  console.log("━".repeat(60));
  console.log("Description:", {json.dumps(description)});
  console.log("Actor Type: ", "{actor_type}");
  console.log("Risk Focus: ", {json.dumps(", ".join(risk_focus))});
  console.log("━".repeat(60) + "\\n");

  // ── Signers ──────────────────────────────────────────────────────
  const signers = await ethers.getSigners();
  const owner   = signers[0];
  const actor   = signers[1];   // primary stress actor
  const victim  = signers[2];   // target for liquidations
  const bystander = signers[3]; // retail observer

  console.log("Owner  :", owner.address);
  console.log("Actor  :", actor.address);
  console.log("Victim :", victim.address);

  // ── Load deploy output ────────────────────────────────────────────
  const deployOutput = require("../deploy_output.json");
  const addrs        = deployOutput.addresses || deployOutput; // support both flat and nested formats
  const VAULT_ADDR   = addrs.VulnerableVault;
  const TOKEN_ADDR   = addrs.MockToken;
  const GATE_ADDR    = addrs.IdentityGatekeeper;

  if (!VAULT_ADDR || !TOKEN_ADDR) {{
    throw new Error("Contract addresses not found in deploy_output.json. Run: npx hardhat run scripts/deploy.js --network localhost");
  }}


  // ── Contract instances ────────────────────────────────────────────
  const VaultFactory = await ethers.getContractFactory("VulnerableVault");
  const TokenFactory = await ethers.getContractFactory("MockToken");

  const vault = VaultFactory.attach(VAULT_ADDR);
  const token = TokenFactory.attach(TOKEN_ADDR);

  console.log("\\n✅ Connected to VulnerableVault at", VAULT_ADDR);
  console.log("✅ Connected to MockToken at", TOKEN_ADDR);

  // ── Snapshot: Before ──────────────────────────────────────────────
  const before = await poolSnapshot(vault);
  console.log("\\n📊 BEFORE snapshot:", JSON.stringify(before, null, 2));

  // ── Fund the actor ────────────────────────────────────────────────
  try {{
    await token.mint(actor.address, ethers.parseEther("10000000")); // 10M tokens
    console.log("\\n💰 Minted 10M tokens to actor");
  }} catch (e) {{
    console.log("⚠️  Could not mint tokens (contract may restrict):", e.message);
  }}

  const gasSummary = [];
  let proposalId = 1n; // placeholder for governance scenarios

  // ════════════════════════════════════════════════════════════════ //
  //  PHASE 1: ACTIONS                                                //
  // ════════════════════════════════════════════════════════════════ //
  console.log("\\n🎬 Phase 1: Executing actor actions...");

{actions_js if actions_js else "  console.log('  ℹ️  No actions defined for this scenario.');"}

  // ════════════════════════════════════════════════════════════════ //
  //  PHASE 2: MARKET EVENTS                                          //
  // ════════════════════════════════════════════════════════════════ //
  console.log("\\n🌪  Phase 2: Applying market events...");

{events_js if events_js else "  console.log('  ℹ️  No market events defined.');"}

  // ── Snapshot: After ───────────────────────────────────────────────
  const after = await poolSnapshot(vault);
  console.log("\\n📊 AFTER snapshot:", JSON.stringify(after, null, 2));

  // ════════════════════════════════════════════════════════════════ //
  //  RISK METRICS                                                    //
  // ════════════════════════════════════════════════════════════════ //
  console.log("\\n🧮 Computing risk metrics...");

  try {{
    const oracleBefore = BigInt(before.oraclePrice ?? "2000000000000000000000");
    const oracleAfter  = BigInt(after.oraclePrice  ?? "2000000000000000000000");
    const priceImpact  = Number((oracleBefore - oracleAfter) * 10000n / oracleBefore) / 100;

    const usdcBefore   = BigInt(before.usdc ?? "0");
    const usdcAfter    = BigInt(after.usdc  ?? "0");
    const liquidityDrain = usdcBefore > 0n
      ? Number((usdcBefore - usdcAfter) * 10000n / usdcBefore) / 100
      : 0;

    // Attack surface score (weighted, 0–100)
    const liquidityRisk  = Math.min(liquidityDrain / 100, 1) * 25;
    const oracleRisk     = Math.min(Math.abs(priceImpact) / 80, 1) * 20;
    const codeVulnScore  = {len(contract_analysis.get('vulnerabilities', []) if contract_analysis else [])} / 8 * 20;
    const flashLoanRisk  = {1 if 'flash_loan' in (scenario.get('risk_focus') or []) else 0} * 15;
    const govRisk        = {1 if 'governance' in (scenario.get('risk_focus') or []) else 0} * 20;
    const attackScore    = Math.min(liquidityRisk + oracleRisk + codeVulnScore + flashLoanRisk + govRisk, 100);

    const insolventProb  = Math.min(Math.abs(liquidityDrain) * 1.2, 100).toFixed(1);

    console.log("\\n┌─────────────────────────────────────────┐");
    console.log("│          STRESS TEST RESULTS            │");
    console.log("├─────────────────────────────────────────┤");
    console.log(`│  Price Impact:       ${{priceImpact.toFixed(2)}}%`);
    console.log(`│  Liquidity Drain:    ${{liquidityDrain.toFixed(2)}}%`);
    console.log(`│  Attack Score:       ${{attackScore.toFixed(1)}} / 100`);
    console.log(`│  Insolvency Prob:    ${{insolventProb}}%`);
    console.log(`│  Risk Focus:         {", ".join(risk_focus) or "general"}`);
    console.log("└─────────────────────────────────────────┘");

    // Emit final metrics as JSON for API capture
    const metrics = {{
      title:            {json.dumps(title)},
      priceImpactPct:   priceImpact,
      liquidityDrainPct: liquidityDrain,
      attackSurfaceScore: parseFloat(attackScore.toFixed(1)),
      insolventProbPct: parseFloat(insolventProb),
      oraclePriceBefore: oracleBefore.toString(),
      oraclePriceAfter:  oracleAfter.toString(),
      liquidityBefore:   usdcBefore.toString(),
      liquidityAfter:    usdcAfter.toString(),
      riskFocus:        {json.dumps(risk_focus)},
      scenarioSteps:    {steps},
    }};
    console.log("\\n__METRICS__:" + JSON.stringify(metrics));

  }} catch (metricsErr) {{
    console.warn("⚠️  Metrics computation partial:", metricsErr.message);
  }}

  console.log("\\n✅ Stress test complete.");
}}

main()
  .then(() => process.exit(0))
  .catch((err) => {{
    console.error("\\n❌ Stress test failed:", err.message);
    process.exit(1);
  }});
'''
    return script


# ─────────────────────────────────────────────────────────────────── #
# CONVENIENCE: generate from scenario_id (historical library)        #
# ─────────────────────────────────────────────────────────────────── #

def generate_from_scenario_id(
    scenario_id: str,
    contract_analysis: Optional[Dict[str, Any]] = None,
) -> str:
    """Generate a Hardhat script from a named historical scenario."""
    from scenarios import get_scenario
    scenario = get_scenario(scenario_id)
    return generate_hardhat_script(scenario, contract_analysis=contract_analysis)