feat(detmir): add rust-first operations tooling
This commit is contained in:
igor04091968
@@ -0,0 +1,387 @@
|
||||
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<String>,
|
||||
|
||||
#[arg(long)]
|
||||
image: Option<PathBuf>,
|
||||
|
||||
#[arg(long)]
|
||||
dictionary_pack: Option<String>,
|
||||
|
||||
#[arg(long)]
|
||||
regex_pack: Option<String>,
|
||||
|
||||
#[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<String>,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
struct Analysis {
|
||||
text: String,
|
||||
dictionary_pack: Option<String>,
|
||||
regex_pack: Option<String>,
|
||||
dictionary_matches: Vec<MatchResult>,
|
||||
regex_matches: Vec<MatchResult>,
|
||||
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::<Vec<_>>();
|
||||
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::<u32>()
|
||||
% 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::<u32>()
|
||||
% 11
|
||||
% 10;
|
||||
let chk12 = nums
|
||||
.iter()
|
||||
.take(11)
|
||||
.zip(c12)
|
||||
.map(|(digit, coef)| digit * coef)
|
||||
.sum::<u32>()
|
||||
% 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::<Vec<_>>();
|
||||
let checksum = nums[9] * 10 + nums[10];
|
||||
let sum = nums
|
||||
.iter()
|
||||
.take(9)
|
||||
.enumerate()
|
||||
.map(|(idx, digit)| digit * (9 - idx as u32))
|
||||
.sum::<u32>();
|
||||
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<Value> {
|
||||
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<PathBuf> {
|
||||
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<Vec<MatchResult>> {
|
||||
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<Vec<MatchResult>> {
|
||||
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<Analysis> {
|
||||
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<Value> {
|
||||
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<Value> {
|
||||
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()
|
||||
)]
|
||||
);
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user