use std::fs; use std::path::{Path, PathBuf}; use std::process::Command; use anyhow::{Context, Result, bail}; use clap::Parser; use regex::Regex; use serde::Serialize; use serde_json::Value; const DEFAULT_BASE_DIR: &str = "/opt/activitywatch/dlp-content-analysis"; #[derive(Debug, Parser)] #[command(about = "Analyze text or screenshot with DLP dictionaries/regex packs.")] struct Cli { #[arg(long)] text: Option, #[arg(long)] image: Option, #[arg(long)] dictionary_pack: Option, #[arg(long)] regex_pack: Option, #[arg(long, default_value = DEFAULT_BASE_DIR)] base_dir: PathBuf, #[arg( long, default_value = "/opt/activitywatch/dlp-content-analysis/.venv/bin/python" )] legacy_python: PathBuf, #[arg( long, default_value = "/opt/activitywatch/dlp-content-analysis/content_analyzer.py" )] legacy_analyzer: PathBuf, } #[derive(Debug, Serialize, PartialEq, Eq)] struct MatchResult { name: String, description: String, value: String, start: usize, end: usize, #[serde(skip_serializing_if = "Option::is_none")] severity: Option, } #[derive(Debug, Serialize)] struct Analysis { text: String, dictionary_pack: Option, regex_pack: Option, dictionary_matches: Vec, regex_matches: Vec, source: String, } fn digits_only(value: &str) -> String { value.chars().filter(|ch| ch.is_ascii_digit()).collect() } fn validate_inn(value: &str) -> bool { let digits = digits_only(value); let nums = digits .bytes() .map(|byte| (byte - b'0') as u32) .collect::>(); match nums.len() { 10 => { let coef = [2, 4, 10, 3, 5, 9, 4, 6, 8]; let chk = nums .iter() .take(9) .zip(coef) .map(|(digit, coef)| digit * coef) .sum::() % 11 % 10; chk == nums[9] } 12 => { let c11 = [7, 2, 4, 10, 3, 5, 9, 4, 6, 8]; let c12 = [3, 7, 2, 4, 10, 3, 5, 9, 4, 6, 8]; let chk11 = nums .iter() .take(10) .zip(c11) .map(|(digit, coef)| digit * coef) .sum::() % 11 % 10; let chk12 = nums .iter() .take(11) .zip(c12) .map(|(digit, coef)| digit * coef) .sum::() % 11 % 10; chk11 == nums[10] && chk12 == nums[11] } _ => false, } } fn validate_snils(value: &str) -> bool { let digits = digits_only(value); if digits.len() != 11 { return false; } let nums = digits .bytes() .map(|byte| (byte - b'0') as u32) .collect::>(); let checksum = nums[9] * 10 + nums[10]; let sum = nums .iter() .take(9) .enumerate() .map(|(idx, digit)| digit * (9 - idx as u32)) .sum::(); let expected = match sum { 0..=99 => sum, 100 | 101 => 0, _ => { let value = sum % 101; if value == 100 { 0 } else { value } } }; checksum == expected } fn validate_passport(value: &str) -> bool { let digits = digits_only(value); digits.len() == 10 && digits != "0000000000" && digits.chars().any(|ch| ch != digits.as_bytes()[0] as char) } fn validate_checksum(kind: &str, value: &str) -> bool { match kind { "inn" => validate_inn(value), "snils" => validate_snils(value), "passport" => validate_passport(value), _ => true, } } fn char_offset(text: &str, byte_offset: usize) -> usize { text[..byte_offset].chars().count() } fn load_json(path: &Path) -> Result { let text = fs::read_to_string(path).with_context(|| format!("read {}", path.display()))?; serde_json::from_str(&text).with_context(|| format!("parse {}", path.display())) } fn resolve_pack(base_dir: &Path, kind: &str, name: Option<&str>) -> Option { let name = name?; let path = base_dir.join(kind).join(format!("{name}.json")); path.exists().then_some(path) } fn match_dictionary(text: &str, path: Option<&Path>) -> Result> { let Some(path) = path else { return Ok(Vec::new()); }; let rules = load_json(path)?; let Some(rules) = rules.as_object() else { return Ok(Vec::new()); }; let mut results = Vec::new(); for (name, rule) in rules { let Some(regex) = rule.get("regex").and_then(Value::as_str) else { continue; }; let regex = Regex::new(regex).with_context(|| format!("compile dictionary regex {name}"))?; let checksum = rule .get("checksum") .and_then(Value::as_str) .unwrap_or("none"); let description = rule .get("description") .and_then(Value::as_str) .unwrap_or(name) .to_string(); for capture in regex.find_iter(text) { let value = capture.as_str(); if validate_checksum(checksum, value) { results.push(MatchResult { name: name.to_string(), description: description.clone(), value: value.to_string(), start: char_offset(text, capture.start()), end: char_offset(text, capture.end()), severity: None, }); } } } results.sort_by_key(|result| (result.start, result.end, result.name.clone())); Ok(results) } fn regex_entries(pack: &Value) -> Vec<(String, String, String, String)> { let mut entries = Vec::new(); if let Some(rules) = pack.get("rules").and_then(Value::as_array) { for rule in rules.iter().filter_map(Value::as_object) { let Some(regex) = rule.get("regex").and_then(Value::as_str) else { continue; }; let id = rule .get("id") .or_else(|| rule.get("name")) .and_then(Value::as_str) .unwrap_or("regex-rule"); let description = rule .get("description") .and_then(Value::as_str) .unwrap_or(id); let severity = rule .get("severity") .and_then(Value::as_str) .unwrap_or("medium"); entries.push(( id.to_string(), description.to_string(), regex.to_string(), severity.to_string(), )); } } else if let Some(patterns) = pack.get("patterns").and_then(Value::as_object) { for (id, rule) in patterns { let Some(rule) = rule.as_object() else { continue; }; let Some(regex) = rule.get("regex").and_then(Value::as_str) else { continue; }; let description = rule .get("description") .and_then(Value::as_str) .unwrap_or(id); let severity = rule .get("severity") .and_then(Value::as_str) .unwrap_or("medium"); entries.push(( id.to_string(), description.to_string(), regex.to_string(), severity.to_string(), )); } } entries } fn match_regex_pack(text: &str, path: Option<&Path>) -> Result> { let Some(path) = path else { return Ok(Vec::new()); }; let pack = load_json(path)?; let mut results = Vec::new(); for (name, description, regex, severity) in regex_entries(&pack) { let regex = Regex::new(®ex).with_context(|| format!("compile regex pack rule {name}"))?; for capture in regex.find_iter(text) { results.push(MatchResult { name: name.clone(), description: description.clone(), value: capture.as_str().to_string(), start: char_offset(text, capture.start()), end: char_offset(text, capture.end()), severity: Some(severity.clone()), }); } } Ok(results) } fn analyze_text(cli: &Cli) -> Result { let text = cli.text.clone().unwrap_or_default(); let dictionary_path = resolve_pack( &cli.base_dir, "dictionaries", cli.dictionary_pack.as_deref(), ); let regex_pack_path = resolve_pack(&cli.base_dir, "regex-packs", cli.regex_pack.as_deref()); Ok(Analysis { dictionary_matches: match_dictionary(&text, dictionary_path.as_deref())?, regex_matches: match_regex_pack(&text, regex_pack_path.as_deref())?, text, dictionary_pack: cli.dictionary_pack.clone(), regex_pack: cli.regex_pack.clone(), source: "text".to_string(), }) } fn run_legacy_image(cli: &Cli, image: &Path) -> Result { if !cli.legacy_python.exists() || !cli.legacy_analyzer.exists() { bail!( "image OCR requires legacy analyzer: {} {}", cli.legacy_python.display(), cli.legacy_analyzer.display() ); } let mut command = Command::new(&cli.legacy_python); command.arg(&cli.legacy_analyzer).arg("--image").arg(image); if let Some(pack) = &cli.dictionary_pack { command.arg("--dictionary-pack").arg(pack); } if let Some(pack) = &cli.regex_pack { command.arg("--regex-pack").arg(pack); } let output = command.output().context("run legacy image analyzer")?; if !output.status.success() { bail!( "legacy image analyzer failed rc={:?}: {}{}", output.status.code(), String::from_utf8_lossy(&output.stdout), String::from_utf8_lossy(&output.stderr) ); } serde_json::from_slice(&output.stdout).context("parse legacy image analyzer JSON") } fn run(cli: &Cli) -> Result { if let Some(image) = &cli.image { return run_legacy_image(cli, image); } Ok(serde_json::to_value(analyze_text(cli)?)?) } fn main() -> Result<()> { let cli = Cli::parse(); let value = run(&cli)?; println!("{}", serde_json::to_string_pretty(&value)?); Ok(()) } #[cfg(test)] mod tests { use super::*; #[test] fn validates_inn_like_python() { assert!(validate_inn("7707083893")); assert!(validate_inn("500100732259")); assert!(!validate_inn("7707083894")); } #[test] fn validates_snils_like_python() { assert!(validate_snils("112-233-445 95")); assert!(!validate_snils("112-233-445 96")); } #[test] fn validates_passport_like_python() { assert!(validate_passport("1234 567890")); assert!(!validate_passport("0000 000000")); assert!(!validate_passport("1111 111111")); } #[test] fn regex_pack_patterns_shape_is_supported() { let pack = serde_json::json!({"patterns": {"secret": {"regex": "token", "description": "Token", "severity": "high"}}}); assert_eq!( regex_entries(&pack), vec![( "secret".to_string(), "Token".to_string(), "token".to_string(), "high".to_string() )] ); } }