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

cc_audit/rules/
custom.rs
