#!/bin/sh
set -eu

usage() {
  echo "usage: $0 <kubeconfig> <context> <session-context>" >&2
  exit 2
}

[ "$#" -eq 3 ] || usage
kubeconfig=$1
context=$2
session_context=$3
proof_namespace=cp9-system
proof_job=cp9-aws-runtime-identity-proof
proof_image=${CP9_AWS_IDENTITY_PROOF_IMAGE:-public.ecr.aws/aws-cli/aws-cli:2.17.47}

unset AWS_ACCESS_KEY_ID AWS_SECRET_ACCESS_KEY AWS_SESSION_TOKEN AWS_SECURITY_TOKEN
unset AWS_PROFILE AWS_DEFAULT_PROFILE AWS_SHARED_CREDENTIALS_FILE AWS_CONFIG_FILE
unset AWS_EC2_METADATA_DISABLED AWS_WEB_IDENTITY_TOKEN_FILE AWS_ROLE_ARN

runtime_json=$(python3 - "$session_context" <<'PY'
import json
import pathlib
import sys

ctx = json.loads(pathlib.Path(sys.argv[1]).read_text())
runtime = (ctx.get("terraform") or {}).get("remote_ec2_runtime") or {}
if (ctx.get("host") or {}).get("mode") != "remote-ec2-k3d" or runtime.get("enabled") is not True:
    raise SystemExit("remote EC2 runtime identity contract is required")
role_arn = str(runtime.get("role_arn") or "").strip()
instance_profile_arn = str(runtime.get("instance_profile_arn") or "").strip()
instance_profile_name = str(runtime.get("instance_profile_name") or "").strip()
evidence_dir = str(ctx.get("evidence_dir") or "").strip()
region = str(((ctx.get("target") or {}).get("region")) or "").strip()
if not role_arn or not instance_profile_arn or not instance_profile_name or not evidence_dir:
    raise SystemExit("remote runtime role/profile identity and evidence_dir are required")
if not region:
    raise SystemExit("remote runtime identity capture requires target.region")
print(json.dumps({
    "evidence_dir": evidence_dir,
    "instance_profile_arn": instance_profile_arn,
    "instance_profile_name": instance_profile_name,
    "region": region,
    "role_arn": role_arn,
}))
PY
)
region=$(python3 -c 'import json,sys; print(json.loads(sys.argv[1])["region"])' "$runtime_json")
evidence_dir=$(python3 -c 'import json,sys; print(json.loads(sys.argv[1])["evidence_dir"])' "$runtime_json")

caller_json=$(mktemp)
proof_pod_json=$(mktemp)
workloads_json=$(mktemp)
upbound_json=$(mktemp)
classic_json=$(mktemp)
secret_names=$(mktemp)
cleanup() {
  KUBECONFIG="$kubeconfig" kubectl --context "$context" --namespace "$proof_namespace" \
    delete job "$proof_job" --ignore-not-found=true >/dev/null 2>&1 || true
  rm -f "$caller_json" "$proof_pod_json" "$workloads_json" "$upbound_json" "$classic_json" "$secret_names"
}
trap cleanup EXIT

KUBECONFIG="$kubeconfig" kubectl --context "$context" --namespace "$proof_namespace" \
  delete job "$proof_job" --ignore-not-found=true >/dev/null
cat <<EOF | KUBECONFIG="$kubeconfig" kubectl --context "$context" apply -f - >/dev/null
apiVersion: batch/v1
kind: Job
metadata:
  name: $proof_job
  namespace: $proof_namespace
spec:
  backoffLimit: 0
  ttlSecondsAfterFinished: 300
  template:
    metadata:
      labels:
        cp9.tidbcloud.com/proof: remote-aws-runtime-identity
    spec:
      restartPolicy: Never
      containers:
        - name: aws-identity
          image: $proof_image
          args: ["sts", "get-caller-identity", "--region", "$region", "--output", "json"]
EOF
KUBECONFIG="$kubeconfig" kubectl --context "$context" --namespace "$proof_namespace" \
  wait --for=condition=complete "job/$proof_job" --timeout=300s >/dev/null
KUBECONFIG="$kubeconfig" kubectl --context "$context" --namespace "$proof_namespace" \
  logs "job/$proof_job" >"$caller_json"
KUBECONFIG="$kubeconfig" kubectl --context "$context" --namespace "$proof_namespace" \
  get pod -l "job-name=$proof_job" -o json >"$proof_pod_json"

KUBECONFIG="$kubeconfig" kubectl --context "$context" get deployment,statefulset -A -o json >"$workloads_json"
printf '%s\n' '{"items":[]}' >"$upbound_json"
printf '%s\n' '{"items":[]}' >"$classic_json"
if KUBECONFIG="$kubeconfig" kubectl --context "$context" get crd providerconfigs.aws.upbound.io >/dev/null 2>&1; then
  KUBECONFIG="$kubeconfig" kubectl --context "$context" get providerconfigs.aws.upbound.io -o json >"$upbound_json"
fi
if KUBECONFIG="$kubeconfig" kubectl --context "$context" get crd providerconfigs.aws.crossplane.io >/dev/null 2>&1; then
  KUBECONFIG="$kubeconfig" kubectl --context "$context" get providerconfigs.aws.crossplane.io -o json >"$classic_json"
fi
KUBECONFIG="$kubeconfig" kubectl --context "$context" get secret -A \
  -o custom-columns=NAMESPACE:.metadata.namespace,NAME:.metadata.name --no-headers >"$secret_names"
if grep -Eq '[[:space:]](aws-creds|cp9-aws-credentials)$' "$secret_names"; then
  echo "remote control plane still contains a human AWS credential Secret" >&2
  exit 1
fi

mkdir -p "$evidence_dir"
python3 - "$runtime_json" "$caller_json" "$proof_pod_json" "$workloads_json" "$upbound_json" "$classic_json" \
  "$evidence_dir/remote-aws-runtime-identity.json" <<'PY'
from datetime import datetime, timezone
import json
import pathlib
import re
import sys

runtime = json.loads(sys.argv[1])
caller = json.loads(pathlib.Path(sys.argv[2]).read_text())
proof_pods = json.loads(pathlib.Path(sys.argv[3]).read_text())
workloads = json.loads(pathlib.Path(sys.argv[4]).read_text())
provider_sets = [
    json.loads(pathlib.Path(sys.argv[5]).read_text()),
    json.loads(pathlib.Path(sys.argv[6]).read_text()),
]
output = pathlib.Path(sys.argv[7])

role_arn = runtime["role_arn"]
match = re.match(r"^arn:aws:iam::(?P<account>\d{12}):role/(?P<role>[^/]+)$", role_arn)
if not match:
    raise SystemExit("remote runtime role ARN is invalid")
caller_arn = str(caller.get("Arn") or "")
prefix = f"arn:aws:sts::{match.group('account')}:assumed-role/{match.group('role')}/"
if not caller_arn.startswith(prefix):
    raise SystemExit(f"in-cluster caller {caller_arn!r} does not match {role_arn!r}")

target_namespaces = {"crossplane-system", "cp9-global", "cp9-premium", "cp9-dedicated"}
human_secrets = {"aws-creds", "cp9-aws-credentials"}
credential_env = {
    "AWS_ACCESS_KEY_ID",
    "AWS_SECRET_ACCESS_KEY",
    "AWS_SESSION_TOKEN",
    "AWS_SHARED_CREDENTIALS_FILE",
}
audited_workloads = []
for item in workloads.get("items") or []:
    metadata = item.get("metadata") or {}
    namespace = str(metadata.get("namespace") or "")
    if namespace not in target_namespaces:
        continue
    kind = str(item.get("kind") or "")
    name = str(metadata.get("name") or "")
    pod_spec = (((item.get("spec") or {}).get("template") or {}).get("spec") or {})
    for volume in pod_spec.get("volumes") or []:
        secret_name = str((volume.get("secret") or {}).get("secretName") or "")
        if secret_name in human_secrets:
            raise SystemExit(f"{namespace}/{kind}/{name} mounts credential Secret {secret_name}")
    for container in (pod_spec.get("initContainers") or []) + (pod_spec.get("containers") or []):
        for mount in container.get("volumeMounts") or []:
            if str(mount.get("name") or "") in human_secrets:
                raise SystemExit(f"{namespace}/{kind}/{name} mounts human credential volume")
        for env in container.get("env") or []:
            env_name = str(env.get("name") or "")
            secret_name = str((((env.get("valueFrom") or {}).get("secretKeyRef") or {}).get("name")) or "")
            if secret_name in human_secrets or (env_name in credential_env and secret_name):
                raise SystemExit(f"{namespace}/{kind}/{name} injects human credential env {env_name}")
    audited_workloads.append(f"{namespace}/{kind}/{name}")

provider_configs = []
for provider_set in provider_sets:
    for item in provider_set.get("items") or []:
        api_version = str(item.get("apiVersion") or "")
        name = str((item.get("metadata") or {}).get("name") or "")
        source = str((((item.get("spec") or {}).get("credentials") or {}).get("source")) or "")
        expected_source = "PodIdentity" if api_version.startswith("aws.upbound.io/") else "InjectedIdentity"
        if source != expected_source:
            raise SystemExit(f"ProviderConfig {name!r} uses credentials.source={source!r}, expected {expected_source!r}")
        provider_configs.append(name)
if not provider_configs:
    raise SystemExit("no AWS ProviderConfigs were available for remote identity proof")
if not audited_workloads:
    raise SystemExit("no control-plane workloads were available for remote identity proof")
proof_pod_uids = sorted(
    str((item.get("metadata") or {}).get("uid") or "")
    for item in proof_pods.get("items") or []
    if str((item.get("metadata") or {}).get("uid") or "")
)
if len(proof_pod_uids) != 1:
    raise SystemExit(f"expected one identity proof pod UID, got {proof_pod_uids!r}")

payload = {
    "schema_version": 1,
    "captured_at": datetime.now(timezone.utc).isoformat().replace("+00:00", "Z"),
    "identity_source": "ec2-instance-profile",
    "instance_profile_arn": runtime["instance_profile_arn"],
    "instance_profile_name": runtime["instance_profile_name"],
    "role_arn": role_arn,
    "caller_arn": caller_arn,
    "caller_account": str(caller.get("Account") or ""),
    "credential_secrets_absent": True,
    "provider_configs_native": sorted(provider_configs),
    "proof_pod_uid": proof_pod_uids[0],
    "audited_workloads": sorted(audited_workloads),
}
output.write_text(json.dumps(payload, indent=2, sort_keys=True) + "\n")
PY

echo "remote AWS runtime identity proof captured: $evidence_dir/remote-aws-runtime-identity.json"