elastree.py

"""
ElasTree — Trace process trees from an Auditbeat Elasticsearch index.

Usage:
    elastree.py init <username> <password>
    elastree.py trace --hostname HOST --name PROC [--time 12h] [--flow parents]
    elastree.py trace --hostname HOST --pid  PID  [--time 6h]  [--flow children]
    elastree.py trace --hostname HOST --name PROC --start "01-01-2024 08:00:00" --end "01-01-2024 10:00:00" [--flow both]
"""

import os
import sys
from collections import defaultdict, deque
from dataclasses import dataclass
from datetime import datetime, timedelta
from zoneinfo import ZoneInfo

import click
from dotenv import dotenv_values, set_key
from elasticsearch import Elasticsearch, TransportError
from rich import print as rprint
from rich.tree import Tree

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------

CONFIG_PATH = os.path.expanduser("~/.elastree.env")
INDEX = "auditbeat-*"
DEFAULT_TZ = "Asia/Kuala_Lumpur"

FIELDS = [
    "@timestamp",
    "process.pid",
    "process.parent.pid",
    "process.ppid",
    "process.name",
    "process.executable",
    "process.args",
    "process.hash.sha1",
    "user.name",
    "event.action",
]

EVENT_ACTIONS = [
    "process_started",
    "executed",
    "acquired-credentials",
    "existing_process",
]

DURATION_MAP: dict[str, timedelta] = {
    "1h": timedelta(hours=1),
    "2h": timedelta(hours=2),
    "3h": timedelta(hours=3),
    "4h": timedelta(hours=4),
    "5h": timedelta(hours=5),
    "6h": timedelta(hours=6),
    "12h": timedelta(hours=12),
    "24h": timedelta(hours=24),
    "2d": timedelta(days=2),
    "3d": timedelta(days=3),
    "6d": timedelta(days=6),
    "30d": timedelta(days=30),
    "60d": timedelta(days=60),
    "90d": timedelta(days=90),
}
DEFAULT_TIMESPAN = "12h"
FLOW_VALUES = ("parents", "children", "both")


# ---------------------------------------------------------------------------
# Config
# ---------------------------------------------------------------------------


def init_config(username: str, password: str) -> None:
    """Write credentials to ~/.elastree.env (plaintext — perms set to 600)."""
    existing = dotenv_values(CONFIG_PATH) if os.path.exists(CONFIG_PATH) else {}
    if os.path.exists(CONFIG_PATH):
        rprint(
            f"[yellow][!][/yellow] Config already exists at [bold]{CONFIG_PATH}[/bold] — updating credentials only."
        )

    for key, val in [
        ("ES_USERNAME", username),
        ("ES_PASSWORD", password),
        ("ES_HOST", str(existing.get("ES_HOST") or "")),  # fill in manually after init
        ("ES_PORT", str(existing.get("ES_PORT") or "")),  # fill in manually after init
        ("ES_VERIFY_CERTS", str(existing.get("ES_VERIFY_CERTS") or "false")),
        ("ES_SSL_SHOW_WARNS", str(existing.get("ES_SSL_SHOW_WARNS") or "true")),
    ]:
        set_key(CONFIG_PATH, key, val, quote_mode="never")

    if os.name != "nt":
        os.chmod(CONFIG_PATH, 0o600)
    rprint(f"[green][+][/green] Config saved to [bold]{CONFIG_PATH}[/bold]")
    if os.name == "nt":
        rprint("[yellow][!][/yellow] Password stored in plaintext — restrict file access manually on Windows.")
    else:
        rprint("[yellow][!][/yellow] Password stored in plaintext — file perms set to 600.")
    rprint(f"[dim]Edit {CONFIG_PATH} and fill in ES_HOST and ES_PORT.[/dim]")


def connect_elasticsearch() -> Elasticsearch:
    """Load ~/.elastree.env and return an authenticated, verified ES client."""
    if not os.path.exists(CONFIG_PATH):
        raise FileNotFoundError(
            f"Config not found: {CONFIG_PATH}\n"
            "Run:  elastree.py init <username> <password>"
        )

    cfg = dotenv_values(CONFIG_PATH)
    host = str(cfg.get("ES_HOST") or "").strip()
    port_str = str(cfg.get("ES_PORT") or "").strip()
    username = str(cfg.get("ES_USERNAME") or "").strip()
    password = str(cfg.get("ES_PASSWORD") or "").strip()

    missing = [
        k
        for k, v in {
            "ES_HOST": host,
            "ES_PORT": port_str,
            "ES_USERNAME": username,
            "ES_PASSWORD": password,
        }.items()
        if not v
    ]
    if missing:
        raise ValueError(
            f"Config incomplete — fill in: {', '.join(missing)} in {CONFIG_PATH}"
        )

    try:
        port = int(port_str)
    except ValueError:
        raise ValueError(f"ES_PORT must be an integer, got: {port_str!r}")

    es = Elasticsearch(
        hosts=[{"host": host, "port": port, "scheme": "https"}],
        verify_certs=str(cfg.get("ES_VERIFY_CERTS") or "false").lower() == "true",
        ssl_show_warn=str(cfg.get("ES_SSL_SHOW_WARNS") or "true").lower() == "true",
        basic_auth=(username, password),
    )
    if not es.ping():
        raise ConnectionError(f"Cannot reach Elasticsearch at {host}:{port}")
    return es


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


@dataclass
class TimeRange:
    start: datetime
    end: datetime


def to_int_pid(value) -> int | None:
    """Safely convert an ES PID (int / float / str) to int."""
    try:
        return int(float(value)) if value is not None else None
    except (TypeError, ValueError):
        return None


def nested(doc: dict, key: str, default=None):
    """Dot-notation access into a nested dict  (e.g. 'process.pid')."""
    for part in key.split("."):
        if isinstance(doc, dict) and part in doc:
            doc = doc[part]
        else:
            return default
    return doc


def ppid_of(proc: dict) -> int | None:
    """Return parent PID, preferring process.parent.pid over process.ppid.

    Bug 1 fix: use `is not None` instead of truthiness so PID 0 is not skipped.
    """
    val = nested(proc, "process.parent.pid")
    if val is None:
        val = nested(proc, "process.ppid")
    return to_int_pid(val)


def parse_dt(s: str) -> str:
    """Parse a dd-mm-yyyy HH:MM:SS string and return a timezone-aware ISO string.

    Improvement D: raises click.UsageError on bad format instead of a raw traceback.
    """
    try:
        return (
            datetime.strptime(s, "%d-%m-%Y %H:%M:%S")
            .replace(tzinfo=local_tz())
            .isoformat()
        )
    except ValueError:
        raise click.UsageError(
            f"Invalid datetime {s!r}. Expected format: dd-mm-yyyy HH:MM:SS"
        )


def parse_dt_obj(s: str) -> datetime:
    """Parse a dd-mm-yyyy HH:MM:SS string and return a timezone-aware datetime."""
    return datetime.fromisoformat(parse_dt(s))


def local_tz() -> ZoneInfo:
    """Return the configured local timezone."""
    return ZoneInfo(os.environ.get("ELASTREE_TZ", DEFAULT_TZ))


# ---------------------------------------------------------------------------
# Query builders
# ---------------------------------------------------------------------------


def _bool_query(
    hostname: str, field: str, value, time_filter: dict, size: int, sort: str
) -> dict:
    """Single shared ES bool query skeleton used by all query builders."""
    host_query = {"term": {"host.name": hostname}}
    value_query = (
        {"term": {field: value}}
        if field in {"process.pid", "process.parent.pid", "process.ppid"}
        else {"match": {field: value}}
    )
    return {
        "size": size,
        "_source": FIELDS,
        "query": {
            "bool": {
                "must": [
                    host_query,
                    value_query,
                    {"terms": {"event.action": EVENT_ACTIONS}},
                ],
                "filter": [{"range": {"@timestamp": time_filter}}],
            }
        },
        "sort": [{"@timestamp": {"order": sort}}],
    }


def make_query(
    hostname, field, value, timespan, tr: TimeRange | None, size=100
) -> dict:
    """Initial search — relative timespan or absolute range."""
    if tr:
        return _bool_query(
            hostname,
            field,
            value,
            {"gte": tr.start.isoformat(), "lte": tr.end.isoformat()},
            size,
            "desc",
        )
    delta = DURATION_MAP.get(timespan, DURATION_MAP[DEFAULT_TIMESPAN])
    return _bool_query(
        hostname,
        field,
        value,
        {"gte": (datetime.now(local_tz()) - delta).isoformat()},
        size,
        "desc",
    )


def query_parent(hostname: str, parent_pid: int, child_ts: str) -> dict:
    """Parent lookup: predates child (lte), wide 90d lower bound to catch long-lived parents."""
    gte = (datetime.now(local_tz()) - DURATION_MAP["90d"]).isoformat()
    return _bool_query(
        hostname, "process.pid", parent_pid, {"gte": gte, "lte": child_ts}, 1, "desc"
    )


def query_children(
    hostname: str, parent_pid: int, parent_ts: str, tr: TimeRange | None
) -> dict:
    """Child lookup: postdates parent (gte), bounded by the user's time window."""
    lte = tr.end.isoformat() if tr else datetime.now(local_tz()).isoformat()
    return {
        "size": 100,
        "_source": FIELDS,
        "query": {
            "bool": {
                "must": [
                    {"term": {"host.name": hostname}},
                    {
                        "bool": {
                            "should": [
                                {"term": {"process.parent.pid": parent_pid}},
                                {"term": {"process.ppid": parent_pid}},
                            ],
                            "minimum_should_match": 1,
                        }
                    },
                    {"terms": {"event.action": EVENT_ACTIONS}},
                ],
                "filter": [{"range": {"@timestamp": {"gte": parent_ts, "lte": lte}}}],
            }
        },
        "sort": [{"@timestamp": {"order": "asc"}}],
    }


# ---------------------------------------------------------------------------
# ES search wrapper
# ---------------------------------------------------------------------------


def search(es: Elasticsearch, query: dict) -> list[dict]:
    """Execute a query and return _source dicts.

    Improvement A: catches ES transport errors and surfaces them clearly
    instead of letting a raw exception bubble up with a confusing traceback.
    """
    try:
        resp = es.search(index=INDEX, **query)
        return [h["_source"] for h in resp.get("hits", {}).get("hits", [])]
    except TransportError as e:
        rprint(f"[red][!][/red] Elasticsearch error: {e}")
        return []


# ---------------------------------------------------------------------------
# Tree tracing
# ---------------------------------------------------------------------------


def trace_parents(
    es: Elasticsearch,
    hostname: str,
    proc: dict,
    seen: set,
    tr: TimeRange | None,
) -> list[dict]:
    """Walk upward to root ancestors. Returns chain ordered root → target."""
    chain, current = [], proc
    while current:
        pid = to_int_pid(nested(current, "process.pid"))
        if pid is None or pid in seen:
            break
        seen.add(pid)
        chain.append(current)

        parent_pid = ppid_of(current)
        child_ts = nested(current, "@timestamp")
        if parent_pid is None or not child_ts:
            break

        hits = search(es, query_parent(hostname, parent_pid, child_ts))
        current = hits[0] if hits else None

    return chain[::-1]  # root → target


def trace_children(
    es: Elasticsearch, hostname: str, root_proc: dict, seen: set, tr: TimeRange | None
) -> list[dict]:
    """BFS downward from root_proc. Returns all descendants.

    Bug 2 fix: uses deque + popleft() (O(1)) instead of list.pop(0) (O(n)).
    Bug 4 fix: removed unused `timespan` parameter.
    """
    root_pid = to_int_pid(nested(root_proc, "process.pid"))
    discovered: list[dict] = []
    queue: deque[dict] = deque([root_proc])

    while queue:
        current = queue.popleft()
        pid = to_int_pid(nested(current, "process.pid"))
        if pid is None:
            continue
        already_seen = pid in seen
        if already_seen and pid != root_pid:
            continue
        seen.add(pid)
        if not already_seen:
            discovered.append(current)

        parent_ts = nested(current, "@timestamp")
        if not parent_ts:
            continue

        for child in search(es, query_children(hostname, pid, parent_ts, tr)):
            if to_int_pid(nested(child, "process.pid")) not in seen:
                queue.append(child)

    return discovered


# ---------------------------------------------------------------------------
# Tree rendering  (rich)
# ---------------------------------------------------------------------------


def proc_label(proc: dict) -> str:
    """Format a single process as a rich-markup string for tree display."""
    pid = nested(proc, "process.pid", "?")
    ppid = ppid_of(proc)
    if ppid is None:
        ppid = "?"
    name = nested(proc, "process.name", "?")
    exe = nested(proc, "process.executable", "")
    args = nested(proc, "process.args", "")
    user = nested(proc, "user.name", "?")
    sha1 = nested(proc, "process.hash.sha1", "")
    sha1_str = f" [dim]{sha1[:12]}…[/dim]" if sha1 else ""
    return (
        f"[cyan]{pid}[/cyan] [magenta]({ppid})[/magenta] "
        f"[green]{name}[/green] [yellow]{exe}[/yellow] "
        f"[blue]{args}[/blue] [red]{user}[/red]{sha1_str}"
    )


def render_tree(procs: list[dict], title: str) -> None:
    """Build and print a rich Tree from a flat list of process dicts."""
    lookup: dict[int, dict] = {}
    adj: dict[int, list] = defaultdict(list)

    for proc in procs:
        pid = to_int_pid(nested(proc, "process.pid"))
        ppid = ppid_of(proc)
        if pid is None:
            continue
        lookup[pid] = proc
        adj[ppid].append(pid)

    roots = sorted(
        pid for pid in lookup if ppid_of(lookup[pid]) not in lookup or ppid_of(lookup[pid]) == pid
    )
    if not roots:
        roots = sorted(lookup)
    root_node = Tree(f"[bold]{title}[/bold]")

    def add_nodes(parent_node: Tree, pid: int, visited: set) -> None:
        if pid in visited:
            return
        visited.add(pid)
        proc = lookup.get(pid)
        if not proc:
            return
        node = parent_node.add(proc_label(proc))
        for kid in sorted(adj.get(pid, [])):
            add_nodes(node, kid, visited)

    visited: set[int] = set()
    for r in roots:
        add_nodes(root_node, r, visited)

    rprint(root_node)


# ---------------------------------------------------------------------------
# CLI
# ---------------------------------------------------------------------------


@click.group()
def cli():
    """ElasTree — Auditbeat process tree tracer."""


@cli.command()
@click.argument("username")
@click.argument("password")
def init(username, password):
    """Initialise the config file with Elasticsearch credentials."""
    init_config(username, password)


@cli.command()
@click.option("--hostname", required=True, help="host.name field value to filter on")
@click.option("--name", default=None, help="Process name to search  (e.g. bash)")
@click.option("--pid", default=None, type=int, help="Process PID to search")
@click.option(
    "--time", default=None, help=f"Relative look-back window: {', '.join(DURATION_MAP)}"
)
@click.option("--start", default=None, help="Range start  (dd-mm-yyyy HH:MM:SS)")
@click.option("--end", default=None, help="Range end    (dd-mm-yyyy HH:MM:SS)")
@click.option(
    "--flow",
    default="parents",
    type=click.Choice(FLOW_VALUES, case_sensitive=False),
    show_default=True,
    help="parents | children | both",
)
def trace(hostname, name, pid, time, start, end, flow):
    """Trace a process tree from the Auditbeat index."""

    # Improvement C: catch --name + --pid used together.
    if name is not None and pid is not None:
        raise click.UsageError("--name and --pid are mutually exclusive.")
    if name is None and pid is None:
        raise click.UsageError("Provide --name or --pid.")
    if time and (start or end):
        raise click.UsageError("Use --time OR --start/--end, not both.")
    if not time and not (start and end):
        raise click.UsageError("Provide --time, or both --start and --end.")
    if time and time not in DURATION_MAP:
        raise click.UsageError(
            f"Unknown --time {time!r}. Valid: {', '.join(DURATION_MAP)}"
        )

    try:
        es = connect_elasticsearch()
    except (FileNotFoundError, ValueError, ConnectionError) as e:
        rprint(f"[red][!][/red] {e}")
        sys.exit(1)

    # Improvement D: parse_dt raises click.UsageError on bad format.
    tr = TimeRange(parse_dt_obj(start), parse_dt_obj(end)) if (start and end) else None
    if tr and tr.start > tr.end:
        raise click.UsageError("--start must be earlier than or equal to --end.")
    field = "process.name" if name else "process.pid"
    value = name or pid

    hits = search(es, make_query(hostname, field, value, time, tr))
    if not hits:
        rprint("[yellow][!][/yellow] No matching processes found.")
        sys.exit(0)

    # De-duplicate by PID — name searches return many events per process.
    initial: dict[int, dict] = {}
    for src in hits:
        p = to_int_pid(nested(src, "process.pid"))
        if p is not None and p not in initial:
            initial[p] = src

    all_procs: list[dict] = []
    seen: set[int] = set()

    for proc in initial.values():
        # --- upward pass ---
        if flow in ("parents", "both"):
            ancestors = trace_parents(es, hostname, proc, seen, tr)
            all_procs.extend(ancestors)
        else:
            ancestors = [proc]

        # --- downward pass ---
        # Bug 3 fix: snapshot the ancestors discovered in THIS iteration only,
        # not all_procs which grows across iterations and causes re-expansion.
        if flow in ("children", "both"):
            expand_from = ancestors if flow == "both" else [proc]
            for ancestor in expand_from:
                # Bug 4 fix: trace_children no longer takes timespan.
                all_procs.extend(trace_children(es, hostname, ancestor, seen, tr))

    if not all_procs:
        rprint("[yellow][!][/yellow] Could not build any process chains.")
        sys.exit(0)

    label = {"parents": "Ancestors", "children": "Descendants", "both": "Full Tree"}
    render_tree(all_procs, f"{label[flow]} — {hostname}")


if __name__ == "__main__":
    cli()