#!/usr/bin/env python3
"""Find, collection: the endpoint evidence collector.

artifact_census.py reads evidence bundles. This is the piece that
produces one on a real machine, using only the standard library and
ordinary user level commands:

    workspace artifacts   copies the small markdown artifacts of
                          any agent workspace found under --root
                          (names and content; these files are the
                          fingerprint, and they are configuration,
                          not credentials)
    processes             ps output, parsed
    schedulers            the user crontab, parsed
    network               a snapshot of established connections
                          from ss; one snapshot is a moment, cadence
                          accrues as snapshots append on later runs

The bundle it writes is exactly the shape artifact_census.py
consumes: <out>/endpoints/<endpoint-id>/...

Production note, stated plainly: a fleet collector would run under
an EDR or MDM with the organization's consent and ship metadata
plus hashes to a collector service. This script is the reference
implementation of the bundle format and the demo collector.

Usage:
    python3 collector.py --root DIR --out BUNDLE_DIR
                         [--endpoint-id NAME]
"""

import argparse
import getpass
import json
import shutil
import socket
import subprocess
from datetime import datetime, timezone
from pathlib import Path

ARTIFACT_NAMES = {"HEARTBEAT.md", "SOUL.md", "IDENTITY.md",
                  "AGENTS.md", "MEMORY.md", "USER.md", "TOOLS.md"}


def now_iso():
    return datetime.now(timezone.utc).isoformat()


def collect_processes(user):
    rows = []
    try:
        out = subprocess.run(["ps", "-eo", "args"],
                             capture_output=True, text=True,
                             timeout=20).stdout
        for line in out.splitlines()[1:]:
            line = line.strip()
            if line:
                rows.append({"process": Path(line.split()[0]).name,
                             "command": line, "user": user})
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass
    return rows


def collect_schedulers():
    rows = []
    try:
        out = subprocess.run(["crontab", "-l"], capture_output=True,
                             text=True, timeout=10)
        for line in (out.stdout or "").splitlines():
            line = line.strip()
            if line and not line.startswith("#"):
                rows.append({"type": "cron", "command": line})
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass
    return rows


def collect_network():
    rows = []
    try:
        out = subprocess.run(["ss", "-tn", "state", "established"],
                             capture_output=True, text=True,
                             timeout=10).stdout
        for line in out.splitlines()[1:]:
            parts = line.split()
            if len(parts) >= 5:
                peer = parts[4]
                host = peer.rsplit(":", 1)[0].strip("[]")
                rows.append({"timestamp": now_iso(), "host": host,
                             "peer": peer, "snapshot": True})
    except (FileNotFoundError, subprocess.TimeoutExpired):
        pass
    return rows


def main():
    parser = argparse.ArgumentParser(
        description="Endpoint evidence collector")
    parser.add_argument("--root", required=True,
                        help="Folder to scan for agent workspaces")
    parser.add_argument("--out", required=True, help="Bundle directory")
    parser.add_argument("--endpoint-id", default=socket.gethostname())
    args = parser.parse_args()

    root = Path(args.root)
    user = getpass.getuser()
    endpoint_dir = Path(args.out) / "endpoints" / args.endpoint_id
    workspace_dir = endpoint_dir / "workspace"
    workspace_dir.mkdir(parents=True, exist_ok=True)

    copied = []
    for path in sorted(root.rglob("*")):
        if path.is_file() and path.name in ARTIFACT_NAMES:
            rel = path.relative_to(root)
            dest = workspace_dir / rel
            dest.parent.mkdir(parents=True, exist_ok=True)
            shutil.copyfile(path, dest)
            copied.append(str(rel))

    (endpoint_dir / "endpoint.json").write_text(json.dumps({
        "endpoint_id": args.endpoint_id,
        "hostname": socket.gethostname(),
        "owner": user, "user": user,
        "collected_at": now_iso()}, indent=2) + "\n")
    for name, rows in (("processes.jsonl", collect_processes(user)),
                       ("schedulers.jsonl", collect_schedulers())):
        with (endpoint_dir / name).open("w") as fh:
            for row in rows:
                fh.write(json.dumps(row) + "\n")
    network = collect_network()
    with (endpoint_dir / "network.jsonl").open("a") as fh:
        for row in network:
            fh.write(json.dumps(row) + "\n")

    print(f"Collector: endpoint {args.endpoint_id}: {len(copied)} "
          f"workspace artifact(s) copied, process snapshot "
          f"{len(collect_processes(user))} rows, "
          f"{len(network)} established connection(s) in this "
          f"snapshot.")
    for rel in copied:
        print(f"  artifact: {rel}")
    print(f"Bundle written under {args.out}")


if __name__ == "__main__":
    main()