619 lines
19 KiB
Rust
619 lines
19 KiB
Rust
use std::collections::BTreeMap;
|
|
use std::fs;
|
|
use std::path::{Path, PathBuf};
|
|
use std::time::Duration;
|
|
|
|
use anyhow::{Context, Result, anyhow, bail};
|
|
use chrono::{DateTime, SecondsFormat, TimeDelta, Utc};
|
|
use clap::Parser;
|
|
use detmir_aw_client::ActivityWatchClient;
|
|
use rusqlite::{Connection, params};
|
|
use serde::{Deserialize, Serialize};
|
|
use serde_json::{Value, json};
|
|
|
|
const DEFAULT_AW_URL: &str = "http://127.0.0.1:5600/api/0";
|
|
const DEFAULT_SQLITE_PATH: &str = "data/dlp-events.sqlite3";
|
|
const DEFAULT_STATE_PATH: &str = "data/dlp-aggregator-state.json";
|
|
const DEFAULT_BUCKET_PREFIXES: &str = "aw-file-operations_,aw-dlp-incidents_";
|
|
|
|
#[derive(Debug, Parser)]
|
|
#[command(about = "Aggregate AWatch-rus DLP buckets into a local warehouse database.")]
|
|
struct Cli {
|
|
#[arg(long, default_value = DEFAULT_AW_URL)]
|
|
aw_url: String,
|
|
|
|
#[arg(long)]
|
|
postgres_dsn: Option<String>,
|
|
|
|
#[arg(long, default_value = DEFAULT_SQLITE_PATH)]
|
|
sqlite_path: PathBuf,
|
|
|
|
#[arg(long, default_value = DEFAULT_STATE_PATH)]
|
|
state_path: PathBuf,
|
|
|
|
#[arg(long, default_value = DEFAULT_BUCKET_PREFIXES)]
|
|
bucket_prefixes: String,
|
|
|
|
#[arg(long)]
|
|
since: Option<String>,
|
|
|
|
#[arg(long, default_value_t = 24)]
|
|
lookback_hours: i64,
|
|
|
|
#[arg(long, default_value_t = 60)]
|
|
overlap_seconds: i64,
|
|
|
|
#[arg(long, default_value_t = 10_000)]
|
|
limit: usize,
|
|
|
|
#[arg(long, default_value_t = 15)]
|
|
timeout: u64,
|
|
|
|
#[arg(long)]
|
|
dry_run: bool,
|
|
}
|
|
|
|
impl Cli {
|
|
fn apply_env(mut self) -> Self {
|
|
self.aw_url = env_value("AW_URL", &self.aw_url);
|
|
self.postgres_dsn = self
|
|
.postgres_dsn
|
|
.or_else(|| std::env::var("DLP_AGGREGATOR_POSTGRES_DSN").ok())
|
|
.filter(|value| !value.is_empty());
|
|
self.sqlite_path = env_path("DLP_AGGREGATOR_SQLITE_PATH", self.sqlite_path);
|
|
self.state_path = env_path("DLP_AGGREGATOR_STATE_PATH", self.state_path);
|
|
self
|
|
}
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize)]
|
|
struct Bucket {
|
|
#[serde(default)]
|
|
r#type: String,
|
|
#[serde(default)]
|
|
hostname: String,
|
|
}
|
|
|
|
#[derive(Debug, Clone, Deserialize)]
|
|
struct RawAwEvent {
|
|
#[serde(default)]
|
|
id: Option<Value>,
|
|
timestamp: String,
|
|
#[serde(default)]
|
|
duration: f64,
|
|
#[serde(default)]
|
|
data: Value,
|
|
}
|
|
|
|
#[derive(Debug, Clone)]
|
|
struct AwEvent {
|
|
bucket_id: String,
|
|
hostname: String,
|
|
stream_type: String,
|
|
event_id: String,
|
|
timestamp: String,
|
|
duration: f64,
|
|
data: Value,
|
|
}
|
|
|
|
#[derive(Debug, Default, Clone, Serialize, Deserialize)]
|
|
struct State {
|
|
#[serde(default)]
|
|
last_end: Option<String>,
|
|
}
|
|
|
|
#[derive(Debug, Serialize)]
|
|
struct DryRunSummary {
|
|
aw_url: String,
|
|
start: String,
|
|
end: String,
|
|
selected_buckets: Vec<String>,
|
|
fetched_events: usize,
|
|
}
|
|
|
|
#[derive(Debug, Serialize)]
|
|
struct RunSummary {
|
|
aw_url: String,
|
|
target: String,
|
|
target_path: String,
|
|
state_path: String,
|
|
start: String,
|
|
end: String,
|
|
selected_buckets: usize,
|
|
fetched_events: usize,
|
|
inserted_events: usize,
|
|
}
|
|
|
|
fn env_value(name: &str, fallback: &str) -> String {
|
|
std::env::var(name)
|
|
.ok()
|
|
.filter(|value| !value.is_empty())
|
|
.unwrap_or_else(|| fallback.to_string())
|
|
}
|
|
|
|
fn env_path(name: &str, fallback: PathBuf) -> PathBuf {
|
|
std::env::var(name)
|
|
.ok()
|
|
.filter(|value| !value.is_empty())
|
|
.map(PathBuf::from)
|
|
.unwrap_or(fallback)
|
|
}
|
|
|
|
fn parse_timestamp(value: &str) -> Result<DateTime<Utc>> {
|
|
let normalized = value.replace('Z', "+00:00");
|
|
Ok(DateTime::parse_from_rfc3339(&normalized)
|
|
.with_context(|| format!("parse timestamp {value}"))?
|
|
.with_timezone(&Utc))
|
|
}
|
|
|
|
fn format_aw_timestamp(value: DateTime<Utc>) -> String {
|
|
value.to_rfc3339_opts(SecondsFormat::Micros, true)
|
|
}
|
|
|
|
fn load_state(path: &Path) -> Result<State> {
|
|
if !path.exists() {
|
|
return Ok(State::default());
|
|
}
|
|
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 save_state(path: &Path, state: &State) -> Result<()> {
|
|
if let Some(parent) = path.parent() {
|
|
fs::create_dir_all(parent).with_context(|| format!("create {}", parent.display()))?;
|
|
let mut temp = tempfile::NamedTempFile::new_in(parent)
|
|
.with_context(|| format!("create temp file in {}", parent.display()))?;
|
|
serde_json::to_writer_pretty(&mut temp, state)?;
|
|
std::io::Write::write_all(&mut temp, b"\n")?;
|
|
temp.persist(path)
|
|
.map_err(|err| anyhow!("persist {}: {}", path.display(), err))?;
|
|
return Ok(());
|
|
}
|
|
fs::write(path, serde_json::to_string_pretty(state)? + "\n")
|
|
.with_context(|| format!("write {}", path.display()))
|
|
}
|
|
|
|
fn get_start_time(cli: &Cli, state: &State) -> Result<DateTime<Utc>> {
|
|
if let Some(since) = &cli.since {
|
|
return parse_timestamp(since);
|
|
}
|
|
if let Some(last_end) = &state.last_end {
|
|
return Ok(parse_timestamp(last_end)? - TimeDelta::seconds(cli.overlap_seconds));
|
|
}
|
|
Ok(Utc::now() - TimeDelta::hours(cli.lookback_hours))
|
|
}
|
|
|
|
fn parse_prefixes(value: &str) -> Result<Vec<String>> {
|
|
let prefixes: Vec<String> = value
|
|
.split(',')
|
|
.map(str::trim)
|
|
.filter(|item| !item.is_empty())
|
|
.map(ToOwned::to_owned)
|
|
.collect();
|
|
if prefixes.is_empty() {
|
|
bail!("at least one bucket prefix is required");
|
|
}
|
|
Ok(prefixes)
|
|
}
|
|
|
|
fn bucket_stream_type(bucket_id: &str, bucket: &Bucket) -> Option<&'static str> {
|
|
if bucket_id.starts_with("aw-file-operations_") || bucket.r#type == "aw.file.operation" {
|
|
return Some("file_operation");
|
|
}
|
|
if bucket_id.starts_with("aw-dlp-incidents_") || bucket.r#type == "aw.dlp.incident" {
|
|
return Some("dlp_incident");
|
|
}
|
|
None
|
|
}
|
|
|
|
fn select_buckets(
|
|
buckets: &BTreeMap<String, Bucket>,
|
|
prefixes: &[String],
|
|
) -> Vec<(String, Bucket, &'static str)> {
|
|
buckets
|
|
.iter()
|
|
.filter_map(|(bucket_id, bucket)| {
|
|
let stream_type = bucket_stream_type(bucket_id, bucket)?;
|
|
if prefixes.iter().any(|prefix| bucket_id.starts_with(prefix)) {
|
|
Some((bucket_id.clone(), bucket.clone(), stream_type))
|
|
} else {
|
|
None
|
|
}
|
|
})
|
|
.collect()
|
|
}
|
|
|
|
fn build_events_path(
|
|
bucket_id: &str,
|
|
start: DateTime<Utc>,
|
|
end: DateTime<Utc>,
|
|
limit: usize,
|
|
) -> String {
|
|
format!(
|
|
"/buckets/{}/events?start={}&end={}&limit={}",
|
|
urlencoding::encode(bucket_id),
|
|
urlencoding::encode(&format_aw_timestamp(start)),
|
|
urlencoding::encode(&format_aw_timestamp(end)),
|
|
limit
|
|
)
|
|
}
|
|
|
|
fn normalized_data(value: Value) -> Value {
|
|
if value.is_object() {
|
|
value
|
|
} else {
|
|
json!({ "raw": value })
|
|
}
|
|
}
|
|
|
|
fn event_key(bucket_id: &str, timestamp: &str, duration: f64, data: &Value) -> String {
|
|
let payload = serde_json::to_string(data).unwrap_or_else(|_| "{}".to_string());
|
|
format!("{bucket_id}|{timestamp}|{duration}|{payload}")
|
|
}
|
|
|
|
fn event_id(bucket_id: &str, event: &RawAwEvent, data: &Value) -> String {
|
|
match &event.id {
|
|
Some(Value::String(value)) if !value.is_empty() => value.clone(),
|
|
Some(Value::Number(value)) => value.to_string(),
|
|
Some(Value::Bool(value)) => value.to_string(),
|
|
_ => event_key(bucket_id, &event.timestamp, event.duration, data),
|
|
}
|
|
}
|
|
|
|
fn fetch_bucket_events(
|
|
client: &ActivityWatchClient,
|
|
bucket_id: &str,
|
|
bucket: &Bucket,
|
|
stream_type: &str,
|
|
start: DateTime<Utc>,
|
|
end: DateTime<Utc>,
|
|
limit: usize,
|
|
) -> Result<Vec<AwEvent>> {
|
|
let path = build_events_path(bucket_id, start, end, limit);
|
|
let raw_events: Vec<RawAwEvent> = client.get_json(&path)?;
|
|
Ok(raw_events
|
|
.into_iter()
|
|
.map(|raw| {
|
|
let data = normalized_data(raw.data.clone());
|
|
let hostname = if bucket.hostname.is_empty() {
|
|
first_string(&data, &["hostname"]).unwrap_or_default()
|
|
} else {
|
|
bucket.hostname.clone()
|
|
};
|
|
AwEvent {
|
|
bucket_id: bucket_id.to_string(),
|
|
hostname,
|
|
stream_type: stream_type.to_string(),
|
|
event_id: event_id(bucket_id, &raw, &data),
|
|
timestamp: raw.timestamp,
|
|
duration: raw.duration,
|
|
data,
|
|
}
|
|
})
|
|
.collect())
|
|
}
|
|
|
|
fn first_string(data: &Value, keys: &[&str]) -> Option<String> {
|
|
let object = data.as_object()?;
|
|
for key in keys {
|
|
if let Some(value) = object.get(*key) {
|
|
match value {
|
|
Value::Null => {}
|
|
Value::String(s) if !s.is_empty() => return Some(s.clone()),
|
|
other => {
|
|
let rendered = other.to_string();
|
|
if !rendered.is_empty() {
|
|
return Some(rendered);
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
None
|
|
}
|
|
|
|
fn bool_as_int(value: Option<&Value>) -> i64 {
|
|
match value {
|
|
Some(Value::Bool(value)) => i64::from(*value),
|
|
Some(Value::String(value)) => i64::from(matches!(
|
|
value.to_ascii_lowercase().as_str(),
|
|
"1" | "true" | "yes" | "y"
|
|
)),
|
|
Some(Value::Number(value)) => i64::from(value.as_i64().unwrap_or(0) != 0),
|
|
Some(Value::Null) | None => 0,
|
|
Some(_) => 1,
|
|
}
|
|
}
|
|
|
|
fn ensure_schema(connection: &Connection) -> Result<()> {
|
|
connection.execute_batch(
|
|
r#"
|
|
create table if not exists dlp_events (
|
|
id integer primary key autoincrement,
|
|
bucket_id text not null,
|
|
event_id text not null,
|
|
stream_type text not null,
|
|
hostname text not null,
|
|
username text,
|
|
event_ts text not null,
|
|
duration real not null default 0,
|
|
operation text,
|
|
file_path text,
|
|
old_file_path text,
|
|
extension text,
|
|
archive_hint integer not null default 0,
|
|
rule_id text,
|
|
action text,
|
|
severity text,
|
|
signal_type text,
|
|
message text,
|
|
source text,
|
|
screenshot_path text,
|
|
raw_json text not null,
|
|
ingested_at text not null,
|
|
unique (bucket_id, event_id)
|
|
);
|
|
create index if not exists idx_dlp_events_event_ts on dlp_events(event_ts);
|
|
create index if not exists idx_dlp_events_host_ts on dlp_events(hostname, event_ts);
|
|
create index if not exists idx_dlp_events_stream_ts on dlp_events(stream_type, event_ts);
|
|
create index if not exists idx_dlp_events_archive on dlp_events(archive_hint, event_ts);
|
|
create index if not exists idx_dlp_events_rule on dlp_events(rule_id, event_ts);
|
|
|
|
create view if not exists dlp_file_operations as
|
|
select *
|
|
from dlp_events
|
|
where stream_type = 'file_operation';
|
|
|
|
create view if not exists dlp_incidents as
|
|
select *
|
|
from dlp_events
|
|
where stream_type = 'dlp_incident';
|
|
"#,
|
|
)?;
|
|
Ok(())
|
|
}
|
|
|
|
fn connect_sqlite(path: &Path) -> Result<Connection> {
|
|
if let Some(parent) = path.parent() {
|
|
fs::create_dir_all(parent).with_context(|| format!("create {}", parent.display()))?;
|
|
}
|
|
let connection = Connection::open(path).with_context(|| format!("open {}", path.display()))?;
|
|
connection.execute_batch(
|
|
r#"
|
|
pragma journal_mode=WAL;
|
|
pragma synchronous=NORMAL;
|
|
pragma foreign_keys=ON;
|
|
"#,
|
|
)?;
|
|
Ok(connection)
|
|
}
|
|
|
|
fn insert_events(connection: &mut Connection, events: &[AwEvent]) -> Result<usize> {
|
|
let ingested_at = format_aw_timestamp(Utc::now());
|
|
let tx = connection.transaction()?;
|
|
let mut inserted = 0;
|
|
{
|
|
let mut statement = tx.prepare(
|
|
r#"
|
|
insert or ignore into dlp_events (
|
|
bucket_id,
|
|
event_id,
|
|
stream_type,
|
|
hostname,
|
|
username,
|
|
event_ts,
|
|
duration,
|
|
operation,
|
|
file_path,
|
|
old_file_path,
|
|
extension,
|
|
archive_hint,
|
|
rule_id,
|
|
action,
|
|
severity,
|
|
signal_type,
|
|
message,
|
|
source,
|
|
screenshot_path,
|
|
raw_json,
|
|
ingested_at
|
|
)
|
|
values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
"#,
|
|
)?;
|
|
for event in events {
|
|
let raw_json = serde_json::to_string(&event.data)?;
|
|
let archive_hint = bool_as_int(event.data.get("archiveHint"));
|
|
let row_count = statement.execute(params![
|
|
event.bucket_id,
|
|
event.event_id,
|
|
event.stream_type,
|
|
event.hostname,
|
|
first_string(&event.data, &["username", "user"]),
|
|
event.timestamp,
|
|
event.duration,
|
|
first_string(&event.data, &["operation"]),
|
|
first_string(&event.data, &["path", "filePath"]),
|
|
first_string(&event.data, &["oldPath", "oldFilePath"]),
|
|
first_string(&event.data, &["extension"]),
|
|
archive_hint,
|
|
first_string(&event.data, &["ruleId", "rule"]),
|
|
first_string(&event.data, &["action"]),
|
|
first_string(&event.data, &["severity"]),
|
|
first_string(&event.data, &["signalType"]),
|
|
first_string(&event.data, &["message"]),
|
|
first_string(&event.data, &["source"]),
|
|
first_string(
|
|
&event.data,
|
|
&["screenshotPath", "capturePath", "artifactPath"]
|
|
),
|
|
raw_json,
|
|
ingested_at,
|
|
])?;
|
|
inserted += usize::from(row_count > 0);
|
|
}
|
|
}
|
|
tx.commit()?;
|
|
Ok(inserted)
|
|
}
|
|
|
|
fn run(cli: &Cli) -> Result<()> {
|
|
if cli.postgres_dsn.is_some() {
|
|
bail!(
|
|
"PostgreSQL mode is not supported by the Rust DLP aggregator yet; use SQLite mode or the legacy Python script"
|
|
);
|
|
}
|
|
let prefixes = parse_prefixes(&cli.bucket_prefixes)?;
|
|
let state = load_state(&cli.state_path)?;
|
|
let start = get_start_time(cli, &state)?;
|
|
let end = Utc::now();
|
|
let client = ActivityWatchClient::new(&cli.aw_url, Duration::from_secs(cli.timeout))?;
|
|
let buckets: BTreeMap<String, Bucket> = client.get_json("/buckets")?;
|
|
let selected = select_buckets(&buckets, &prefixes);
|
|
let mut events = Vec::new();
|
|
for (bucket_id, bucket, stream_type) in &selected {
|
|
events.extend(fetch_bucket_events(
|
|
&client,
|
|
bucket_id,
|
|
bucket,
|
|
stream_type,
|
|
start,
|
|
end,
|
|
cli.limit,
|
|
)?);
|
|
}
|
|
|
|
if cli.dry_run {
|
|
let summary = DryRunSummary {
|
|
aw_url: cli.aw_url.clone(),
|
|
start: format_aw_timestamp(start),
|
|
end: format_aw_timestamp(end),
|
|
selected_buckets: selected
|
|
.iter()
|
|
.map(|(bucket_id, _bucket, _stream_type)| bucket_id.clone())
|
|
.collect(),
|
|
fetched_events: events.len(),
|
|
};
|
|
println!("{}", serde_json::to_string_pretty(&summary)?);
|
|
return Ok(());
|
|
}
|
|
|
|
let mut connection = connect_sqlite(&cli.sqlite_path)?;
|
|
ensure_schema(&connection)?;
|
|
let inserted = insert_events(&mut connection, &events)?;
|
|
save_state(
|
|
&cli.state_path,
|
|
&State {
|
|
last_end: Some(format_aw_timestamp(end)),
|
|
},
|
|
)?;
|
|
let summary = RunSummary {
|
|
aw_url: cli.aw_url.clone(),
|
|
target: "sqlite".to_string(),
|
|
target_path: cli.sqlite_path.display().to_string(),
|
|
state_path: cli.state_path.display().to_string(),
|
|
start: format_aw_timestamp(start),
|
|
end: format_aw_timestamp(end),
|
|
selected_buckets: selected.len(),
|
|
fetched_events: events.len(),
|
|
inserted_events: inserted,
|
|
};
|
|
println!("{}", serde_json::to_string_pretty(&summary)?);
|
|
Ok(())
|
|
}
|
|
|
|
fn main() -> Result<()> {
|
|
let cli = Cli::parse().apply_env();
|
|
run(&cli)
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
|
|
#[test]
|
|
fn selects_expected_bucket_prefixes() {
|
|
let buckets = BTreeMap::from([
|
|
(
|
|
"aw-file-operations_HOST".to_string(),
|
|
Bucket {
|
|
r#type: "aw.file.operation".to_string(),
|
|
hostname: "HOST".to_string(),
|
|
},
|
|
),
|
|
(
|
|
"aw-dlp-incidents_HOST".to_string(),
|
|
Bucket {
|
|
r#type: "aw.dlp.incident".to_string(),
|
|
hostname: "HOST".to_string(),
|
|
},
|
|
),
|
|
(
|
|
"aw-watcher-window_HOST".to_string(),
|
|
Bucket {
|
|
r#type: "currentwindow".to_string(),
|
|
hostname: "HOST".to_string(),
|
|
},
|
|
),
|
|
]);
|
|
let selected = select_buckets(
|
|
&buckets,
|
|
&parse_prefixes("aw-file-operations_,aw-dlp-incidents_").unwrap(),
|
|
);
|
|
assert_eq!(selected.len(), 2);
|
|
assert_eq!(selected[0].2, "dlp_incident");
|
|
assert_eq!(selected[1].2, "file_operation");
|
|
}
|
|
|
|
#[test]
|
|
fn event_row_helpers_match_python_contract() {
|
|
let data = json!({
|
|
"username": "igor",
|
|
"path": "C:/x.zip",
|
|
"oldFilePath": "C:/x.tmp",
|
|
"archiveHint": "true",
|
|
"rule": "archive",
|
|
"signalType": "file_operation"
|
|
});
|
|
assert_eq!(
|
|
first_string(&data, &["username", "user"]).as_deref(),
|
|
Some("igor")
|
|
);
|
|
assert_eq!(
|
|
first_string(&data, &["oldPath", "oldFilePath"]).as_deref(),
|
|
Some("C:/x.tmp")
|
|
);
|
|
assert_eq!(bool_as_int(data.get("archiveHint")), 1);
|
|
assert_eq!(
|
|
first_string(&data, &["ruleId", "rule"]).as_deref(),
|
|
Some("archive")
|
|
);
|
|
}
|
|
|
|
#[test]
|
|
fn inserts_events_once() {
|
|
let temp = tempfile::tempdir().unwrap();
|
|
let path = temp.path().join("dlp.sqlite");
|
|
let mut connection = connect_sqlite(&path).unwrap();
|
|
ensure_schema(&connection).unwrap();
|
|
let event = AwEvent {
|
|
bucket_id: "aw-file-operations_HOST".to_string(),
|
|
hostname: "HOST".to_string(),
|
|
stream_type: "file_operation".to_string(),
|
|
event_id: "1".to_string(),
|
|
timestamp: "2026-05-31T10:00:00Z".to_string(),
|
|
duration: 0.0,
|
|
data: json!({"username":"igor","archiveHint":true}),
|
|
};
|
|
assert_eq!(
|
|
insert_events(&mut connection, std::slice::from_ref(&event)).unwrap(),
|
|
1
|
|
);
|
|
assert_eq!(insert_events(&mut connection, &[event]).unwrap(), 0);
|
|
let count: i64 = connection
|
|
.query_row("select count(*) from dlp_events", [], |row| row.get(0))
|
|
.unwrap();
|
|
assert_eq!(count, 1);
|
|
}
|
|
}
|