527 lines
17 KiB
Python
Executable File
527 lines
17 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
import argparse
|
|
import json
|
|
import os
|
|
import sqlite3
|
|
import sys
|
|
import urllib.error
|
|
import urllib.parse
|
|
import urllib.request
|
|
from collections.abc import Iterable
|
|
from dataclasses import dataclass
|
|
from datetime import UTC, datetime, timedelta
|
|
from pathlib import Path
|
|
from typing import Protocol, TypeAlias
|
|
|
|
|
|
JsonScalar: TypeAlias = str | int | float | bool | None
|
|
JsonValue: TypeAlias = JsonScalar | list["JsonValue"] | dict[str, "JsonValue"]
|
|
|
|
|
|
DEFAULT_BUCKET_PREFIXES = (
|
|
"aw-file-operations_",
|
|
"aw-dlp-incidents_",
|
|
"aw-dlp-endpoint-signals_",
|
|
"aw-email-monitor_",
|
|
)
|
|
DEFAULT_SQLITE_PATH = "data/dlp-events.sqlite3"
|
|
EVENT_COLUMNS = (
|
|
"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",
|
|
)
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Bucket:
|
|
id: str
|
|
type: str
|
|
client: str
|
|
hostname: str
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class AwEvent:
|
|
bucket_id: str
|
|
hostname: str
|
|
stream_type: str
|
|
event_id: str
|
|
timestamp: str
|
|
duration: float
|
|
data: dict[str, JsonValue]
|
|
|
|
|
|
class PsycopgConnection(Protocol):
|
|
def cursor(self):
|
|
...
|
|
|
|
def commit(self) -> None:
|
|
...
|
|
|
|
|
|
def utc_now() -> datetime:
|
|
return datetime.now(tz=UTC)
|
|
|
|
|
|
def parse_timestamp(value: str) -> datetime:
|
|
normalized = value.replace("Z", "+00:00")
|
|
parsed = datetime.fromisoformat(normalized)
|
|
if parsed.tzinfo is None:
|
|
return parsed.replace(tzinfo=UTC)
|
|
return parsed.astimezone(UTC)
|
|
|
|
|
|
def format_aw_timestamp(value: datetime) -> str:
|
|
return value.astimezone(UTC).isoformat().replace("+00:00", "Z")
|
|
|
|
|
|
def load_state(path: Path) -> dict[str, str]:
|
|
if not path.exists():
|
|
return {}
|
|
return json.loads(path.read_text(encoding="utf-8"))
|
|
|
|
|
|
def save_state(path: Path, state: dict[str, str]) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(json.dumps(state, ensure_ascii=False, indent=2, sort_keys=True) + "\n", encoding="utf-8")
|
|
|
|
|
|
def normalize_base_url(base_url: str) -> str:
|
|
return base_url.rstrip("/")
|
|
|
|
|
|
def aw_get_json(base_url: str, path: str, timeout: int) -> JsonValue:
|
|
url = normalize_base_url(base_url) + path
|
|
request = urllib.request.Request(url, headers={"Accept": "application/json"})
|
|
with urllib.request.urlopen(request, timeout=timeout) as response:
|
|
return json.loads(response.read().decode("utf-8"))
|
|
|
|
|
|
def list_buckets(base_url: str, timeout: int) -> list[Bucket]:
|
|
payload = aw_get_json(base_url, "/buckets", timeout)
|
|
if not isinstance(payload, dict):
|
|
raise ValueError("ActivityWatch /buckets response must be a JSON object")
|
|
buckets: list[Bucket] = []
|
|
for bucket_id, bucket_data in payload.items():
|
|
if not isinstance(bucket_data, dict):
|
|
continue
|
|
buckets.append(
|
|
Bucket(
|
|
id=str(bucket_id),
|
|
type=str(bucket_data.get("type", "")),
|
|
client=str(bucket_data.get("client", "")),
|
|
hostname=str(bucket_data.get("hostname", "")),
|
|
)
|
|
)
|
|
return buckets
|
|
|
|
|
|
def bucket_stream_type(bucket: Bucket) -> str | None:
|
|
if bucket.id.startswith("aw-file-operations_") or bucket.type == "aw.file.operation":
|
|
return "file_operation"
|
|
if bucket.id.startswith("aw-dlp-incidents_") or bucket.type == "aw.dlp.incident":
|
|
return "dlp_incident"
|
|
if bucket.id.startswith("aw-dlp-endpoint-signals_") or bucket.type == "aw.dlp.endpoint.signal":
|
|
return "dlp_endpoint_signal"
|
|
if bucket.id.startswith("aw-email-monitor_") or bucket.type == "aw.email.signal":
|
|
return "email_monitor"
|
|
return None
|
|
|
|
|
|
def select_buckets(buckets: Iterable[Bucket], prefixes: tuple[str, ...]) -> list[tuple[Bucket, str]]:
|
|
selected: list[tuple[Bucket, str]] = []
|
|
for bucket in buckets:
|
|
stream_type = bucket_stream_type(bucket)
|
|
if stream_type and any(bucket.id.startswith(prefix) for prefix in prefixes):
|
|
selected.append((bucket, stream_type))
|
|
return selected
|
|
|
|
|
|
def build_events_path(bucket_id: str, start: datetime, end: datetime, limit: int) -> str:
|
|
query = urllib.parse.urlencode(
|
|
{
|
|
"start": format_aw_timestamp(start),
|
|
"end": format_aw_timestamp(end),
|
|
"limit": str(limit),
|
|
}
|
|
)
|
|
return f"/buckets/{urllib.parse.quote(bucket_id, safe='')}/events?{query}"
|
|
|
|
|
|
def event_key(bucket_id: str, timestamp: str, duration: float, data: dict[str, JsonValue]) -> str:
|
|
payload = json.dumps(data, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
|
|
return f"{bucket_id}|{timestamp}|{duration}|{payload}"
|
|
|
|
|
|
def fetch_bucket_events(
|
|
base_url: str,
|
|
bucket: Bucket,
|
|
stream_type: str,
|
|
start: datetime,
|
|
end: datetime,
|
|
limit: int,
|
|
timeout: int,
|
|
) -> list[AwEvent]:
|
|
payload = aw_get_json(base_url, build_events_path(bucket.id, start, end, limit), timeout)
|
|
if not isinstance(payload, list):
|
|
raise ValueError(f"ActivityWatch events response for {bucket.id} must be a JSON array")
|
|
events: list[AwEvent] = []
|
|
for item in payload:
|
|
if not isinstance(item, dict):
|
|
continue
|
|
timestamp = str(item["timestamp"])
|
|
duration = float(item.get("duration", 0) or 0)
|
|
data = item.get("data") or {}
|
|
if not isinstance(data, dict):
|
|
data = {"raw": data}
|
|
item_id = str(item.get("id") or event_key(bucket.id, timestamp, duration, data))
|
|
events.append(
|
|
AwEvent(
|
|
bucket_id=bucket.id,
|
|
hostname=bucket.hostname or str(data.get("hostname") or ""),
|
|
stream_type=stream_type,
|
|
event_id=item_id,
|
|
timestamp=timestamp,
|
|
duration=duration,
|
|
data=data,
|
|
)
|
|
)
|
|
return events
|
|
|
|
|
|
def connect_sqlite(path: Path) -> sqlite3.Connection:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
connection = sqlite3.connect(str(path))
|
|
connection.execute("PRAGMA journal_mode=WAL")
|
|
connection.execute("PRAGMA synchronous=NORMAL")
|
|
connection.execute("PRAGMA foreign_keys=ON")
|
|
return connection
|
|
|
|
|
|
def ensure_schema(connection: sqlite3.Connection) -> None:
|
|
connection.executescript(
|
|
"""
|
|
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';
|
|
"""
|
|
)
|
|
connection.commit()
|
|
|
|
|
|
def ensure_postgres_schema(connection: PsycopgConnection) -> None:
|
|
with connection.cursor() as cursor:
|
|
cursor.execute(
|
|
"""
|
|
create table if not exists dlp_events (
|
|
id bigserial primary key,
|
|
bucket_id text not null,
|
|
event_id text not null,
|
|
stream_type text not null,
|
|
hostname text not null,
|
|
username text,
|
|
event_ts timestamptz not null,
|
|
duration double precision not null default 0,
|
|
operation text,
|
|
file_path text,
|
|
old_file_path text,
|
|
extension text,
|
|
archive_hint boolean not null default false,
|
|
rule_id text,
|
|
action text,
|
|
severity text,
|
|
signal_type text,
|
|
message text,
|
|
source text,
|
|
screenshot_path text,
|
|
raw_json jsonb not null,
|
|
ingested_at timestamptz 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 or replace view dlp_file_operations as
|
|
select *
|
|
from dlp_events
|
|
where stream_type = 'file_operation';
|
|
|
|
create or replace view dlp_incidents as
|
|
select *
|
|
from dlp_events
|
|
where stream_type = 'dlp_incident';
|
|
"""
|
|
)
|
|
connection.commit()
|
|
|
|
|
|
def first_string(data: dict[str, JsonValue], keys: tuple[str, ...]) -> str | None:
|
|
for key in keys:
|
|
value = data.get(key)
|
|
if value is not None and str(value) != "":
|
|
return str(value)
|
|
return None
|
|
|
|
|
|
def bool_as_int(value: JsonValue) -> int:
|
|
if isinstance(value, bool):
|
|
return int(value)
|
|
if isinstance(value, str):
|
|
return int(value.lower() in {"1", "true", "yes", "y"})
|
|
return int(bool(value))
|
|
|
|
|
|
def event_row(event: AwEvent, ingested_at: str) -> tuple[JsonValue, ...]:
|
|
data = event.data
|
|
event_id = event.event_id or event_key(event.bucket_id, event.timestamp, event.duration, data)
|
|
return (
|
|
event.bucket_id,
|
|
event_id,
|
|
event.stream_type,
|
|
event.hostname,
|
|
first_string(data, ("username", "user")),
|
|
event.timestamp,
|
|
event.duration,
|
|
first_string(data, ("operation",)),
|
|
first_string(data, ("path", "filePath")),
|
|
first_string(data, ("oldPath", "oldFilePath")),
|
|
first_string(data, ("extension",)),
|
|
bool_as_int(data.get("archiveHint")),
|
|
first_string(data, ("ruleId", "rule")),
|
|
first_string(data, ("action",)),
|
|
first_string(data, ("severity",)),
|
|
first_string(data, ("signalType",)),
|
|
first_string(data, ("message",)),
|
|
first_string(data, ("source",)),
|
|
first_string(data, ("screenshotPath", "capturePath", "artifactPath")),
|
|
json.dumps(data, ensure_ascii=False, sort_keys=True),
|
|
ingested_at,
|
|
)
|
|
|
|
|
|
def insert_events(connection: sqlite3.Connection, events: Iterable[AwEvent]) -> int:
|
|
inserted = 0
|
|
now = format_aw_timestamp(utc_now())
|
|
for event in events:
|
|
cursor = connection.execute(
|
|
"""
|
|
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 (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
|
|
""",
|
|
event_row(event, now),
|
|
)
|
|
inserted += int(cursor.rowcount > 0)
|
|
connection.commit()
|
|
return inserted
|
|
|
|
|
|
def insert_postgres_events(dsn: str, events: Iterable[AwEvent]) -> int:
|
|
try:
|
|
import psycopg
|
|
except ImportError as exc:
|
|
raise SystemExit("PostgreSQL mode requires psycopg: python3 -m pip install 'psycopg[binary]'") from exc
|
|
|
|
inserted = 0
|
|
now = format_aw_timestamp(utc_now())
|
|
columns = ", ".join(EVENT_COLUMNS)
|
|
placeholders = ", ".join(["%s"] * len(EVENT_COLUMNS))
|
|
sql = f"""
|
|
insert into dlp_events ({columns})
|
|
values ({placeholders})
|
|
on conflict (bucket_id, event_id) do nothing
|
|
"""
|
|
with psycopg.connect(dsn) as connection:
|
|
ensure_postgres_schema(connection)
|
|
with connection.cursor() as cursor:
|
|
for event in events:
|
|
row = list(event_row(event, now))
|
|
row[EVENT_COLUMNS.index("archive_hint")] = bool(row[EVENT_COLUMNS.index("archive_hint")])
|
|
cursor.execute(sql, row)
|
|
inserted += int(cursor.rowcount > 0)
|
|
connection.commit()
|
|
return inserted
|
|
|
|
|
|
def get_start_time(args: argparse.Namespace, state: dict[str, str]) -> datetime:
|
|
if args.since:
|
|
return parse_timestamp(args.since)
|
|
if state.get("last_end"):
|
|
return parse_timestamp(state["last_end"]) - timedelta(seconds=args.overlap_seconds)
|
|
return utc_now() - timedelta(hours=args.lookback_hours)
|
|
|
|
|
|
def parse_prefixes(value: str) -> tuple[str, ...]:
|
|
prefixes = tuple(item.strip() for item in value.split(",") if item.strip())
|
|
if not prefixes:
|
|
raise argparse.ArgumentTypeError("at least one bucket prefix is required")
|
|
return prefixes
|
|
|
|
|
|
def build_parser() -> argparse.ArgumentParser:
|
|
parser = argparse.ArgumentParser(description="Aggregate AWatch-rus DLP buckets into a local warehouse database.")
|
|
parser.add_argument("--aw-url", default=os.environ.get("AW_URL", "http://127.0.0.1:5600/api/0"))
|
|
parser.add_argument("--postgres-dsn", default=os.environ.get("DLP_AGGREGATOR_POSTGRES_DSN"))
|
|
parser.add_argument("--sqlite-path", default=os.environ.get("DLP_AGGREGATOR_SQLITE_PATH", DEFAULT_SQLITE_PATH))
|
|
parser.add_argument("--state-path", default=os.environ.get("DLP_AGGREGATOR_STATE_PATH", "data/dlp-aggregator-state.json"))
|
|
parser.add_argument("--bucket-prefixes", type=parse_prefixes, default=DEFAULT_BUCKET_PREFIXES)
|
|
parser.add_argument("--since", help="UTC ISO timestamp. Overrides saved state, for example 2026-05-02T00:00:00Z.")
|
|
parser.add_argument("--lookback-hours", type=int, default=24)
|
|
parser.add_argument("--overlap-seconds", type=int, default=60)
|
|
parser.add_argument("--limit", type=int, default=10000)
|
|
parser.add_argument("--timeout", type=int, default=15)
|
|
parser.add_argument("--dry-run", action="store_true")
|
|
return parser
|
|
|
|
|
|
def main() -> int:
|
|
args = build_parser().parse_args()
|
|
state_path = Path(args.state_path)
|
|
state = load_state(state_path)
|
|
start = get_start_time(args, state)
|
|
end = utc_now()
|
|
|
|
buckets = select_buckets(list_buckets(args.aw_url, args.timeout), args.bucket_prefixes)
|
|
all_events: list[AwEvent] = []
|
|
for bucket, stream_type in buckets:
|
|
all_events.extend(fetch_bucket_events(args.aw_url, bucket, stream_type, start, end, args.limit, args.timeout))
|
|
|
|
if args.dry_run:
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"aw_url": args.aw_url,
|
|
"start": format_aw_timestamp(start),
|
|
"end": format_aw_timestamp(end),
|
|
"selected_buckets": [bucket.id for bucket, _stream_type in buckets],
|
|
"fetched_events": len(all_events),
|
|
},
|
|
ensure_ascii=False,
|
|
indent=2,
|
|
)
|
|
)
|
|
return 0
|
|
|
|
if args.postgres_dsn:
|
|
target = "postgres"
|
|
target_path = args.postgres_dsn.split("@")[-1]
|
|
inserted = insert_postgres_events(args.postgres_dsn, all_events)
|
|
else:
|
|
target = "sqlite"
|
|
sqlite_path = Path(args.sqlite_path)
|
|
target_path = str(sqlite_path)
|
|
connection = connect_sqlite(sqlite_path)
|
|
try:
|
|
ensure_schema(connection)
|
|
inserted = insert_events(connection, all_events)
|
|
finally:
|
|
connection.close()
|
|
|
|
state["last_end"] = format_aw_timestamp(end)
|
|
save_state(state_path, state)
|
|
print(
|
|
json.dumps(
|
|
{
|
|
"aw_url": args.aw_url,
|
|
"target": target,
|
|
"target_path": target_path,
|
|
"state_path": str(state_path),
|
|
"start": format_aw_timestamp(start),
|
|
"end": format_aw_timestamp(end),
|
|
"selected_buckets": len(buckets),
|
|
"fetched_events": len(all_events),
|
|
"inserted_events": inserted,
|
|
},
|
|
ensure_ascii=False,
|
|
indent=2,
|
|
)
|
|
)
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
try:
|
|
raise SystemExit(main())
|
|
except urllib.error.URLError as exc:
|
|
print(f"ActivityWatch API request failed: {exc}", file=sys.stderr)
|
|
raise SystemExit(2)
|