envmapper.py

#!/usr/bin/env python3
"""
Netographer (LLM Edition)
=========================

A drop-in, LLM-assisted rewrite of the original `netographer` script that grouped
user activity from Windows `tasklist.exe` exports and Enum4Linux output.

What changed vs. the original
-----------------------------
1) **Same inputs**: still supports parsing Windows Tasklist text dumps and Enum4Linux output.
2) **Richer features**: builds a per-user "activity profile" (process inventory + frequencies + AD groups).
3) **Embeddings-powered grouping**: converts each user profile into a text description and embeds it with a
   pluggable provider (OpenAI by default), then clusters in vector space (HDBSCAN → KMeans fallback).
4) **LLM labeling**: auto-generates human-readable labels/descriptions + risk for each cluster.
5) **Reports**: dumps machine-readable JSON and a Markdown report you can hand to an analyst.

Quick start
-----------
```bash
pip install --upgrade pandas numpy scikit-learn hdbscan tenacity python-dateutil tqdm tabulate
# (Optional) OpenAI provider
pip install --upgrade openai

export OPENAI_API_KEY=...  # Or swap providers by implementing LLMClient below

python netographer_llm.py cluster \
  --tasklist-glob "/path/to/tasklists/*.txt" \
  --enum4linux "/path/to/enum4linux.txt" \
  --json-out clusters.json \
  --md-out clusters.md
```

Inputs
------
- **Tasklist files**: Plain-text output of `tasklist.exe /V` redirected to a file. Multiple files allowed.
- **Enum4Linux file**: Typical enum4linux-ng output, used to enrich users with AD groups.

Outputs
-------
- **clusters.json**: Per-user cluster assignments + cluster labels/risks.
- **clusters.md**: Markdown report summarizing clusters (top processes, sample users, LLM label/description).

Notes
-----
- Keep PII out of prompts. The script masks emails and IPs in process lines.
- If you cannot use OpenAI, replace `LLMClient`'s `.chat()` and `.embed()` with your provider.
- This code is designed to be a clean replacement; command names mirror the original where useful.
"""
from __future__ import annotations

import argparse
import json
import os
import re
from collections import Counter, OrderedDict, defaultdict
from dataclasses import dataclass
from pathlib import Path
from typing import Any, Dict, Iterable, List, Optional, Tuple

import numpy as np
import pandas as pd
from tenacity import retry, stop_after_attempt, wait_exponential
from tqdm import tqdm

# ML
try:
    import hdbscan  # type: ignore
except Exception:
    hdbscan = None

from sklearn.cluster import KMeans
from sklearn.metrics import silhouette_score

# Optional provider (OpenAI)
try:
    from openai import OpenAI  # pip install openai
except Exception:
    OpenAI = None

# ---------------------------- LLM client ----------------------------
@dataclass
class LLMConfig:
    chat_model: str = os.getenv("LLM_CHAT_MODEL", "gpt-4o-mini")
    embed_model: str = os.getenv("LLM_EMBED_MODEL", "text-embedding-3-large")
    temperature: float = float(os.getenv("LLM_TEMPERATURE", "0.1"))
    max_tokens: int = int(os.getenv("LLM_MAX_TOKENS", "600"))


class LLMClient:
    """Swap this out to use Anthropic, Azure OpenAI, etc."""

    def __init__(self, cfg: Optional[LLMConfig] = None):
        self.cfg = cfg or LLMConfig()
        self._oai = None
        if OpenAI is not None and os.getenv("OPENAI_API_KEY"):
            self._oai = OpenAI()

    @retry(wait=wait_exponential(min=1, max=20), stop=stop_after_attempt(5), reraise=True)
    def chat(self, system: str, user: str) -> str:
        if self._oai is None:
            raise RuntimeError("No LLM provider configured. Set OPENAI_API_KEY or implement LLMClient.chat().")
        resp = self._oai.chat.completions.create(
            model=self.cfg.chat_model,
            temperature=self.cfg.temperature,
            max_tokens=self.cfg.max_tokens,
            messages=[{"role": "system", "content": system}, {"role": "user", "content": user}],
        )
        return resp.choices[0].message.content.strip()

    @retry(wait=wait_exponential(min=1, max=20), stop=stop_after_attempt(5), reraise=True)
    def embed(self, texts: List[str]) -> List[List[float]]:
        if self._oai is None:
            raise RuntimeError("No embedding provider configured. Set OPENAI_API_KEY or implement LLMClient.embed().")
        out: List[List[float]] = []
        batch = 256
        for i in range(0, len(texts), batch):
            sub = texts[i : i + batch]
            resp = self._oai.embeddings.create(model=self.cfg.embed_model, input=sub)
            out.extend([d.embedding for d in resp.data])
        return out

# ---------------------------- Parsers ----------------------------
EMAIL_RE = re.compile(r"[A-Za-z0-9._%+-]+@[A-Za-z0-9.-]+\.[A-Za-z]{2,}")
IP_RE = re.compile(r"\b\d{1,3}(?:\.\d{1,3}){3}\b")


def _mask(s: str) -> str:
    s = EMAIL_RE.sub("<EMAIL>", s)
    s = IP_RE.sub("<IP>", s)
    return s


class TasklistFile:
    """Parses plain-text `tasklist.exe /V` output.

    We auto-detect the header and fixed-width columns.
    Returns per-user -> set(process names).
    """

    def __init__(self, path: Path):
        self.path = Path(path)
        self.headers: List[str] = []
        self.rows: List[Dict[str, str]] = []

    def _split_columns(self, line: str, widths: List[int]) -> List[str]:
        out = []
        i = 0
        for w in widths:
            out.append(line[i : i + w])
            i += w
        out[-1] = out[-1].rstrip("\n")
        return out

    def load(self) -> "TasklistFile":
        with self.path.open("r", encoding="utf-8", errors="ignore") as f:
            header_line = None
            widths: List[int] = []
            for raw in f:
                line = raw.rstrip("\n")
                if not self.headers:
                    if not line.strip():
                        continue  # skip leading blanks
                    if header_line is None:
                        header_line = line
                        continue
                    # widths line (==== blocks separated by space)
                    widths = [len(x) + 1 for x in line.rstrip().split(" ")]
                    self.headers = [h.strip() for h in self._split_columns(header_line, widths)]
                    continue
                # data rows
                cols = self._split_columns(line, widths)
                row = {h: c.strip() for h, c in zip(self.headers, cols)}
                self.rows.append(row)
        return self

    def user_processes(self) -> Dict[str, set[str]]:
        if not self.rows:
            return {}
        if not {"User Name", "Image Name"}.issubset(self.headers):
            raise ValueError(f"{self.path} missing required columns 'User Name'/'Image Name'")
        users = set(r.get("User Name", "") for r in self.rows)
        up: Dict[str, set[str]] = {u: set() for u in users}
        for r in self.rows:
            u = r.get("User Name", "")
            p = r.get("Image Name", "")
            if u and p:
                up[u].add(_mask(p))
        return up


class Enum4LinuxFile:
    """Minimal parser that extracts users and groups from enum4linux output."""

    def __init__(self, path: Path):
        self.path = Path(path)
        self.domain = ""
        self.users: Dict[str, List[str]] = {}
        self.groups: Dict[str, List[str]] = {}

    def load(self) -> "Enum4LinuxFile":
        section = None
        with self.path.open("r", encoding="utf-8", errors="ignore") as f:
            for raw in f:
                line = raw.rstrip("\n")
                if line.startswith("Domain Name:"):
                    self.domain = line.split(":", 1)[1].strip()
                if line.startswith("Users on"):
                    section = "users"
                    continue
                if line.startswith("Groups on"):
                    section = "groups"
                    continue
                if not section:
                    continue
                if section == "users" and line.startswith("user:"):
                    user, rid = line.split(" rid:")
                    user = f"{self.domain}\\{user[6:-1]}"  # user: 'NAME'
                    self.users[user] = []
                if section == "groups":
                    if line.startswith("group:"):
                        group, rid = line.split(" rid:")
                        group = f"{self.domain}\\{group[7:-1]}"  # group: 'NAME'
                        self.groups[group] = []
                    elif line.startswith("Group '") and " has member: " in line and not line.endswith("Couldn't lookup SIDs"):
                        group, member = line.split(" has member: ")
                        group = f"{self.domain}\\{group[6:].split(" (RID: ")[0][1:-1]}"
                        member = f"{self.domain}\\{member.strip()}"
                        if group in self.groups and member in self.users:
                            self.groups[group].append(member)
                            self.users[member].append(group)
        return self

# ---------------------------- Feature engineering ----------------------------

def merge_user_processes(tasklist_paths: List[Path]) -> Dict[str, List[str]]:
    combined: Dict[str, set[str]] = defaultdict(set)
    for p in tasklist_paths:
        try:
            up = TasklistFile(p).load().user_processes()
        except Exception:
            continue
        for user, procs in up.items():
            combined[user].update(procs)
    return {u: sorted(list(v)) for u, v in combined.items()}


def process_frequencies(user_processes: Dict[str, List[str]]) -> Counter[str]:
    c: Counter[str] = Counter()
    for procs in user_processes.values():
        c.update(procs)
    return c


def build_user_profile_text(user: str, procs: List[str], groups: Optional[Dict[str, List[str]]] = None) -> str:
    groups_str = ", ".join(sorted(groups.get(user, []))) if groups else ""
    top = ", ".join(procs[:50])  # limit prompt length
    return (
        f"User: {user}\n"
        f"AD Groups: {groups_str}\n"
        f"Processes: {top}\n"
        "Summarize the likely role, workflows, and risk-relevant behaviors based on these processes and groups."
    )

# ---------------------------- Clustering ----------------------------

def cluster_embeddings(X: np.ndarray, k_hint: Optional[int] = None) -> Tuple[np.ndarray, Dict[str, Any]]:
    info: Dict[str, Any] = {}
    labels: Optional[np.ndarray] = None
    if hdbscan is not None:
        clusterer = hdbscan.HDBSCAN(min_cluster_size=max(3, int(len(X) * 0.03)), min_samples=1)
        labels = clusterer.fit_predict(X)
        info["algo"] = "hdbscan"
        info["clusters"] = int(len(set(labels)) - (1 if -1 in labels else 0))
    if labels is None or (labels >= 0).sum() < max(2, int(0.5 * len(X))):
        # Fallback to KMeans with silhouette selection
        best_score = -1
        best = None
        k_range = range(2, min(12, len(X)))
        for k in k_range:
            km = KMeans(n_clusters=k, n_init="auto", random_state=42)
            y = km.fit_predict(X)
            score = silhouette_score(X, y) if len(set(y)) > 1 else -1
            if score > best_score:
                best_score, best = score, (y, {"algo": "kmeans", "clusters": k, "silhouette": float(score)})
        labels, info = best
    return labels, info

# ---------------------------- LLM prompts ----------------------------
SESSION_SYSTEM = (
    "You are a seasoned SOC analyst. Given per-user process inventories and optional group memberships, "
    "write a terse JSON object capturing: summary, likely_role, top_workflows (list), risk ('low'|'medium'|'high'). "
    "Be concise. Output ONLY JSON."
)

CLUSTER_SYSTEM = (
    "You are classifying users by behavior. Given several short user summaries from the same cluster, "
    "produce clustering metadata as JSON with keys: label (<=5 words), description (1-2 sentences), "
    "risk ('low'|'medium'|'high')."
)

# ---------------------------- Commands ----------------------------

def cmd_users(args: argparse.Namespace) -> None:
    enum = Enum4LinuxFile(Path(args.enum4linux)).load()
    print(json.dumps(enum.users, indent=2))


def cmd_groups(args: argparse.Namespace) -> None:
    enum = Enum4LinuxFile(Path(args.enum4linux)).load()
    print(json.dumps(enum.groups, indent=2))


def _load_inputs(tasklist_glob: str, enum_path: Optional[str]) -> Tuple[Dict[str, List[str]], Optional[Enum4LinuxFile]]:
    tl_paths = sorted([Path(p) for p in Path().glob(tasklist_glob)]) if any(ch in tasklist_glob for ch in "*?[]") else [Path(tasklist_glob)]
    tl_paths = [p for p in tl_paths if p.exists()]
    if not tl_paths:
        raise SystemExit(f"No tasklist files matched: {tasklist_glob}")
    up = merge_user_processes(tl_paths)
    enum = Enum4LinuxFile(Path(enum_path)).load() if enum_path else None
    return up, enum


def cmd_cluster(args: argparse.Namespace) -> None:
    llm = LLMClient()
    user_procs, enum = _load_inputs(args.tasklist_glob, args.enum4linux)

    # 1) Per-user LLM summaries → text blocks → embeddings
    users = sorted(user_procs.keys())
    summaries: Dict[str, Dict[str, Any]] = {}
    texts: List[str] = []

    for u in tqdm(users, desc="Summarizing users"):
        prompt = build_user_profile_text(u, user_procs[u], enum.groups if enum else None)
        raw = llm.chat(SESSION_SYSTEM, prompt)
        try:
            data = json.loads(raw)
        except Exception:
            data = {"summary": raw, "likely_role": "unknown", "top_workflows": [], "risk": "low"}
        summaries[u] = data
        # Combine fields for embedding
        text = f"User {u}: {data.get('summary','')}. Role: {data.get('likely_role','')}. Workflows: {', '.join(data.get('top_workflows', []))}. Risk: {data.get('risk','')}"
        texts.append(text)

    vecs = np.array(llm.embed(texts))

    # 2) Cluster
    labels, info = cluster_embeddings(vecs)

    # 3) LLM label per cluster
    clusters: Dict[int, Dict[str, Any]] = {}
    for cid in sorted(set(labels)):
        members = [u for u, y in zip(users, labels) if y == cid]
        # pick up to N sample summaries to keep prompt small
        examples = [summaries[u].get("summary", "") for u in members[: min(15, len(members))]]
        user_prompt = "\n".join(f"- {s}" for s in examples if s)
        raw = llm.chat(CLUSTER_SYSTEM, user_prompt or "- Typical office user")
        try:
            meta = json.loads(raw)
        except Exception:
            meta = {"label": f"Cluster {cid}", "description": raw[:200], "risk": "low"}
        clusters[int(cid)] = {
            "label": meta.get("label", f"Cluster {cid}"),
            "description": meta.get("description", ""),
            "risk": meta.get("risk", "low"),
            "members": members,
        }

    # 4) Build JSON output
    out = {
        "algo": info.get("algo"),
        "clusters": clusters,
        "users": {
            u: {
                "cluster": int(lab),
                "risk": summaries[u].get("risk", "low"),
                "summary": summaries[u].get("summary", ""),
                "likely_role": summaries[u].get("likely_role", ""),
                "top_workflows": summaries[u].get("top_workflows", []),
                "processes": user_procs[u],
                "groups": (enum.users.get(u, []) if enum else []),
            }
            for u, lab in zip(users, labels)
        },
    }

    if args.json_out:
        Path(args.json_out).write_text(json.dumps(out, indent=2), encoding="utf-8")
        print(f"Wrote {args.json_out}")

    # 5) Markdown report
    if args.md_out:
        lines: List[str] = []
        lines.append(f"# Netographer (LLM Edition)\n")
        lines.append(f"Algorithm: {out['algo']} — Clusters: {len(clusters)}\n")
        for cid, meta in clusters.items():
            lines.append(f"\n## Cluster {cid}: {meta['label']} ({meta['risk']})\n")
            lines.append(f"{meta['description']}\n")
            sample = meta["members"][: min(10, len(meta["members"]))]
            lines.append("**Sample users:** " + ", ".join(sample) + "\n")
            # top processes in cluster
            proc_counter = Counter()
            for u in meta["members"]:
                proc_counter.update(user_procs[u])
            top = ", ".join([f"{p}×{c}" for p, c in proc_counter.most_common(15)])
            lines.append("**Top processes:** " + top + "\n")
        Path(args.md_out).write_text("\n".join(lines), encoding="utf-8")
        print(f"Wrote {args.md_out}")

    # Also print a small summary to stdout
    print(json.dumps({"algo": out["algo"], "clusters": {cid: {"label": c["label"], "risk": c["risk"], "size": len(c["members"]) } for cid, c in clusters.items()}}, indent=2))

# ---------------------------- CLI wiring ----------------------------

def build_parser() -> argparse.ArgumentParser:
    p = argparse.ArgumentParser(description="LLM-assisted user activity grouping from tasklist + enum4linux.")
    sub = p.add_subparsers(dest="command", required=True)

    p_users = sub.add_parser("users", help="Dump users parsed from enum4linux output")
    p_users.add_argument("--enum4linux", required=True, help="Path to enum4linux output file")
    p_users.set_defaults(func=cmd_users)

    p_groups = sub.add_parser("groups", help="Dump groups parsed from enum4linux output")
    p_groups.add_argument("--enum4linux", required=True, help="Path to enum4linux output file")
    p_groups.set_defaults(func=cmd_groups)

    p_cluster = sub.add_parser("cluster", help="Cluster users with embeddings + label with LLM")
    p_cluster.add_argument("--tasklist-glob", required=True, help="Glob or path to tasklist text files")
    p_cluster.add_argument("--enum4linux", default=None, help="Optional path to enum4linux output")
    p_cluster.add_argument("--json-out", default=None, help="Write JSON results here")
    p_cluster.add_argument("--md-out", default=None, help="Write Markdown report here")
    p_cluster.set_defaults(func=cmd_cluster)

    return p


def main(argv: Optional[List[str]] = None) -> None:
    args = build_parser().parse_args(argv)
    args.func(args)


if __name__ == "__main__":
    main()