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

cc_audit/rules/
custom.rs
