#!/usr/bin/env python3
"""
aggregate_privileged_accounts.py

Crawls a directory of privileged-account discovery CSVs (as produced by
Discover-PrivilegedAccounts.ps1 and discover_privileged_accounts.sh), merges
them, de-duplicates accounts across repeated uploads, and writes a single JSON
summary with the total count of privileged accounts plus breakdowns.

Standard library only (no pip dependencies). Python 3.6+.

Usage:
    python3 aggregate_privileged_accounts.py \
        --input-dir ./output \
        --output ./aggregate/privileged_accounts_summary.json

    # counts only, omit the per-account inventory:
    python3 aggregate_privileged_accounts.py --no-inventory
"""
import argparse
import csv
import glob
import json
import os
import sys
from collections import Counter
from datetime import datetime, timezone

# Columns every valid CSV must contain (order-independent).
EXPECTED_COLUMNS = [
    "CollectionTimestamp", "Hostname", "IPAddress", "OSPlatform", "OSVersion",
    "AccountName", "AccountDomain", "AccountType", "AccountID", "PrivilegeSource",
    "GroupMemberships", "Enabled", "IsBuiltIn", "PasswordLastSet",
    "PasswordNeverExpires", "PasswordRequired", "LastLogon", "AccountExpires",
]

# Leading characters the collectors prefix with a single quote to defuse
# spreadsheet formula injection. We reverse exactly that escaping on read.
_FORMULA_LEADERS = ("=", "+", "-", "@")


def unescape_cell(value):
    """Reverse the collectors' CSV-injection guard: a leading ' followed by a
    formula character was added by us, so strip that single quote back off.
    Legitimate values that merely start with ' are left untouched."""
    if value is None:
        return ""
    if len(value) >= 2 and value[0] == "'" and value[1] in _FORMULA_LEADERS:
        return value[1:]
    return value


def parse_timestamp(value):
    """Parse 'YYYY-MM-DD HH:MM:SS'; return a comparable datetime (min on fail)."""
    if not value:
        return datetime.min
    try:
        return datetime.strptime(value.strip(), "%Y-%m-%d %H:%M:%S")
    except (ValueError, AttributeError):
        return datetime.min


def to_bool(value):
    """Normalize the string booleans the collectors emit ('true'/'false')."""
    if value is None:
        return None
    v = value.strip().lower()
    if v in ("true", "1", "yes"):
        return True
    if v in ("false", "0", "no"):
        return False
    return None


def read_csv_file(path):
    """Yield dict rows from one CSV, cells unescaped. Raises ValueError if the
    file's header does not contain the expected columns."""
    # utf-8-sig transparently strips the BOM PowerShell's Export-Csv may add.
    with open(path, "r", encoding="utf-8-sig", newline="") as fh:
        # restkey collects any overflow cells from ragged rows so they don't
        # leak into the record under a None key; we ignore it below.
        reader = csv.DictReader(fh, restkey="_overflow", restval="")
        header = reader.fieldnames or []
        missing = [c for c in EXPECTED_COLUMNS if c not in header]
        if missing:
            raise ValueError("missing columns: %s" % ", ".join(missing))
        for row in reader:
            # Project strictly onto the known columns: drops overflow, and
            # coerces missing cells (short rows -> None) to empty strings.
            yield {col: unescape_cell(row.get(col) or "") for col in EXPECTED_COLUMNS}


def dedup_key(row, strict_identity):
    """Identity used to collapse duplicate uploads of the same account.

    Default: (host, account). With strict_identity, also include the account's
    domain and ID (SID on Windows / UID on Linux) so that same-named but
    distinct accounts (e.g. WEB01\\admin vs CORP\\admin) are counted apart."""
    host = (row.get("Hostname") or "").strip().lower()
    account = (row.get("AccountName") or "").strip().lower()
    if strict_identity:
        domain = (row.get("AccountDomain") or "").strip().lower()
        acct_id = (row.get("AccountID") or "").strip().lower()
        return (host, account, domain, acct_id)
    return (host, account)


def trusted_host_from_path(path, input_dir):
    """In a per-host chrooted layout (/srv/sftp/<host>/incoming/file.csv) the
    first path segment under --input-dir is the SFTP account name, which the
    server enforces and a host cannot forge. Return it, or None when the file
    sits directly in input_dir (nothing to verify against)."""
    rel = os.path.relpath(path, input_dir)
    parts = [p for p in rel.split(os.sep) if p not in ("", ".")]
    if len(parts) >= 2:            # <host>/.../file.csv
        return parts[0]
    return None                    # file directly under input_dir


def normalize_hostname(name):
    """Fold a hostname to its comparison key: lowercase short (pre-dot) form,
    so 'WEB01', 'web01', and 'web01.corp.example.com' all compare equal."""
    return (name or "").strip().lower().split(".")[0]


def hostnames_match(trusted, claimed):
    """Compare the trusted directory segment with the CSV's self-reported
    Hostname, case-insensitively and on the short (pre-dot) form, so a directory
    'web01' matches a claimed 'WEB01' or 'web01.corp.example.com'."""
    t = normalize_hostname(trusted)
    return bool(t) and t == normalize_hostname(claimed)


def summarize_records(records):
    """Count a de-duplicated record set into the standard breakdown block.
    Returns (block, by_host_counter) — the counter is reused for roster
    reconciliation. Shared by the whole-run summary and each monthly bucket."""
    by_platform = Counter()
    by_host = Counter()
    by_source = Counter()
    enabled_count = disabled_count = enabled_unknown = 0

    for rec in records:
        by_platform[(rec.get("OSPlatform") or "Unknown").strip()] += 1
        by_host[(rec.get("Hostname") or "Unknown").strip()] += 1
        # PrivilegeSource is a "; "-joined list; count each distinct reason.
        for src in (rec.get("PrivilegeSource") or "").split(";"):
            src = src.strip()
            if src:
                # Normalize "Group: sudo (primary)" and "Group: sudo" together.
                label = src.split(":")[0].strip() if ":" in src else src
                by_source[label] += 1
        state = to_bool(rec.get("Enabled"))
        if state is True:
            enabled_count += 1
        elif state is False:
            disabled_count += 1
        else:
            enabled_unknown += 1

    block = {
        "total_privileged_accounts": len(records),
        "unique_hosts": len(by_host),
        "counts": {
            "by_platform": dict(by_platform.most_common()),
            "by_privilege_source": dict(by_source.most_common()),
            "by_host": dict(by_host.most_common()),
            "enabled": enabled_count,
            "disabled": disabled_count,
            "enabled_unknown": enabled_unknown,
        },
    }
    return block, by_host


def reconcile_roster(roster, by_host):
    """Compare the expected-host roster against the hosts that reported.
    Comparison is on the normalized short form; original spellings are reported.
    `by_host` is a counter/iterable of reported hostnames."""
    roster_norm = {}
    for name in roster:
        n = normalize_hostname(name)
        if n:
            roster_norm.setdefault(n, name.strip())
    reporting_norm = {}
    for h in by_host:
        n = normalize_hostname(h)
        if n:
            reporting_norm.setdefault(n, h)

    missing = sorted(orig for n, orig in roster_norm.items()
                     if n not in reporting_norm)          # expected, no data
    unexpected = sorted(orig for n, orig in reporting_norm.items()
                        if n not in roster_norm)          # reported, off-roster
    present = sum(1 for n in roster_norm if n in reporting_norm)
    coverage = round(100.0 * present / len(roster_norm), 1) if roster_norm else 0.0

    return {
        "expected_hosts": len(roster_norm),
        "reporting_hosts": len(reporting_norm),
        "covered_hosts": present,
        "coverage_pct": coverage,
        "missing_from_reports": missing,      # in roster but did NOT report
        "unexpected_reporters": unexpected,   # reported but NOT in roster
    }


def _sorted_accounts(records):
    """Stable ordering so diffs between runs are meaningful."""
    return sorted(records, key=lambda r: ((r.get("Hostname") or "").lower(),
                                          (r.get("AccountName") or "").lower()))


def aggregate(input_dir, include_inventory=True, strict_identity=False,
              max_age_days=None, recursive=False,
              verify_hostname=False, drop_mismatched=False,
              month=None, roster=None, group_by_month=False):
    if recursive:
        # Walk per-host subdirectories, e.g. a chrooted SFTP layout where each
        # host uploads into /srv/sftp/<host>/incoming/.
        csv_paths = sorted(glob.glob(os.path.join(input_dir, "**", "*.csv"),
                                     recursive=True))
    else:
        csv_paths = sorted(glob.glob(os.path.join(input_dir, "*.csv")))

    files_processed = []
    files_skipped = []

    # Optional recency filter: drop records collected more than N days ago so
    # the total reflects only currently-reporting hosts. Collectors write local
    # time, so compare against a naive local 'now'. Records whose timestamp
    # cannot be parsed are treated as stale and dropped when this is enabled.
    cutoff = None
    if max_age_days is not None:
        from datetime import timedelta
        cutoff = datetime.now() - timedelta(days=max_age_days)

    # Optional month filter: keep only records whose CollectionTimestamp falls
    # in the given calendar month (YYYY-MM). Lets you produce point-in-time
    # summaries per month from a retained CSV archive.
    month_ym = None
    if month:
        month_ym = (int(month[:4]), int(month[5:7]))

    # When the same identity shows up in multiple files (e.g. daily uploads),
    # the record with the newest CollectionTimestamp wins. In group-by-month
    # mode we keep one such map per calendar month instead of a single one.
    latest = {}                  # key -> (timestamp, record)   [single mode]
    month_buckets = {}           # "YYYY-MM" -> {key: (ts, record)}  [group mode]
    total_rows = 0
    aged_out = 0
    month_excluded = 0
    undated_rows = 0             # rows with unparseable timestamp (group mode)

    # Hostname verification (only meaningful over a per-host directory tree).
    hostname_mismatches = []     # detail of rows whose Hostname != source dir
    mismatch_dropped = 0
    unverifiable_files = []      # files with no host segment to check against

    for path in csv_paths:
        trusted_host = trusted_host_from_path(path, input_dir) if verify_hostname else None
        if verify_hostname and trusted_host is None:
            unverifiable_files.append(os.path.basename(path))
        try:
            row_count = 0
            for row in read_csv_file(path):
                row_count += 1
                total_rows += 1
                host = (row.get("Hostname") or "").strip()
                account = (row.get("AccountName") or "").strip()
                if not host and not account:
                    continue

                # Cross-check the self-reported Hostname against the trustworthy
                # source-directory name. A mismatch means a host uploaded a row
                # claiming to be a different host (or a misconfigured account).
                if verify_hostname and trusted_host is not None and not hostnames_match(trusted_host, host):
                    hostname_mismatches.append({
                        "file": os.path.relpath(path, input_dir),
                        "source_directory": trusted_host,
                        "claimed_hostname": host,
                        "account": account,
                    })
                    if drop_mismatched:
                        mismatch_dropped += 1
                        continue

                ts = parse_timestamp(row.get("CollectionTimestamp"))
                if cutoff is not None and ts < cutoff:
                    aged_out += 1
                    continue
                if month_ym is not None and (ts == datetime.min
                                             or (ts.year, ts.month) != month_ym):
                    month_excluded += 1
                    continue
                key = dedup_key(row, strict_identity)
                if group_by_month:
                    if ts == datetime.min:
                        undated_rows += 1        # cannot bucket without a date
                        continue
                    mk = "%04d-%02d" % (ts.year, ts.month)
                    bucket = month_buckets.setdefault(mk, {})
                    if key not in bucket or ts >= bucket[key][0]:
                        bucket[key] = (ts, row)
                elif key not in latest or ts >= latest[key][0]:
                    latest[key] = (ts, row)
            files_processed.append({"file": os.path.basename(path), "rows": row_count})
        except (ValueError, OSError, UnicodeDecodeError) as exc:
            files_skipped.append({"file": os.path.basename(path), "reason": str(exc)})

    # --- Common run metadata -------------------------------------------------
    summary = {
        "generated_at": datetime.now(timezone.utc).strftime("%Y-%m-%dT%H:%M:%SZ"),
        "source_directory": os.path.abspath(input_dir),
        "files_processed": len(files_processed),
        "files_skipped": files_skipped,
        "rows_read": total_rows,
        "rows_aged_out": aged_out,
        "rows_excluded_by_month": month_excluded,
        "month_filter": month,
        "dedup_identity": "host+account+domain+id" if strict_identity else "host+account",
        "max_age_days": max_age_days,
    }

    if verify_hostname:
        summary["hostname_verification"] = {
            "enabled": True,
            "policy": "drop" if drop_mismatched else "report",
            "mismatch_count": len(hostname_mismatches),
            "rows_dropped": mismatch_dropped,
            "unverifiable_files": unverifiable_files,
            "mismatches": hostname_mismatches,
        }

    if group_by_month:
        # One de-duplicated summary per calendar month, plus a compact trend.
        summary["group_by_month"] = True
        summary["undated_rows"] = undated_rows
        months = {}
        trend = []
        for mk in sorted(month_buckets):
            recs = [rec for (_, rec) in month_buckets[mk].values()]
            block, by_host = summarize_records(recs)
            if roster is not None:
                block["roster_reconciliation"] = reconcile_roster(roster, by_host)
            if include_inventory:
                block["accounts"] = _sorted_accounts(recs)
            months[mk] = block
            point = {
                "month": mk,
                "total_privileged_accounts": block["total_privileged_accounts"],
                "unique_hosts": block["unique_hosts"],
            }
            if roster is not None:
                point["coverage_pct"] = block["roster_reconciliation"]["coverage_pct"]
            trend.append(point)
        summary["months"] = months
        summary["trend"] = trend
        summary["processed_file_detail"] = files_processed
        return summary

    # --- Single (whole-run) summary ------------------------------------------
    records = [rec for (_, rec) in latest.values()]
    block, by_host = summarize_records(records)
    summary.update(block)
    summary["processed_file_detail"] = files_processed

    if roster is not None:
        summary["roster_reconciliation"] = reconcile_roster(roster, by_host)
    if include_inventory:
        summary["accounts"] = _sorted_accounts(records)

    return summary


def main(argv=None):
    parser = argparse.ArgumentParser(
        description="Aggregate privileged-account discovery CSVs into one JSON summary.")
    parser.add_argument("--input-dir", default="./output",
                        help="Directory to crawl for *.csv (default: ./output)")
    parser.add_argument("--output", default="./aggregate/privileged_accounts_summary.json",
                        help="Output JSON path (default: ./aggregate/privileged_accounts_summary.json)")
    parser.add_argument("--no-inventory", action="store_true",
                        help="Emit counts only; omit the per-account list.")
    parser.add_argument("--strict-identity", action="store_true",
                        help="De-duplicate on host+account+domain+ID (SID/UID) so "
                             "same-named but distinct accounts count separately.")
    parser.add_argument("--max-age-days", type=int, default=None,
                        help="Ignore records collected more than N days ago, so the "
                             "total reflects only currently-reporting hosts.")
    parser.add_argument("--recursive", action="store_true",
                        help="Recurse into subdirectories (e.g. a per-host chrooted "
                             "SFTP layout: /srv/sftp/<host>/incoming/*.csv).")
    parser.add_argument("--verify-hostname", action="store_true",
                        help="Cross-check each row's self-reported Hostname against "
                             "the trustworthy source-directory name (the first path "
                             "segment under --input-dir). Reports mismatches. Use with "
                             "--recursive over a per-host layout.")
    parser.add_argument("--drop-mismatched", action="store_true",
                        help="With --verify-hostname, exclude mismatched rows from the "
                             "totals instead of only reporting them.")
    parser.add_argument("--roster", default=None, metavar="FILE",
                        help="A text file of expected hostnames (one per line, '#' "
                             "comments allowed). Reports which expected hosts did not "
                             "report (coverage gaps) and which reporters are off-roster.")
    parser.add_argument("--month", default=None, metavar="YYYY-MM",
                        help="Only count records whose CollectionTimestamp falls in "
                             "this calendar month, for a point-in-time monthly summary.")
    parser.add_argument("--group-by-month", action="store_true",
                        help="Emit a per-month breakdown and trend in one pass "
                             "(a 'months' map plus a 'trend' array). Cannot be "
                             "combined with --month.")
    args = parser.parse_args(argv)

    if not os.path.isdir(args.input_dir):
        print("Input directory not found: %s" % args.input_dir, file=sys.stderr)
        return 1

    if args.group_by_month and args.month is not None:
        print("--group-by-month and --month are mutually exclusive.", file=sys.stderr)
        return 1

    if args.max_age_days is not None and args.max_age_days < 0:
        print("--max-age-days must be >= 0", file=sys.stderr)
        return 1

    if args.month is not None:
        m = args.month
        well_formed = (len(m) == 7 and m[4] == "-"
                       and m[:4].isdigit() and m[5:].isdigit())
        try:
            if not well_formed:
                raise ValueError
            datetime.strptime(m, "%Y-%m")   # also range-checks month 01-12
        except ValueError:
            print("--month must be formatted YYYY-MM with a zero-padded month "
                  "(e.g. 2026-07)", file=sys.stderr)
            return 1

    roster = None
    if args.roster is not None:
        if not os.path.isfile(args.roster):
            print("Roster file not found: %s" % args.roster, file=sys.stderr)
            return 1
        with open(args.roster, "r", encoding="utf-8-sig") as fh:
            roster = [ln.strip() for ln in fh
                      if ln.strip() and not ln.lstrip().startswith("#")]

    # --drop-mismatched only makes sense alongside verification; imply it on.
    verify_hostname = args.verify_hostname or args.drop_mismatched
    if verify_hostname and not args.recursive:
        print("NOTE: --verify-hostname needs a per-host directory tree; add "
              "--recursive (files directly in --input-dir cannot be verified).",
              file=sys.stderr)

    summary = aggregate(args.input_dir,
                        include_inventory=not args.no_inventory,
                        strict_identity=args.strict_identity,
                        max_age_days=args.max_age_days,
                        recursive=args.recursive,
                        verify_hostname=verify_hostname,
                        drop_mismatched=args.drop_mismatched,
                        month=args.month,
                        roster=roster,
                        group_by_month=args.group_by_month)

    out_dir = os.path.dirname(os.path.abspath(args.output))
    os.makedirs(out_dir, exist_ok=True)
    # Restrict the aggregate (it is a whole-estate privileged inventory).
    fd = os.open(args.output, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
    with os.fdopen(fd, "w", encoding="utf-8") as fh:
        json.dump(summary, fh, indent=2, ensure_ascii=False)
        fh.write("\n")

    if summary.get("group_by_month"):
        print("Aggregated %d month(s) from %d file(s) -> %s"
              % (len(summary["months"]), summary["files_processed"], args.output))
        for point in summary["trend"]:
            cov = (" (%.1f%% coverage)" % point["coverage_pct"]
                   if "coverage_pct" in point else "")
            print("  %s: %d accounts across %d host(s)%s"
                  % (point["month"], point["total_privileged_accounts"],
                     point["unique_hosts"], cov))
    else:
        print("Aggregated %d unique privileged accounts from %d file(s) -> %s"
              % (summary["total_privileged_accounts"],
                 summary["files_processed"], args.output))

    if summary["files_skipped"]:
        print("WARNING: skipped %d file(s):" % len(summary["files_skipped"]),
              file=sys.stderr)
        for s in summary["files_skipped"]:
            print("  - %s: %s" % (s["file"], s["reason"]), file=sys.stderr)

    hv = summary.get("hostname_verification")
    if hv and hv["mismatch_count"]:
        verb = "dropped" if hv["policy"] == "drop" else "flagged"
        print("WARNING: hostname verification %s %d row(s):"
              % (verb, hv["mismatch_count"]), file=sys.stderr)
        for m in hv["mismatches"]:
            print("  - %s: dir '%s' vs claimed '%s' (account %s)"
                  % (m["file"], m["source_directory"], m["claimed_hostname"],
                     m["account"]), file=sys.stderr)

    rr = summary.get("roster_reconciliation")
    if rr:
        print("Roster coverage: %d/%d hosts (%.1f%%) reported."
              % (rr["covered_hosts"], rr["expected_hosts"], rr["coverage_pct"]))
        if rr["missing_from_reports"]:
            print("WARNING: %d expected host(s) did NOT report: %s"
                  % (len(rr["missing_from_reports"]),
                     ", ".join(rr["missing_from_reports"])), file=sys.stderr)
        if rr["unexpected_reporters"]:
            print("WARNING: %d off-roster host(s) reported: %s"
                  % (len(rr["unexpected_reporters"]),
                     ", ".join(rr["unexpected_reporters"])), file=sys.stderr)
    return 0


if __name__ == "__main__":
    sys.exit(main())