#!/usr/bin/env bash
# Acquire a client-credentials bearer token for a ComputerGuru MSP app tier.
# Usage: get-token.sh <tenant-id-or-domain> <tier>
#
# Tiers and their app + resource scope:
#   investigator       ComputerGuru Security Investigator -> Graph API  (read-only breach checks)
#   investigator-exo   ComputerGuru Security Investigator -> Exchange Online  (EXO read: Get-InboxRule, Get-Mailbox)
#   exchange-op        ComputerGuru Exchange Operator     -> Exchange Online  (EXO write: Set-Mailbox, Remove-InboxRule, revoke sessions)
#   user-manager       ComputerGuru User Manager          -> Graph API  (user create/update/disable, license assign, MFA reset)
#   tenant-admin       ComputerGuru Tenant Admin          -> Graph API  (app roles, CA policy, directory write — high privilege)
#   defender           ComputerGuru Defender Add-on       -> Defender ATP API  (MDE-licensed tenants only)
#
# Authentication: certificate-based (client_assertion JWT, RS256) is preferred
# when cert_thumbprint_b64url + cert_private_key_pem_b64 are present in the vault
# entry's credentials block. Falls back to client_secret otherwise.
#
# Override via env: REMEDIATION_AUTH=cert | secret  (default: auto, prefer cert)
#   - REMEDIATION_AUTH=cert   forces cert; errors if cert fields missing (no fallback).
#   - REMEDIATION_AUTH=secret forces client_secret; errors if secret missing.
#   - unset / any other value: auto-detect (cert preferred, secret fallback).
#
# Cert fields read from the vault entry (under credentials:):
#   cert_thumbprint_b64url   x5t header value (base64url SHA1, no padding)
#   cert_private_key_pem_b64 base64-encoded RSA private key PEM
#
# Output (stdout): bearer token. Exits 0 on success.
# Cache: /tmp/remediation-tool/{tenant-id}/{tier}.jwt  TTL 55 minutes.
set -euo pipefail

TARGET="${1:?usage: get-token.sh <tenant-id|domain> <tier>}"
TIER="${2:?usage: get-token.sh <tenant-id|domain> <tier>}"

# Resolve domain to tenant GUID if needed
if [[ "$TARGET" =~ ^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$ ]]; then
  TENANT_ID="$TARGET"
else
  SCRIPT_DIR="$(dirname "${BASH_SOURCE[0]}")"
  TENANT_ID=$("$SCRIPT_DIR/resolve-tenant.sh" "$TARGET")
fi

# Map tier -> client_id, vault SOPS path, resource scope
case "$TIER" in
  investigator)
    CLIENT_ID="bfbc12a4-f0dd-4e12-b06d-997e7271e10c"
    VAULT_PATH="msp-tools/computerguru-security-investigator.sops.yaml"
    SCOPE_URL="https://graph.microsoft.com/.default"
    ;;
  investigator-exo)
    CLIENT_ID="bfbc12a4-f0dd-4e12-b06d-997e7271e10c"
    VAULT_PATH="msp-tools/computerguru-security-investigator.sops.yaml"
    SCOPE_URL="https://outlook.office365.com/.default"
    ;;
  exchange-op)
    CLIENT_ID="b43e7342-5b4b-492f-890f-bb5a4f7f40e9"
    VAULT_PATH="msp-tools/computerguru-exchange-operator.sops.yaml"
    SCOPE_URL="https://outlook.office365.com/.default"
    ;;
  user-manager)
    CLIENT_ID="64fac46b-8b44-41ad-93ee-7da03927576c"
    VAULT_PATH="msp-tools/computerguru-user-manager.sops.yaml"
    SCOPE_URL="https://graph.microsoft.com/.default"
    ;;
  tenant-admin)
    CLIENT_ID="709e6eed-0711-4875-9c44-2d3518c47063"
    VAULT_PATH="msp-tools/computerguru-tenant-admin.sops.yaml"
    SCOPE_URL="https://graph.microsoft.com/.default"
    ;;
  tenant-admin-onboard)
    # Same app as tenant-admin; this alias signals the token is being used
    # during initial tenant onboarding (clearer error messages on failure).
    CLIENT_ID="709e6eed-0711-4875-9c44-2d3518c47063"
    VAULT_PATH="msp-tools/computerguru-tenant-admin.sops.yaml"
    SCOPE_URL="https://graph.microsoft.com/.default"
    ;;
  defender)
    CLIENT_ID="dbf8ad1a-54f4-4bb8-8a9e-ea5b9634635b"
    VAULT_PATH="msp-tools/computerguru-defender-addon.sops.yaml"
    SCOPE_URL="https://api.securitycenter.microsoft.com/.default"
    ;;
  intune-manager)
    CLIENT_ID="46986910-aa47-4e5e-b596-f65c6b485abb"
    VAULT_PATH="msp-tools/computerguru-intune-manager.sops.yaml"
    SCOPE_URL="https://graph.microsoft.com/.default"
    ;;
  *)
    echo "ERROR: unknown tier '$TIER'." >&2
    echo "Valid tiers: investigator | investigator-exo | exchange-op | user-manager | tenant-admin | tenant-admin-onboard | defender | intune-manager" >&2
    exit 2
    ;;
esac

CACHE_DIR="/tmp/remediation-tool/$TENANT_ID"
mkdir -p "$CACHE_DIR"
CACHE_FILE="$CACHE_DIR/${TIER}.jwt"

# Return cached token if < 55 minutes old
if [[ -f "$CACHE_FILE" ]] && [[ $(find "$CACHE_FILE" -mmin -55 2>/dev/null) ]]; then
  cat "$CACHE_FILE"
  exit 0
fi

# Locate vault repo via .claude/identity.json (per-machine, gitignored).
# Falls back to VAULT_PATH env var if set.
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
CLAUDETOOLS_ROOT="$(cd "$SCRIPT_DIR/../../../.." && pwd)"
IDENTITY_FILE="$CLAUDETOOLS_ROOT/.claude/identity.json"

VAULT_ROOT="${VAULT_ROOT_ENV:-}"
if [[ -z "$VAULT_ROOT" && -f "$IDENTITY_FILE" ]]; then
  if command -v jq >/dev/null 2>&1; then
    VAULT_ROOT=$(jq -r '.vault_path // empty' "$IDENTITY_FILE" 2>/dev/null)
  fi
  if [[ -z "$VAULT_ROOT" ]]; then
    IDENTITY_FILE_WIN=$(cygpath -w "$IDENTITY_FILE" 2>/dev/null || echo "$IDENTITY_FILE")
    for py in py python3 python; do
      if command -v "$py" >/dev/null 2>&1; then
        VAULT_ROOT=$("$py" -c "import json; print(json.load(open(r'${IDENTITY_FILE_WIN}')).get('vault_path',''))" 2>/dev/null) && break
      fi
    done
  fi
fi
[[ -z "$VAULT_ROOT" ]] && { echo "ERROR: vault_path not set in $IDENTITY_FILE and VAULT_ROOT_ENV env var not set" >&2; exit 3; }
[[ ! -d "$VAULT_ROOT" ]] && { echo "ERROR: vault not found at $VAULT_ROOT (check vault_path in $IDENTITY_FILE)" >&2; exit 3; }

SOPS_FILE="$VAULT_ROOT/$VAULT_PATH"
[[ ! -f "$SOPS_FILE" ]] && { echo "ERROR: vault file not found: $SOPS_FILE" >&2; exit 3; }

# Pick a Python interpreter once for all helpers.
PYTHON_BIN=""
for p in py python python3; do command -v "$p" >/dev/null 2>&1 && PYTHON_BIN="$p" && break; done

# Decrypt the SOPS file once and cache plaintext in a memory-mapped tmp file.
# All field reads below pull from this single decrypt; saves multiple sops invocations
# (each is ~1s on Windows due to age key load) and avoids piping ~2KB key material
# through subshells repeatedly.
PLAINTEXT_FILE="$(mktemp -t sops-plain.XXXXXX)"
trap 'rm -f "$PLAINTEXT_FILE"' EXIT

# Try vault.sh first (handles yq vs python-yaml selection internally).
# If that fails (e.g. PyYAML missing on this host), do a raw sops decrypt.
SOPS_OK=0
if [[ -f "$VAULT_ROOT/scripts/vault.sh" ]]; then
  if bash "$VAULT_ROOT/scripts/vault.sh" get "$VAULT_PATH" > "$PLAINTEXT_FILE" 2>/dev/null && [[ -s "$PLAINTEXT_FILE" ]]; then
    SOPS_OK=1
  fi
fi
if [[ $SOPS_OK -eq 0 ]]; then
  command -v sops >/dev/null 2>&1 || { echo "ERROR: vault.sh failed and sops not on PATH (needed for fallback decrypt)" >&2; exit 3; }
  if ! sops -d "$SOPS_FILE" > "$PLAINTEXT_FILE" 2>/dev/null; then
    echo "ERROR: sops decrypt failed for $SOPS_FILE" >&2
    exit 3
  fi
fi

# read_field <field-name> -> stdout (empty if missing)
# Tries PyYAML for correctness; falls back to a tolerant flat-key regex that handles
# the actual vault file shape (single-line scalars under a 4-space-indented credentials:
# block — including 2KB+ values like cert_private_key_pem_b64).
read_field() {
  local field="$1"
  local val=""
  if [[ -n "$PYTHON_BIN" ]]; then
    val=$("$PYTHON_BIN" - "$PLAINTEXT_FILE" "$field" <<'PY' 2>/dev/null || true
import sys
path, field = sys.argv[1], sys.argv[2]
try:
    import yaml
except ImportError:
    sys.exit(1)
try:
    with open(path, "r", encoding="utf-8") as f:
        doc = yaml.safe_load(f) or {}
except Exception:
    sys.exit(1)
creds = (doc.get("credentials") or {})
v = creds.get(field)
if v is None:
    sys.exit(0)
print(v)
PY
)
    val="${val%$'\r'}"
    if [[ -n "$val" ]]; then
      printf '%s' "$val"
      return 0
    fi
  fi
  # Regex fallback: flat keys under credentials:, indent 2 or 4 spaces, value on same line.
  # Strips surrounding quotes if present. Stops at first match. Works for the actual
  # vault format because cert_private_key_pem_b64 is emitted as a single long line.
  if [[ -n "$PYTHON_BIN" ]]; then
    val=$("$PYTHON_BIN" - "$PLAINTEXT_FILE" "$field" <<'PY' 2>/dev/null || true
import sys, re
path, field = sys.argv[1], sys.argv[2]
try:
    with open(path, "r", encoding="utf-8") as f:
        text = f.read()
except Exception:
    sys.exit(1)
m = re.search(r'(?ms)^credentials:\s*\n((?:[ \t]+.*\n?)+)', text)
if not m:
    sys.exit(0)
block = m.group(1)
pat = re.compile(r'^[ \t]+' + re.escape(field) + r'\s*:\s*(.*?)\s*$', re.MULTILINE)
fm = pat.search(block)
if not fm:
    sys.exit(0)
v = fm.group(1)
if (v.startswith('"') and v.endswith('"')) or (v.startswith("'") and v.endswith("'")):
    v = v[1:-1]
print(v)
PY
)
    val="${val%$'\r'}"
    printf '%s' "$val"
  fi
}

# Auth-method selection: env override > auto-detect.
AUTH_OVERRIDE="${REMEDIATION_AUTH:-}"
AUTH_METHOD=""
CERT_X5T=""
CERT_KEY_B64=""
CLIENT_SECRET=""

case "$AUTH_OVERRIDE" in
  cert)
    CERT_X5T=$(read_field cert_thumbprint_b64url | tr -d '\r\n')
    CERT_KEY_B64=$(read_field cert_private_key_pem_b64 | tr -d '\r\n')
    if [[ -z "$CERT_X5T" || -z "$CERT_KEY_B64" ]]; then
      echo "ERROR: REMEDIATION_AUTH=cert but cert fields missing in vault ($VAULT_PATH)" >&2
      echo "  Required fields under credentials: cert_thumbprint_b64url, cert_private_key_pem_b64" >&2
      exit 4
    fi
    AUTH_METHOD="cert"
    ;;
  secret)
    CLIENT_SECRET=$(read_field client_secret | tr -d '\r\n')
    if [[ -z "$CLIENT_SECRET" ]]; then
      CLIENT_SECRET=$(read_field credential | tr -d '\r\n')
    fi
    if [[ -z "$CLIENT_SECRET" ]]; then
      echo "ERROR: REMEDIATION_AUTH=secret but client_secret missing in vault ($VAULT_PATH)" >&2
      echo "  Check field: credentials.client_secret  (or credentials.credential for older entries)" >&2
      exit 4
    fi
    AUTH_METHOD="secret"
    ;;
  *)
    CERT_X5T=$(read_field cert_thumbprint_b64url | tr -d '\r\n')
    CERT_KEY_B64=$(read_field cert_private_key_pem_b64 | tr -d '\r\n')
    if [[ -n "$CERT_X5T" && -n "$CERT_KEY_B64" ]]; then
      AUTH_METHOD="cert"
    else
      CLIENT_SECRET=$(read_field client_secret | tr -d '\r\n')
      if [[ -z "$CLIENT_SECRET" ]]; then
        CLIENT_SECRET=$(read_field credential | tr -d '\r\n')
      fi
      if [[ -z "$CLIENT_SECRET" ]]; then
        echo "ERROR: no usable credential found in $VAULT_PATH" >&2
        echo "  Need either credentials.cert_thumbprint_b64url + credentials.cert_private_key_pem_b64," >&2
        echo "  or credentials.client_secret (or legacy credentials.credential)." >&2
        exit 4
      fi
      AUTH_METHOD="secret"
    fi
    ;;
esac

echo "[INFO] auth=$AUTH_METHOD" >&2

# Build request and POST.
TOKEN_URL="https://login.microsoftonline.com/${TENANT_ID}/oauth2/v2.0/token"

if [[ "$AUTH_METHOD" == "cert" ]]; then
  [[ -z "$PYTHON_BIN" ]] && { echo "ERROR: cert auth requires python (py/python/python3) for JWT signing" >&2; exit 4; }

  CLIENT_ASSERTION=$(CERT_X5T_ENV="$CERT_X5T" \
                     CERT_KEY_B64_ENV="$CERT_KEY_B64" \
                     CLIENT_ID_ENV="$CLIENT_ID" \
                     TOKEN_URL_ENV="$TOKEN_URL" \
                     "$PYTHON_BIN" - <<'PY' 2>&1
import os, sys, time, uuid, base64
try:
    import jwt
except ImportError:
    sys.stderr.write("ERROR: PyJWT not installed (pip install PyJWT cryptography)\n")
    sys.exit(2)
try:
    from cryptography.hazmat.primitives.serialization import load_pem_private_key
except ImportError:
    sys.stderr.write("ERROR: cryptography not installed (pip install cryptography)\n")
    sys.exit(2)

x5t = os.environ["CERT_X5T_ENV"]
key_b64 = os.environ["CERT_KEY_B64_ENV"]
client_id = os.environ["CLIENT_ID_ENV"]
aud = os.environ["TOKEN_URL_ENV"]

try:
    key_pem = base64.b64decode(key_b64)
except Exception as e:
    sys.stderr.write(f"ERROR: cert_private_key_pem_b64 is not valid base64: {e}\n")
    sys.exit(3)

# Validate the PEM parses before handing to PyJWT, for a clearer error.
try:
    load_pem_private_key(key_pem, password=None)
except Exception as e:
    sys.stderr.write(f"ERROR: cert_private_key_pem_b64 did not decode to a valid PEM private key: {e}\n")
    sys.exit(3)

now = int(time.time())
payload = {
    "aud": aud,
    "iss": client_id,
    "sub": client_id,
    "jti": str(uuid.uuid4()),
    "exp": now + 300,
    "nbf": now,
}
token = jwt.encode(payload, key_pem, algorithm="RS256", headers={"x5t": x5t})
sys.stdout.write(token)
PY
)
  ASSERT_RC=$?
  if [[ $ASSERT_RC -ne 0 || -z "$CLIENT_ASSERTION" ]]; then
    echo "ERROR: failed to build client_assertion JWT" >&2
    [[ -n "$CLIENT_ASSERTION" ]] && echo "$CLIENT_ASSERTION" >&2
    exit 4
  fi

  RESP=$(curl -s --max-time 15 -X POST "$TOKEN_URL" \
    --data-urlencode "client_id=${CLIENT_ID}" \
    --data-urlencode "client_assertion_type=urn:ietf:params:oauth:client-assertion-type:jwt-bearer" \
    --data-urlencode "client_assertion=${CLIENT_ASSERTION}" \
    --data-urlencode "scope=${SCOPE_URL}" \
    --data-urlencode "grant_type=client_credentials")
else
  RESP=$(curl -s --max-time 15 -X POST "$TOKEN_URL" \
    --data-urlencode "client_id=${CLIENT_ID}" \
    --data-urlencode "client_secret=${CLIENT_SECRET}" \
    --data-urlencode "scope=${SCOPE_URL}" \
    --data-urlencode "grant_type=client_credentials")
fi

TOKEN=$(echo "$RESP" | jq -r '.access_token // empty')

if [[ -z "$TOKEN" ]]; then
  ERROR_CODE=$(echo "$RESP" | jq -r '.error_codes[0] // empty' 2>/dev/null || true)
  ERROR_DESC=$(echo "$RESP" | jq -r '.error_description // empty' 2>/dev/null || true)

  # AADSTS7000229 — service principal not found in tenant (not consented)
  if echo "$ERROR_DESC" | grep -qi "7000229\|AADSTS7000229" || [[ "$ERROR_CODE" == "7000229" ]]; then
    echo "ERROR: AADSTS7000229 — app not consented in tenant $TENANT_ID (tier=$TIER, auth=$AUTH_METHOD)" >&2
    echo "" >&2
    echo "  The '${TIER}' service principal has not been authorized in this tenant." >&2
    echo "  Send this consent URL to the customer Global Admin:" >&2
    echo "" >&2
    echo "  https://login.microsoftonline.com/${TENANT_ID}/adminconsent?client_id=${CLIENT_ID}&redirect_uri=https://azcomputerguru.com&prompt=consent" >&2
    echo "" >&2
    echo "  After the admin accepts, run onboard-tenant.sh to assign required directory roles:" >&2
    SCRIPT_DIR_ERR="$(dirname "${BASH_SOURCE[0]}")"
    echo "  bash ${SCRIPT_DIR_ERR}/onboard-tenant.sh ${TARGET}" >&2
    exit 5
  fi

  echo "ERROR: token request failed (tenant=$TENANT_ID tier=$TIER auth=$AUTH_METHOD)" >&2
  echo "$RESP" >&2
  exit 5
fi

echo "$TOKEN" > "$CACHE_FILE"
chmod 600 "$CACHE_FILE" 2>/dev/null || true
echo "$TOKEN"