mirror of
https://github.com/marcredhat/kql
synced 2026-06-08 21:27:09 +00:00
marc
135 lines
4.5 KiB
Python
135 lines
4.5 KiB
Python
"""SentinelOne SDL client (uses `requests` for reliable I/O)."""
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import time
|
|
from pathlib import Path
|
|
|
|
import requests
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
CFG = json.loads((ROOT / "config.json").read_text())
|
|
|
|
import os, uuid
|
|
|
|
BASE = CFG["base_url"].rstrip("/")
|
|
WRITE_KEY = CFG["log_write_key"]
|
|
READ_KEY = CFG["log_read_key"]
|
|
# Make the session unique per *process* so SDL never dedupes re-runs of the
|
|
# same payload (SDL hashes session+ts on the server side and silently drops
|
|
# events whose (session, ts) tuple was already accepted -> bytesCharged=0).
|
|
SESSION = os.environ.get("KQL_PROOF_SESSION") or f"kql-proof-{uuid.uuid4()}"
|
|
VERIFY = CFG.get("verify_tls", True)
|
|
TIMEOUT = CFG.get("timeout_seconds", 120)
|
|
print(f"[sdl_client] session = {SESSION}")
|
|
|
|
|
|
def _post(path: str, body: dict, token: str, timeout: int | None = None) -> dict:
|
|
url = f"{BASE}{path}"
|
|
r = requests.post(
|
|
url,
|
|
json=body,
|
|
headers={"Content-Type": "application/json",
|
|
"Authorization": f"Bearer {token}"},
|
|
timeout=timeout or TIMEOUT,
|
|
verify=VERIFY,
|
|
)
|
|
try:
|
|
return r.json()
|
|
except ValueError:
|
|
return {"status": "error", "http_status": r.status_code, "raw": r.text[:500]}
|
|
|
|
|
|
# --- addEvents -------------------------------------------------------------
|
|
def add_events(events: list[dict], session_info: dict | None = None) -> dict:
|
|
payload = {
|
|
"session": SESSION,
|
|
"sessionInfo": session_info or {
|
|
"serverHost": "kql-proof",
|
|
"logfile": "kql-proof.jsonl",
|
|
"parser": "json",
|
|
},
|
|
"events": events,
|
|
"threads": [{"id": "T1", "name": "kql-proof"}],
|
|
}
|
|
return _post("/api/addEvents", payload, WRITE_KEY)
|
|
|
|
|
|
def _clean_attrs(rec: dict) -> dict:
|
|
"""SDL silently rejects events that contain `null` attribute values
|
|
(the call returns status=success but bytesCharged=0 and the event is
|
|
not queryable). Strip them, and coerce everything else to JSON-safe
|
|
primitives that SDL's parser indexes correctly."""
|
|
out: dict = {}
|
|
for k, v in rec.items():
|
|
if v is None:
|
|
continue
|
|
if isinstance(v, bool):
|
|
out[k] = str(v).lower() # SDL stores bools as strings reliably
|
|
elif isinstance(v, (int, float, str)):
|
|
out[k] = v
|
|
else:
|
|
# dict/list -> JSON string
|
|
out[k] = json.dumps(v, default=str)
|
|
return out
|
|
|
|
|
|
def upload_logs(body: str, server_host: str = "kql-proof",
|
|
logfile: str = "kql-proof.jsonl",
|
|
parser: str = "json") -> dict:
|
|
"""POST /api/uploadLogs. Body is raw text; SDL applies the named parser."""
|
|
url = f"{BASE}/api/uploadLogs"
|
|
headers = {
|
|
"Authorization": f"Bearer {WRITE_KEY}",
|
|
"Content-Type": "text/plain",
|
|
"parser": parser,
|
|
"server-host": server_host,
|
|
"logfile": logfile,
|
|
}
|
|
r = requests.post(url, data=body.encode(), headers=headers,
|
|
timeout=TIMEOUT, verify=VERIFY)
|
|
try:
|
|
return r.json()
|
|
except ValueError:
|
|
return {"status": "error", "http_status": r.status_code, "raw": r.text[:500]}
|
|
|
|
|
|
def ingest_jsonl(jsonl_path: Path, run_id: str | None = None,
|
|
batch_lines: int = 2000) -> tuple[int, str]:
|
|
"""Ingest the entire JSONL via uploadLogs. Stamps every event with the
|
|
given `run_id` (or a fresh uuid) so subsequent PowerQueries can scope to
|
|
a single run. Returns (events_sent, run_id)."""
|
|
run_id = run_id or f"run-{uuid.uuid4().hex[:10]}"
|
|
sent = 0
|
|
buf: list[str] = []
|
|
|
|
def flush():
|
|
nonlocal sent
|
|
if not buf:
|
|
return
|
|
r = upload_logs("\n".join(buf))
|
|
if r.get("status") != "success":
|
|
raise RuntimeError(f"uploadLogs rejected batch: {r}")
|
|
sent += len(buf); buf.clear()
|
|
|
|
for line in jsonl_path.read_text().splitlines():
|
|
if not line.strip():
|
|
continue
|
|
rec = json.loads(line)
|
|
rec["proof_run_id"] = run_id
|
|
buf.append(json.dumps(rec, default=str))
|
|
if len(buf) >= batch_lines:
|
|
flush()
|
|
flush()
|
|
return sent, run_id
|
|
|
|
|
|
# --- powerQuery ------------------------------------------------------------
|
|
def power_query(query: str,
|
|
start_time: str | int = "7d",
|
|
end_time: str | int | None = None) -> dict:
|
|
body: dict = {"query": query, "startTime": str(start_time)}
|
|
if end_time is not None:
|
|
body["endTime"] = str(end_time)
|
|
return _post("/api/powerQuery", body, READ_KEY)
|