cc_audit/rules/
custom.rs

1use crate::rules::types::{Category, Confidence, Finding, Location, Severity};
2use regex::Regex;
3use serde::{Deserialize, Serialize};
4use std::path::Path;
5
6/// A dynamic rule loaded from YAML configuration.
7#[derive(Debug, Clone)]
8pub struct DynamicRule {
9    pub id: String,
10    pub name: String,
11    pub description: String,
12    pub severity: Severity,
13    pub category: Category,
14    pub confidence: Confidence,
15    pub patterns: Vec<Regex>,
16    pub exclusions: Vec<Regex>,
17    pub message: String,
18    pub recommendation: String,
19    pub fix_hint: Option<String>,
20    pub cwe_ids: Vec<String>,
21}
22
23impl DynamicRule {
24    /// Check if a line matches this rule.
25    pub fn matches(&self, line: &str) -> bool {
26        let pattern_match = self.patterns.iter().any(|p| p.is_match(line));
27        let excluded = self.exclusions.iter().any(|e| e.is_match(line));
28        pattern_match && !excluded
29    }
30
31    /// Create a Finding from this rule.
32    pub fn create_finding(&self, location: Location, code: String) -> Finding {
33        Finding {
34            id: self.id.clone(),
35            severity: self.severity,
36            category: self.category,
37            confidence: self.confidence,
38            name: self.name.clone(),
39            location,
40            code,
41            message: self.message.clone(),
42            recommendation: self.recommendation.clone(),
43            fix_hint: self.fix_hint.clone(),
44            cwe_ids: self.cwe_ids.clone(),
45            rule_severity: None,
46        }
47    }
48}
49
50/// YAML schema for custom rules file.
51#[derive(Debug, Serialize, Deserialize)]
52pub struct CustomRulesConfig {
53    pub version: String,
54    pub rules: Vec<YamlRule>,
55}
56
57#[derive(Debug, Clone, Serialize, Deserialize)]
58pub struct YamlRule {
59    pub id: String,
60    pub name: String,
61    #[serde(default)]
62    pub description: String,
63    pub severity: String,
64    pub category: String,
65    #[serde(default = "default_confidence")]
66    pub confidence: String,
67    pub patterns: Vec<String>,
68    #[serde(default)]
69    pub exclusions: Vec<String>,
70    pub message: String,
71    #[serde(default)]
72    pub recommendation: String,
73    #[serde(default)]
74    pub fix_hint: Option<String>,
75    #[serde(default)]
76    pub cwe: Vec<String>,
77}
78
79fn default_confidence() -> String {
80    "firm".to_string()
81}
82
83/// Error type for custom rule loading.
84#[derive(Debug)]
85pub enum CustomRuleError {
86    IoError(std::io::Error),
87    ParseError(serde_yaml::Error),
88    InvalidPattern {
89        rule_id: String,
90        pattern: String,
91        error: regex::Error,
92    },
93    InvalidSeverity {
94        rule_id: String,
95        value: String,
96    },
97    InvalidCategory {
98        rule_id: String,
99        value: String,
100    },
101    InvalidConfidence {
102        rule_id: String,
103        value: String,
104    },
105}
106
107impl std::fmt::Display for CustomRuleError {
108    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
109        match self {
110            Self::IoError(e) => write!(f, "Failed to read custom rules file: {}", e),
111            Self::ParseError(e) => write!(f, "Failed to parse custom rules YAML: {}", e),
112            Self::InvalidPattern {
113                rule_id,
114                pattern,
115                error,
116            } => {
117                write!(
118                    f,
119                    "Invalid regex pattern '{}' in rule {}: {}",
120                    pattern, rule_id, error
121                )
122            }
123            Self::InvalidSeverity { rule_id, value } => {
124                write!(
125                    f,
126                    "Invalid severity '{}' in rule {}. Expected: critical, high, medium, low",
127                    value, rule_id
128                )
129            }
130            Self::InvalidCategory { rule_id, value } => {
131                write!(
132                    f,
133                    "Invalid category '{}' in rule {}. Expected: exfiltration, privilege-escalation, persistence, prompt-injection, overpermission, obfuscation, supply-chain, secret-leak",
134                    value, rule_id
135                )
136            }
137            Self::InvalidConfidence { rule_id, value } => {
138                write!(
139                    f,
140                    "Invalid confidence '{}' in rule {}. Expected: certain, firm, tentative",
141                    value, rule_id
142                )
143            }
144        }
145    }
146}
147
148impl std::error::Error for CustomRuleError {}
149
150/// Loads custom rules from a YAML file.
151pub struct CustomRuleLoader;
152
153impl CustomRuleLoader {
154    /// Load rules from a YAML file path.
155    pub fn load_from_file(path: &Path) -> Result<Vec<DynamicRule>, CustomRuleError> {
156        let content = std::fs::read_to_string(path).map_err(CustomRuleError::IoError)?;
157        Self::load_from_string(&content)
158    }
159
160    /// Load rules from a YAML string.
161    pub fn load_from_string(content: &str) -> Result<Vec<DynamicRule>, CustomRuleError> {
162        let config: CustomRulesConfig =
163            serde_yaml::from_str(content).map_err(CustomRuleError::ParseError)?;
164
165        let mut rules = Vec::new();
166        for yaml_rule in config.rules {
167            let rule = Self::convert_yaml_rule(yaml_rule)?;
168            rules.push(rule);
169        }
170        Ok(rules)
171    }
172
173    /// Convert a vector of YamlRules to DynamicRules.
174    pub fn convert_yaml_rules(rules: Vec<YamlRule>) -> Result<Vec<DynamicRule>, CustomRuleError> {
175        rules.into_iter().map(Self::convert_yaml_rule).collect()
176    }
177
178    /// Convert a single YamlRule to a DynamicRule.
179    pub fn convert_yaml_rule(yaml: YamlRule) -> Result<DynamicRule, CustomRuleError> {
180        let severity = Self::parse_severity(&yaml.id, &yaml.severity)?;
181        let category = Self::parse_category(&yaml.id, &yaml.category)?;
182        let confidence = Self::parse_confidence(&yaml.id, &yaml.confidence)?;
183
184        let patterns = yaml
185            .patterns
186            .iter()
187            .map(|p| {
188                Regex::new(p).map_err(|e| CustomRuleError::InvalidPattern {
189                    rule_id: yaml.id.clone(),
190                    pattern: p.clone(),
191                    error: e,
192                })
193            })
194            .collect::<Result<Vec<_>, _>>()?;
195
196        let exclusions = yaml
197            .exclusions
198            .iter()
199            .map(|p| {
200                Regex::new(p).map_err(|e| CustomRuleError::InvalidPattern {
201                    rule_id: yaml.id.clone(),
202                    pattern: p.clone(),
203                    error: e,
204                })
205            })
206            .collect::<Result<Vec<_>, _>>()?;
207
208        Ok(DynamicRule {
209            id: yaml.id,
210            name: yaml.name,
211            description: yaml.description,
212            severity,
213            category,
214            confidence,
215            patterns,
216            exclusions,
217            message: yaml.message,
218            recommendation: yaml.recommendation,
219            fix_hint: yaml.fix_hint,
220            cwe_ids: yaml.cwe,
221        })
222    }
223
224    fn parse_severity(rule_id: &str, value: &str) -> Result<Severity, CustomRuleError> {
225        match value.to_lowercase().as_str() {
226            "critical" => Ok(Severity::Critical),
227            "high" => Ok(Severity::High),
228            "medium" => Ok(Severity::Medium),
229            "low" => Ok(Severity::Low),
230            _ => Err(CustomRuleError::InvalidSeverity {
231                rule_id: rule_id.to_string(),
232                value: value.to_string(),
233            }),
234        }
235    }
236
237    fn parse_category(rule_id: &str, value: &str) -> Result<Category, CustomRuleError> {
238        match value.to_lowercase().replace('_', "-").as_str() {
239            "exfiltration" | "data-exfiltration" => Ok(Category::Exfiltration),
240            "privilege-escalation" | "privilege" => Ok(Category::PrivilegeEscalation),
241            "persistence" => Ok(Category::Persistence),
242            "prompt-injection" | "injection" => Ok(Category::PromptInjection),
243            "overpermission" | "permission" => Ok(Category::Overpermission),
244            "obfuscation" => Ok(Category::Obfuscation),
245            "supply-chain" | "supplychain" => Ok(Category::SupplyChain),
246            "secret-leak" | "secrets" | "secretleak" => Ok(Category::SecretLeak),
247            _ => Err(CustomRuleError::InvalidCategory {
248                rule_id: rule_id.to_string(),
249                value: value.to_string(),
250            }),
251        }
252    }
253
254    fn parse_confidence(rule_id: &str, value: &str) -> Result<Confidence, CustomRuleError> {
255        match value.to_lowercase().as_str() {
256            "certain" => Ok(Confidence::Certain),
257            "firm" => Ok(Confidence::Firm),
258            "tentative" => Ok(Confidence::Tentative),
259            _ => Err(CustomRuleError::InvalidConfidence {
260                rule_id: rule_id.to_string(),
261                value: value.to_string(),
262            }),
263        }
264    }
265}
266
267#[cfg(test)]
268mod tests {
269    use super::*;
270
271    #[test]
272    fn test_load_valid_yaml() {
273        let yaml = r#"
274version: "1"
275rules:
276  - id: "CUSTOM-001"
277    name: "Internal API access"
278    description: "Detects access to internal APIs"
279    severity: "high"
280    category: "exfiltration"
281    confidence: "firm"
282    patterns:
283      - 'https?://internal\.'
284    exclusions:
285      - 'localhost'
286    message: "Internal API access detected"
287    recommendation: "Review if this is intended"
288    cwe:
289      - "CWE-200"
290"#;
291        let rules = CustomRuleLoader::load_from_string(yaml).unwrap();
292        assert_eq!(rules.len(), 1);
293        assert_eq!(rules[0].id, "CUSTOM-001");
294        assert_eq!(rules[0].name, "Internal API access");
295        assert_eq!(rules[0].severity, Severity::High);
296        assert_eq!(rules[0].category, Category::Exfiltration);
297        assert_eq!(rules[0].confidence, Confidence::Firm);
298        assert_eq!(rules[0].cwe_ids, vec!["CWE-200"]);
299    }
300
301    #[test]
302    fn test_load_multiple_rules() {
303        let yaml = r#"
304version: "1"
305rules:
306  - id: "CUSTOM-001"
307    name: "Rule One"
308    severity: "critical"
309    category: "exfiltration"
310    patterns:
311      - 'pattern1'
312    message: "Message 1"
313  - id: "CUSTOM-002"
314    name: "Rule Two"
315    severity: "low"
316    category: "obfuscation"
317    patterns:
318      - 'pattern2'
319    message: "Message 2"
320"#;
321        let rules = CustomRuleLoader::load_from_string(yaml).unwrap();
322        assert_eq!(rules.len(), 2);
323        assert_eq!(rules[0].id, "CUSTOM-001");
324        assert_eq!(rules[1].id, "CUSTOM-002");
325    }
326
327    #[test]
328    fn test_invalid_severity() {
329        let yaml = r#"
330version: "1"
331rules:
332  - id: "CUSTOM-001"
333    name: "Test"
334    severity: "invalid"
335    category: "exfiltration"
336    patterns:
337      - 'test'
338    message: "Test"
339"#;
340        let result = CustomRuleLoader::load_from_string(yaml);
341        assert!(result.is_err());
342        let err = result.unwrap_err();
343        assert!(matches!(err, CustomRuleError::InvalidSeverity { .. }));
344    }
345
346    #[test]
347    fn test_invalid_category() {
348        let yaml = r#"
349version: "1"
350rules:
351  - id: "CUSTOM-001"
352    name: "Test"
353    severity: "high"
354    category: "invalid"
355    patterns:
356      - 'test'
357    message: "Test"
358"#;
359        let result = CustomRuleLoader::load_from_string(yaml);
360        assert!(result.is_err());
361        let err = result.unwrap_err();
362        assert!(matches!(err, CustomRuleError::InvalidCategory { .. }));
363    }
364
365    #[test]
366    fn test_invalid_regex() {
367        let yaml = r#"
368version: "1"
369rules:
370  - id: "CUSTOM-001"
371    name: "Test"
372    severity: "high"
373    category: "exfiltration"
374    patterns:
375      - '[invalid('
376    message: "Test"
377"#;
378        let result = CustomRuleLoader::load_from_string(yaml);
379        assert!(result.is_err());
380        let err = result.unwrap_err();
381        assert!(matches!(err, CustomRuleError::InvalidPattern { .. }));
382    }
383
384    #[test]
385    fn test_default_confidence() {
386        let yaml = r#"
387version: "1"
388rules:
389  - id: "CUSTOM-001"
390    name: "Test"
391    severity: "high"
392    category: "exfiltration"
393    patterns:
394      - 'test'
395    message: "Test"
396"#;
397        let rules = CustomRuleLoader::load_from_string(yaml).unwrap();
398        assert_eq!(rules[0].confidence, Confidence::Firm);
399    }
400
401    #[test]
402    fn test_rule_matches() {
403        let yaml = r#"
404version: "1"
405rules:
406  - id: "CUSTOM-001"
407    name: "API Key Pattern"
408    severity: "high"
409    category: "secret-leak"
410    patterns:
411      - 'API_KEY\s*=\s*"[^"]+"'
412    exclusions:
413      - 'test'
414      - 'example'
415    message: "API key detected"
416"#;
417        let rules = CustomRuleLoader::load_from_string(yaml).unwrap();
418        let rule = &rules[0];
419
420        // Should match
421        assert!(rule.matches(r#"API_KEY = "secret123""#));
422
423        // Should not match (exclusion)
424        assert!(!rule.matches(r#"test API_KEY = "secret123""#));
425
426        // Should not match (no pattern match)
427        assert!(!rule.matches("random text"));
428    }
429
430    #[test]
431    fn test_create_finding() {
432        let yaml = r#"
433version: "1"
434rules:
435  - id: "CUSTOM-001"
436    name: "Test Rule"
437    severity: "critical"
438    category: "exfiltration"
439    confidence: "certain"
440    patterns:
441      - 'test'
442    message: "Test message"
443    recommendation: "Fix it"
444    fix_hint: "Do this"
445    cwe:
446      - "CWE-200"
447      - "CWE-319"
448"#;
449        let rules = CustomRuleLoader::load_from_string(yaml).unwrap();
450        let rule = &rules[0];
451
452        let location = Location {
453            file: "test.txt".to_string(),
454            line: 10,
455            column: None,
456        };
457        let finding = rule.create_finding(location, "test code".to_string());
458
459        assert_eq!(finding.id, "CUSTOM-001");
460        assert_eq!(finding.name, "Test Rule");
461        assert_eq!(finding.severity, Severity::Critical);
462        assert_eq!(finding.category, Category::Exfiltration);
463        assert_eq!(finding.confidence, Confidence::Certain);
464        assert_eq!(finding.message, "Test message");
465        assert_eq!(finding.recommendation, "Fix it");
466        assert_eq!(finding.fix_hint, Some("Do this".to_string()));
467        assert_eq!(finding.cwe_ids, vec!["CWE-200", "CWE-319"]);
468    }
469
470    #[test]
471    fn test_category_variations() {
472        let test_cases = vec![
473            ("exfiltration", Category::Exfiltration),
474            ("data-exfiltration", Category::Exfiltration),
475            ("privilege-escalation", Category::PrivilegeEscalation),
476            ("privilege", Category::PrivilegeEscalation),
477            ("persistence", Category::Persistence),
478            ("prompt-injection", Category::PromptInjection),
479            ("injection", Category::PromptInjection),
480            ("overpermission", Category::Overpermission),
481            ("permission", Category::Overpermission),
482            ("obfuscation", Category::Obfuscation),
483            ("supply-chain", Category::SupplyChain),
484            ("supplychain", Category::SupplyChain),
485            ("secret-leak", Category::SecretLeak),
486            ("secrets", Category::SecretLeak),
487            ("secretleak", Category::SecretLeak),
488        ];
489
490        for (input, expected) in test_cases {
491            let result = CustomRuleLoader::parse_category("test", input);
492            assert_eq!(result.unwrap(), expected, "Failed for input: {}", input);
493        }
494    }
495
496    #[test]
497    fn test_error_display() {
498        let io_err = CustomRuleError::IoError(std::io::Error::new(
499            std::io::ErrorKind::NotFound,
500            "file not found",
501        ));
502        assert!(io_err.to_string().contains("Failed to read"));
503
504        let severity_err = CustomRuleError::InvalidSeverity {
505            rule_id: "TEST".to_string(),
506            value: "bad".to_string(),
507        };
508        assert!(severity_err.to_string().contains("Invalid severity"));
509
510        let category_err = CustomRuleError::InvalidCategory {
511            rule_id: "TEST".to_string(),
512            value: "bad".to_string(),
513        };
514        assert!(category_err.to_string().contains("Invalid category"));
515
516        let confidence_err = CustomRuleError::InvalidConfidence {
517            rule_id: "TEST".to_string(),
518            value: "bad".to_string(),
519        };
520        assert!(confidence_err.to_string().contains("Invalid confidence"));
521    }
522
523    #[test]
524    fn test_load_from_file() {
525        use std::fs;
526        use tempfile::TempDir;
527
528        let dir = TempDir::new().unwrap();
529        let file_path = dir.path().join("rules.yaml");
530        fs::write(
531            &file_path,
532            r#"
533version: "1"
534rules:
535  - id: "FILE-001"
536    name: "File Test"
537    severity: "high"
538    category: "exfiltration"
539    patterns:
540      - 'test_from_file'
541    message: "Test"
542"#,
543        )
544        .unwrap();
545
546        let rules = CustomRuleLoader::load_from_file(&file_path).unwrap();
547        assert_eq!(rules.len(), 1);
548        assert_eq!(rules[0].id, "FILE-001");
549    }
550
551    #[test]
552    fn test_load_from_file_not_found() {
553        let result =
554            CustomRuleLoader::load_from_file(std::path::Path::new("/nonexistent/file.yaml"));
555        assert!(result.is_err());
556        assert!(matches!(result, Err(CustomRuleError::IoError(_))));
557    }
558
559    #[test]
560    fn test_convert_yaml_rules() {
561        let yaml_rules = vec![YamlRule {
562            id: "CONV-001".to_string(),
563            name: "Convert Test".to_string(),
564            description: "Test".to_string(),
565            severity: "high".to_string(),
566            category: "exfiltration".to_string(),
567            confidence: "firm".to_string(),
568            patterns: vec!["test".to_string()],
569            exclusions: vec![],
570            message: "Test".to_string(),
571            recommendation: "".to_string(),
572            fix_hint: None,
573            cwe: vec![],
574        }];
575
576        let rules = CustomRuleLoader::convert_yaml_rules(yaml_rules).unwrap();
577        assert_eq!(rules.len(), 1);
578        assert_eq!(rules[0].id, "CONV-001");
579    }
580
581    #[test]
582    fn test_invalid_confidence() {
583        let yaml = r#"
584version: "1"
585rules:
586  - id: "CUSTOM-001"
587    name: "Test"
588    severity: "high"
589    category: "exfiltration"
590    confidence: "invalid"
591    patterns:
592      - 'test'
593    message: "Test"
594"#;
595        let result = CustomRuleLoader::load_from_string(yaml);
596        assert!(result.is_err());
597        let err = result.unwrap_err();
598        assert!(matches!(err, CustomRuleError::InvalidConfidence { .. }));
599    }
600
601    #[test]
602    fn test_parse_error_display() {
603        let invalid_yaml = "invalid: yaml: [";
604        let result: Result<CustomRulesConfig, _> = serde_yaml::from_str(invalid_yaml);
605        let yaml_err = result.unwrap_err();
606        let err = CustomRuleError::ParseError(yaml_err);
607        assert!(err.to_string().contains("Failed to parse"));
608    }
609
610    #[test]
611    #[allow(clippy::invalid_regex)]
612    fn test_invalid_pattern_error_display() {
613        let regex_err = Regex::new("[invalid(").unwrap_err();
614        let err = CustomRuleError::InvalidPattern {
615            rule_id: "TEST-001".to_string(),
616            pattern: "[invalid(".to_string(),
617            error: regex_err,
618        };
619        let display = err.to_string();
620        assert!(display.contains("Invalid regex pattern"));
621        assert!(display.contains("[invalid("));
622        assert!(display.contains("TEST-001"));
623    }
624
625    #[test]
626    fn test_invalid_exclusion_regex() {
627        let yaml = r#"
628version: "1"
629rules:
630  - id: "CUSTOM-001"
631    name: "Test"
632    severity: "high"
633    category: "exfiltration"
634    patterns:
635      - 'valid_pattern'
636    exclusions:
637      - '[invalid('
638    message: "Test"
639"#;
640        let result = CustomRuleLoader::load_from_string(yaml);
641        assert!(result.is_err());
642        let err = result.unwrap_err();
643        assert!(matches!(err, CustomRuleError::InvalidPattern { .. }));
644    }
645}
cc_audit/rules/custom.rs

cc_audit/rules/
custom.rs
