contract_analyzer.py

"""
contract_analyzer.py
====================
Static analysis of uploaded Solidity smart contracts.
Extracts ABI, functions, state variables, modifiers, events,
and builds a contract interaction map without requiring solc.
"""

import re
import json
from pathlib import Path
from typing import Dict, List, Any, Optional


# ─────────────────────────────────────────────────────────────────── #
# SOLIDITY PARSER  (regex-based — zero external deps)                 #
# ─────────────────────────────────────────────────────────────────── #

_UINT_RE   = re.compile(r'\buint\d*\b')
_INT_RE    = re.compile(r'\bint\d*\b')
_ADDR_RE   = re.compile(r'\baddress(?:\s+payable)?\b')
_BOOL_RE   = re.compile(r'\bbool\b')
_BYTES_RE  = re.compile(r'\bbytes\d*\b')
_STRING_RE = re.compile(r'\bstring\b')


def _parse_type(tok: str) -> str:
    """Normalise a Solidity type token."""
    tok = tok.strip()
    for pat in [_UINT_RE, _INT_RE, _ADDR_RE, _BOOL_RE, _BYTES_RE, _STRING_RE]:
        if pat.match(tok):
            return tok
    if tok.startswith("mapping") or tok.startswith("(") or "[" in tok:
        return tok
    return tok or "uint256"


def _strip_comments(src: str) -> str:
    """Remove // and /* */ comments from Solidity source."""
    # Block comments
    src = re.sub(r'/\*.*?\*/', ' ', src, flags=re.DOTALL)
    # Line comments
    src = re.sub(r'//[^\n]*', '', src)
    return src


def _parse_params(raw: str) -> List[Dict[str, str]]:
    """Parse a parameter list string into list of {name, type}."""
    raw = raw.strip()
    if not raw:
        return []
    params = []
    depth = 0
    current = ""
    for ch in raw:
        if ch in "(<":
            depth += 1
            current += ch
        elif ch in ")>":
            depth -= 1
            current += ch
        elif ch == "," and depth == 0:
            params.append(current.strip())
            current = ""
        else:
            current += ch
    if current.strip():
        params.append(current.strip())

    result = []
    for p in params:
        parts = p.split()
        if not parts:
            continue
        typ = parts[0]
        # Handle memory/calldata/storage keywords
        filtered = [x for x in parts if x not in ("memory", "calldata", "storage", "payable")]
        if len(filtered) >= 2:
            result.append({"type": filtered[0], "name": filtered[-1]})
        elif len(filtered) == 1:
            result.append({"type": filtered[0], "name": ""})
    return result


class ContractAnalyzer:
    """
    Parses a Solidity source file and extracts:
    - Contract names
    - Functions (with visibility, mutability, params, returns)
    - State variables
    - Modifiers
    - Events
    - Identified vulnerabilities (static heuristics)
    - Interaction map (which functions call which)
    """

    def __init__(self, source: str, filename: str = "contract.sol"):
        self.source   = source
        self.filename = filename
        self._clean   = _strip_comments(source)
        self.results: Dict[str, Any] = {}

    # ── Public API ──────────────────────────────────────────────────

    def analyze(self) -> Dict[str, Any]:
        contracts   = self._extract_contracts()
        functions   = self._extract_functions()
        state_vars  = self._extract_state_vars()
        modifiers   = self._extract_modifiers()
        events      = self._extract_events()
        vulns       = self._detect_vulnerabilities(functions)
        interaction = self._build_interaction_map(functions)

        self.results = {
            "filename":         self.filename,
            "contracts":        contracts,
            "functions":        functions,
            "state_variables":  state_vars,
            "modifiers":        modifiers,
            "events":           events,
            "vulnerabilities":  vulns,
            "interaction_map":  interaction,
            "stats": {
                "total_contracts": len(contracts),
                "total_functions": len(functions),
                "total_state_vars": len(state_vars),
                "total_events":    len(events),
                "total_modifiers": len(modifiers),
                "vuln_count":      len(vulns),
            }
        }
        return self.results

    # ── Extractors ───────────────────────────────────────────────────

    def _extract_contracts(self) -> List[Dict[str, Any]]:
        pattern = re.compile(
            r'\b(contract|interface|library|abstract\s+contract)\s+(\w+)'
            r'(?:\s+is\s+([\w\s,]+?))?(?:\s*\{)',
            re.MULTILINE
        )
        results = []
        for m in pattern.finditer(self._clean):
            kind      = m.group(1).strip().split()[0]
            name      = m.group(2)
            inherits  = [x.strip() for x in (m.group(3) or "").split(",") if x.strip()]
            results.append({"kind": kind, "name": name, "inherits": inherits})
        return results

    def _extract_functions(self) -> List[Dict[str, Any]]:
        # Match: function name(params) visibility mutability modifiers returns(...)
        pattern = re.compile(
            r'\bfunction\s+(\w+)\s*\(([^)]*)\)'
            r'((?:\s+(?:public|private|internal|external|pure|view|payable|virtual|override|'
            r'nonReentrant|\w+))*)'
            r'(?:\s+returns\s*\(([^)]*)\))?'
            r'\s*(?:\{|;)',
            re.MULTILINE
        )
        results = []
        for m in pattern.finditer(self._clean):
            name        = m.group(1)
            params_raw  = m.group(2)
            qualifiers  = m.group(3) or ""
            returns_raw = m.group(4) or ""

            quals_list = qualifiers.split()

            visibility = "internal"
            for v in ("public", "private", "internal", "external"):
                if v in quals_list:
                    visibility = v
                    break

            mutability = "nonpayable"
            for mm in ("pure", "view", "payable"):
                if mm in quals_list:
                    mutability = mm
                    break

            custom_mods = [
                q for q in quals_list
                if q not in ("public", "private", "internal", "external",
                             "pure", "view", "payable", "virtual", "override",
                             "nonReentrant", "returns")
            ]

            results.append({
                "name":       name,
                "params":     _parse_params(params_raw),
                "returns":    _parse_params(returns_raw),
                "visibility": visibility,
                "mutability": mutability,
                "modifiers":  custom_mods,
            })
        return results

    def _extract_state_vars(self) -> List[Dict[str, str]]:
        pattern = re.compile(
            r'^\s*(uint\d*|int\d*|address(?:\s+payable)?|bool|bytes\d*|string|'
            r'mapping\([^)]+\)|IERC20|IERC721|I\w+)\s+'
            r'(public|private|internal|immutable|constant|\s)*'
            r'(\w+)\s*(?:=\s*[^;]+)?;',
            re.MULTILINE
        )
        results = []
        seen = set()
        for m in pattern.finditer(self._clean):
            typ  = m.group(1).strip()
            name = m.group(3).strip()
            qual = (m.group(2) or "").strip()
            if name not in seen and name not in ("function", "event", "modifier"):
                seen.add(name)
                results.append({"type": typ, "name": name, "visibility": qual or "internal"})
        return results

    def _extract_modifiers(self) -> List[Dict[str, Any]]:
        pattern = re.compile(r'\bmodifier\s+(\w+)\s*\(([^)]*)\)', re.MULTILINE)
        results = []
        for m in pattern.finditer(self._clean):
            results.append({
                "name":   m.group(1),
                "params": _parse_params(m.group(2)),
            })
        return results

    def _extract_events(self) -> List[Dict[str, Any]]:
        pattern = re.compile(r'\bevent\s+(\w+)\s*\(([^)]*)\)', re.MULTILINE)
        results = []
        for m in pattern.finditer(self._clean):
            results.append({
                "name":   m.group(1),
                "params": _parse_params(m.group(2)),
            })
        return results

    # ── Vulnerability Heuristics ─────────────────────────────────────

    def _detect_vulnerabilities(self, functions: List[Dict]) -> List[Dict[str, str]]:
        vulns = []
        src   = self._clean

        # Reentrancy: external call before state update
        if re.search(r'\.call\{value', src) and re.search(r'balance\s*=', src):
            vulns.append({
                "type": "Reentrancy",
                "severity": "CRITICAL",
                "description": "External call with value found. Ensure state updates precede external calls (Checks-Effects-Interactions).",
            })

        # Unprotected oracle setter
        for fn in functions:
            if ("oracle" in fn["name"].lower() or "price" in fn["name"].lower() or "set" in fn["name"].lower()):
                if fn["visibility"] in ("public", "external") and not fn["modifiers"]:
                    vulns.append({
                        "type": "Unprotected Oracle / Price Manipulation",
                        "severity": "HIGH",
                        "description": f"Function `{fn['name']}` is {fn['visibility']} with no access-control modifier.",
                        "function": fn["name"],
                    })

        # tx.origin auth
        if "tx.origin" in src:
            vulns.append({
                "type": "tx.origin Authentication",
                "severity": "HIGH",
                "description": "tx.origin used for authentication — vulnerable to phishing attacks.",
            })

        # Integer overflow (pre-0.8)
        if "pragma solidity ^0.7" in src or "pragma solidity 0.6" in src:
            vulns.append({
                "type": "Integer Overflow/Underflow",
                "severity": "HIGH",
                "description": "Solidity <0.8 detected — arithmetic is unchecked; use SafeMath.",
            })

        # Timestamp dependence
        if "block.timestamp" in src:
            vulns.append({
                "type": "Timestamp Dependence",
                "severity": "LOW",
                "description": "block.timestamp used — miners can manipulate slightly.",
            })

        # Delegatecall
        if "delegatecall" in src:
            vulns.append({
                "type": "Delegatecall Risk",
                "severity": "HIGH",
                "description": "delegatecall found — ensure implementation contracts are trusted and immutable.",
            })

        # selfdestruct
        if "selfdestruct" in src:
            vulns.append({
                "type": "Selfdestruct",
                "severity": "MEDIUM",
                "description": "selfdestruct present — can permanently destroy the contract and send ETH.",
            })

        # Flash loan surface
        if re.search(r'flashLoan|FlashLoan|flash_loan', src):
            vulns.append({
                "type": "Flash Loan Attack Surface",
                "severity": "MEDIUM",
                "description": "Flash loan functionality detected — ensure reentrancy guards and price oracle integrity.",
            })

        return vulns

    # ── Interaction Map ──────────────────────────────────────────────

    def _build_interaction_map(self, functions: List[Dict]) -> Dict[str, List[str]]:
        """
        For each function, find which other contract functions it references.
        """
        fn_names = {f["name"] for f in functions}
        interaction_map = {}
        # Split source into rough blocks by function
        source = self._clean
        for fn in functions:
            calls = set()
            name  = fn["name"]
            # Find the block after this function's signature
            pat = re.compile(
                r'\bfunction\s+' + re.escape(name) + r'\s*\([^{]*\{',
                re.DOTALL
            )
            m = pat.search(source)
            if m:
                start = m.end()
                # Find next function or end of source
                next_fn = re.search(r'\bfunction\s+\w+\s*\(', source[start:])
                block  = source[start: start + (next_fn.start() if next_fn else len(source[start:]))]
                for other in fn_names:
                    if other != name and re.search(r'\b' + re.escape(other) + r'\s*\(', block):
                        calls.add(other)
            interaction_map[name] = sorted(calls)
        return interaction_map


# ─────────────────────────────────────────────────────────────────── #
# HELPER: load from file or string                                    #
# ─────────────────────────────────────────────────────────────────── #

def analyze_source(source: str, filename: str = "contract.sol") -> Dict[str, Any]:
    """Public entry point — analyze Solidity source code."""
    return ContractAnalyzer(source, filename).analyze()


def analyze_file(path: str) -> Dict[str, Any]:
    """Load a .sol file and analyze it."""
    p = Path(path)
    if not p.exists() or not p.suffix == ".sol":
        raise ValueError(f"Invalid Solidity file: {path}")
    return analyze_source(p.read_text(encoding="utf-8"), filename=p.name)