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

cc_audit/rules/
custom.rs
