forge_guard/templates/

mod.rs

//! Template library — pre-built audit profiles for common contract types.
//!
//! Templates adjust enabled checks, severity thresholds, and scoring weights
//! to focus on the contract's specific risk profile.

use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use std::path::PathBuf;

/// An audit template configuration.
#[derive(Debug, Clone, Default, Serialize, Deserialize)]
pub struct AuditTemplate {
    /// Template name identifier.
    pub name: String,
    /// Human-readable description.
    pub description: String,
    /// Checks to enable (check IDs or categories).
    #[serde(default)]
    pub enabled_checks: Vec<String>,
    /// Checks to disable (check IDs or categories).
    #[serde(default)]
    pub disabled_checks: Vec<String>,
    /// Severity thresholds per category (category -> min score).
    #[serde(default)]
    pub min_scores: HashMap<String, u8>,
    /// Focus areas — categories weighted higher in scoring.
    #[serde(default)]
    pub focus_areas: Vec<String>,
}

impl AuditTemplate {
    /// Create a new template.
    pub fn new(name: &str, description: &str) -> Self {
        Self {
            name: name.to_string(),
            description: description.to_string(),
            enabled_checks: Vec::new(),
            disabled_checks: Vec::new(),
            min_scores: HashMap::new(),
            focus_areas: Vec::new(),
        }
    }
}

/// Built-in template definitions.
fn builtin_templates() -> Vec<AuditTemplate> {
    vec![
        AuditTemplate {
            name: "erc20".into(),
            description: "Focuses on ERC20 compliance, approval bugs, permit replay, and token-specific vulnerabilities".into(),
            enabled_checks: vec![
                "FA-H-014".into(),  // ERC20 Issues
                "FA-M-005".into(),  // Unsafe Events
            ],
            disabled_checks: vec![],
            min_scores: HashMap::from([
                ("access_control".into(), 85u8),
                ("dependencies".into(), 80u8),
            ]),
            focus_areas: vec!["access_control".into(), "dependencies".into()],
        },
        AuditTemplate {
            name: "erc721".into(),
            description: "Focuses on NFT-specific patterns, royalties, metadata, and collection integrity".into(),
            enabled_checks: vec![
                "FA-M-005".into(),  // Unsafe Events (metadata)
            ],
            disabled_checks: vec![],
            min_scores: HashMap::from([
                ("access_control".into(), 85u8),
            ]),
            focus_areas: vec!["access_control".into(), "security".into()],
        },
        AuditTemplate {
            name: "defi".into(),
            description: "Focuses on lending, AMM, oracle dependencies, flash loans, and MEV resistance".into(),
            enabled_checks: vec![
                "FA-H-011".into(),  // Oracle Manipulation
                "FA-H-016".into(),  // Flash Loan Issues
                "FA-H-017".into(),  // MEV Issues
                "FA-M-003".into(),  // Timestamp Manipulation
            ],
            disabled_checks: vec![],
            min_scores: HashMap::from([
                ("exploit_resistance".into(), 80u8),
                ("security".into(), 75u8),
            ]),
            focus_areas: vec!["exploit_resistance".into(), "security".into(), "chain_compatibility".into()],
        },
        AuditTemplate {
            name: "bridge".into(),
            description: "Focuses on cross-chain messaging, validator sets, replay attacks, and bridge-specific patterns".into(),
            enabled_checks: vec![
                "FA-H-015".into(),  // Bridge Vulnerabilities
                "FA-H-018".into(),  // Cross-Chain Issues
                "FA-H-013".into(),  // Replay Attacks
            ],
            disabled_checks: vec![],
            min_scores: HashMap::from([
                ("chain_compatibility".into(), 90u8),
                ("security".into(), 80u8),
            ]),
            focus_areas: vec!["chain_compatibility".into(), "security".into()],
        },
        AuditTemplate {
            name: "upgradeable".into(),
            description: "Focuses on proxy patterns, storage layouts, initializers, and upgrade path safety".into(),
            enabled_checks: vec![
                "FA-H-007".into(),  // Storage Collision
                "FA-H-010".into(),  // Proxy Vulnerabilities
                "FA-H-022".into(),  // Unsafe Upgrade Paths
                "FA-H-023".into(),  // Clone Vulnerabilities
            ],
            disabled_checks: vec![],
            min_scores: HashMap::from([
                ("upgradeability".into(), 90u8),
                ("proxy_safety".into(), 90u8),
            ]),
            focus_areas: vec!["upgradeability".into(), "proxy_safety".into(), "architecture".into()],
        },
    ]
}

/// Load user-custom templates from `~/.forge-guard/templates/*.json`.
fn load_user_templates() -> Vec<AuditTemplate> {
    let mut templates = Vec::new();

    if let Some(user_dir) = user_template_dir() {
        if user_dir.exists() {
            if let Ok(entries) = std::fs::read_dir(&user_dir) {
                for entry in entries.flatten() {
                    let path = entry.path();
                    if path.extension().is_some_and(|e| e == "json") {
                        if let Ok(content) = std::fs::read_to_string(&path) {
                            if let Ok(template) = serde_json::from_str::<AuditTemplate>(&content) {
                                templates.push(template);
                            }
                        }
                    }
                }
            }
        }
    }

    templates
}

/// Get a template by name (built-in or user-custom).
/// User templates take precedence over built-ins with the same name.
pub fn get_template(name: &str) -> Option<AuditTemplate> {
    // User templates override built-ins
    if let Some(t) = load_user_templates().into_iter().find(|t| t.name == name) {
        return Some(t);
    }
    builtin_templates().into_iter().find(|t| t.name == name)
}

/// List all available templates (built-in + user-custom).
/// User templates with the same name as a built-in replace it in the list,
/// matching what `get_template` will actually return.
pub fn list_templates() -> Vec<AuditTemplate> {
    let mut templates = load_user_templates();

    // Add built-ins that aren't shadowed by a user template
    for template in builtin_templates() {
        if !templates.iter().any(|t| t.name == template.name) {
            templates.push(template);
        }
    }

    templates
}

/// Apply an audit template to a project config.
///
/// Adjusts the security engine's enabled/disabled check lists so the audit
/// focuses on the template's risk profile.
pub fn apply_to_config(config: &mut crate::core::ProjectConfig, template: &AuditTemplate) {
    config.security.enabled_checks = template.enabled_checks.clone();
    config.security.disabled_checks = template.disabled_checks.clone();
}

/// Get the user template directory path.
fn user_template_dir() -> Option<PathBuf> {
    let home = if cfg!(target_os = "windows") {
        std::env::var("USERPROFILE").ok()
    } else {
        std::env::var("HOME").ok()
    };
    home.map(|h| PathBuf::from(h).join(".forge-guard").join("templates"))
}

/// Print template info for the terminal.
pub fn print_template(t: &AuditTemplate) {
    use colored::*;
    println!("  {} — {}", t.name.bold().cyan(), t.description);
    if !t.focus_areas.is_empty() {
        println!(
            "      {} {}",
            "Focus areas:".dimmed(),
            t.focus_areas.join(", ")
        );
    }
    if !t.enabled_checks.is_empty() {
        println!(
            "      {} {}",
            "Extra checks:".dimmed(),
            t.enabled_checks.join(", ")
        );
    }
    println!();
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_builtin_templates_exist() {
        let templates = builtin_templates();
        assert_eq!(templates.len(), 5);
        assert!(templates.iter().any(|t| t.name == "erc20"));
        assert!(templates.iter().any(|t| t.name == "erc721"));
        assert!(templates.iter().any(|t| t.name == "defi"));
        assert!(templates.iter().any(|t| t.name == "bridge"));
        assert!(templates.iter().any(|t| t.name == "upgradeable"));
    }

    #[test]
    fn test_get_template_found() {
        let t = get_template("defi").unwrap();
        assert_eq!(t.name, "defi");
        assert!(t.enabled_checks.contains(&"FA-H-011".to_string()));
        assert!(t.focus_areas.contains(&"exploit_resistance".to_string()));
    }

    #[test]
    fn test_get_template_not_found() {
        assert!(get_template("nonexistent").is_none());
    }

    #[test]
    fn test_list_templates_includes_builtin() {
        let templates = list_templates();
        assert!(templates.iter().any(|t| t.name == "erc20"));
    }

    #[test]
    fn test_template_has_descriptions() {
        for t in builtin_templates() {
            assert!(
                !t.description.is_empty(),
                "Template {} has no description",
                t.name
            );
        }
    }

    #[test]
    fn test_erc20_template_checks() {
        let t = get_template("erc20").unwrap();
        assert!(t.enabled_checks.contains(&"FA-H-014".to_string()));
        assert!(t.focus_areas.contains(&"access_control".to_string()));
    }

    #[test]
    fn test_bridge_template_checks() {
        let t = get_template("bridge").unwrap();
        assert!(t.enabled_checks.contains(&"FA-H-015".to_string()));
        assert!(t.enabled_checks.contains(&"FA-H-018".to_string()));
        assert!(t.focus_areas.contains(&"chain_compatibility".to_string()));
    }

    #[test]
    fn test_upgradeable_template_checks() {
        let t = get_template("upgradeable").unwrap();
        assert!(t.enabled_checks.contains(&"FA-H-007".to_string()));
        assert!(t.enabled_checks.contains(&"FA-H-010".to_string()));
        assert!(t.focus_areas.contains(&"upgradeability".to_string()));
    }

    #[test]
    fn test_template_serialization() {
        let t = get_template("defi").unwrap();
        let json = serde_json::to_string_pretty(&t).unwrap();
        assert!(json.contains("defi"));
        assert!(json.contains("oracle"));

        let deserialized: AuditTemplate = serde_json::from_str(&json).unwrap();
        assert_eq!(deserialized.name, "defi");
        assert_eq!(deserialized.focus_areas.len(), 3);
    }

    #[test]
    fn test_template_new_constructor() {
        let t = AuditTemplate::new("custom", "A custom template");
        assert_eq!(t.name, "custom");
        assert_eq!(t.description, "A custom template");
        assert!(t.enabled_checks.is_empty());
        assert!(t.focus_areas.is_empty());
    }

    #[test]
    fn test_user_template_dir_returns_some() {
        // Should return Some path (the actual dir may not exist)
        let dir = user_template_dir();
        assert!(dir.is_some());
    }

    #[test]
    fn test_print_template_does_not_panic() {
        let t = get_template("erc20").unwrap();
        // Should not panic
        print_template(&t);
    }

    #[test]
    fn test_apply_to_config_sets_checks() {
        let mut config = crate::core::ProjectConfig::default();
        let t = get_template("defi").unwrap();

        crate::templates::apply_to_config(&mut config, &t);

        assert!(config
            .security
            .enabled_checks
            .contains(&"FA-H-011".to_string()));
        assert!(config
            .security
            .enabled_checks
            .contains(&"FA-H-016".to_string()));
        assert!(config.security.disabled_checks.is_empty());
    }

    #[test]
    fn test_apply_to_config_with_disabled_checks() {
        let mut config = crate::core::ProjectConfig::default();
        let mut t = AuditTemplate::new("custom", "Custom template");
        t.enabled_checks = vec!["FA-H-001".into()];
        t.disabled_checks = vec!["FA-H-014".into()];

        crate::templates::apply_to_config(&mut config, &t);

        assert_eq!(config.security.enabled_checks, vec!["FA-H-001".to_string()]);
        assert_eq!(
            config.security.disabled_checks,
            vec!["FA-H-014".to_string()]
        );
    }

    #[test]
    fn test_get_template_user_override() {
        // Write a user template that shadows a built-in name
        let dir = user_template_dir().unwrap();
        let _ = std::fs::remove_dir_all(&dir);
        std::fs::create_dir_all(&dir).unwrap();

        let custom = AuditTemplate {
            name: "erc20".into(),
            description: "User override".into(),
            enabled_checks: vec!["FA-H-001".into()],
            ..Default::default()
        };
        std::fs::write(
            dir.join("erc20.json"),
            serde_json::to_string(&custom).unwrap(),
        )
        .unwrap();

        let t = get_template("erc20").unwrap();
        assert_eq!(t.description, "User override");
        assert_eq!(t.enabled_checks, vec!["FA-H-001".to_string()]);

        let _ = std::fs::remove_dir_all(&dir);
    }

    #[test]
    fn test_list_templates_dedupes_user_override() {
        let templates = list_templates();
        let count = templates.iter().filter(|t| t.name == "erc20").count();
        assert_eq!(count, 1, "erc20 should appear exactly once");
    }
}