ai_coding_shield/catalog/

mod.rs

pub mod threats;

use anyhow::Result;
use colored::*;
use serde::{Deserialize, Serialize};
use std::fs;
use std::path::PathBuf;

pub use threats::*;

const CATALOG_FILE: &str = "threats.yaml";

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct ThreatCatalog {
    pub version: String,
    pub last_updated: String,
    pub categories: Vec<ThreatCategory>,
    #[serde(default)]
    pub trusted_authors: Vec<String>,
    #[serde(default)]
    pub trusted_domains: Vec<String>,
}

impl ThreatCatalog {
    /// Load threat catalog from embedded or local file
    pub fn load() -> Result<Self> {
        // Try to load from local file first
        let local_path = Self::local_catalog_path();
        
        if local_path.exists() {
            let content = fs::read_to_string(&local_path)?;
            let catalog: ThreatCatalog = serde_yaml::from_str(&content)?;
            return Ok(catalog);
        }

        // Fall back to embedded catalog
        let embedded = include_str!("../../config/threats.yaml");
        let catalog: ThreatCatalog = serde_yaml::from_str(embedded)?;
        Ok(catalog)
    }

    fn local_catalog_path() -> PathBuf {
        dirs::config_dir()
            .unwrap_or_else(|| PathBuf::from("."))
            .join("ai-coding-shield")
            .join("config")
            .join(CATALOG_FILE)
    }

    /// Get all rules across all categories
    pub fn all_rules(&self) -> Vec<&ThreatRule> {
        self.categories
            .iter()
            .flat_map(|cat| cat.rules.iter())
            .collect()
    }

    /// Find a specific threat by ID
    pub fn find_threat(&self, threat_id: &str) -> Option<(&ThreatCategory, &ThreatRule)> {
        for category in &self.categories {
            for rule in &category.rules {
                if rule.id == threat_id {
                    return Some((category, rule));
                }
            }
        }
        None
    }

    /// Show detailed information about a threat
    pub fn show_threat_info(&self, threat_id: &str) -> Result<()> {
        match self.find_threat(threat_id) {
            Some((category, rule)) => {
                println!("{}", format!("🔍 Threat: {}", rule.id).bold().cyan());
                println!("{}", "━".repeat(50).cyan());
                println!();
                println!("{}: {}", "Name".bold(), rule.description);
                println!("{}: {}", "Category".bold(), category.name);
                println!("{}: {}", "Severity".bold(), rule.severity.as_str().red());
                println!();
                
                if let Some(mitre) = &category.mitre_id {
                    println!("{}: {}", "MITRE ATT&CK".bold(), mitre);
                }
                if let Some(cwe) = &category.cwe_id {
                    println!("{}: {}", "CWE".bold(), cwe);
                }
                println!();
                
                println!("{}", "Pattern:".bold());
                println!("  {}", rule.pattern);
                println!();
                
                if !rule.examples.is_empty() {
                    println!("{}", "Examples:".bold());
                    for example in &rule.examples {
                        println!("  • {}", example.dimmed());
                    }
                    println!();
                }
                
                println!("{}", "Remediation:".bold());
                println!("  {}", rule.remediation);
                
                Ok(())
            }
            None => {
                eprintln!("{}", format!("❌ Threat ID '{}' not found", threat_id).red());
                Err(anyhow::anyhow!("Threat not found"))
            }
        }
    }

    /// Save current catalog to local config file
    pub fn save(&self) -> Result<()> {
        let local_path = Self::local_catalog_path();
        
        if let Some(parent) = local_path.parent() {
            fs::create_dir_all(parent)?;
        }
        
        let content = serde_yaml::to_string(self)?;
        fs::write(local_path, content)?;
        Ok(())
    }

    pub fn add_trusted_author(&mut self, author: String) {
        if !self.trusted_authors.contains(&author) {
            self.trusted_authors.push(author);
        }
    }

    pub fn remove_trusted_author(&mut self, author: &str) -> bool {
        if let Some(pos) = self.trusted_authors.iter().position(|x| x == author) {
            self.trusted_authors.remove(pos);
            true
        } else {
            false
        }
    }

    pub fn add_trusted_domain(&mut self, domain: String) {
        if !self.trusted_domains.contains(&domain) {
            self.trusted_domains.push(domain);
        }
    }

    pub fn remove_trusted_domain(&mut self, domain: &str) -> bool {
        if let Some(pos) = self.trusted_domains.iter().position(|x| x == domain) {
            self.trusted_domains.remove(pos);
            true
        } else {
            false
        }
    }

    /// List all threats, optionally filtered
    pub fn list_threats(&self, category_filter: Option<&str>, severity_filter: Option<&str>) -> Result<()> {
        println!("{}", "📋 Threat Catalog".bold().cyan());
        println!("{}", "━".repeat(50).cyan());
        println!();

        for category in &self.categories {
            if let Some(filter) = category_filter {
                if !category.id.to_lowercase().contains(&filter.to_lowercase()) {
                    continue;
                }
            }

            println!("{}", format!("▸ {}", category.name).bold());
            
            for rule in &category.rules {
                if let Some(filter) = severity_filter {
                    if rule.severity.as_str().to_lowercase() != filter.to_lowercase() {
                        continue;
                    }
                }

                let severity_color = match rule.severity {
                    RiskLevel::Critical => "CRITICAL".red().bold(),
                    RiskLevel::High => "HIGH".yellow().bold(),
                    RiskLevel::Medium => "MEDIUM".blue(),
                    RiskLevel::Low => "LOW".green(),
                };

                println!("  {} [{}] {}", 
                    rule.id.cyan(), 
                    severity_color,
                    rule.description.dimmed()
                );
            }
            println!();
        }

        Ok(())
    }
}


pub fn update_catalog(_force: bool) -> Result<()> {
    // TODO: Implement fetching from MITRE, CWE, etc.
    // For now, just a placeholder
    println!("⚠️  Threat catalog update not yet implemented");
    println!("   This will fetch from:");
    println!("   • MITRE ATT&CK");
    println!("   • CWE Database");
    println!("   • GitHub Security Advisories");
    Ok(())
}

use crate::types::RiskLevel;