"""
Vulnerability RAG (Retrieval-Augmented Generation) for MIESC
=============================================================

RAG system that enhances vulnerability analysis with contextual
information from a knowledge base of known vulnerabilities.

Based on SmartLLM (2024) - using RAG for improved detection accuracy.

Features:
- Embedding-based similarity search
- SWC Registry integration
- Real-world exploit examples
- Contextual analysis enhancement

Author: Fernando Boiero <fboiero@frvm.utn.edu.ar>
Institution: UNDEF - IUA
Date: January 2026
"""

import json
import logging
import hashlib
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Tuple
from pathlib import Path

logger = logging.getLogger(__name__)


@dataclass
class VulnerabilityExample:
    """A known vulnerability example for RAG context."""
    id: str
    swc_id: Optional[str]
    cwe_id: Optional[str]
    title: str
    description: str
    vulnerable_code: str
    fixed_code: Optional[str]
    attack_scenario: Optional[str]
    severity: str
    category: str
    real_exploit: Optional[str] = None
    tags: List[str] = field(default_factory=list)
    embedding: Optional[List[float]] = None


@dataclass
class SimilarVuln:
    """A similar vulnerability from the knowledge base."""
    example: VulnerabilityExample
    similarity_score: float
    relevance_reason: str


@dataclass
class EnhancedFinding:
    """A finding enhanced with RAG context."""
    original_finding: Dict[str, Any]
    similar_vulns: List[SimilarVuln]
    enhanced_description: str
    exploit_scenario: Optional[str]
    remediation_code: Optional[str]
    confidence_adjustment: float
    additional_references: List[str]


# =============================================================================
# SWC REGISTRY - Smart Contract Weakness Classification
# =============================================================================

SWC_REGISTRY: Dict[str, VulnerabilityExample] = {
    "SWC-107": VulnerabilityExample(
        id="SWC-107",
        swc_id="SWC-107",
        cwe_id="CWE-841",
        title="Reentrancy",
        description=(
            "One of the major dangers of calling external contracts is that they can "
            "take over the control flow. In the reentrancy attack, a malicious contract "
            "calls back into the calling contract before the first invocation of the "
            "function is finished."
        ),
        vulnerable_code="""
function withdraw(uint256 _amount) public {
    require(balances[msg.sender] >= _amount);
    (bool success,) = msg.sender.call{value: _amount}("");
    require(success);
    balances[msg.sender] -= _amount;  // State update AFTER external call
}
""",
        fixed_code="""
function withdraw(uint256 _amount) public nonReentrant {
    require(balances[msg.sender] >= _amount);
    balances[msg.sender] -= _amount;  // State update BEFORE external call
    (bool success,) = msg.sender.call{value: _amount}("");
    require(success);
}
""",
        attack_scenario=(
            "Attacker deploys contract with fallback that calls withdraw() again. "
            "First call: balance check passes, ETH sent, fallback triggers. "
            "Second call: balance still not updated, check passes again. "
            "Repeat until contract drained."
        ),
        severity="high",
        category="reentrancy",
        real_exploit="The DAO Hack ($60M, 2016)",
        tags=["reentrancy", "external-call", "state-update"],
    ),

    "SWC-105": VulnerabilityExample(
        id="SWC-105",
        swc_id="SWC-105",
        cwe_id="CWE-284",
        title="Unprotected Ether Withdrawal",
        description=(
            "Due to missing or insufficient access controls, malicious parties can "
            "withdraw some or all Ether from the contract account."
        ),
        vulnerable_code="""
function withdrawAll() public {
    // No access control!
    payable(msg.sender).transfer(address(this).balance);
}
""",
        fixed_code="""
function withdrawAll() public onlyOwner {
    payable(msg.sender).transfer(address(this).balance);
}
""",
        attack_scenario=(
            "Anyone can call withdrawAll() and drain all ETH from the contract. "
            "No authentication or authorization check is performed."
        ),
        severity="critical",
        category="access-control",
        real_exploit="Parity Multisig Hack ($150M frozen, 2017)",
        tags=["access-control", "withdrawal", "authorization"],
    ),

    "SWC-106": VulnerabilityExample(
        id="SWC-106",
        swc_id="SWC-106",
        cwe_id="CWE-284",
        title="Unprotected SELFDESTRUCT Instruction",
        description=(
            "A contract without access control on selfdestruct can be destroyed by "
            "anyone, permanently losing all funds and functionality."
        ),
        vulnerable_code="""
function kill() public {
    selfdestruct(payable(msg.sender));
}
""",
        fixed_code="""
function kill() public onlyOwner {
    selfdestruct(payable(owner));
}
""",
        attack_scenario=(
            "Attacker calls kill() function. Contract is destroyed. "
            "All funds are sent to attacker. Contract code is erased from blockchain."
        ),
        severity="critical",
        category="access-control",
        real_exploit="Parity Kill Bug ($150M frozen, 2017)",
        tags=["selfdestruct", "access-control", "permanent-loss"],
    ),

    "SWC-104": VulnerabilityExample(
        id="SWC-104",
        swc_id="SWC-104",
        cwe_id="CWE-252",
        title="Unchecked Call Return Value",
        description=(
            "The return value of a message call is not checked. Execution will resume "
            "even if the called contract throws an exception."
        ),
        vulnerable_code="""
function sendEther(address payable _to, uint256 _amount) public {
    _to.call{value: _amount}("");  // Return value not checked!
    // Continues even if call fails
}
""",
        fixed_code="""
function sendEther(address payable _to, uint256 _amount) public {
    (bool success,) = _to.call{value: _amount}("");
    require(success, "Transfer failed");
}
""",
        attack_scenario=(
            "External call fails silently. Contract state may become inconsistent. "
            "User thinks operation succeeded but it didn't."
        ),
        severity="medium",
        category="unchecked-call",
        tags=["unchecked-return", "low-level-call", "error-handling"],
    ),

    "SWC-115": VulnerabilityExample(
        id="SWC-115",
        swc_id="SWC-115",
        cwe_id="CWE-477",
        title="Authorization through tx.origin",
        description=(
            "tx.origin is a global variable in Solidity which returns the address of "
            "the account that sent the transaction. Using tx.origin for authorization "
            "could make a contract vulnerable if an authorized account calls a malicious contract."
        ),
        vulnerable_code="""
function transferTo(address payable dest, uint amount) public {
    require(tx.origin == owner);  // BAD: uses tx.origin
    dest.transfer(amount);
}
""",
        fixed_code="""
function transferTo(address payable dest, uint amount) public {
    require(msg.sender == owner);  // GOOD: uses msg.sender
    dest.transfer(amount);
}
""",
        attack_scenario=(
            "Attacker creates malicious contract with legitimate-looking function. "
            "Owner interacts with it. Malicious contract calls transferTo(). "
            "tx.origin is still owner, authorization passes, funds stolen."
        ),
        severity="high",
        category="access-control",
        tags=["tx-origin", "phishing", "authorization"],
    ),

    "SWC-116": VulnerabilityExample(
        id="SWC-116",
        swc_id="SWC-116",
        cwe_id="CWE-829",
        title="Block values as a proxy for time",
        description=(
            "Contracts often need access to time values to perform certain types of "
            "functionality. Values such as block.timestamp and block.number can be "
            "manipulated by miners to some degree."
        ),
        vulnerable_code="""
function play() public payable {
    require(msg.value >= 1 ether);
    if (block.timestamp % 15 == 0) {  // Miner can manipulate
        payable(msg.sender).transfer(address(this).balance);
    }
}
""",
        fixed_code="""
// Use Chainlink VRF for randomness instead
function play() public payable {
    require(msg.value >= 1 ether);
    requestRandomness();  // Use verifiable random function
}
""",
        attack_scenario=(
            "Miner submits transaction with manipulated timestamp. "
            "Wins lottery by ensuring condition is met. "
            "Block.timestamp can be adjusted by ~15 seconds."
        ),
        severity="medium",
        category="randomness",
        tags=["timestamp", "block-values", "randomness", "miner-manipulation"],
    ),

    "SWC-101": VulnerabilityExample(
        id="SWC-101",
        swc_id="SWC-101",
        cwe_id="CWE-190",
        title="Integer Overflow and Underflow",
        description=(
            "An overflow/underflow happens when an arithmetic operation reaches the "
            "maximum or minimum size of a type. For instance if a number is stored "
            "in the uint8 type, it means that the number is stored in a 8 bits unsigned "
            "number ranging from 0 to 2^8-1."
        ),
        vulnerable_code="""
// Solidity < 0.8.0
function transfer(address _to, uint256 _amount) public {
    require(balances[msg.sender] >= _amount);
    balances[msg.sender] -= _amount;  // Can underflow!
    balances[_to] += _amount;  // Can overflow!
}
""",
        fixed_code="""
// Solidity 0.8.0+: Built-in overflow checks
// Or use SafeMath for older versions
using SafeMath for uint256;

function transfer(address _to, uint256 _amount) public {
    require(balances[msg.sender] >= _amount);
    balances[msg.sender] = balances[msg.sender].sub(_amount);
    balances[_to] = balances[_to].add(_amount);
}
""",
        attack_scenario=(
            "Attacker transfers more than they have, causing underflow. "
            "Balance wraps around to max uint256. "
            "Attacker now has ~2^256 tokens."
        ),
        severity="high",
        category="arithmetic",
        real_exploit="BEC Token ($900M theoretically, 2018)",
        tags=["overflow", "underflow", "arithmetic", "safemath"],
    ),

    "SWC-112": VulnerabilityExample(
        id="SWC-112",
        swc_id="SWC-112",
        cwe_id="CWE-829",
        title="Delegatecall to Untrusted Callee",
        description=(
            "Delegatecall is a special variant of message call that allows the callee "
            "to modify the caller's state. If the target of a delegatecall is controlled "
            "by an attacker, they can execute arbitrary code in the context of the calling contract."
        ),
        vulnerable_code="""
function execute(address _target, bytes memory _data) public {
    // Arbitrary delegatecall!
    _target.delegatecall(_data);
}
""",
        fixed_code="""
mapping(address => bool) public allowedTargets;

function execute(address _target, bytes memory _data) public onlyOwner {
    require(allowedTargets[_target], "Target not allowed");
    _target.delegatecall(_data);
}
""",
        attack_scenario=(
            "Attacker provides malicious contract as target. "
            "Delegatecall executes attacker's code with victim's storage. "
            "Attacker can modify any state variable, including owner."
        ),
        severity="critical",
        category="delegatecall",
        real_exploit="Parity Wallet ($30M, 2017)",
        tags=["delegatecall", "proxy", "code-injection"],
    ),

    "SWC-133": VulnerabilityExample(
        id="SWC-133",
        swc_id="SWC-133",
        cwe_id="CWE-829",
        title="Hash Collisions With Multiple Variable Length Arguments",
        description=(
            "Using abi.encodePacked() with multiple variable length arguments can lead to "
            "hash collisions. This is because abi.encodePacked concatenates the arguments "
            "directly without length prefixes."
        ),
        vulnerable_code="""
function collideHash(string memory a, string memory b) public pure returns (bytes32) {
    // "ab" + "c" produces same hash as "a" + "bc"
    return keccak256(abi.encodePacked(a, b));
}
""",
        fixed_code="""
function safeHash(string memory a, string memory b) public pure returns (bytes32) {
    // Use abi.encode which adds length prefixes
    return keccak256(abi.encode(a, b));
}
""",
        attack_scenario=(
            "Attacker finds inputs that produce hash collision. "
            "Bypasses signature verification or access control. "
            "Example: ('a', 'bc') and ('ab', 'c') produce same hash."
        ),
        severity="medium",
        category="hash-collision",
        tags=["encodepacked", "hash", "collision", "signature"],
    ),

    # =========================================================================
    # ADDITIONAL SWC ENTRIES (v4.4.0)
    # =========================================================================

    "SWC-100": VulnerabilityExample(
        id="SWC-100",
        swc_id="SWC-100",
        cwe_id="CWE-710",
        title="Function Default Visibility",
        description=(
            "Functions that do not have a function visibility type specified are public "
            "by default. This can lead to a vulnerability if a developer forgot to set "
            "the visibility and a malicious user is able to make unauthorized state changes."
        ),
        vulnerable_code="""
function setOwner(address _owner) {  // No visibility specified - defaults to public!
    owner = _owner;
}
""",
        fixed_code="""
function setOwner(address _owner) internal {  // Explicitly set visibility
    owner = _owner;
}
""",
        attack_scenario=(
            "Function without explicit visibility defaults to public. "
            "Anyone can call it and modify contract state. "
            "Attacker calls setOwner() and becomes the owner."
        ),
        severity="high",
        category="access-control",
        tags=["visibility", "public", "access-control", "default"],
    ),

    "SWC-102": VulnerabilityExample(
        id="SWC-102",
        swc_id="SWC-102",
        cwe_id="CWE-937",
        title="Outdated Compiler Version",
        description=(
            "Using an outdated compiler version can be problematic especially if there "
            "are publicly disclosed bugs and issues that affect the current compiler version."
        ),
        vulnerable_code="""
pragma solidity 0.4.24;  // Old version with known vulnerabilities
""",
        fixed_code="""
pragma solidity 0.8.20;  // Use recent stable version
""",
        attack_scenario=(
            "Old compiler has known bugs (e.g., optimizer bugs, ABIEncoderV2 issues). "
            "Attacker exploits known compiler vulnerability. "
            "Contract behaves unexpectedly due to compiler bug."
        ),
        severity="medium",
        category="compiler",
        tags=["compiler", "version", "outdated", "pragma"],
    ),

    "SWC-103": VulnerabilityExample(
        id="SWC-103",
        swc_id="SWC-103",
        cwe_id="CWE-664",
        title="Floating Pragma",
        description=(
            "Contracts should be deployed with the same compiler version and flags "
            "that they have been tested with. Locking the pragma helps ensure that "
            "contracts do not accidentally get deployed using a different compiler version."
        ),
        vulnerable_code="""
pragma solidity ^0.8.0;  // Any 0.8.x version - inconsistent behavior possible
""",
        fixed_code="""
pragma solidity 0.8.20;  // Locked to specific tested version
""",
        attack_scenario=(
            "Contract deployed with different compiler version than tested. "
            "Optimizer or language behavior differs between versions. "
            "Unexpected vulnerabilities introduced by version mismatch."
        ),
        severity="low",
        category="compiler",
        tags=["pragma", "floating", "compiler", "version"],
    ),

    "SWC-108": VulnerabilityExample(
        id="SWC-108",
        swc_id="SWC-108",
        cwe_id="CWE-710",
        title="State Variable Default Visibility",
        description=(
            "Labeling the visibility of state variables explicitly makes it easier to "
            "catch incorrect assumptions about who can access the variable."
        ),
        vulnerable_code="""
uint256 secretNumber;  // Default internal, but intention unclear
address owner;  // Is this meant to be public or internal?
""",
        fixed_code="""
uint256 private secretNumber;  // Explicitly private
address public owner;  // Explicitly public
""",
        attack_scenario=(
            "Developer assumes variable is private but it's internal. "
            "Derived contracts can access and modify sensitive data. "
            "Security assumptions violated due to unclear visibility."
        ),
        severity="low",
        category="visibility",
        tags=["visibility", "state-variable", "internal", "default"],
    ),

    "SWC-109": VulnerabilityExample(
        id="SWC-109",
        swc_id="SWC-109",
        cwe_id="CWE-824",
        title="Uninitialized Storage Pointer",
        description=(
            "Uninitialized local storage variables can point to unexpected storage "
            "locations in the contract, which can lead to vulnerabilities."
        ),
        vulnerable_code="""
function register(string name) public {
    User user;  // Uninitialized storage pointer!
    user.name = name;  // Overwrites slot 0 (owner!)
    users[msg.sender] = user;
}
""",
        fixed_code="""
function register(string memory name) public {
    User storage user = users[msg.sender];  // Explicit storage location
    user.name = name;
}
""",
        attack_scenario=(
            "Uninitialized storage pointer defaults to slot 0. "
            "Writing to it overwrites critical state (e.g., owner). "
            "Attacker takes control of contract."
        ),
        severity="critical",
        category="storage",
        real_exploit="Multiple contracts affected pre-0.5.0",
        tags=["storage", "uninitialized", "pointer", "slot0"],
    ),

    "SWC-110": VulnerabilityExample(
        id="SWC-110",
        swc_id="SWC-110",
        cwe_id="CWE-617",
        title="Assert Violation",
        description=(
            "The assert() function is meant to assert invariants. Properly functioning "
            "code should never reach a failing assert statement. A reachable assert "
            "indicates a bug in the contract or invalid state."
        ),
        vulnerable_code="""
function withdraw(uint256 amount) public {
    balances[msg.sender] -= amount;  // Can underflow in old Solidity
    assert(balances[msg.sender] >= 0);  // Always true for uint, useless
}
""",
        fixed_code="""
function withdraw(uint256 amount) public {
    require(balances[msg.sender] >= amount, "Insufficient balance");
    balances[msg.sender] -= amount;
}
""",
        attack_scenario=(
            "Assert used for input validation instead of require. "
            "Failed assert consumes all gas (unlike require). "
            "Misuse indicates logic errors in contract."
        ),
        severity="medium",
        category="assertion",
        tags=["assert", "invariant", "gas", "require"],
    ),

    "SWC-111": VulnerabilityExample(
        id="SWC-111",
        swc_id="SWC-111",
        cwe_id="CWE-477",
        title="Use of Deprecated Solidity Functions",
        description=(
            "Several functions and operators in Solidity are deprecated. Using them "
            "leads to reduced code quality and potential vulnerabilities."
        ),
        vulnerable_code="""
function hashData(bytes data) public returns (bytes32) {
    return sha3(data);  // Deprecated, use keccak256
}

function destroy() public {
    suicide(owner);  // Deprecated, use selfdestruct
}

function send(address to) public {
    to.callcode(data);  // Deprecated, use delegatecall
}
""",
        fixed_code="""
function hashData(bytes memory data) public pure returns (bytes32) {
    return keccak256(data);
}

function destroy() public onlyOwner {
    selfdestruct(payable(owner));
}

function execute(address to, bytes memory data) public {
    to.delegatecall(data);
}
""",
        attack_scenario=(
            "Deprecated functions may have unexpected behavior. "
            "sha3 vs keccak256 confusion possible. "
            "Future compiler versions may remove support."
        ),
        severity="low",
        category="deprecated",
        tags=["deprecated", "sha3", "suicide", "callcode"],
    ),

    "SWC-113": VulnerabilityExample(
        id="SWC-113",
        swc_id="SWC-113",
        cwe_id="CWE-703",
        title="DoS with Failed Call",
        description=(
            "External calls can fail accidentally or deliberately. A DoS vulnerability "
            "can occur when the contract handles call failures improperly."
        ),
        vulnerable_code="""
function refundAll() public {
    for (uint i = 0; i < refundAddresses.length; i++) {
        // If one transfer fails, all refunds fail!
        require(refundAddresses[i].send(refunds[refundAddresses[i]]));
    }
}
""",
        fixed_code="""
mapping(address => uint256) public pendingWithdrawals;

function withdraw() public {
    uint256 amount = pendingWithdrawals[msg.sender];
    pendingWithdrawals[msg.sender] = 0;
    (bool success,) = msg.sender.call{value: amount}("");
    if (!success) {
        pendingWithdrawals[msg.sender] = amount;
    }
}
""",
        attack_scenario=(
            "Malicious contract reverts on receive. "
            "Entire batch operation fails due to one revert. "
            "Legitimate users cannot receive refunds."
        ),
        severity="high",
        category="dos",
        tags=["dos", "failed-call", "batch", "pull-payment"],
    ),

    "SWC-114": VulnerabilityExample(
        id="SWC-114",
        swc_id="SWC-114",
        cwe_id="CWE-362",
        title="Transaction Order Dependence",
        description=(
            "The outcome of a transaction can depend on the order of transactions in "
            "the block. Miners/validators can exploit this by reordering transactions."
        ),
        vulnerable_code="""
function solve(bytes32 solution) public {
    // Front-runnable: attacker sees solution in mempool
    require(keccak256(solution) == puzzle);
    winner = msg.sender;
    msg.sender.transfer(prize);
}
""",
        fixed_code="""
mapping(address => bytes32) public commits;

function commit(bytes32 hash) public {
    commits[msg.sender] = hash;
}

function reveal(bytes32 solution, bytes32 secret) public {
    require(commits[msg.sender] == keccak256(abi.encodePacked(solution, secret)));
    require(keccak256(abi.encodePacked(solution)) == puzzle);
    winner = msg.sender;
}
""",
        attack_scenario=(
            "User submits solution to puzzle. "
            "Attacker sees tx in mempool, copies solution. "
            "Attacker front-runs with higher gas price, wins prize."
        ),
        severity="medium",
        category="front-running",
        real_exploit="Common in DeFi, NFT mints, puzzle games",
        tags=["front-running", "mev", "mempool", "ordering"],
    ),

    "SWC-117": VulnerabilityExample(
        id="SWC-117",
        swc_id="SWC-117",
        cwe_id="CWE-347",
        title="Signature Malleability",
        description=(
            "ECDSA signatures are malleable: for every valid signature, there exists "
            "a different valid signature for the same message. This can break assumptions "
            "about signature uniqueness."
        ),
        vulnerable_code="""
mapping(bytes32 => bool) usedSignatures;

function claimWithSig(uint256 amount, bytes memory signature) public {
    bytes32 hash = keccak256(abi.encodePacked(msg.sender, amount));
    require(!usedSignatures[keccak256(signature)]);  // Can be bypassed!
    require(ecrecover(hash, v, r, s) == trustedSigner);
    usedSignatures[keccak256(signature)] = true;
    _transfer(msg.sender, amount);
}
""",
        fixed_code="""
mapping(bytes32 => bool) usedHashes;

function claimWithSig(uint256 amount, bytes memory signature) public {
    bytes32 hash = keccak256(abi.encodePacked(msg.sender, amount));
    require(!usedHashes[hash]);  // Track message hash, not signature
    require(ECDSA.recover(hash, signature) == trustedSigner);  // Use OZ ECDSA
    usedHashes[hash] = true;
    _transfer(msg.sender, amount);
}
""",
        attack_scenario=(
            "User claims with valid signature. "
            "Attacker computes malleable signature (s' = n - s). "
            "Attacker claims again with same message but different signature."
        ),
        severity="high",
        category="signature",
        tags=["ecdsa", "malleability", "signature", "replay"],
    ),

    "SWC-118": VulnerabilityExample(
        id="SWC-118",
        swc_id="SWC-118",
        cwe_id="CWE-665",
        title="Incorrect Constructor Name",
        description=(
            "Before Solidity 0.4.22, constructors were defined as functions with the "
            "same name as the contract. A misspelled constructor becomes a regular "
            "function that anyone can call."
        ),
        vulnerable_code="""
contract Owned {
    address public owner;

    function owned() public {  // Misspelled! Should be 'Owned'
        owner = msg.sender;
    }
}
""",
        fixed_code="""
contract Owned {
    address public owner;

    constructor() {  // Use constructor keyword (0.4.22+)
        owner = msg.sender;
    }
}
""",
        attack_scenario=(
            "Constructor name misspelled, becomes public function. "
            "Anyone can call it and become owner. "
            "Attacker calls 'owned()' and takes ownership."
        ),
        severity="critical",
        category="constructor",
        real_exploit="Rubixi ($100K+, 2016)",
        tags=["constructor", "typo", "ownership", "legacy"],
    ),

    "SWC-119": VulnerabilityExample(
        id="SWC-119",
        swc_id="SWC-119",
        cwe_id="CWE-710",
        title="Shadowing State Variables",
        description=(
            "Solidity allows for ambiguous naming of state variables when inheritance "
            "is used. A derived contract that defines a variable with the same name "
            "shadows the parent's variable."
        ),
        vulnerable_code="""
contract Parent {
    uint256 public value = 100;
}

contract Child is Parent {
    uint256 public value = 200;  // Shadows Parent.value!

    function getParentValue() public view returns (uint256) {
        return value;  // Returns 200, not 100
    }
}
""",
        fixed_code="""
contract Parent {
    uint256 public value = 100;
}

contract Child is Parent {
    uint256 public childValue = 200;  // Different name

    function getParentValue() public view returns (uint256) {
        return Parent.value;  // Explicit reference
    }
}
""",
        attack_scenario=(
            "Developer expects to use parent variable. "
            "Child variable shadows it unexpectedly. "
            "Logic uses wrong value, leading to vulnerabilities."
        ),
        severity="medium",
        category="inheritance",
        tags=["shadowing", "inheritance", "variable", "state"],
    ),

    "SWC-120": VulnerabilityExample(
        id="SWC-120",
        swc_id="SWC-120",
        cwe_id="CWE-330",
        title="Weak Sources of Randomness from Chain Attributes",
        description=(
            "Using block variables like block.timestamp, block.difficulty (now prevrandao), "
            "or blockhash for randomness is insecure. Miners can influence these values."
        ),
        vulnerable_code="""
function random() private view returns (uint256) {
    return uint256(keccak256(abi.encodePacked(
        block.timestamp,
        block.difficulty,
        msg.sender
    ))) % 100;
}

function bet(uint256 guess) public payable {
    if (guess == random()) {
        msg.sender.transfer(address(this).balance);
    }
}
""",
        fixed_code="""
import "@chainlink/contracts/src/v0.8/VRFConsumerBase.sol";

function requestRandomness() public returns (bytes32) {
    return requestRandomness(keyHash, fee);
}

function fulfillRandomness(bytes32 requestId, uint256 randomness) internal override {
    randomResult = randomness;
}
""",
        attack_scenario=(
            "Attacker analyzes block parameters before transaction. "
            "Miner manipulates timestamp to win lottery. "
            "Validator with stake can predict prevrandao."
        ),
        severity="high",
        category="randomness",
        real_exploit="SmartBillions lottery ($400K, 2017)",
        tags=["randomness", "block", "miner", "vrf", "chainlink"],
    ),

    "SWC-121": VulnerabilityExample(
        id="SWC-121",
        swc_id="SWC-121",
        cwe_id="CWE-294",
        title="Missing Protection against Signature Replay Attacks",
        description=(
            "Signed messages that do not include a nonce or expiration can be replayed. "
            "Cross-chain replay is also possible if chain ID is not included."
        ),
        vulnerable_code="""
function executeWithSig(address to, uint256 amount, bytes memory sig) public {
    bytes32 hash = keccak256(abi.encodePacked(to, amount));
    require(ecrecover(hash, v, r, s) == owner);
    payable(to).transfer(amount);  // Can be replayed!
}
""",
        fixed_code="""
mapping(uint256 => bool) public usedNonces;

function executeWithSig(
    address to,
    uint256 amount,
    uint256 nonce,
    uint256 deadline,
    bytes memory sig
) public {
    require(block.timestamp <= deadline, "Expired");
    require(!usedNonces[nonce], "Nonce used");
    bytes32 hash = keccak256(abi.encodePacked(
        to, amount, nonce, deadline, block.chainid, address(this)
    ));
    require(ECDSA.recover(hash, sig) == owner);
    usedNonces[nonce] = true;
    payable(to).transfer(amount);
}
""",
        attack_scenario=(
            "Owner signs message to transfer 1 ETH. "
            "Transaction executes successfully. "
            "Attacker replays same signature, drains 1 ETH again."
        ),
        severity="high",
        category="signature",
        tags=["replay", "nonce", "signature", "chainid"],
    ),

    "SWC-122": VulnerabilityExample(
        id="SWC-122",
        swc_id="SWC-122",
        cwe_id="CWE-347",
        title="Lack of Proper Signature Verification",
        description=(
            "Failing to properly verify signatures can allow unauthorized actions. "
            "Common issues: not checking return value of ecrecover, missing zero-address check."
        ),
        vulnerable_code="""
function verify(bytes32 hash, uint8 v, bytes32 r, bytes32 s) public view returns (bool) {
    address signer = ecrecover(hash, v, r, s);
    return signer == trustedSigner;  // ecrecover can return 0x0!
}
""",
        fixed_code="""
function verify(bytes32 hash, bytes memory signature) public view returns (bool) {
    address signer = ECDSA.recover(hash, signature);  // Reverts on invalid sig
    require(signer != address(0), "Invalid signature");
    return signer == trustedSigner;
}
""",
        attack_scenario=(
            "Invalid signature causes ecrecover to return 0x0. "
            "If trustedSigner is uninitialized (0x0), check passes. "
            "Attacker executes actions with invalid signature."
        ),
        severity="high",
        category="signature",
        tags=["ecrecover", "zero-address", "signature", "validation"],
    ),

    "SWC-123": VulnerabilityExample(
        id="SWC-123",
        swc_id="SWC-123",
        cwe_id="CWE-573",
        title="Requirement Violation",
        description=(
            "A require statement should be used for input validation. If a require fails "
            "due to improper usage, it indicates either invalid user input or a bug."
        ),
        vulnerable_code="""
function divide(uint256 a, uint256 b) public pure returns (uint256) {
    require(b > 0);  // Missing error message
    return a / b;
}

function withdraw(uint256 amount) public {
    require(amount > 0 && balances[msg.sender] >= amount);  // Complex require
    // Logic error: what if amount == 0 is valid?
}
""",
        fixed_code="""
function divide(uint256 a, uint256 b) public pure returns (uint256) {
    require(b > 0, "Division by zero");
    return a / b;
}

function withdraw(uint256 amount) public {
    require(amount > 0, "Amount must be positive");
    require(balances[msg.sender] >= amount, "Insufficient balance");
    balances[msg.sender] -= amount;
}
""",
        attack_scenario=(
            "Missing or unclear require conditions. "
            "Users cannot understand why transaction reverts. "
            "Complex conditions may have logic errors."
        ),
        severity="low",
        category="validation",
        tags=["require", "validation", "error-message", "input"],
    ),

    "SWC-124": VulnerabilityExample(
        id="SWC-124",
        swc_id="SWC-124",
        cwe_id="CWE-123",
        title="Write to Arbitrary Storage Location",
        description=(
            "A smart contract's storage can be written to by manipulating array indices "
            "or using assembly. This can lead to overwriting critical state variables."
        ),
        vulnerable_code="""
uint256[] public data;

function write(uint256 index, uint256 value) public {
    // No bounds check! Can write anywhere in storage
    data[index] = value;
}
""",
        fixed_code="""
uint256[] public data;

function write(uint256 index, uint256 value) public {
    require(index < data.length, "Index out of bounds");
    data[index] = value;
}
""",
        attack_scenario=(
            "Attacker calculates storage slot of owner variable. "
            "Calls write() with index that maps to owner slot. "
            "Overwrites owner address with attacker's address."
        ),
        severity="critical",
        category="storage",
        tags=["storage", "arbitrary-write", "array", "overflow"],
    ),

    "SWC-125": VulnerabilityExample(
        id="SWC-125",
        swc_id="SWC-125",
        cwe_id="CWE-696",
        title="Incorrect Inheritance Order",
        description=(
            "Solidity uses C3 linearization for multiple inheritance. Incorrect ordering "
            "can lead to unexpected behavior when contracts override the same function."
        ),
        vulnerable_code="""
contract A {
    function foo() public virtual returns (string memory) { return "A"; }
}
contract B is A {
    function foo() public virtual override returns (string memory) { return "B"; }
}
contract C is A {
    function foo() public virtual override returns (string memory) { return "C"; }
}
contract D is B, C {  // Ambiguous: which foo() is called?
    function foo() public override(B, C) returns (string memory) {
        return super.foo();  // Returns "C" due to C3 linearization
    }
}
""",
        fixed_code="""
contract D is C, B {  // Explicit order: B.foo() will be called
    function foo() public override(B, C) returns (string memory) {
        return super.foo();  // Returns "B"
    }
}
""",
        attack_scenario=(
            "Developer expects parent A's function to be called. "
            "Due to C3 linearization, different parent is used. "
            "Security checks in expected parent are bypassed."
        ),
        severity="medium",
        category="inheritance",
        tags=["inheritance", "c3-linearization", "override", "diamond"],
    ),

    "SWC-126": VulnerabilityExample(
        id="SWC-126",
        swc_id="SWC-126",
        cwe_id="CWE-400",
        title="Insufficient Gas Griefing",
        description=(
            "When a contract makes an external call with limited gas, the called contract "
            "may run out of gas and revert, causing the original transaction to fail or "
            "behave unexpectedly."
        ),
        vulnerable_code="""
function relay(address to, bytes memory data) public {
    // Only 2300 gas forwarded - not enough for complex operations
    (bool success,) = to.call{gas: 2300}(data);
    require(success);
}
""",
        fixed_code="""
function relay(address to, bytes memory data, uint256 gasLimit) public {
    require(gasLimit >= 50000, "Minimum gas required");
    (bool success,) = to.call{gas: gasLimit}(data);
    // Handle failure gracefully
    if (!success) {
        emit RelayFailed(to, data);
    }
}
""",
        attack_scenario=(
            "Caller provides just enough gas for outer call. "
            "Inner call gets 1/64th gas due to EIP-150. "
            "Inner call fails, outer call may succeed with wrong state."
        ),
        severity="medium",
        category="gas",
        tags=["gas", "griefing", "external-call", "eip-150"],
    ),

    "SWC-127": VulnerabilityExample(
        id="SWC-127",
        swc_id="SWC-127",
        cwe_id="CWE-695",
        title="Arbitrary Jump with Function Type Variable",
        description=(
            "In older Solidity versions, function type variables could be manipulated "
            "to jump to arbitrary locations in the bytecode."
        ),
        vulnerable_code="""
// Solidity < 0.5.0
function() internal stored;

function set(function() internal f) public {
    stored = f;
}

function execute() public {
    stored();  // Could jump anywhere if 'stored' is manipulated
}
""",
        fixed_code="""
// Use proper function dispatch or interface
interface ICallback {
    function callback() external;
}

ICallback public stored;

function set(ICallback _callback) public onlyOwner {
    stored = _callback;
}

function execute() public {
    stored.callback();
}
""",
        attack_scenario=(
            "Attacker finds way to set function pointer to arbitrary bytecode offset. "
            "Execute jumps to attacker-controlled code. "
            "Bypasses access control, steals funds."
        ),
        severity="critical",
        category="control-flow",
        tags=["function-pointer", "jump", "bytecode", "legacy"],
    ),

    "SWC-128": VulnerabilityExample(
        id="SWC-128",
        swc_id="SWC-128",
        cwe_id="CWE-400",
        title="DoS With Block Gas Limit",
        description=(
            "Operations that iterate over unbounded arrays can exceed the block gas limit, "
            "making the function impossible to execute."
        ),
        vulnerable_code="""
address[] public participants;

function distribute() public {
    for (uint i = 0; i < participants.length; i++) {
        // If participants grows too large, this exceeds block gas limit
        participants[i].transfer(1 ether);
    }
}
""",
        fixed_code="""
mapping(address => uint256) public pendingWithdrawals;

function distribute(uint256 start, uint256 end) public {
    require(end <= participants.length);
    require(end - start <= 100, "Batch too large");
    for (uint i = start; i < end; i++) {
        pendingWithdrawals[participants[i]] += 1 ether;
    }
}

function withdraw() public {
    uint256 amount = pendingWithdrawals[msg.sender];
    pendingWithdrawals[msg.sender] = 0;
    payable(msg.sender).transfer(amount);
}
""",
        attack_scenario=(
            "Array grows over time with users. "
            "Eventually exceeds block gas limit. "
            "Function becomes permanently uncallable."
        ),
        severity="high",
        category="dos",
        real_exploit="GovernMental Ponzi ($1M stuck, 2016)",
        tags=["dos", "gas-limit", "unbounded", "loop"],
    ),

    "SWC-129": VulnerabilityExample(
        id="SWC-129",
        swc_id="SWC-129",
        cwe_id="CWE-480",
        title="Typographical Error",
        description=(
            "A typo or incorrect operator can lead to unintended behavior. "
            "Common examples: using = instead of ==, + instead of -, wrong variable name."
        ),
        vulnerable_code="""
function transfer(address to, uint256 amount) public {
    require(balances[msg.sender] =+ amount);  // Typo: =+ instead of >=
    balances[msg.sender] -= amount;
    balances[to] += amount;
}
""",
        fixed_code="""
function transfer(address to, uint256 amount) public {
    require(balances[msg.sender] >= amount, "Insufficient balance");
    balances[msg.sender] -= amount;
    balances[to] += amount;
}
""",
        attack_scenario=(
            "Typo in operator changes logic completely. "
            "=+ sets value instead of comparing. "
            "Security checks bypassed due to typo."
        ),
        severity="high",
        category="typo",
        tags=["typo", "operator", "comparison", "assignment"],
    ),

    "SWC-130": VulnerabilityExample(
        id="SWC-130",
        swc_id="SWC-130",
        cwe_id="CWE-451",
        title="Right-To-Left-Override Control Character (U+202E)",
        description=(
            "Malicious actors can use the Right-To-Left-Override Unicode character "
            "to make code appear different than it actually is, hiding malicious logic."
        ),
        vulnerable_code="""
// Looks like: require(msg.sender == owner)
// Actually is: require(owner == msg.sender) with RTLO hiding swap
// The ‮ character reverses text display
function withdraw() public {
    require(msg.sender ==/*‮*/ owner);
    payable(msg.sender).transfer(address(this).balance);
}
""",
        fixed_code="""
// Scan code for Unicode control characters
// Use tools like slither to detect RTLO
function withdraw() public onlyOwner {
    payable(msg.sender).transfer(address(this).balance);
}
""",
        attack_scenario=(
            "Attacker embeds RTLO character in code. "
            "Code review shows correct logic visually. "
            "Actual bytecode executes different logic."
        ),
        severity="high",
        category="unicode",
        tags=["rtlo", "unicode", "obfuscation", "visual"],
    ),

    "SWC-131": VulnerabilityExample(
        id="SWC-131",
        swc_id="SWC-131",
        cwe_id="CWE-1164",
        title="Presence of Unused Variables",
        description=(
            "Unused variables indicate dead code or potential bugs. They can also lead "
            "to confusion and increased gas costs for storage variables."
        ),
        vulnerable_code="""
contract Example {
    uint256 public unusedVar;  // Never used
    uint256 public fee = 100;

    function calculate(uint256 amount) public view returns (uint256) {
        uint256 result = amount * 2;
        uint256 temp = amount + 1;  // Unused!
        return result;
    }
}
""",
        fixed_code="""
contract Example {
    uint256 public fee = 100;

    function calculate(uint256 amount) public pure returns (uint256) {
        return amount * 2;
    }
}
""",
        attack_scenario=(
            "Unused variable indicates developer mistake. "
            "Logic that should use variable doesn't. "
            "Business logic vulnerability due to omission."
        ),
        severity="low",
        category="code-quality",
        tags=["unused", "variable", "dead-code", "gas"],
    ),

    "SWC-132": VulnerabilityExample(
        id="SWC-132",
        swc_id="SWC-132",
        cwe_id="CWE-667",
        title="Unexpected Ether Balance",
        description=(
            "Contracts can receive Ether through selfdestruct or pre-calculated addresses. "
            "Logic that relies on exact balance checks can be manipulated."
        ),
        vulnerable_code="""
function claimPrize() public {
    require(address(this).balance == 10 ether);  // Can be manipulated!
    payable(msg.sender).transfer(address(this).balance);
}
""",
        fixed_code="""
uint256 public deposits;

function deposit() public payable {
    deposits += msg.value;
}

function claimPrize() public {
    require(deposits >= 10 ether);  // Track deposits, not balance
    uint256 prize = deposits;
    deposits = 0;
    payable(msg.sender).transfer(prize);
}
""",
        attack_scenario=(
            "Contract expects exactly 10 ETH to trigger payout. "
            "Attacker selfdestructs contract sending 0.1 ETH. "
            "Balance becomes 10.1 ETH, condition never met, funds stuck."
        ),
        severity="medium",
        category="balance",
        tags=["balance", "selfdestruct", "force-send", "invariant"],
    ),

    "SWC-134": VulnerabilityExample(
        id="SWC-134",
        swc_id="SWC-134",
        cwe_id="CWE-665",
        title="Message Call with Hardcoded Gas Amount",
        description=(
            "Hardcoding gas in calls can cause failures when EVM gas costs change, "
            "as happened with EIP-1884. transfer() and send() use 2300 gas stipend."
        ),
        vulnerable_code="""
function withdraw() public {
    // transfer uses 2300 gas - not enough for contracts with receive()
    payable(msg.sender).transfer(address(this).balance);
}

function relay(address to) public {
    // Hardcoded gas may be insufficient after gas cost changes
    to.call{gas: 2300}("");
}
""",
        fixed_code="""
function withdraw() public {
    // Forward all available gas
    (bool success,) = msg.sender.call{value: address(this).balance}("");
    require(success, "Transfer failed");
}

function relay(address to) public {
    // Let EVM manage gas
    (bool success,) = to.call("");
    require(success, "Call failed");
}
""",
        attack_scenario=(
            "Contract uses transfer() to send ETH. "
            "After EIP-1884, SLOAD costs 800 gas (was 200). "
            "Recipient contracts with storage operations fail to receive."
        ),
        severity="medium",
        category="gas",
        real_exploit="Multiple contracts broken by EIP-1884 (2019)",
        tags=["gas", "transfer", "2300", "eip-1884"],
    ),

    "SWC-135": VulnerabilityExample(
        id="SWC-135",
        swc_id="SWC-135",
        cwe_id="CWE-1164",
        title="Code With No Effects",
        description=(
            "Code that does not affect state or return values indicates a bug. "
            "Common examples: standalone expression statements, assignments to self."
        ),
        vulnerable_code="""
function process(uint256 x) public {
    x + 1;  // No effect! Result discarded
    x;      // No effect! Just an expression
    x == 10;  // No effect! Comparison result unused
}
""",
        fixed_code="""
function process(uint256 x) public pure returns (uint256) {
    return x + 1;
}
""",
        attack_scenario=(
            "Developer intends to modify value but forgets assignment. "
            "x + 1 should be x = x + 1. "
            "Critical calculation has no effect."
        ),
        severity="medium",
        category="code-quality",
        tags=["no-effect", "dead-code", "expression", "bug"],
    ),

    "SWC-136": VulnerabilityExample(
        id="SWC-136",
        swc_id="SWC-136",
        cwe_id="CWE-311",
        title="Unencrypted Private Data On-Chain",
        description=(
            "Private variables are not truly private - they're just not directly accessible. "
            "Anyone can read storage slots directly from the blockchain."
        ),
        vulnerable_code="""
contract SecretStorage {
    bytes32 private secretPassword;  // NOT ACTUALLY SECRET!

    function setPassword(bytes32 _password) public {
        secretPassword = _password;
    }

    function checkPassword(bytes32 _password) public view returns (bool) {
        return secretPassword == _password;
    }
}
""",
        fixed_code="""
contract SecretStorage {
    bytes32 public passwordHash;  // Store hash, not plaintext

    function setPassword(bytes32 _password) public {
        passwordHash = keccak256(abi.encodePacked(_password, msg.sender));
    }

    function checkPassword(bytes32 _password) public view returns (bool) {
        return passwordHash == keccak256(abi.encodePacked(_password, msg.sender));
    }
}
""",
        attack_scenario=(
            "Contract stores 'private' password. "
            "Attacker reads storage slot 0 using web3.eth.getStorageAt(). "
            "Password revealed, authentication bypassed."
        ),
        severity="high",
        category="privacy",
        tags=["private", "storage", "encryption", "blockchain-transparency"],
    ),
}


# Additional examples from real exploits
EXPLOIT_EXAMPLES: Dict[str, VulnerabilityExample] = {
    "dao-hack": VulnerabilityExample(
        id="dao-hack",
        swc_id="SWC-107",
        cwe_id="CWE-841",
        title="The DAO Reentrancy Attack",
        description=(
            "The attacker exploited the splitDAO function which called the recipient's "
            "fallback function before updating state. By repeatedly calling splitDAO "
            "from the fallback, they drained ~$60M in ETH."
        ),
        vulnerable_code="""
function splitDAO(uint _proposalID, address _newCurator) returns (bool) {
    ...
    // Transfer happens before state update
    rewardToken[msg.sender] += reward;
    if (!p.creator.send(p.amount)) throw;
    balances[msg.sender] = 0;  // Too late!
    ...
}
""",
        fixed_code="""
function splitDAO(uint _proposalID, address _newCurator) nonReentrant returns (bool) {
    ...
    // Update state first
    uint balanceToTransfer = balances[msg.sender];
    balances[msg.sender] = 0;
    rewardToken[msg.sender] += reward;
    if (!p.creator.send(p.amount)) throw;
    ...
}
""",
        attack_scenario=(
            "1. Attacker creates malicious contract\n"
            "2. Calls splitDAO with malicious contract as recipient\n"
            "3. splitDAO sends ETH, triggers fallback\n"
            "4. Fallback calls splitDAO again\n"
            "5. Balance not yet updated, passes check\n"
            "6. Repeat until drained"
        ),
        severity="critical",
        category="reentrancy",
        real_exploit="The DAO ($60M, June 2016)",
        tags=["reentrancy", "historic", "ethereum-fork"],
    ),

    "beanstalk-governance": VulnerabilityExample(
        id="beanstalk-governance",
        swc_id=None,
        cwe_id="CWE-284",
        title="Beanstalk Governance Flash Loan Attack",
        description=(
            "Attacker used flash loan to borrow governance tokens, immediately "
            "passed a malicious proposal, and drained the protocol's funds."
        ),
        vulnerable_code="""
function propose(address[] memory targets, bytes[] memory calldatas) public {
    require(getVotes(msg.sender) >= proposalThreshold());
    // No time delay or snapshot before proposal
    ...
}
""",
        fixed_code="""
function propose(address[] memory targets, bytes[] memory calldatas) public {
    // Snapshot votes BEFORE proposal
    uint256 proposerVotes = getPastVotes(msg.sender, block.number - 1);
    require(proposerVotes >= proposalThreshold());
    // Timelock execution
    ...
}
""",
        attack_scenario=(
            "1. Flash loan $1B in governance tokens\n"
            "2. Propose malicious governance action\n"
            "3. Vote with borrowed tokens (>67% voting power)\n"
            "4. Execute proposal immediately\n"
            "5. Drain treasury ($182M)\n"
            "6. Repay flash loan"
        ),
        severity="critical",
        category="governance",
        real_exploit="Beanstalk ($182M, April 2022)",
        tags=["governance", "flash-loan", "defi"],
    ),

    "wormhole-bridge": VulnerabilityExample(
        id="wormhole-bridge",
        swc_id=None,
        cwe_id="CWE-347",
        title="Wormhole Bridge Signature Verification Bypass",
        description=(
            "Attacker exploited a vulnerability in the signature verification "
            "to mint 120,000 wETH without backing."
        ),
        vulnerable_code="""
function complete_transfer(bytes memory vaa) public {
    // Insufficient signature verification
    (IWormhole.Signature[] memory signatures, bytes memory body) = parseVAA(vaa);
    require(signatures.length >= guardian_set.num_guardians / 3);
    // Attacker bypassed this check
    ...
}
""",
        fixed_code="""
function complete_transfer(bytes memory vaa) public {
    // Proper signature verification
    require(verifyVM(vaa).valid, "Invalid VAA");
    require(!processedVAAs[vaa.hash], "Already processed");
    processedVAAs[vaa.hash] = true;
    ...
}
""",
        attack_scenario=(
            "1. Attacker found bug in signature verification\n"
            "2. Crafted fake VAA (Verified Action Approval)\n"
            "3. Bridge accepted forged message\n"
            "4. Minted 120,000 wETH ($326M) on Solana\n"
            "5. Bridged some to Ethereum and sold"
        ),
        severity="critical",
        category="bridge",
        real_exploit="Wormhole ($326M, February 2022)",
        tags=["bridge", "signature", "cross-chain"],
    ),

    # =========================================================================
    # 2023-2025 EXPLOIT EXAMPLES (v4.4.0)
    # =========================================================================

    "euler-finance-2023": VulnerabilityExample(
        id="euler-finance-2023",
        swc_id="SWC-107",
        cwe_id="CWE-841",
        title="Euler Finance Flash Loan + Donation Attack",
        description=(
            "Attacker exploited a missing health check in the donateToReserves function. "
            "By donating collateral to reserves, they could manipulate their own debt position "
            "and trigger a liquidation that profited them."
        ),
        vulnerable_code="""
function donateToReserves(address subAccount, uint amount) external {
    // Missing: health check after donation!
    AssetStorage storage assetStorage = eTokenLookup[eTokenAddress];
    assetStorage.users[subAccount].balance -= amount;
    assetStorage.reserveBalance += amount;
    // Attacker's health factor drops below 1, but no check!
}
""",
        fixed_code="""
function donateToReserves(address subAccount, uint amount) external {
    AssetStorage storage assetStorage = eTokenLookup[eTokenAddress];
    assetStorage.users[subAccount].balance -= amount;
    assetStorage.reserveBalance += amount;
    // Check health factor after state change
    require(checkLiquidity(subAccount), "Health check failed");
}
""",
        attack_scenario=(
            "1. Flash loan $30M DAI from Aave\n"
            "2. Deposit DAI as collateral in Euler\n"
            "3. Mint eDAI (2x leverage)\n"
            "4. Donate eDAI to reserves (no health check)\n"
            "5. Self-liquidate at profit\n"
            "6. Repeat with larger amounts\n"
            "7. Drain $197M total"
        ),
        severity="critical",
        category="flash-loan",
        real_exploit="Euler Finance ($197M, March 2023)",
        tags=["flash-loan", "donation", "health-check", "liquidation", "defi"],
    ),

    "curve-vyper-reentrancy-2023": VulnerabilityExample(
        id="curve-vyper-reentrancy-2023",
        swc_id="SWC-107",
        cwe_id="CWE-841",
        title="Curve Pool Vyper Compiler Reentrancy",
        description=(
            "A bug in Vyper compiler versions 0.2.15, 0.2.16, and 0.3.0 caused the "
            "@nonreentrant decorator to malfunction, allowing reentrancy attacks on "
            "Curve pools using those compiler versions."
        ),
        vulnerable_code="""
# Vyper 0.2.15-0.3.0 - @nonreentrant was broken!
@nonreentrant('lock')
@external
def remove_liquidity(amount: uint256) -> uint256[2]:
    # Reentrancy guard didn't work properly
    self._burn(msg.sender, amount)
    raw_call(msg.sender, b"", value=amount)  # Callback possible!
    return [amount, amount]
""",
        fixed_code="""
# Vyper 0.3.1+ - fixed nonreentrant
@nonreentrant('lock')
@external
def remove_liquidity(amount: uint256) -> uint256[2]:
    # Update state before external call
    self._burn(msg.sender, amount)
    self._update_reserves()
    raw_call(msg.sender, b"", value=amount)
    return [amount, amount]
""",
        attack_scenario=(
            "1. Vyper compiler bug: @nonreentrant didn't work\n"
            "2. Attacker calls remove_liquidity\n"
            "3. During ETH transfer, callback triggers\n"
            "4. Callback calls remove_liquidity again\n"
            "5. Pool state not updated, double withdrawal\n"
            "6. Multiple pools affected: alETH, msETH, pETH"
        ),
        severity="critical",
        category="reentrancy",
        real_exploit="Curve/Vyper ($70M, July 2023)",
        tags=["vyper", "reentrancy", "compiler-bug", "curve", "defi"],
    ),

    "sentiment-readonly-reentrancy-2023": VulnerabilityExample(
        id="sentiment-readonly-reentrancy-2023",
        swc_id="SWC-107",
        cwe_id="CWE-841",
        title="Sentiment Protocol Read-Only Reentrancy",
        description=(
            "Attacker exploited read-only reentrancy in Balancer pools. During a join/exit, "
            "the pool's rate is temporarily incorrect. Sentiment read this stale rate "
            "during the callback, allowing the attacker to borrow more than they should."
        ),
        vulnerable_code="""
function getPrice() external view returns (uint256) {
    // Gets rate from Balancer pool
    // During join/exit callback, this rate is STALE!
    return balancerPool.getRate();
}

function borrow(uint256 amount) external {
    uint256 collateralValue = getCollateralValue();  // Uses stale price!
    require(collateralValue >= amount * collateralRatio);
    // Borrow allowed with inflated collateral value
}
""",
        fixed_code="""
function getPrice() external view returns (uint256) {
    // Check if pool is in a callback (reentrancy)
    (,, uint256 lastJoinExitAmp,) = balancerPool.getLastJoinExitData();
    require(block.number > lastJoinExitBlock, "Pool in transition");
    return balancerPool.getRate();
}
""",
        attack_scenario=(
            "1. Attacker joins Balancer pool with callback\n"
            "2. During join, pool.getRate() returns stale value\n"
            "3. In callback, attacker calls Sentiment.borrow()\n"
            "4. Sentiment uses inflated collateral value\n"
            "5. Attacker borrows more than allowed\n"
            "6. Exit Balancer callback, rate normalizes\n"
            "7. Attacker has undercollateralized loan"
        ),
        severity="high",
        category="reentrancy",
        real_exploit="Sentiment ($1M, April 2023)",
        tags=["read-only-reentrancy", "balancer", "oracle", "defi"],
    ),

    "bonqdao-oracle-2023": VulnerabilityExample(
        id="bonqdao-oracle-2023",
        swc_id=None,
        cwe_id="CWE-20",
        title="BonqDAO Oracle Manipulation Attack",
        description=(
            "Attacker manipulated the price oracle used by BonqDAO by exploiting low "
            "liquidity in the Tellor oracle. They staked TRB and submitted a false price "
            "to liquidate other users' positions."
        ),
        vulnerable_code="""
function getPrice(address token) external view returns (uint256) {
    // Uses Tellor oracle without TWAP or validation
    (bool success, bytes memory data,) = tellor.getDataBefore(
        queryId, block.timestamp - 15 minutes
    );
    return abi.decode(data, (uint256));
    // No check for price deviation!
}
""",
        fixed_code="""
function getPrice(address token) external view returns (uint256) {
    uint256 price = tellor.getDataBefore(queryId, block.timestamp - 15 minutes);
    uint256 twapPrice = getTWAP(token, 1 hours);
    // Check deviation
    require(
        price >= twapPrice * 95 / 100 && price <= twapPrice * 105 / 100,
        "Price deviation too high"
    );
    return price;
}
""",
        attack_scenario=(
            "1. Stake 10 TRB in Tellor oracle\n"
            "2. Submit false WALBT price (100x actual)\n"
            "3. Wait for oracle update period\n"
            "4. Mint BEUR stablecoin against inflated collateral\n"
            "5. Submit false price (near zero)\n"
            "6. Liquidate other users' positions\n"
            "7. Profit $120M in BEUR and WALBT"
        ),
        severity="critical",
        category="oracle",
        real_exploit="BonqDAO ($120M, February 2023)",
        tags=["oracle", "tellor", "price-manipulation", "liquidation"],
    ),

    "hundred-finance-2023": VulnerabilityExample(
        id="hundred-finance-2023",
        swc_id="SWC-107",
        cwe_id="CWE-841",
        title="Hundred Finance Flash Loan Reentrancy",
        description=(
            "Attacker exploited reentrancy in Hundred Finance's Compound fork on Optimism. "
            "The redeemUnderlying function was vulnerable when using tokens with callbacks."
        ),
        vulnerable_code="""
function redeemUnderlying(uint redeemAmount) external returns (uint) {
    // Transfer underlying before updating state
    doTransferOut(msg.sender, redeemAmount);  // ERC-677 callback!
    // State update after transfer
    totalSupply = totalSupply - redeemTokens;
    accountTokens[msg.sender] = accountTokens[msg.sender] - redeemTokens;
}
""",
        fixed_code="""
function redeemUnderlying(uint redeemAmount) external nonReentrant returns (uint) {
    // Update state before transfer
    totalSupply = totalSupply - redeemTokens;
    accountTokens[msg.sender] = accountTokens[msg.sender] - redeemTokens;
    // Transfer after state update
    doTransferOut(msg.sender, redeemAmount);
}
""",
        attack_scenario=(
            "1. Flash loan WBTC\n"
            "2. Deposit as collateral in Hundred\n"
            "3. Borrow hWBTC\n"
            "4. Call redeemUnderlying with ERC-677 token\n"
            "5. In callback, borrow more (state not updated)\n"
            "6. Repeat until pool drained\n"
            "7. Repay flash loan, keep profit"
        ),
        severity="critical",
        category="reentrancy",
        real_exploit="Hundred Finance ($7M, April 2023)",
        tags=["reentrancy", "compound-fork", "flash-loan", "optimism"],
    ),

    "platypus-2023": VulnerabilityExample(
        id="platypus-2023",
        swc_id="SWC-107",
        cwe_id="CWE-841",
        title="Platypus Finance Flash Loan Attack",
        description=(
            "Attacker exploited a logic flaw in Platypus's MasterPlatypusV4 contract. "
            "The emergencyWithdraw function didn't properly check position solvency, "
            "allowing borrowers to withdraw collateral without repaying debt."
        ),
        vulnerable_code="""
function emergencyWithdraw(uint256 _pid) public {
    UserInfo storage user = userInfo[_pid][msg.sender];
    // No solvency check! User can have debt
    uint256 amount = user.amount;
    user.amount = 0;
    lpToken.safeTransfer(msg.sender, amount);
}
""",
        fixed_code="""
function emergencyWithdraw(uint256 _pid) public {
    UserInfo storage user = userInfo[_pid][msg.sender];
    // Check solvency before withdrawal
    require(getUserDebt(msg.sender) == 0, "Must repay debt first");
    uint256 amount = user.amount;
    user.amount = 0;
    lpToken.safeTransfer(msg.sender, amount);
}
""",
        attack_scenario=(
            "1. Flash loan 44M USDC\n"
            "2. Deposit USDC, get LP tokens\n"
            "3. Stake LP in MasterPlatypusV4\n"
            "4. Borrow USP stablecoin (max)\n"
            "5. emergencyWithdraw (keeps borrowed USP!)\n"
            "6. Swap USP for USDC\n"
            "7. Repay flash loan, profit $8.5M"
        ),
        severity="critical",
        category="flash-loan",
        real_exploit="Platypus Finance ($8.5M, February 2023)",
        tags=["flash-loan", "emergency-withdraw", "solvency", "avalanche"],
    ),

    "socket-bungee-2024": VulnerabilityExample(
        id="socket-bungee-2024",
        swc_id="SWC-112",
        cwe_id="CWE-20",
        title="Socket/Bungee Arbitrary Call Vulnerability",
        description=(
            "Attacker exploited insufficient validation in Socket's bridge aggregator. "
            "The performAction function allowed arbitrary calls to any address with "
            "user-controlled data, enabling token theft."
        ),
        vulnerable_code="""
function performAction(
    address target,
    bytes calldata data,
    uint256 value
) external {
    // No validation of target or data!
    (bool success,) = target.call{value: value}(data);
    require(success, "Action failed");
}
""",
        fixed_code="""
mapping(address => bool) public whitelistedTargets;

function performAction(
    address target,
    bytes calldata data,
    uint256 value
) external {
    require(whitelistedTargets[target], "Target not allowed");
    require(validateCalldata(data), "Invalid calldata");
    (bool success,) = target.call{value: value}(data);
    require(success, "Action failed");
}
""",
        attack_scenario=(
            "1. Identify Socket contracts with user approvals\n"
            "2. Craft calldata for transferFrom on approved tokens\n"
            "3. Call performAction(tokenAddress, transferFromData)\n"
            "4. Socket executes transferFrom as itself\n"
            "5. Steal approved tokens from users\n"
            "6. Repeat for all approved addresses"
        ),
        severity="critical",
        category="arbitrary-call",
        real_exploit="Socket/Bungee ($3.3M, January 2024)",
        tags=["arbitrary-call", "bridge", "approval", "transferFrom"],
    ),

    "lifi-2024": VulnerabilityExample(
        id="lifi-2024",
        swc_id="SWC-112",
        cwe_id="CWE-20",
        title="LI.FI Protocol Arbitrary Call Attack",
        description=(
            "Similar to Socket, LI.FI's bridge aggregator had an arbitrary call vulnerability. "
            "Attacker could call any contract with any data through the bridge interface, "
            "stealing tokens that users had approved to the LI.FI contract."
        ),
        vulnerable_code="""
struct SwapData {
    address callTo;      // Arbitrary address!
    bytes callData;      // Arbitrary data!
    address receivingAsset;
}

function _executeSwaps(SwapData[] calldata swaps) internal {
    for (uint i = 0; i < swaps.length; i++) {
        // No validation - can call anything!
        (bool success,) = swaps[i].callTo.call(swaps[i].callData);
        require(success);
    }
}
""",
        fixed_code="""
mapping(address => bool) public allowedDEXs;

function _executeSwaps(SwapData[] calldata swaps) internal {
    for (uint i = 0; i < swaps.length; i++) {
        require(allowedDEXs[swaps[i].callTo], "DEX not allowed");
        require(_isValidSwapSelector(swaps[i].callData), "Invalid function");
        (bool success,) = swaps[i].callTo.call(swaps[i].callData);
        require(success);
    }
}
""",
        attack_scenario=(
            "1. Find wallets with token approvals to LI.FI\n"
            "2. Craft SwapData with callTo=token, callData=transferFrom\n"
            "3. Call LI.FI swap function with malicious SwapData\n"
            "4. LI.FI executes transferFrom, stealing approved tokens\n"
            "5. $12M stolen from ~200 wallets"
        ),
        severity="critical",
        category="arbitrary-call",
        real_exploit="LI.FI ($12M, July 2024)",
        tags=["arbitrary-call", "bridge", "aggregator", "approval"],
    ),

    "penpie-2024": VulnerabilityExample(
        id="penpie-2024",
        swc_id="SWC-107",
        cwe_id="CWE-841",
        title="Penpie Reentrancy Attack",
        description=(
            "Attacker exploited reentrancy in Penpie's staking mechanism. During reward "
            "distribution, the contract made external calls before updating state, "
            "allowing multiple claims of the same rewards."
        ),
        vulnerable_code="""
function claimRewards(address _user) external {
    uint256 pending = pendingRewards[_user];
    // External call before state update!
    rewardToken.safeTransfer(_user, pending);
    // State update too late
    pendingRewards[_user] = 0;
    lastClaimTime[_user] = block.timestamp;
}
""",
        fixed_code="""
function claimRewards(address _user) external nonReentrant {
    uint256 pending = pendingRewards[_user];
    // State update first
    pendingRewards[_user] = 0;
    lastClaimTime[_user] = block.timestamp;
    // External call last
    rewardToken.safeTransfer(_user, pending);
}
""",
        attack_scenario=(
            "1. Stake tokens in Penpie\n"
            "2. Wait for rewards to accumulate\n"
            "3. Call claimRewards\n"
            "4. In token's callback, call claimRewards again\n"
            "5. Rewards claimed multiple times\n"
            "6. Drain $27M in rewards"
        ),
        severity="critical",
        category="reentrancy",
        real_exploit="Penpie ($27M, September 2024)",
        tags=["reentrancy", "staking", "rewards", "pendle"],
    ),

    "uwu-lend-2024": VulnerabilityExample(
        id="uwu-lend-2024",
        swc_id=None,
        cwe_id="CWE-20",
        title="UwU Lend Oracle Manipulation",
        description=(
            "Attacker manipulated Curve pool prices used by UwU Lend's oracle. "
            "By flash loaning and swapping large amounts, they temporarily moved "
            "the price to liquidate positions and borrow against inflated collateral."
        ),
        vulnerable_code="""
function getPrice(address token) external view returns (uint256) {
    // Uses spot price from Curve pool
    uint256 price = curvePool.get_virtual_price();
    // No TWAP, no deviation check
    return price;
}
""",
        fixed_code="""
function getPrice(address token) external view returns (uint256) {
    uint256 spotPrice = curvePool.get_virtual_price();
    uint256 twapPrice = curveTWAP.consult(token, 30 minutes);
    // Max 5% deviation allowed
    require(
        spotPrice >= twapPrice * 95 / 100 &&
        spotPrice <= twapPrice * 105 / 100,
        "Price manipulation detected"
    );
    return twapPrice;
}
""",
        attack_scenario=(
            "1. Flash loan $40M in stablecoins\n"
            "2. Swap to manipulate Curve pool price\n"
            "3. sUSDe price inflated by 20%\n"
            "4. Deposit sUSDe as collateral at inflated value\n"
            "5. Borrow stablecoins (over-borrow)\n"
            "6. Swap back, repay flash loan\n"
            "7. Keep undercollateralized loans worth $20M"
        ),
        severity="critical",
        category="oracle",
        real_exploit="UwU Lend ($20M, June 2024)",
        tags=["oracle", "price-manipulation", "curve", "flash-loan"],
    ),

    "velocore-2024": VulnerabilityExample(
        id="velocore-2024",
        swc_id="SWC-101",
        cwe_id="CWE-682",
        title="Velocore Rounding Error Attack",
        description=(
            "Attacker exploited a precision loss vulnerability in Velocore's AMM math. "
            "By carefully crafting swap amounts, they could extract more tokens than "
            "they should due to rounding in their favor."
        ),
        vulnerable_code="""
function swap(uint256 amountIn) external returns (uint256 amountOut) {
    // Precision loss: multiply then divide
    uint256 fee = amountIn * feeRate / 10000;
    uint256 amountInAfterFee = amountIn - fee;
    // More precision loss in output calculation
    amountOut = reserve1 * amountInAfterFee / (reserve0 + amountInAfterFee);
    // Attacker exploits cumulative rounding errors
}
""",
        fixed_code="""
function swap(uint256 amountIn) external returns (uint256 amountOut) {
    // Use higher precision for intermediate calculations
    uint256 amountInScaled = amountIn * PRECISION;
    uint256 feeScaled = amountInScaled * feeRate / 10000;
    uint256 amountInAfterFee = (amountInScaled - feeScaled) / PRECISION;
    // Round down output (favor protocol)
    amountOut = reserve1 * amountInAfterFee / (reserve0 + amountInAfterFee);
}
""",
        attack_scenario=(
            "1. Analyze AMM math for precision loss\n"
            "2. Calculate optimal amounts that maximize rounding benefit\n"
            "3. Execute many small swaps\n"
            "4. Each swap extracts tiny extra amount\n"
            "5. Aggregate to $6.8M profit"
        ),
        severity="high",
        category="precision",
        real_exploit="Velocore ($6.8M, June 2024)",
        tags=["precision", "rounding", "amm", "math", "zkSync"],
    ),

    "radiant-capital-2024": VulnerabilityExample(
        id="radiant-capital-2024",
        swc_id="SWC-107",
        cwe_id="CWE-841",
        title="Radiant Capital Flash Loan Attack",
        description=(
            "Attacker exploited a known issue in Compound/Aave forks where new markets "
            "can be manipulated. By being the first depositor in a new market, they "
            "inflated the exchange rate to steal from later depositors."
        ),
        vulnerable_code="""
function mint(uint256 amount) external returns (uint256 shares) {
    if (totalSupply == 0) {
        // First depositor gets 1:1 shares
        shares = amount;
    } else {
        // Later depositors calculated based on exchange rate
        shares = amount * totalSupply / totalAssets();
        // Can round to 0 if exchange rate is high!
    }
    _mint(msg.sender, shares);
}
""",
        fixed_code="""
uint256 constant INITIAL_SHARES = 1000;

function mint(uint256 amount) external returns (uint256 shares) {
    if (totalSupply == 0) {
        // Mint dead shares to prevent manipulation
        _mint(address(0), INITIAL_SHARES);
        shares = amount;
    } else {
        shares = amount * totalSupply / totalAssets();
        require(shares > 0, "Zero shares");
    }
    _mint(msg.sender, shares);
}
""",
        attack_scenario=(
            "1. New RDNT-WETH market deployed\n"
            "2. Attacker deposits 1 wei, gets 1 share\n"
            "3. Donate large amount directly to market\n"
            "4. Exchange rate becomes huge\n"
            "5. Victim deposits, gets 0 shares (rounds down)\n"
            "6. Attacker withdraws, gets victim's funds"
        ),
        severity="critical",
        category="flash-loan",
        real_exploit="Radiant Capital ($4.5M, January 2024)",
        tags=["first-depositor", "exchange-rate", "compound-fork", "arbitrum"],
    ),

    "gamma-strategies-2024": VulnerabilityExample(
        id="gamma-strategies-2024",
        swc_id=None,
        cwe_id="CWE-20",
        title="Gamma Strategies Price Manipulation",
        description=(
            "Attacker manipulated Uniswap V3 concentrated liquidity positions managed by "
            "Gamma Strategies. By moving the price, they caused the vault to rebalance "
            "at unfavorable prices, extracting value."
        ),
        vulnerable_code="""
function rebalance() external {
    // Get current price from pool
    (uint160 sqrtPriceX96,,,,,,) = pool.slot0();
    // Rebalance based on current (manipulated) price
    if (isOutOfRange(sqrtPriceX96)) {
        _removeLiquidity();
        _addLiquidity(newLower, newUpper);  // At bad price!
    }
}
""",
        fixed_code="""
function rebalance() external onlyKeeper {
    (uint160 sqrtPriceX96,,,,,,) = pool.slot0();
    uint256 twapPrice = getTWAP(30 minutes);
    // Check price deviation
    require(
        _priceDeviation(sqrtPriceX96, twapPrice) < MAX_DEVIATION,
        "Price manipulated"
    );
    if (isOutOfRange(sqrtPriceX96)) {
        _removeLiquidity();
        _addLiquidity(newLower, newUpper);
    }
}
""",
        attack_scenario=(
            "1. Flash loan large amount of token0\n"
            "2. Swap to move Uniswap V3 price\n"
            "3. Trigger Gamma vault rebalance\n"
            "4. Vault sells at manipulated price\n"
            "5. Swap back to original price\n"
            "6. Vault now has less value\n"
            "7. Profit from arbitrage"
        ),
        severity="high",
        category="price-manipulation",
        real_exploit="Gamma Strategies ($6M, January 2024)",
        tags=["uniswap-v3", "concentrated-liquidity", "rebalance", "manipulation"],
    ),

    "abracadabra-2024": VulnerabilityExample(
        id="abracadabra-2024",
        swc_id="SWC-101",
        cwe_id="CWE-682",
        title="Abracadabra Cauldron Rounding Error",
        description=(
            "Attacker exploited rounding errors in Abracadabra's Cauldron contracts. "
            "The borrow and repay functions had precision issues that could be "
            "exploited through repeated small transactions."
        ),
        vulnerable_code="""
function borrow(address to, uint256 amount) public returns (uint256 part) {
    // Rounding in favor of borrower
    uint256 feeAmount = amount * BORROW_FEE / PRECISION;
    totalBorrow.elastic += amount + feeAmount;
    totalBorrow.base += part;
    // part calculation has precision loss
    part = totalBorrow.base == 0 ? amount : amount * totalBorrow.base / totalBorrow.elastic;
}
""",
        fixed_code="""
function borrow(address to, uint256 amount) public returns (uint256 part) {
    uint256 feeAmount = amount * BORROW_FEE / PRECISION;
    // Round up (favor protocol) for borrow
    part = totalBorrow.base == 0
        ? amount
        : (amount * totalBorrow.base + totalBorrow.elastic - 1) / totalBorrow.elastic;
    totalBorrow.elastic += amount + feeAmount;
    totalBorrow.base += part;
}
""",
        attack_scenario=(
            "1. Identify cauldron with rounding vulnerability\n"
            "2. Execute many small borrows/repays\n"
            "3. Each transaction rounds in attacker's favor\n"
            "4. Accumulate small gains into significant profit\n"
            "5. Extract $6.5M from affected cauldrons"
        ),
        severity="high",
        category="precision",
        real_exploit="Abracadabra ($6.5M, January 2024)",
        tags=["rounding", "precision", "cauldron", "mim"],
    ),

    "prisma-finance-2024": VulnerabilityExample(
        id="prisma-finance-2024",
        swc_id="SWC-107",
        cwe_id="CWE-841",
        title="Prisma Finance Flash Loan Attack",
        description=(
            "Attacker exploited a vulnerability in Prisma's MigrateTroveZap contract. "
            "The contract allowed migrating troves without proper validation, enabling "
            "flash loan attacks to extract value."
        ),
        vulnerable_code="""
function migrateTrove(
    address fromTroveManager,
    address toTroveManager,
    uint256 maxFeePercentage
) external {
    // Close old trove
    ITroveManager(fromTroveManager).closeTrove(msg.sender);
    // Open new trove with same collateral
    // Vulnerability: no validation between close and open
    ITroveManager(toTroveManager).openTrove(...);
}
""",
        fixed_code="""
function migrateTrove(
    address fromTroveManager,
    address toTroveManager,
    uint256 maxFeePercentage
) external nonReentrant {
    uint256 collateral = getTroveCollateral(fromTroveManager, msg.sender);
    uint256 debt = getTroveDebt(fromTroveManager, msg.sender);
    ITroveManager(fromTroveManager).closeTrove(msg.sender);
    require(
        collateral > 0 && debt > 0,
        "Invalid trove state"
    );
    ITroveManager(toTroveManager).openTrove(collateral, debt, ...);
}
""",
        attack_scenario=(
            "1. Flash loan large amount of stETH\n"
            "2. Open trove in Prisma with collateral\n"
            "3. Borrow mkUSD against collateral\n"
            "4. Use MigrateTroveZap to manipulate trove\n"
            "5. Extract extra collateral during migration\n"
            "6. Repay flash loan, keep $12M profit"
        ),
        severity="critical",
        category="flash-loan",
        real_exploit="Prisma Finance ($12M, March 2024)",
        tags=["flash-loan", "trove", "migration", "prisma"],
    ),

    "thala-2024": VulnerabilityExample(
        id="thala-2024",
        swc_id=None,
        cwe_id="CWE-284",
        title="Thala Governance Attack",
        description=(
            "Attacker exploited Thala's governance mechanism by accumulating voting power "
            "through flash loans and manipulating proposal execution timing."
        ),
        vulnerable_code="""
function executeProposal(uint256 proposalId) external {
    Proposal storage proposal = proposals[proposalId];
    require(proposal.forVotes > proposal.againstVotes, "Not passed");
    require(block.timestamp > proposal.endTime, "Voting ongoing");
    // No timelock! Executes immediately
    (bool success,) = proposal.target.call(proposal.callData);
    require(success);
}
""",
        fixed_code="""
uint256 constant EXECUTION_DELAY = 2 days;

function executeProposal(uint256 proposalId) external {
    Proposal storage proposal = proposals[proposalId];
    require(proposal.forVotes > proposal.againstVotes, "Not passed");
    require(block.timestamp > proposal.endTime, "Voting ongoing");
    require(
        block.timestamp > proposal.endTime + EXECUTION_DELAY,
        "Timelock not expired"
    );
    (bool success,) = proposal.target.call(proposal.callData);
    require(success);
}
""",
        attack_scenario=(
            "1. Acquire significant voting power\n"
            "2. Create malicious governance proposal\n"
            "3. Vote for proposal with accumulated power\n"
            "4. Execute proposal immediately (no timelock)\n"
            "5. Drain treasury of $25.5M"
        ),
        severity="critical",
        category="governance",
        real_exploit="Thala ($25.5M, November 2024)",
        tags=["governance", "timelock", "treasury", "aptos"],
    ),

    "dough-finance-2024": VulnerabilityExample(
        id="dough-finance-2024",
        swc_id="SWC-107",
        cwe_id="CWE-841",
        title="Dough Finance Flash Loan Attack",
        description=(
            "Attacker exploited Dough Finance's ConnectorDeleverageParaswap contract "
            "which lacked proper validation of external call data, allowing "
            "manipulation of the deleverage process."
        ),
        vulnerable_code="""
function deleverage(
    address tokenIn,
    uint256 amount,
    bytes calldata paraswapData
) external {
    // Withdraw collateral
    aave.withdraw(tokenIn, amount, address(this));
    // Execute arbitrary swap via Paraswap
    // No validation of paraswapData!
    (bool success,) = PARASWAP.call(paraswapData);
    require(success);
    // Repay debt
    aave.repay(debtToken, receivedAmount, 2, msg.sender);
}
""",
        fixed_code="""
function deleverage(
    address tokenIn,
    uint256 amount,
    bytes calldata paraswapData
) external nonReentrant {
    require(validateParaswapData(paraswapData, tokenIn), "Invalid swap");
    uint256 balanceBefore = IERC20(tokenOut).balanceOf(address(this));
    aave.withdraw(tokenIn, amount, address(this));
    (bool success,) = PARASWAP.call(paraswapData);
    require(success);
    uint256 received = IERC20(tokenOut).balanceOf(address(this)) - balanceBefore;
    require(received >= minExpected, "Slippage too high");
    aave.repay(debtToken, received, 2, msg.sender);
}
""",
        attack_scenario=(
            "1. Flash loan USDC\n"
            "2. Deposit as collateral in Aave via Dough\n"
            "3. Borrow against collateral\n"
            "4. Call deleverage with malicious paraswapData\n"
            "5. ParaswapData actually calls transferFrom on approved tokens\n"
            "6. Steal approved tokens from contract\n"
            "7. Repay flash loan, profit $2M"
        ),
        severity="critical",
        category="flash-loan",
        real_exploit="Dough Finance ($2M, July 2024)",
        tags=["flash-loan", "arbitrary-call", "paraswap", "defi"],
    ),

    "bedrock-2024": VulnerabilityExample(
        id="bedrock-2024",
        swc_id="SWC-107",
        cwe_id="CWE-841",
        title="Bedrock Reentrancy Attack",
        description=(
            "Attacker exploited a reentrancy vulnerability in Bedrock's uniBTC vault. "
            "The mint function was vulnerable to reentrancy through the underlying "
            "token's transfer callback."
        ),
        vulnerable_code="""
function mint(uint256 amount) external returns (uint256 shares) {
    // Transfer tokens in (callback possible!)
    IERC20(underlying).safeTransferFrom(msg.sender, address(this), amount);
    // Calculate shares based on current state
    shares = _calculateShares(amount);
    // Mint shares
    _mint(msg.sender, shares);
}
""",
        fixed_code="""
function mint(uint256 amount) external nonReentrant returns (uint256 shares) {
    uint256 balanceBefore = IERC20(underlying).balanceOf(address(this));
    IERC20(underlying).safeTransferFrom(msg.sender, address(this), amount);
    uint256 received = IERC20(underlying).balanceOf(address(this)) - balanceBefore;
    shares = _calculateShares(received);
    _mint(msg.sender, shares);
}
""",
        attack_scenario=(
            "1. Deploy malicious token with callback\n"
            "2. Call mint on Bedrock vault\n"
            "3. During transfer callback, call mint again\n"
            "4. State not updated, mint extra shares\n"
            "5. Repeat to drain $2M"
        ),
        severity="critical",
        category="reentrancy",
        real_exploit="Bedrock ($2M, September 2024)",
        tags=["reentrancy", "vault", "unibtc", "mint"],
    ),

    "woofi-2024": VulnerabilityExample(
        id="woofi-2024",
        swc_id=None,
        cwe_id="CWE-20",
        title="WooFi Price Manipulation Attack",
        description=(
            "Attacker manipulated WooFi's custom oracle by exploiting low liquidity "
            "in the reference pools. The oracle used prices from pools that could be "
            "easily moved with flash loans."
        ),
        vulnerable_code="""
function getPrice(address token) external view returns (uint256) {
    // Uses spot price from WooPP pool
    uint256 price = wooPP.querySellQuote(baseToken, 1e18);
    // No TWAP, no deviation check, easily manipulated
    return price;
}
""",
        fixed_code="""
function getPrice(address token) external view returns (uint256) {
    uint256 wooPrice = wooPP.querySellQuote(baseToken, 1e18);
    uint256 chainlinkPrice = getChainlinkPrice(token);
    // Max 3% deviation from Chainlink
    require(
        wooPrice >= chainlinkPrice * 97 / 100 &&
        wooPrice <= chainlinkPrice * 103 / 100,
        "Price deviation"
    );
    return wooPrice;
}
""",
        attack_scenario=(
            "1. Flash loan $10M in WOO tokens\n"
            "2. Swap to move WooPP price significantly\n"
            "3. Borrow against manipulated collateral value\n"
            "4. Swap back, repay flash loan\n"
            "5. Keep $8.75M in undercollateralized loans"
        ),
        severity="critical",
        category="oracle",
        real_exploit="WooFi ($8.75M, March 2024)",
        tags=["oracle", "price-manipulation", "flash-loan", "arbitrum"],
    ),

    "balancer-2023": VulnerabilityExample(
        id="balancer-2023",
        swc_id=None,
        cwe_id="CWE-682",
        title="Balancer Flash Loan Rate Manipulation",
        description=(
            "Attacker exploited Balancer's rate provider mechanism. By manipulating "
            "the rate during a flash loan callback, they could affect the pool's "
            "internal accounting and extract value."
        ),
        vulnerable_code="""
function getRate() external view returns (uint256) {
    // Rate can be stale during flash loan callback
    return _cachedRate;
}

function joinPool(uint256[] memory amounts) external {
    uint256 rate = rateProvider.getRate();  // Stale during callback!
    uint256 bptOut = calculateBPT(amounts, rate);
    _mint(msg.sender, bptOut);
}
""",
        fixed_code="""
function getRate() external view returns (uint256) {
    // Prevent rate reads during flash loan
    require(!_inFlashLoan(), "Rate unavailable during flash loan");
    return _cachedRate;
}

function joinPool(uint256[] memory amounts) external {
    require(!_inFlashLoan(), "Cannot join during flash loan");
    uint256 rate = rateProvider.getRate();
    uint256 bptOut = calculateBPT(amounts, rate);
    _mint(msg.sender, bptOut);
}
""",
        attack_scenario=(
            "1. Take flash loan from Balancer pool\n"
            "2. In callback, rate is temporarily stale\n"
            "3. Join pool using stale (incorrect) rate\n"
            "4. Get more BPT than should\n"
            "5. Exit flash loan, rate normalizes\n"
            "6. Exit pool at correct rate, profit"
        ),
        severity="high",
        category="flash-loan",
        real_exploit="Balancer ($2M, August 2023)",
        tags=["balancer", "flash-loan", "rate-manipulation", "read-only-reentrancy"],
    ),
}


class VulnerabilityRAG:
    """
    Retrieval-Augmented Generation for vulnerability analysis.

    Uses a knowledge base of known vulnerabilities to enhance
    detection and provide contextual information.
    """

    def __init__(
        self,
        ollama_base_url: str = "http://localhost:11434",
        model: str = "deepseek-coder:6.7b",
        top_k: int = 3,
    ):
        """
        Initialize the RAG system.

        Args:
            ollama_base_url: Ollama API URL
            model: Model for analysis
            top_k: Number of similar vulns to retrieve
        """
        self.base_url = ollama_base_url
        self.model = model
        self.top_k = top_k

        # Combine all knowledge bases
        self.knowledge_base: Dict[str, VulnerabilityExample] = {}
        self.knowledge_base.update(SWC_REGISTRY)
        self.knowledge_base.update(EXPLOIT_EXAMPLES)

        # Build tag index for fast lookup
        self._tag_index: Dict[str, List[str]] = {}
        self._build_tag_index()

        logger.info(f"VulnerabilityRAG initialized with {len(self.knowledge_base)} examples")

    def _build_tag_index(self) -> None:
        """Build inverted index from tags to example IDs."""
        for example_id, example in self.knowledge_base.items():
            for tag in example.tags:
                tag_lower = tag.lower()
                if tag_lower not in self._tag_index:
                    self._tag_index[tag_lower] = []
                self._tag_index[tag_lower].append(example_id)

            # Also index by category
            if example.category:
                cat_lower = example.category.lower()
                if cat_lower not in self._tag_index:
                    self._tag_index[cat_lower] = []
                self._tag_index[cat_lower].append(example_id)

    def search_similar(
        self,
        finding: Dict[str, Any],
        code_context: str = "",
    ) -> List[SimilarVuln]:
        """
        Find similar vulnerabilities in the knowledge base.

        Args:
            finding: The vulnerability finding to search for
            code_context: Code context for better matching

        Returns:
            List of similar vulnerabilities
        """
        vuln_type = finding.get("type", "").lower()
        description = finding.get("description", "").lower()

        candidates = set()
        scores: Dict[str, float] = {}

        # Search by type/category
        if vuln_type in self._tag_index:
            for example_id in self._tag_index[vuln_type]:
                candidates.add(example_id)
                scores[example_id] = scores.get(example_id, 0) + 0.5

        # Search by keywords in description
        keywords = self._extract_keywords(description)
        for keyword in keywords:
            if keyword in self._tag_index:
                for example_id in self._tag_index[keyword]:
                    candidates.add(example_id)
                    scores[example_id] = scores.get(example_id, 0) + 0.3

        # Search by SWC ID if present
        swc_id = finding.get("swc_id") or finding.get("swc")
        if swc_id:
            swc_upper = swc_id.upper()
            if swc_upper in self.knowledge_base:
                candidates.add(swc_upper)
                scores[swc_upper] = scores.get(swc_upper, 0) + 0.8

        # Score candidates and get top-k
        results = []
        for example_id in candidates:
            example = self.knowledge_base.get(example_id)
            if not example:
                continue

            score = scores.get(example_id, 0)

            # Boost for code similarity
            if code_context and example.vulnerable_code:
                code_sim = self._code_similarity(code_context, example.vulnerable_code)
                score += code_sim * 0.4

            # Normalize score
            score = min(1.0, score)

            results.append(SimilarVuln(
                example=example,
                similarity_score=score,
                relevance_reason=self._get_relevance_reason(finding, example),
            ))

        # Sort by score and return top-k
        results.sort(key=lambda x: -x.similarity_score)
        return results[:self.top_k]

    async def analyze_with_context(
        self,
        finding: Dict[str, Any],
        code: str,
    ) -> EnhancedFinding:
        """
        Enhance a finding with RAG context.

        Args:
            finding: The vulnerability finding
            code: Contract source code

        Returns:
            Enhanced finding with context
        """
        # Find similar vulnerabilities
        similar = self.search_similar(finding, code)

        # Build context for LLM
        context_parts = []
        for sim in similar:
            context_parts.append(
                f"**Similar Vulnerability: {sim.example.title}**\n"
                f"- Description: {sim.example.description[:200]}...\n"
                f"- Real Exploit: {sim.example.real_exploit or 'N/A'}\n"
                f"- Relevance: {sim.relevance_reason}"
            )

        context = "\n\n".join(context_parts) if context_parts else "No similar vulnerabilities found."

        # Query LLM for enhanced analysis
        enhanced = await self._enhance_with_llm(finding, code, context)

        # Build references
        references = []
        for sim in similar:
            if sim.example.swc_id:
                references.append(f"https://swcregistry.io/docs/{sim.example.swc_id}")
            if sim.example.real_exploit:
                references.append(sim.example.real_exploit)

        # Calculate confidence adjustment
        confidence_adjustment = 0.0
        if similar:
            # Boost confidence if similar known vulnerabilities exist
            confidence_adjustment = min(0.2, similar[0].similarity_score * 0.25)

        return EnhancedFinding(
            original_finding=finding,
            similar_vulns=similar,
            enhanced_description=enhanced.get("description", finding.get("description", "")),
            exploit_scenario=enhanced.get("exploit_scenario"),
            remediation_code=enhanced.get("remediation_code"),
            confidence_adjustment=confidence_adjustment,
            additional_references=references,
        )

    async def _enhance_with_llm(
        self,
        finding: Dict[str, Any],
        code: str,
        context: str,
    ) -> Dict[str, Any]:
        """Use LLM to enhance finding with RAG context."""
        import aiohttp

        prompt = f"""Based on similar known vulnerabilities, enhance this security finding:

## Original Finding
- Type: {finding.get('type', 'unknown')}
- Severity: {finding.get('severity', 'unknown')}
- Description: {finding.get('description', '')}

## Similar Known Vulnerabilities
{context}

## Code Context
```solidity
{code[:2000]}
```

Provide:
1. Enhanced description with more technical detail
2. Specific exploit scenario for this contract
3. Remediation code example

Respond in JSON:
{{
    "description": "Enhanced description",
    "exploit_scenario": "Step-by-step exploit",
    "remediation_code": "Fixed code example"
}}
"""

        try:
            async with aiohttp.ClientSession() as session:
                async with session.post(
                    f"{self.base_url}/api/chat",
                    json={
                        "model": self.model,
                        "messages": [
                            {"role": "system", "content": "You are an expert smart contract security auditor."},
                            {"role": "user", "content": prompt}
                        ],
                        "stream": False,
                        "options": {"temperature": 0.1},
                    },
                    timeout=aiohttp.ClientTimeout(total=60),
                ) as resp:
                    if resp.status == 200:
                        data = await resp.json()
                        content = data.get("message", {}).get("content", "")
                        return self._parse_json(content)
        except Exception as e:
            logger.warning(f"LLM enhancement failed: {e}")

        return {}

    def _extract_keywords(self, text: str) -> List[str]:
        """Extract relevant keywords from text."""
        # Common vulnerability-related keywords
        relevant_keywords = [
            "reentrancy", "reentrant", "overflow", "underflow", "arithmetic",
            "access", "control", "owner", "admin", "delegatecall", "proxy",
            "selfdestruct", "suicide", "timestamp", "block", "random",
            "oracle", "price", "flash", "loan", "front-running", "sandwich",
            "signature", "ecrecover", "hash", "collision", "transfer",
            "call", "send", "withdraw", "deposit", "balance", "token",
            "erc20", "erc721", "approval", "allowance", "governance", "vote",
        ]

        found = []
        text_lower = text.lower()
        for keyword in relevant_keywords:
            if keyword in text_lower:
                found.append(keyword)

        return found[:5]  # Limit to top 5

    def _code_similarity(self, code1: str, code2: str) -> float:
        """Calculate simple code similarity."""
        # Extract function calls and keywords
        import re

        def extract_tokens(code: str) -> set:
            # Extract function calls
            calls = set(re.findall(r'\b(\w+)\s*\(', code.lower()))
            # Extract keywords
            keywords = set(re.findall(r'\b(require|if|revert|transfer|send|call|delegatecall)\b', code.lower()))
            return calls.union(keywords)

        tokens1 = extract_tokens(code1)
        tokens2 = extract_tokens(code2)

        if not tokens1 or not tokens2:
            return 0.0

        intersection = tokens1.intersection(tokens2)
        union = tokens1.union(tokens2)

        return len(intersection) / len(union) if union else 0.0

    def _get_relevance_reason(
        self,
        finding: Dict[str, Any],
        example: VulnerabilityExample,
    ) -> str:
        """Explain why an example is relevant."""
        reasons = []

        vuln_type = finding.get("type", "").lower()
        if example.category and vuln_type in example.category.lower():
            reasons.append(f"Same vulnerability category ({example.category})")

        if example.real_exploit:
            reasons.append(f"Known real-world exploit: {example.real_exploit}")

        if example.swc_id:
            reasons.append(f"SWC Registry: {example.swc_id}")

        return "; ".join(reasons) if reasons else "Pattern similarity"

    def _parse_json(self, content: str) -> Dict[str, Any]:
        """Parse JSON from response."""
        try:
            start = content.find('{')
            end = content.rfind('}') + 1
            if start >= 0 and end > start:
                return json.loads(content[start:end])
        except json.JSONDecodeError:
            pass
        return {}

    def get_swc_info(self, swc_id: str) -> Optional[VulnerabilityExample]:
        """Get SWC Registry information by ID."""
        return self.knowledge_base.get(swc_id.upper())

    def get_all_categories(self) -> List[str]:
        """Get all vulnerability categories in the knowledge base."""
        categories = set()
        for example in self.knowledge_base.values():
            if example.category:
                categories.add(example.category)
        return sorted(categories)


# Export
__all__ = [
    "VulnerabilityRAG",
    "VulnerabilityExample",
    "SimilarVuln",
    "EnhancedFinding",
    "SWC_REGISTRY",
    "EXPLOIT_EXAMPLES",
]
