Skip to main content

tirith_core/rules/
custom.rs

1use regex::Regex;
2
3use crate::extract::ScanContext;
4use crate::policy::CustomRule;
5use crate::verdict::{Evidence, Finding, RuleId, Severity};
6
7/// A compiled custom rule ready for matching.
8pub struct CompiledCustomRule {
9    pub id: String,
10    pub regex: Regex,
11    pub contexts: Vec<ScanContext>,
12    pub severity: Severity,
13    pub title: String,
14    pub description: String,
15}
16
17/// Compile custom rules from policy. Invalid regexes are logged and skipped.
18pub fn compile_rules(rules: &[CustomRule]) -> Vec<CompiledCustomRule> {
19    let mut compiled = Vec::new();
20    for rule in rules {
21        if rule.pattern.len() > 1024 {
22            eprintln!(
23                "tirith: custom rule '{}' pattern too long ({} chars), skipping",
24                rule.id,
25                rule.pattern.len()
26            );
27            continue;
28        }
29        let regex = match Regex::new(&rule.pattern) {
30            Ok(r) => r,
31            Err(e) => {
32                eprintln!(
33                    "tirith: warning: custom rule '{}' has invalid regex: {e}",
34                    rule.id
35                );
36                continue;
37            }
38        };
39
40        let contexts: Vec<ScanContext> = rule
41            .context
42            .iter()
43            .filter_map(|c| match c.as_str() {
44                "exec" => Some(ScanContext::Exec),
45                "paste" => Some(ScanContext::Paste),
46                "file" => Some(ScanContext::FileScan),
47                other => {
48                    eprintln!(
49                        "tirith: warning: custom rule '{}' has unknown context: {other}",
50                        rule.id
51                    );
52                    None
53                }
54            })
55            .collect();
56
57        if contexts.is_empty() {
58            eprintln!(
59                "tirith: warning: custom rule '{}' has no valid contexts, skipping",
60                rule.id
61            );
62            continue;
63        }
64
65        compiled.push(CompiledCustomRule {
66            id: rule.id.clone(),
67            regex,
68            contexts,
69            severity: rule.severity,
70            title: rule.title.clone(),
71            description: rule.description.clone(),
72        });
73    }
74    compiled
75}
76
77/// Check input against compiled custom rules for a given context.
78pub fn check(input: &str, context: ScanContext, compiled: &[CompiledCustomRule]) -> Vec<Finding> {
79    let mut findings = Vec::new();
80
81    for rule in compiled {
82        if !rule.contexts.contains(&context) {
83            continue;
84        }
85
86        if let Some(m) = rule.regex.find(input) {
87            let matched_text = m.as_str();
88            let preview: String = matched_text.chars().take(100).collect();
89
90            findings.push(Finding {
91                rule_id: RuleId::CustomRuleMatch,
92                severity: rule.severity,
93                title: rule.title.clone(),
94                description: if rule.description.is_empty() {
95                    format!("Custom rule '{}' matched", rule.id)
96                } else {
97                    rule.description.clone()
98                },
99                evidence: vec![Evidence::Text {
100                    detail: format!("Matched: \"{preview}\""),
101                }],
102                human_view: None,
103                agent_view: None,
104                mitre_id: None,
105                custom_rule_id: Some(rule.id.clone()),
106            });
107        }
108    }
109
110    findings
111}
112
113#[cfg(test)]
114mod tests {
115    use super::*;
116
117    fn make_rule(id: &str, pattern: &str, contexts: &[&str]) -> CustomRule {
118        CustomRule {
119            id: id.to_string(),
120            pattern: pattern.to_string(),
121            context: contexts.iter().map(|s| s.to_string()).collect(),
122            severity: Severity::High,
123            title: format!("Test rule: {id}"),
124            description: String::new(),
125        }
126    }
127
128    #[test]
129    fn test_compile_valid_rule() {
130        let rules = vec![make_rule("test1", r"internal\.corp", &["exec"])];
131        let compiled = compile_rules(&rules);
132        assert_eq!(compiled.len(), 1);
133        assert_eq!(compiled[0].id, "test1");
134    }
135
136    #[test]
137    fn test_compile_invalid_regex_skipped() {
138        let rules = vec![make_rule("bad", r"(unclosed", &["exec"])];
139        let compiled = compile_rules(&rules);
140        assert_eq!(compiled.len(), 0);
141    }
142
143    #[test]
144    fn test_check_matches_in_context() {
145        let rules = vec![make_rule(
146            "corp",
147            r"internal\.corp\.example\.com",
148            &["exec"],
149        )];
150        let compiled = compile_rules(&rules);
151
152        let findings = check(
153            "curl https://internal.corp.example.com/api",
154            ScanContext::Exec,
155            &compiled,
156        );
157        assert_eq!(findings.len(), 1);
158        assert_eq!(findings[0].rule_id, RuleId::CustomRuleMatch);
159        assert_eq!(findings[0].custom_rule_id.as_deref(), Some("corp"));
160    }
161
162    #[test]
163    fn test_check_no_match_wrong_context() {
164        let rules = vec![make_rule("corp", r"internal\.corp", &["exec"])];
165        let compiled = compile_rules(&rules);
166
167        let findings = check("internal.corp.example.com", ScanContext::Paste, &compiled);
168        assert_eq!(findings.len(), 0);
169    }
170
171    #[test]
172    fn test_check_no_match_when_pattern_absent() {
173        let rules = vec![make_rule("corp", r"internal\.corp", &["exec"])];
174        let compiled = compile_rules(&rules);
175
176        let findings = check("curl https://example.com", ScanContext::Exec, &compiled);
177        assert_eq!(findings.len(), 0);
178    }
179}