#!/usr/bin/env python3
"""
Compare all users across the Orchard and Delphi Snowflake accounts.

Produces a report of:
  - Users present in both accounts with differing attributes
  - Users only in Delphi

Uses real-time SHOW USERS, INFORMATION_SCHEMA.POLICY_REFERENCES and DESCRIBE
USER rather than SNOWFLAKE.ACCOUNT_USAGE.{USERS, POLICY_REFERENCES}, which has
a latency of up to ~2 hours.

Usage:
  pip install snowflake-connector-python
  python3 compare-snowflake-users.py [--csv output.csv]

Authentication:
  Uses key-pair auth if SNOWFLAKE_PRIVATE_KEY_PATH is set, otherwise falls back
  to externalbrowser SSO. Per-account overrides via ORCHARD_* / DELPHI_* prefixes.

Required env vars:
  SNOWFLAKE_USER   - your Snowflake username

Optional env vars:
  SNOWFLAKE_PRIVATE_KEY_PATH         - path to PEM private key file
  SNOWFLAKE_PRIVATE_KEY_PASSPHRASE   - passphrase if key is encrypted
  SNOWFLAKE_ROLE                     - role to use (default: SECURITYADMIN)
  ORCHARD_SNOWFLAKE_USER / DELPHI_SNOWFLAKE_USER   - per-account user override
  ORCHARD_SNOWFLAKE_ROLE / DELPHI_SNOWFLAKE_ROLE   - per-account role override
  ORCHARD_ACCOUNT / DELPHI_ACCOUNT                 - account identifier overrides
"""

import argparse
import csv
import os
import sys
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any

try:
    import snowflake.connector
except ImportError:
    sys.exit(
        "snowflake-connector-python is required.\n"
        "Install it with: pip install snowflake-connector-python"
    )

# ---------------------------------------------------------------------------
# Attributes compared between accounts.
# All attrs except RSA fingerprints come from SHOW USERS (real-time).
# NETWORK_POLICY / AUTHENTICATION_POLICY are pulled per-user from
# INFORMATION_SCHEMA.POLICY_REFERENCES (real-time table function).
# RSA key fingerprints are pulled per-user from DESCRIBE USER.
# ---------------------------------------------------------------------------
COMPARE_ATTRS = [
    "LOGIN_NAME",
    "DISPLAY_NAME",
    "FIRST_NAME",
    "LAST_NAME",
    "EMAIL",
    "TYPE",
    "DISABLED",
    # "DEFAULT_WAREHOUSE",
    # "DEFAULT_NAMESPACE",
    # "DEFAULT_ROLE",
    "DEFAULT_SECONDARY_ROLE",
    "COMMENT",
    # "OWNER",
    # "HAS_PASSWORD",
    # "MUST_CHANGE_PASSWORD",
    # "HAS_MFA",
    "HAS_PAT",
    "HAS_WORKLOAD_IDENTITY",
    # "EXT_AUTHN_DUO",
    # "EXT_AUTHN_UID",
    "SNOWFLAKE_LOCK",
    "IS_FROM_ORGANIZATION_USER",
    # Fetched via parallel POLICY_REFERENCES calls
    "NETWORK_POLICY",
    "AUTHENTICATION_POLICY",
    # Fetched via parallel DESCRIBE USER calls
    "RSA_PUBLIC_KEY_FP",
    "RSA_PUBLIC_KEY_2_FP",
]

# SHOW USERS column → COMPARE_ATTRS key. SHOW USERS uses lowercase; the
# `default_secondary_roles` column is plural in SHOW USERS but exposed as
# DEFAULT_SECONDARY_ROLE (singular) in ACCOUNT_USAGE — keep the singular form
# in the report for back-compat.
_SHOW_USERS_MAP = {
    "login_name": "LOGIN_NAME",
    "display_name": "DISPLAY_NAME",
    "first_name": "FIRST_NAME",
    "last_name": "LAST_NAME",
    "email": "EMAIL",
    "type": "TYPE",
    "disabled": "DISABLED",
    "default_secondary_roles": "DEFAULT_SECONDARY_ROLE",
    "comment": "COMMENT",
    "has_pat": "HAS_PAT",
    "has_workload_identity": "HAS_WORKLOAD_IDENTITY",
    "snowflake_lock": "SNOWFLAKE_LOCK",
    "is_from_organization_user": "IS_FROM_ORGANIZATION_USER",
}


def _private_key_bytes() -> bytes | None:
    """Load a private key from disk if configured."""
    path = os.environ.get("SNOWFLAKE_PRIVATE_KEY_PATH")
    if not path:
        return None
    from cryptography.hazmat.backends import default_backend
    from cryptography.hazmat.primitives import serialization

    passphrase_str = os.environ.get("SNOWFLAKE_PRIVATE_KEY_PASSPHRASE")
    passphrase = passphrase_str.encode() if passphrase_str else None
    with open(path, "rb") as fh:
        private_key = serialization.load_pem_private_key(
            fh.read(), password=passphrase, backend=default_backend()
        )
    return private_key.private_bytes(
        encoding=serialization.Encoding.DER,
        format=serialization.PrivateFormat.PKCS8,
        encryption_algorithm=serialization.NoEncryption(),
    )


def _conn_params(account_name: str, env_prefix: str, default_warehouse: str) -> dict[str, Any]:
    """Build snowflake.connector connection params for one account."""

    def env(key: str) -> str | None:
        return os.environ.get(f"{env_prefix}{key}") or os.environ.get(key)

    params: dict[str, Any] = {
        "account": account_name,
        "user": env("SNOWFLAKE_USER"),
        "role": env("SNOWFLAKE_ROLE") or "SECURITYADMIN",
        "warehouse": env("SNOWFLAKE_WAREHOUSE") or default_warehouse,
    }

    pk_bytes = _private_key_bytes()
    if pk_bytes:
        params["private_key"] = pk_bytes
        params["authenticator"] = "snowflake_jwt"
    else:
        params["authenticator"] = "externalbrowser"

    return params


def _ci(row: dict, key: str) -> Any:
    """Case-insensitive lookup — SHOW commands return lowercase keys, queries return upper."""
    return row.get(key) if key in row else row.get(key.lower())


def _describe_user_fps(
    conn: "snowflake.connector.SnowflakeConnection", name: str
) -> tuple[str | None, str | None]:
    """Return (rsa_public_key_fp, rsa_public_key_2_fp) for a single user."""
    cur = conn.cursor(snowflake.connector.DictCursor)
    cur.execute(f'DESCRIBE USER "{name}"')
    rows = {r["property"]: r["value"] for r in cur.fetchall()}
    return rows.get("RSA_PUBLIC_KEY_FP"), rows.get("RSA_PUBLIC_KEY_2_FP")


def _user_policy_refs(
    conn: "snowflake.connector.SnowflakeConnection", name: str
) -> tuple[str | None, str | None]:
    """Return (network_policy, authentication_policy) for a single user."""
    cur = conn.cursor(snowflake.connector.DictCursor)
    # Use the SNOWFLAKE shared db's INFORMATION_SCHEMA — it's available to any
    # role that can read ACCOUNT_USAGE (which we already require).
    cur.execute(
        "SELECT POLICY_KIND, POLICY_NAME "
        "FROM TABLE(SNOWFLAKE.INFORMATION_SCHEMA.POLICY_REFERENCES("
        "REF_ENTITY_NAME => %(name)s, REF_ENTITY_DOMAIN => 'USER'))",
        {"name": name},
    )
    network_policy: str | None = None
    auth_policy: str | None = None
    for row in cur.fetchall():
        kind = _ci(row, "POLICY_KIND")
        policy_name = _ci(row, "POLICY_NAME")
        if kind == "NETWORK_POLICY":
            network_policy = policy_name
        elif kind == "AUTHENTICATION_POLICY":
            auth_policy = policy_name
    return network_policy, auth_policy


def _enrich_users(
    conn_params: dict[str, Any],
    users: dict[str, dict],
    names_with_key: set[str],
    label: str,
) -> None:
    """
    Per-user enrichment: POLICY_REFERENCES for every user, plus DESCRIBE USER
    for users that have an RSA key set. Mutates `users` in place.
    """
    names = sorted(users)
    if not names:
        return

    print(
        f"  Fetching policy refs (+ RSA fingerprints for {len(names_with_key):,}) "
        f"for {len(names):,} users in {label}…",
        flush=True,
    )

    def _fetch_one(name: str) -> tuple[str, str | None, str | None, str | None, str | None]:
        # Each thread needs its own connection — cursors aren't thread-safe.
        c = snowflake.connector.connect(**conn_params)
        try:
            np, ap = _user_policy_refs(c, name)
            if name in names_with_key:
                fp1, fp2 = _describe_user_fps(c, name)
            else:
                fp1, fp2 = None, None
        finally:
            c.close()
        return name, np, ap, fp1, fp2

    with ThreadPoolExecutor(max_workers=20) as pool:
        futures = {pool.submit(_fetch_one, n): n for n in names}
        done = 0
        for future in as_completed(futures):
            name, np, ap, fp1, fp2 = future.result()
            users[name]["NETWORK_POLICY"] = np
            users[name]["AUTHENTICATION_POLICY"] = ap
            users[name]["RSA_PUBLIC_KEY_FP"] = fp1
            users[name]["RSA_PUBLIC_KEY_2_FP"] = fp2
            done += 1
            if done % 100 == 0:
                print(f"    …{done}/{len(names)}", flush=True)


def fetch_users(account_name: str, env_prefix: str, label: str, default_warehouse: str) -> dict[str, dict]:
    """
    Connect to `account_name` and return {username: {attr: value}} dict.

    Uses real-time SHOW USERS for the bulk of attributes, plus per-user
    POLICY_REFERENCES + DESCRIBE USER calls. The previous version queried
    SNOWFLAKE.ACCOUNT_USAGE.{USERS, POLICY_REFERENCES}, which lag by ~2 hours.
    """
    print(f"Connecting to {label} ({account_name})…", flush=True)
    params = _conn_params(account_name, env_prefix, default_warehouse)
    conn = snowflake.connector.connect(**params)
    try:
        cur = conn.cursor(snowflake.connector.DictCursor)
        cur.execute("SHOW USERS")
        rows = cur.fetchall()
    finally:
        conn.close()

    users: dict[str, dict] = {}
    names_with_key: set[str] = set()
    for row in rows:
        name = _ci(row, "NAME")
        record: dict[str, Any] = {attr: None for attr in COMPARE_ATTRS}
        for show_col, attr in _SHOW_USERS_MAP.items():
            record[attr] = row.get(show_col)
        users[name] = record
        # SHOW USERS returns has_rsa_public_key as the string "true"/"false".
        if str(row.get("has_rsa_public_key", "")).lower() == "true":
            names_with_key.add(name)

    print(
        f"  → {len(users):,} users found in {label} "
        f"({len(names_with_key):,} with RSA keys)",
        flush=True,
    )

    _enrich_users(params, users, names_with_key, label)

    return users


def compare(
        orchard: dict[str, dict], delphi: dict[str, dict]
) -> tuple[list, list]:
    """
    Returns:
        common_diffs  - list of (name, list_of_(attr, orchard_val, delphi_val))
        delphi_only   - list of names
    """
    orchard_names = set(orchard)
    delphi_names = set(delphi)

    delphi_only = sorted(delphi_names - orchard_names)

    common_diffs = []
    for name in sorted(orchard_names & delphi_names):
        diffs = []
        for attr in COMPARE_ATTRS:
            o_val = orchard[name].get(attr)
            d_val = delphi[name].get(attr)

            if attr == "DEFAULT_SECONDARY_ROLE":
                # ["ALL"] and [ALL] are equivalent representations (quotes optional)
                o_norm = str(o_val).replace('"', '').upper()
                d_norm = str(d_val).replace('"', '').upper()
                if o_norm == d_norm:
                    continue

            if attr == "TYPE":
                # PERSON and blank/None are equivalent (some accounts omit the default)
                o_type = (str(o_val).upper() if o_val else "")
                d_type = (str(d_val).upper() if d_val else "")
                if {o_type, d_type} <= {"PERSON", "NONE", ""}:
                    continue

            if attr in ("RSA_PUBLIC_KEY_FP", "RSA_PUBLIC_KEY_2_FP"):
                # DESCRIBE USER returns the literal string "null" when no key is
                # set; users we never DESCRIBE'd carry a Python None. Treat both
                # forms (and the empty string) as equivalent.
                if _is_blank_fp(o_val) and _is_blank_fp(d_val):
                    continue

            if str(o_val) != str(d_val):
                diffs.append((attr, o_val, d_val))
        if diffs:
            common_diffs.append((name, diffs))

    return common_diffs, delphi_only


def print_report(common_diffs, delphi_only):
    sep = "=" * 72

    print(f"\n{sep}")
    print("SNOWFLAKE USER COMPARISON: ORCHARD vs DELPHI")
    print(sep)

    print(f"\nUSERS WITH ATTRIBUTE DIFFERENCES ({len(common_diffs)} user(s)):\n")
    for name, diffs in common_diffs:
        print(f"  {name}")
        print(f"    {'ATTRIBUTE':<40} {'ORCHARD':<35} {'DELPHI'}")
        print(f"    {'-'*40} {'-'*35} {'-'*35}")
        for attr, o_val, d_val in diffs:
            o_str = _truncate(str(o_val), 34)
            d_str = _truncate(str(d_val), 34)
            print(f"    {attr:<40} {o_str:<35} {d_str}")
        print()

    print(f"USERS ONLY IN DELPHI ({len(delphi_only)} user(s)):")
    if delphi_only:
        for chunk in _chunks(delphi_only, 6):
            print("  " + "  ".join(chunk))
    else:
        print("  (none)")

    print(f"\n{sep}\n")


def write_csv(path: str, common_diffs, delphi_only):
    with open(path, "w", newline="") as fh:
        writer = csv.writer(fh)
        writer.writerow(
            ["username", "category", "attribute", "orchard_value", "delphi_value"]
        )
        for name, diffs in common_diffs:
            for attr, o_val, d_val in diffs:
                writer.writerow([name, "DIFF", attr, o_val, d_val])
        for name in delphi_only:
            writer.writerow([name, "DELPHI_ONLY", "", "", ""])
    print(f"CSV written to: {path}")


def _is_blank_fp(value: Any) -> bool:
    """True for the values DESCRIBE USER and SHOW USERS use to mean 'no key set'."""
    return value is None or str(value).strip().lower() in {"", "null"}


def _truncate(s: str, maxlen: int) -> str:
    return s if len(s) <= maxlen else s[: maxlen - 1] + "…"


def _chunks(lst, n):
    for i in range(0, len(lst), n):
        yield lst[i : i + n]


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--csv", metavar="FILE", help="also write results to CSV")
    parser.add_argument(
        "--orchard-account",
        default=os.environ.get("ORCHARD_ACCOUNT", "sme-orchard"),
        help="Orchard account identifier (default: sme-orchard)",
    )
    parser.add_argument(
        "--delphi-account",
        default=os.environ.get("DELPHI_ACCOUNT", "sme-delphi"),
        help="Delphi account identifier (default: sme-delphi)",
    )
    args = parser.parse_args()

    orchard_users = fetch_users(args.orchard_account, "ORCHARD_", "Orchard", "DEV_OWS_WAREHOUSE")
    delphi_users = fetch_users(args.delphi_account, "DELPHI_", "Delphi", "DEV_OWS_WH")

    common_diffs, delphi_only = compare(orchard_users, delphi_users)
    print_report(common_diffs, delphi_only)

    if args.csv:
        write_csv(args.csv, common_diffs, delphi_only)


if __name__ == "__main__":
    main()