# { "Depends": "py-genlayer:test" } import json from dataclasses import dataclass from genlayer import * ERROR_EXPECTED = "[EXPECTED]" ERROR_LLM = "[LLM_ERROR]" SEVERITY_TOLERANCE = 1 @allow_storage @dataclass class Audit: audit_id: u256 submitter: Address code_or_address: str chain: str risk_level: str vulnerability_count: u256 findings_json: str summary: str def _risk_from_count(critical: int, high: int) -> str: if critical > 0: return "CRITICAL" if high > 0: return "HIGH" if high == 0 and critical == 0: return "LOW" return "MEDIUM" def _audit_prompt(code_or_address: str, chain: str) -> str: return f"""You are AutoAudit, an AI smart contract security auditor. Analyze the following smart contract code or address for vulnerabilities. Code/Address: \"\"\"{code_or_address}\"\"\" Chain: {chain} Perform a comprehensive security audit checking for: 1. Reentrancy attacks 2. Integer overflow/underflow 3. Access control issues 4. Front-running vulnerabilities 5. Flash loan attack vectors 6. Unchecked external calls 7. Gas optimization issues 8. Logic errors 9. Centralization risks 10. Oracle manipulation Return ONLY valid JSON: {{ "risk_level": "CRITICAL|HIGH|MEDIUM|LOW", "vulnerabilities": [ {{ "severity": "CRITICAL|HIGH|MEDIUM|LOW|INFO", "title": "", "description": "", "location": "

",
      "recommendation": ""
    }}
  ],
  "summary": "<2-3 sentence overall assessment>",
  "gas_optimizations": ["", ""],
  "best_practices_missing": ["", ""]
}}
"""


class AutoAudit(gl.Contract):
    owner: Address
    audit_count: u256
    audits: TreeMap[u256, Audit]
    audit_ids: DynArray[u256]

    def __init__(self):
        self.owner = gl.message.sender_address
        self.audit_count = u256(0)

    def _run_audit(self, code_or_address: str, chain: str) -> dict:
        prompt = _audit_prompt(code_or_address, chain)

        def leader_fn():
            result = gl.nondet.exec_prompt(prompt, response_format="json")
            if not isinstance(result, dict):
                raise gl.vm.UserError(f"{ERROR_LLM} Invalid response")
            vulns = result.get("vulnerabilities", [])
            critical = sum(1 for v in vulns if v.get("severity") == "CRITICAL")
            high = sum(1 for v in vulns if v.get("severity") == "HIGH")
            result["risk_level"] = _risk_from_count(critical, high)
            result["vulnerability_count"] = len(vulns)
            return result

        def validator_fn(leaders_res):
            if not isinstance(leaders_res, gl.vm.Return):
                return False
            v = leader_fn()
            leader_count = leaders_res.calldata.get("vulnerability_count", 0)
            validator_count = v.get("vulnerability_count", 0)
            if abs(leader_count - validator_count) > SEVERITY_TOLERANCE:
                return False
            return leaders_res.calldata.get("risk_level") == v.get("risk_level")

        return gl.vm.run_nondet_unsafe(leader_fn, validator_fn)

    @gl.public.write
    def submit_audit(self, code_or_address: str, chain: str) -> u256:
        cleaned = code_or_address.strip()
        if len(cleaned) < 10:
            raise gl.vm.UserError(f"{ERROR_EXPECTED} Code or address too short")
        if len(cleaned) > 50000:
            raise gl.vm.UserError(f"{ERROR_EXPECTED} Code exceeds maximum length")

        analysis = self._run_audit(cleaned, chain.strip())

        audit_id = self.audit_count
        findings_json = json.dumps(analysis.get("vulnerabilities", []), sort_keys=True)

        self.audits[audit_id] = Audit(
            audit_id=audit_id,
            submitter=gl.message.sender_address,
            code_or_address=cleaned[:500],
            chain=chain.strip(),
            risk_level=str(analysis.get("risk_level", "UNKNOWN")),
            vulnerability_count=int(analysis.get("vulnerability_count", 0)),
            findings_json=findings_json,
            summary=str(analysis.get("summary", ""))[:500],
        )
        self.audit_ids.append(audit_id)
        self.audit_count = audit_id + 1

        return audit_id

    @gl.public.view
    def get_audit(self, audit_id: u256) -> dict:
        if audit_id not in self.audits:
            raise gl.vm.UserError(f"{ERROR_EXPECTED} Audit not found")
        a = self.audits[audit_id]
        return {
            "audit_id": a.audit_id,
            "submitter": a.submitter.as_hex,
            "code_or_address": a.code_or_address,
            "chain": a.chain,
            "risk_level": a.risk_level,
            "vulnerability_count": a.vulnerability_count,
            "findings": json.loads(a.findings_json),
            "summary": a.summary,
        }

    @gl.public.view
    def get_all_audits(self) -> list:
        results = []
        count = int(self.audit_count)
        i = count - 1
        while i >= 0 and len(results) < 50:
            a = self.audits[i]
            preview = a.code_or_address
            if len(preview) > 80:
                preview = preview[:80] + "..."
            results.append({
                "audit_id": a.audit_id,
                "submitter": a.submitter.as_hex,
                "code_preview": preview,
                "chain": a.chain,
                "risk_level": a.risk_level,
                "vulnerability_count": a.vulnerability_count,
            })
            i -= 1
        return results