#!/usr/bin/env zsh
# Novesai Copilot Spend Snapshot collector (read-only), macOS/Linux.
#
# Reads Microsoft 365 Copilot Cost Management data from Microsoft Graph /beta
# and writes a JSON file in the novesai.copilot-snapshot 1.0 format.
#
# Read-only guarantee: every request goes through graph_get, which runs
# `curl --get` only, against an allowlist of read paths on graph.microsoft.com.
# Anything else exits. Every call is recorded in the output "calls" list.
#
# Sign-in: device code flow with the Microsoft Graph Command Line Tools app
# (delegated), or an existing token in NCS_ACCESS_TOKEN.
# Needs: zsh, curl, jq.
#
# Options:
#   --tenant <id|domain>     tenant (default: organizations / home tenant)
#   --user-scope <s>         PolicyMembers (default) | AllUsers | None
#   --max-users <n>          cap for per-user lookups (default 500)
#   --output <file>          output path
#   --minimal                balance + policies only, fewest scopes, no per-user data
#   --anonymize              hash names, UPNs, ids, group names and recipients
#
# The output file is sensitive. Admin consent for Microsoft Graph Command Line
# Tools stays in your tenant; remove it in Entra > Enterprise applications >
# Microsoft Graph Command Line Tools > Permissions when you are done.
#
# MIT License. (c) 2026 Novesai. https://noves.ai/copilot-snapshot.html

emulate -L zsh
setopt err_return no_unset pipe_fail

readonly COLLECTOR_VERSION="0.1.0"
readonly GRAPH="https://graph.microsoft.com"
readonly COST="/beta/copilot/costManagement"
readonly CLIENT_ID="14d82eec-204b-4c2f-b7e8-296a70dab67e"   # Microsoft Graph Command Line Tools
readonly FULL_SCOPES="CopilotCostManagement.Read.All CopilotCostManagement-Policy.Read.All CopilotCostManagement-Assignment.Read.All CopilotCostManagement-UserData.Read.All GroupMember.Read.All User.ReadBasic.All User.Read"
readonly MIN_SCOPES="CopilotCostManagement.Read.All CopilotCostManagement-Policy.Read.All User.Read"

tenant="organizations"; user_scope="PolicyMembers"; max_users=500; output=""; minimal=0; anonymize=0

usage() { sed -n '2,26p' "$0" | sed 's/^# \{0,1\}//'; }

while (( $# )); do
  case "$1" in
    --tenant) tenant="$2"; shift 2 ;;
    --user-scope) user_scope="$2"; shift 2 ;;
    --max-users) max_users="$2"; shift 2 ;;
    --output) output="$2"; shift 2 ;;
    --minimal) minimal=1; shift ;;
    --anonymize) anonymize=1; shift ;;
    -h|--help) usage; exit 0 ;;
    *) printf 'Unknown option: %s\n' "$1" >&2; usage; exit 2 ;;
  esac
done

case "$user_scope" in PolicyMembers|AllUsers|None) ;; *) printf '%s\n' "Invalid --user-scope" >&2; exit 2 ;; esac
if ! [[ "$max_users" =~ ^[0-9]+$ ]]; then printf '%s\n' "--max-users must be a number" >&2; exit 2; fi
for bin in curl jq; do command -v "$bin" >/dev/null || { printf 'Missing %s\n' "$bin" >&2; exit 1; }; done

scopes="$FULL_SCOPES"; mode="full"
if (( minimal )); then scopes="$MIN_SCOPES"; mode="minimal"; user_scope="None"; fi

work="$(mktemp -d)"; trap 'rm -rf "$work"' EXIT
calls="$work/calls.jsonl"; warnings="$work/warnings.jsonl"; : > "$calls"; : > "$warnings"
warn() { jq -nc --arg w "$1" '$w' >> "$warnings"; }

# ---- sign in (device code) -------------------------------------------------
token="${NCS_ACCESS_TOKEN:-}"
if [[ -z "$token" ]]; then
  login="https://login.microsoftonline.com/$tenant/oauth2/v2.0"
  dc="$(curl -sS -X POST "$login/devicecode" --data-urlencode "client_id=$CLIENT_ID" --data-urlencode "scope=$scopes offline_access")"
  printf '%s\n' "$(jq -r '.message' <<< "$dc")"
  device_code="$(jq -r '.device_code' <<< "$dc")"; interval="$(jq -r '.interval // 5' <<< "$dc")"
  while true; do
    sleep "$interval"
    tr="$(curl -sS -X POST "$login/token" --data-urlencode "grant_type=urn:ietf:params:oauth:grant-type:device_code" \
      --data-urlencode "client_id=$CLIENT_ID" --data-urlencode "device_code=$device_code")"
    token="$(jq -r '.access_token // empty' <<< "$tr")"
    [[ -n "$token" ]] && break
    case "$(jq -r '.error' <<< "$tr")" in
      authorization_pending) ;;
      slow_down) (( interval += 5 )) ;;
      *) printf 'Sign-in failed: %s\n' "$(jq -r '.error_description' <<< "$tr")" >&2; exit 1 ;;
    esac
  done
fi

# ---- read-only GET with allowlist ------------------------------------------
readonly RE_ALLOW_1='^/beta/copilot/costManagement/(getTenantCreditBalance|spendingPolicies)$'
readonly RE_ALLOW_2='^/beta/copilot/costManagement/spendingPolicies/[^/]+/(serviceSettings|emailThresholdSettings|emailThresholdSettings/recipients|microsoft\.graph\.selectedGroupsSpendingPolicy/assignedGroups)$'
readonly RE_ALLOW_3='^/beta/copilot/costManagement/userBalances/[^/]+$'
readonly RE_ALLOW_4='^/v1\.0/(users|organization|groups/[^/]+/transitiveMembers/microsoft\.graph\.user)$'
allowed() {
  local p="${1%%\?*}"
  if [[ "$p" =~ $RE_ALLOW_1 || "$p" =~ $RE_ALLOW_2 || "$p" =~ $RE_ALLOW_3 || "$p" =~ $RE_ALLOW_4 ]]; then return 0; fi
  return 1
}

# Note: zsh ties $path to $PATH, so never name a variable "path".
# graph_get <path-or-nextLink> <outfile> [all]   returns 0 on success, 1 on HTTP error
graph_get() {
  local target="$1" out="$2" all="${3:-}" url rpath code body ms attempt res
  : > "$out.items"
  url="$target"; [[ "$url" == http* ]] || url="$GRAPH$target"
  while [[ -n "$url" ]]; do
    [[ "$url" == "$GRAPH/"* ]] || { printf 'Blocked: %s is not graph.microsoft.com\n' "$url" >&2; exit 3; }
    rpath="${url#"$GRAPH"}"
    allowed "$rpath" || { printf 'Blocked: %s is not on the read-only allowlist\n' "${rpath%%\?*}" >&2; exit 3; }
    for attempt in 1 2 3 4; do
      res="$(curl -sS --get -o "$work/body" -w '%{http_code} %{time_total}' -H "Authorization: Bearer $token" -H 'Accept: application/json' "$url")" || res="0 0"
      code="${res%% *}"; ms="$(awk -v t="${res##* }" 'BEGIN{printf "%d", t*1000}')"
      if [[ "$code" == 429 || "$code" == 5* ]]; then sleep $(( 2 ** attempt )); else break; fi
    done
    jq -nc --arg p "$(sed -E 's#/userBalances/[^/?]+#/userBalances/{id}#' <<< "$rpath")" --argjson s "${code:-0}" --argjson d "$ms" \
      '{method:"GET",path:$p,status:$s,durationMs:$d}' >> "$calls"
    if [[ "$code" != 2* ]]; then
      body="$(jq -r '.error.message // "request failed"' "$work/body" 2>/dev/null || printf 'request failed')"
      warn "GET ${rpath%%\?*} failed ($code): $body"
      return 1
    fi
    if [[ -n "$all" ]]; then
      jq -c '.value[]?' "$work/body" >> "$out.items"
      url="$(jq -r '."@odata.nextLink" // empty' "$work/body")"
    else
      cp "$work/body" "$out"; url=""
    fi
  done
  if [[ -n "$all" ]]; then jq -s '.' "$out.items" > "$out"; fi
  return 0
}

printf '%s\n' "Novesai Copilot Spend Snapshot $COLLECTOR_VERSION (read-only, GET only). Mode: $mode."

# shellcheck disable=SC2016  # $select is an OData parameter, not a shell variable
graph_get '/v1.0/organization?$select=id,displayName,verifiedDomains' "$work/org" || printf '%s' '{"value":[]}' > "$work/org"
graph_get "$COST/getTenantCreditBalance" "$work/balance" || printf '%s' 'null' > "$work/balance"
graph_get "$COST/spendingPolicies" "$work/policies" all || printf '%s' '[]' > "$work/policies"

printf '[]' > "$work/entries"
while IFS= read -r pid; do
  [[ -z "$pid" ]] && continue
  enc="$(jq -rn --arg v "$pid" '$v|@uri')"
  graph_get "$COST/spendingPolicies/$enc/serviceSettings" "$work/svc" all || printf '%s' '[]' > "$work/svc"
  graph_get "$COST/spendingPolicies/$enc/emailThresholdSettings" "$work/alert" || printf '%s' 'null' > "$work/alert"
  printf '[]' > "$work/recip"; printf '[]' > "$work/groups"
  if (( ! minimal )); then
    if [[ "$(jq -r 'if . == null then "no" else "yes" end' "$work/alert")" == yes ]]; then
      graph_get "$COST/spendingPolicies/$enc/emailThresholdSettings/recipients" "$work/recip" all || printf '%s' '[]' > "$work/recip"
    fi
    if jq -e --arg id "$pid" '.[] | select(.id == $id) | ."@odata.type" | test("selectedGroups")' "$work/policies" >/dev/null; then
      graph_get "$COST/spendingPolicies/$enc/microsoft.graph.selectedGroupsSpendingPolicy/assignedGroups" "$work/groups" all || printf '%s' '[]' > "$work/groups"
    fi
  fi
  jq --arg id "$pid" --slurpfile p "$work/policies" --slurpfile s "$work/svc" --slurpfile a "$work/alert" \
     --slurpfile r "$work/recip" --slurpfile g "$work/groups" \
     '. + [{policy: ($p[0][] | select(.id == $id)), serviceSettings: $s[0], alertSetting: $a[0], alertRecipients: $r[0], assignedGroups: $g[0]}]' \
     "$work/entries" > "$work/entries.tmp" && mv "$work/entries.tmp" "$work/entries"
done < <(jq -r '.[].id' "$work/policies")

# ---- users -----------------------------------------------------------------
: > "$work/userlist"
# shellcheck disable=SC2016  # OData parameters, not shell variables
sel='$select=id,displayName,userPrincipalName&$top=999'
has_all_users="$(jq -r '[.[] | select((."@odata.type" // "" | test("allUsers")) and .status == "enabled")] | length' "$work/policies")"
if [[ "$user_scope" == AllUsers || ( "$user_scope" == PolicyMembers && "$has_all_users" != 0 ) ]]; then
  if [[ "$user_scope" == PolicyMembers ]]; then warn "An enabled all-users policy exists, so PolicyMembers reads users up to --max-users."; fi
  graph_get "/v1.0/users?$sel" "$work/u" all && jq -c '.[]' "$work/u" >> "$work/userlist"
fi
if [[ "$user_scope" == PolicyMembers ]]; then
  while IFS= read -r gid; do
    [[ -z "$gid" ]] && continue
    graph_get "/v1.0/groups/$(jq -rn --arg v "$gid" '$v|@uri')/transitiveMembers/microsoft.graph.user?$sel" "$work/gm" all && jq -c '.[]' "$work/gm" >> "$work/userlist"
  done < <(jq -r '.[].assignedGroups[]?.id' "$work/entries")
fi
jq -s --argjson m "$max_users" 'unique_by(.id) | .[:$m]' "$work/userlist" > "$work/users_sel"
total="$(jq 'length' "$work/users_sel")"
if (( total >= max_users && max_users > 0 )); then warn "User lookups capped at --max-users $max_users."; fi

: > "$work/users.jsonl"; i=0
while IFS= read -r row; do
  [[ -z "$row" ]] && continue
  i=$((i + 1)); printf '\rReading per-user balances %d/%d' "$i" "$total"
  uid="$(jq -r '.id' <<< "$row")"
  if graph_get "$COST/userBalances/$(jq -rn --arg v "$uid" '$v|@uri')?\$expand=serviceBalances" "$work/ub"; then b="$(cat "$work/ub")"; else b=null; fi
  jq -c --argjson b "$b" '{id, userPrincipalName, displayName, balance: $b}' <<< "$row" >> "$work/users.jsonl"
done < <(jq -c '.[]' "$work/users_sel")
if (( total > 0 )); then printf '\n'; fi
jq -s '.' "$work/users.jsonl" > "$work/users"

# ---- assemble --------------------------------------------------------------
salt="$(head -c 16 /dev/urandom | od -An -tx1 | tr -d ' \n')"
[[ -n "$output" ]] || output="copilot-snapshot-$(jq -r '.value[0].verifiedDomains[]? | select(.isDefault) | .name' "$work/org" | head -1 | sed 's/^$/tenant/')-$(date +%Y%m%d-%H%M).json"

jq -n \
  --arg ver "$COLLECTOR_VERSION" --arg mode "$mode" --arg us "$user_scope" --argjson mu "$max_users" \
  --argjson anon "$(if (( anonymize )); then printf true; else printf false; fi)" --arg scopes "$scopes" --arg salt "$salt" \
  --arg now "$(date -u +%Y-%m-%dT%H:%M:%SZ)" \
  --slurpfile org "$work/org" --slurpfile bal "$work/balance" --slurpfile ent "$work/entries" \
  --slurpfile users "$work/users" --slurpfile calls <(jq -s 'group_by(.path + "|" + (.status|tostring)) | map(.[0] + {durationMs: (map(.durationMs) | add), count: length})' "$calls") \
  --slurpfile warn <(jq -s '.' "$warnings") '
  def h: "anon-" + ($salt + "|" + tostring | @base64 | gsub("[^A-Za-z0-9]"; "") | .[-12:]);
  ($org[0].value[0] // {}) as $o
  | {
      "$schema": "https://noves.ai/schema/copilot-snapshot-1.0.schema.json",
      format: "novesai.copilot-snapshot", formatVersion: "1.0", generatedAt: $now,
      collector: {name: "novesai-copilot-snapshot", version: $ver, platform: "zsh"},
      tenant: ({id: $o.id, displayName: $o.displayName, domain: ([$o.verifiedDomains[]? | select(.isDefault) | .name][0])} | with_entries(select(.value != null))),
      options: {mode: $mode, userScope: $us, maxUsers: $mu, anonymized: $anon, scopes: ($scopes | split(" "))},
      creditBalance: $bal[0],
      spendingPolicies: $ent[0],
      users: $users[0],
      calls: $calls[0],
      warnings: $warn[0]
    }
  | if $anon then
      .users |= map((.id | h) as $x | .id = $x | .userPrincipalName = $x | .displayName = $x | (if .balance then .balance.id = $x else . end))
      | .spendingPolicies |= map(.alertRecipients |= map(with_entries(if .key == "displayName" or .key == "mail" or .key == "userPrincipalName" then .value |= h else . end))
                                 | .assignedGroups |= map(if .displayName then .displayName |= h else . end))
    else . end' > "$output"

printf '%s\n' "Wrote $output ($total users, $(jq '.spendingPolicies|length' "$output") policies)."
printf '%s\n' "This file is sensitive. Open it at https://noves.ai/copilot-snapshot.html; it is read in your browser and never uploaded."
printf '%s\n' "Done with Graph Command Line Tools? Remove its consent in Entra > Enterprise applications > Microsoft Graph Command Line Tools > Permissions."