yara_forge/utils/
mod.rs

1//! Utility functions for YARA rule generation
2//! Provides helper functions for common tasks
3
4use serde_json::{json, Value};
5use std::fs;
6use std::io::{self, Read, Write};
7use std::path::Path;
8
9/// Export a rule to JSON format
10pub fn export_rule_to_json(rule: &crate::Rule) -> Value {
11    json!({
12        "name": rule.name,
13        "tags": rule.tags,
14        "metadata": rule.metadata,
15        "strings": rule.strings,
16        "condition": rule.condition
17    })
18}
19
20/// Import a rule from JSON format
21pub fn import_rule_from_json(json: Value) -> Result<crate::Rule, crate::YaraError> {
22    let name = json["name"]
23        .as_str()
24        .ok_or_else(|| crate::YaraError::MissingField("name".to_string()))?
25        .to_string();
26
27    let tags = json["tags"]
28        .as_array()
29        .map(|arr| {
30            arr.iter()
31                .filter_map(|v| v.as_str())
32                .map(|s| s.to_string())
33                .collect()
34        })
35        .unwrap_or_default();
36
37    let metadata = json["metadata"]
38        .as_object()
39        .map(|obj| {
40            obj.iter()
41                .filter_map(|(k, v)| v.as_str().map(|s| (k.clone(), s.to_string())))
42                .collect()
43        })
44        .unwrap_or_default();
45
46    let strings = json["strings"]
47        .as_array()
48        .ok_or_else(|| crate::YaraError::MissingField("strings".to_string()))?
49        .iter()
50        .map(|s| crate::StringDefinition {
51            identifier: s["identifier"].as_str().unwrap_or("").to_string(),
52            pattern: s["pattern"].as_str().unwrap_or("").to_string(),
53            is_hex: s["is_hex"].as_bool().unwrap_or(false),
54            modifiers: s["modifiers"]
55                .as_array()
56                .map(|arr| {
57                    arr.iter()
58                        .filter_map(|v| v.as_str())
59                        .map(|s| s.to_string())
60                        .collect()
61                })
62                .unwrap_or_default(),
63        })
64        .collect();
65
66    let condition = json["condition"]
67        .as_str()
68        .ok_or_else(|| crate::YaraError::MissingField("condition".to_string()))?
69        .to_string();
70
71    Ok(crate::Rule {
72        name,
73        tags,
74        metadata,
75        strings,
76        condition,
77    })
78}
79
80/// Save a rule to a file
81pub fn save_rule_to_file(rule: &crate::Rule, path: impl AsRef<Path>) -> io::Result<()> {
82    let mut file = fs::File::create(path)?;
83    file.write_all(rule.to_string().as_bytes())?;
84    Ok(())
85}
86
87/// Load a rule from a file
88pub fn load_rule_from_file(path: impl AsRef<Path>) -> io::Result<String> {
89    let mut file = fs::File::open(path)?;
90    let mut contents = String::new();
91    file.read_to_string(&mut contents)?;
92    Ok(contents)
93}
94
95/// Escape special characters in strings
96pub fn escape_string(s: &str) -> String {
97    s.replace('\\', "\\\\")
98        .replace('"', "\\\"")
99        .replace('\n', "\\n")
100        .replace('\r', "\\r")
101        .replace('\t', "\\t")
102}
103
104/// Convert bytes to hex string
105pub fn bytes_to_hex(bytes: &[u8]) -> String {
106    bytes
107        .iter()
108        .map(|b| format!("{:02X}", b))
109        .collect::<Vec<_>>()
110        .join(" ")
111}
112
113/// Generate a unique rule name
114pub fn generate_unique_rule_name(prefix: &str) -> String {
115    use chrono::Utc;
116    let timestamp = Utc::now().timestamp();
117    format!("{}_{}", prefix, timestamp)
118}
119
120/// Calculate Shannon entropy of data
121pub fn calculate_entropy(data: &[u8]) -> f64 {
122    let mut counts = [0u32; 256];
123    for &byte in data {
124        counts[byte as usize] += 1;
125    }
126
127    let len = data.len() as f64;
128    let mut entropy = 0.0;
129
130    for &count in &counts {
131        if count > 0 {
132            let p = count as f64 / len;
133            entropy -= p * p.log2();
134        }
135    }
136
137    entropy
138}
139
140#[cfg(test)]
141mod tests {
142    use super::*;
143
144    #[test]
145    fn test_rule_json_export() {
146        let rule = crate::RuleBuilder::new("test_rule")
147            .with_tag("test")
148            .with_metadata("author", "test")
149            .with_string("$test", "test")
150            .unwrap()
151            .with_condition("$test")
152            .build()
153            .unwrap();
154
155        let json = export_rule_to_json(&rule);
156        assert_eq!(json["name"], "test_rule");
157    }
158
159    #[test]
160    fn test_string_escaping() {
161        let input = "test\"string\nwith\tspecial\\chars";
162        let escaped = escape_string(input);
163        assert_eq!(escaped, "test\\\"string\\nwith\\tspecial\\\\chars");
164    }
165
166    #[test]
167    fn test_bytes_to_hex() {
168        let bytes = &[0x48, 0x45, 0x4C, 0x4C, 0x4F];
169        let hex = bytes_to_hex(bytes);
170        assert_eq!(hex, "48 45 4C 4C 4F");
171    }
172
173    #[test]
174    fn test_entropy_calculation() {
175        let data = b"AAAAABBBCC";
176        let entropy = calculate_entropy(data);
177        assert!(entropy > 0.0 && entropy < 8.0);
178    }
179}
yara_forge/utils/mod.rs

yara_forge/utils/
mod.rs
