forge-guard 0.3.6

//! Auditor agents — AI-powered audit specialists for Solidity code.
//!
//! Each auditor implements the [`AuditorAgent`] trait and targets a specific
//! domain (security, gas, business logic). Auditors construct a domain-specific
//! prompt, send it to a configured [`LlmProvider`], and parse the structured
//! JSON response into a list of [`AuditorFinding`] values.

use crate::ai::providers::LlmProvider;
use crate::core::{ForgeGuardError, Severity};
use std::sync::Arc;

use super::{AuditContext, AuditorFinding};

// ─────────────────────────────────────────────────────────────
// AuditorAgent trait
// ─────────────────────────────────────────────────────────────

/// An AI auditor that analyzes Solidity source code for a specific domain.
pub trait AuditorAgent: Send + Sync {
    /// Human-readable auditor name (e.g. "security-auditor").
    fn name(&self) -> &'static str;

    /// The domain this auditor specializes in.
    fn domain(&self) -> &'static str;

    /// Analyze source code and return a list of findings.
    fn analyze(&self, context: &AuditContext) -> Result<Vec<AuditorFinding>, ForgeGuardError>;
}

// ─────────────────────────────────────────────────────────────
// Security Auditor
// ─────────────────────────────────────────────────────────────

/// Security auditor — detects vulnerabilities, access control issues,
/// reentrancy, oracle manipulation, etc.
pub struct SecurityAuditor {
    provider: Arc<dyn LlmProvider>,
    max_chunk_size: usize,
}

impl SecurityAuditor {
    /// Create a new security auditor backed by the given LLM provider.
    ///
    /// `max_chunk_size` controls source-code chunking (in bytes) for
    /// contracts that exceed the context window. Defaults to 10 000.
    pub fn new(provider: Arc<dyn LlmProvider>) -> Self {
        Self {
            provider,
            max_chunk_size: 10_000,
        }
    }

    /// Build the system prompt for security audit.
    fn system_prompt() -> &'static str {
        r#"You are a world-class Solidity security auditor. Analyze the provided Solidity source code for security vulnerabilities.

Return a JSON object with the following structure — no markdown, no code fences:
{
  "findings": [
    {
      "title": "Short vulnerability title",
      "description": "Detailed description of the issue, where it occurs, and why it is a problem",
      "severity": "critical|high|medium|low|info",
      "line_numbers": [45, 47],
      "recommendation": "Specific actionable fix recommendation",
      "category": "Reentrancy|AccessControl|Oracle|Arithmetic|DeFi|Upgradeability|Gas|Logic|Cryptography|Compliance|Other"
    }
  ]
}

Check for these vulnerability categories (in priority order):
1. **Reentrancy** — external calls before state updates, missing CEI pattern, missing reentrancy guards
2. **Access Control** — unprotected sensitive functions, missing onlyOwner/onlyRole modifiers, tx.origin usage, improper initialization
3. **Oracle Manipulation** — single-source price feeds, missing TWAP, unchecked oracle return values
4. **Arithmetic Issues** — uncheckable overflow/underflow (post Solidity 0.8), unsafe casting
5. **DeFi Logic** — flash loan attacks, sandwich attacks, liquidity manipulation, incorrect fee calculations
6. **Upgradeability** — storage collision, missing __gap, unsafe delegatecall, initializer front-running
7. **Cryptography** — weak signature schemes, missing nonce/replay protection, ecrecover pitfalls
8. **Compliance / Business Logic** — logic errors, race conditions, incorrect state transitions

Be thorough but precise. Only report genuine issues with high confidence. For each finding, assign the correct severity according to real-world impact potential. Do NOT include commentary outside the JSON object."#
    }

    /// Chunk source code by contract boundaries, keeping each chunk under `max_chunk_size`.
    fn chunk_source<'a>(&self, source: &'a str) -> Vec<&'a str> {
        if source.len() <= self.max_chunk_size {
            return vec![source];
        }

        let mut chunks = Vec::new();
        let mut remaining = source;

        while !remaining.is_empty() {
            if remaining.len() <= self.max_chunk_size {
                chunks.push(remaining);
                break;
            }

            // Try to split at a contract boundary within the chunk
            let chunk_end = remaining[..self.max_chunk_size]
                .rfind("\ncontract ")
                .or_else(|| {
                    remaining[..self.max_chunk_size]
                        .rfind("\nlibrary ")
                        .or_else(|| {
                            remaining[..self.max_chunk_size]
                                .rfind("\ninterface ")
                                .or_else(|| remaining[..self.max_chunk_size].rfind('\n'))
                        })
                })
                .map(|i| i + 1) // include the newline
                .unwrap_or(self.max_chunk_size);

            chunks.push(&remaining[..chunk_end]);
            remaining = &remaining[chunk_end..];
        }

        chunks
    }
}

impl AuditorAgent for SecurityAuditor {
    fn name(&self) -> &'static str {
        "security-auditor"
    }

    fn domain(&self) -> &'static str {
        "Security"
    }

    fn analyze(&self, context: &AuditContext) -> Result<Vec<AuditorFinding>, ForgeGuardError> {
        let mut all_findings = Vec::new();
        let chunks = self.chunk_source(&context.source_code);

        for chunk in &chunks {
            let user_prompt = format!(
                r#"Analyze this Solidity file for security vulnerabilities.

File: {}
Compiler: {}

```solidity
{}
```"#,
                context.file_name, context.compiler_version, chunk
            );

            let response = self.provider.call(Self::system_prompt(), &user_prompt)?;
            let findings = parse_findings_json(&response, self.name())?;
            all_findings.extend(findings);
        }

        Ok(all_findings)
    }
}

// ─────────────────────────────────────────────────────────────
// Gas Auditor
// ─────────────────────────────────────────────────────────────

/// Gas auditor — detects expensive patterns and suggests optimizations.
pub struct GasAuditor {
    provider: Arc<dyn LlmProvider>,
}

impl GasAuditor {
    pub fn new(provider: Arc<dyn LlmProvider>) -> Self {
        Self { provider }
    }

    fn system_prompt() -> &'static str {
        r#"You are a Solidity gas optimization expert. Analyze the provided Solidity code and suggest gas optimizations.

Return a JSON object with this structure:
{
  "findings": [
    {
      "title": "Gas optimization title",
      "description": "How much gas is wasted and the pattern causing it",
      "severity": "medium|low|info",
      "line_numbers": [12],
      "recommendation": "Specific gas-saving rewrite",
      "category": "Gas"
    }
  ]
}

Look for:
1. Storage vs memory: reading/writing storage repeatedly in loops
2. Unchecked arithmetic blocks for gas savings pre-Solidity 0.8
3. Redundant state reads
4. Packing structs for tighter storage (use uint128/uint64 instead of uint256)
5. Using 'delete' vs zero-assignment
6. Short-circuiting require() statements
7. Using calldata instead of memory for read-only params
8. Pre-increment vs post-increment (++i vs i++)
9. Caching array length in for loops
10. Unnecessary intermediate variables"#
    }
}

impl AuditorAgent for GasAuditor {
    fn name(&self) -> &'static str {
        "gas-auditor"
    }

    fn domain(&self) -> &'static str {
        "Gas"
    }

    fn analyze(&self, context: &AuditContext) -> Result<Vec<AuditorFinding>, ForgeGuardError> {
        let user_prompt = format!(
            r#"Analyze this Solidity file for gas optimization opportunities.

File: {}

```solidity
{}
```"#,
            context.file_name, context.source_code
        );

        let response = self.provider.call(Self::system_prompt(), &user_prompt)?;
        parse_findings_json(&response, self.name())
    }
}

// ─────────────────────────────────────────────────────────────
// Logic Auditor
// ─────────────────────────────────────────────────────────────

/// Logic auditor — verifies business logic correctness.
pub struct LogicAuditor {
    provider: Arc<dyn LlmProvider>,
}

impl LogicAuditor {
    pub fn new(provider: Arc<dyn LlmProvider>) -> Self {
        Self { provider }
    }

    fn system_prompt() -> &'static str {
        r#"You are a Solidity business logic auditor. Analyze the provided Solidity code for logic errors, race conditions, and incorrect state transitions.

Return a JSON object with this structure:
{
  "findings": [
    {
      "title": "Logic issue title",
      "description": "How the logic is incorrect and what the consequences could be",
      "severity": "critical|high|medium|low|info",
      "line_numbers": [22, 25],
      "recommendation": "Specific fix for the logic error",
      "category": "Logic|Compliance|DeFi"
    }
  ]
}

Check for:
1. State transition correctness: can the contract reach an invalid state?
2. Off-by-one errors in comparisons (>= vs >)
3. Rounding errors: division before multiplication truncating to zero
4. Shares/asset calculation correctness
5. Deadline / timelock logic
6. Pausable logic: can funds get stuck?
7. Fee calculation correctness
8. Cross-contract invariant consistency
9. Withdrawal logic: anyone can call, double-claim prevention
10. Approval / allowance logic correctness
11. Balance calculation and totalSupply tracking"#
    }
}

impl AuditorAgent for LogicAuditor {
    fn name(&self) -> &'static str {
        "logic-auditor"
    }

    fn domain(&self) -> &'static str {
        "Business Logic"
    }

    fn analyze(&self, context: &AuditContext) -> Result<Vec<AuditorFinding>, ForgeGuardError> {
        let user_prompt = format!(
            r#"Analyze this Solidity file for business logic errors.

File: {}

```solidity
{}
```"#,
            context.file_name, context.source_code
        );

        let response = self.provider.call(Self::system_prompt(), &user_prompt)?;
        parse_findings_json(&response, self.name())
    }
}

// ─────────────────────────────────────────────────────────────
// Response Parsing
// ─────────────────────────────────────────────────────────────

/// Parse a JSON response from an LLM into structured findings.
///
/// The response is expected to match:
/// ```json
/// { "findings": [{ "title": "...", "description": "...", ... }] }
/// ```
pub fn parse_findings_json(
    response: &str,
    auditor_name: &str,
) -> Result<Vec<AuditorFinding>, ForgeGuardError> {
    // Strip markdown code fences if present
    let cleaned = response
        .trim()
        .trim_start_matches("```json")
        .trim_start_matches("```")
        .trim_end_matches("```")
        .trim();

    let parsed: serde_json::Value = serde_json::from_str(cleaned).map_err(|e| {
        ForgeGuardError::Parse(format!(
            "[{auditor_name}] Failed to parse response as JSON: {e}\nRaw response (first 200 chars): {}",
            &response[..response.len().min(200)]
        ))
    })?;

    let findings_array = parsed["findings"].as_array().ok_or_else(|| {
        ForgeGuardError::Parse(format!(
            "[{auditor_name}] Response missing 'findings' array. Keys: {:?}",
            parsed
                .as_object()
                .map(|o| o.keys().cloned().collect::<Vec<_>>())
        ))
    })?;

    let mut findings = Vec::with_capacity(findings_array.len());
    for entry in findings_array {
        let title = entry["title"]
            .as_str()
            .unwrap_or("Unknown issue")
            .to_owned();
        let description = entry["description"].as_str().unwrap_or("").to_owned();
        let severity_str = entry["severity"].as_str().unwrap_or("medium");
        let recommendation = entry["recommendation"].as_str().unwrap_or("").to_owned();
        let category = entry["category"].as_str().unwrap_or("Other").to_owned();

        let line_numbers: Vec<usize> = entry["line_numbers"]
            .as_array()
            .map(|arr| {
                arr.iter()
                    .filter_map(|v| v.as_u64().map(|n| n as usize))
                    .collect()
            })
            .unwrap_or_default();

        let confidence = entry["confidence"].as_f64().unwrap_or(0.85);

        let severity = match severity_str.to_lowercase().as_str() {
            "critical" => Severity::Critical,
            "high" => Severity::High,
            "medium" => Severity::Medium,
            "low" => Severity::Low,
            "info" | "informational" => Severity::Informational,
            _ => Severity::Medium,
        };

        findings.push(AuditorFinding {
            title,
            description,
            confidence,
            severity,
            suggestion: recommendation,
            line_numbers,
            category,
        });
    }

    Ok(findings)
}

/// Build an `AuditContext` from a source file and parsed Solidity data.
pub fn build_audit_context(source_code: &str, file_name: &str) -> AuditContext {
    // Extract compiler version from pragma
    let compiler_version = extract_compiler_version(source_code);

    AuditContext {
        source_code: source_code.to_owned(),
        file_name: file_name.to_owned(),
        compiler_version,
        additional: Default::default(),
    }
}

/// Try to extract the Solidity compiler version pragma.
fn extract_compiler_version(source: &str) -> String {
    for line in source.lines() {
        let trimmed = line.trim();
        if trimmed.starts_with("pragma solidity") {
            return trimmed
                .trim_start_matches("pragma solidity")
                .trim()
                .trim_end_matches(';')
                .to_owned();
        }
    }
    "unknown".into()
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::ai::providers::OllamaProvider;

    // ── Parse findings ──────────────────────────────────────

    #[test]
    fn test_parse_empty_findings() {
        let findings = parse_findings_json(r#"{"findings": []}"#, "test").unwrap();
        assert!(findings.is_empty());
    }

    #[test]
    fn test_parse_single_finding() {
        let json = r#"{
            "findings": [{
                "title": "Reentrancy",
                "description": "External call before state update",
                "severity": "high",
                "line_numbers": [15, 17],
                "recommendation": "Apply CEI pattern",
                "category": "Reentrancy",
                "confidence": 0.95
            }]
        }"#;
        let findings = parse_findings_json(json, "test").unwrap();
        assert_eq!(findings.len(), 1);
        assert_eq!(findings[0].title, "Reentrancy");
        assert_eq!(findings[0].severity, Severity::High);
        assert_eq!(findings[0].line_numbers, vec![15, 17]);
        assert!((findings[0].confidence - 0.95).abs() < 0.01);
    }

    #[test]
    fn test_parse_multiple_findings() {
        let json = r#"{
            "findings": [
                {
                    "title": "Access Control",
                    "description": "Missing onlyOwner",
                    "severity": "high",
                    "line_numbers": [5],
                    "recommendation": "Add modifier",
                    "category": "AccessControl"
                },
                {
                    "title": "Gas",
                    "description": "Loop gas waste",
                    "severity": "low",
                    "line_numbers": [12],
                    "recommendation": "Cache length",
                    "category": "Gas"
                }
            ]
        }"#;
        let findings = parse_findings_json(json, "test").unwrap();
        assert_eq!(findings.len(), 2);
        assert_eq!(findings[0].title, "Access Control");
        assert_eq!(findings[0].severity, Severity::High);
        assert_eq!(findings[1].title, "Gas");
        assert_eq!(findings[1].severity, Severity::Low);
    }

    #[test]
    fn test_parse_strips_markdown_fences() {
        let json = "```json\n{\"findings\": [{\"title\": \"Test\", \"description\": \"desc\", \"severity\": \"medium\", \"line_numbers\": [], \"recommendation\": \"fix\", \"category\": \"Other\"}]}\n```";
        let findings = parse_findings_json(json, "test").unwrap();
        assert_eq!(findings.len(), 1);
        assert_eq!(findings[0].title, "Test");
    }

    #[test]
    fn test_parse_missing_findings_key() {
        let result = parse_findings_json(r#"{"error": "something"}"#, "test");
        assert!(result.is_err());
        assert!(result
            .unwrap_err()
            .to_string()
            .contains("missing 'findings'"));
    }

    #[test]
    fn test_parse_invalid_json() {
        let result = parse_findings_json("not json at all", "test");
        assert!(result.is_err());
    }

    #[test]
    fn test_parse_severity_mapping() {
        let test_cases = [
            ("critical", Severity::Critical),
            ("high", Severity::High),
            ("medium", Severity::Medium),
            ("low", Severity::Low),
            ("info", Severity::Informational),
            ("unknown", Severity::Medium),
        ];
        for (input, expected) in &test_cases {
            let json = format!(
                r#"{{"findings": [{{"title":"T","description":"d","severity":"{input}","line_numbers":[],"recommendation":"r","category":"C"}}]}}"#
            );
            let findings = parse_findings_json(&json, "test").unwrap();
            assert_eq!(
                findings[0].severity, *expected,
                "severity '{input}' should map to {expected:?}"
            );
        }
    }

    // ── Compiler version extraction ─────────────────────────

    #[test]
    fn test_extract_compiler_version() {
        let src = "// SPDX\npragma solidity ^0.8.20;\ncontract C {}\n";
        assert_eq!(extract_compiler_version(src), "^0.8.20");
    }

    #[test]
    fn test_extract_compiler_version_none() {
        let src = "contract C {}\n";
        assert_eq!(extract_compiler_version(src), "unknown");
    }

    // ── Chunking ────────────────────────────────────────────

    #[test]
    fn test_chunk_source_small() {
        let auditor = SecurityAuditor::new(Arc::new(
            // Dummy provider — won't be called
            OllamaProvider::new(Some("http://127.0.0.1:1".into()), "test", 0.0),
        ));
        let src = "contract Small {}";
        let chunks = auditor.chunk_source(src);
        assert_eq!(chunks.len(), 1);
        assert_eq!(chunks[0], src);
    }

    #[test]
    fn test_chunk_source_splits_on_contract() {
        // Set max_chunk_size small to force splitting
        let mut auditor = SecurityAuditor::new(Arc::new(OllamaProvider::new(
            Some("http://127.0.0.1:1".into()),
            "test",
            0.0,
        )));
        auditor.max_chunk_size = 50;

        let src = "contract First {\n  uint x;\n}\n\ncontract Second {\n  uint y;\n}\n";
        let chunks = auditor.chunk_source(src);
        assert!(
            chunks.len() >= 2,
            "should split into at least 2 chunks, got {}",
            chunks.len()
        );
        assert!(chunks[0].contains("contract First"));
        assert!(chunks.last().unwrap().contains("contract Second"));
    }

    // ── Provider error propagation ──────────────────────────

    #[test]
    fn test_provider_call_error() {
        // Pointing at a port that won't respond
        let provider = Arc::new(OllamaProvider::new(
            Some("http://127.0.0.1:1".into()),
            "test",
            0.0,
        ));
        let auditor = SecurityAuditor::new(provider);
        let ctx = AuditContext {
            source_code: "contract C {}".into(),
            file_name: "test.sol".into(),
            compiler_version: "0.8.20".into(),
            additional: Default::default(),
        };
        let result = auditor.analyze(&ctx);
        assert!(result.is_err());
    }
}