#!/usr/bin/env python3
"""Emit the dry-run bootstrap and revocation plan for GitHub-to-GCP IAP SSH."""

from __future__ import annotations

import argparse
import json
import shlex

PROJECT = "teleo-501523"
PROJECT_NUMBER = "785938879453"
GITHUB_REPOSITORY = "living-ip/teleo-infrastructure"
WORKFLOW_PATH = ".github/workflows/gcp-iap-operator.yml"
WORKFLOW_REF = f"{GITHUB_REPOSITORY}/{WORKFLOW_PATH}@refs/heads/main"
POOL = "github-actions"
PROVIDER = "teleo-iap-operator"
PROVIDER_RESOURCE = f"projects/{PROJECT_NUMBER}/locations/global/workloadIdentityPools/{POOL}/providers/{PROVIDER}"

ZONE = "europe-west6-a"
INSTANCE = "teleo-prod-1"
INSTANCE_INTERNAL_IP = "10.60.0.3"
NETWORK = "teleo-staging-net"
TARGET_TAG = "teleo-prod-ssh"
VM_SERVICE_ACCOUNT = f"sa-teleo-prod-vm@{PROJECT}.iam.gserviceaccount.com"
STATUS_SERVICE_ACCOUNT_ID = "sa-teleo-iap-status"
STATUS_SERVICE_ACCOUNT = f"{STATUS_SERVICE_ACCOUNT_ID}@{PROJECT}.iam.gserviceaccount.com"
CLONE_SERVICE_ACCOUNT_ID = "sa-teleo-iap-clone-operator"
CLONE_SERVICE_ACCOUNT = f"{CLONE_SERVICE_ACCOUNT_ID}@{PROJECT}.iam.gserviceaccount.com"

CUSTOM_ROLE_ID = "teleoIapSshDiscovery"
CUSTOM_ROLE = f"projects/{PROJECT}/roles/{CUSTOM_ROLE_ID}"
CUSTOM_PERMISSIONS = [
    "compute.instances.get",
    "compute.instances.list",
    "compute.projects.get",
]
IAP_ROLE = "roles/iap.tunnelResourceAccessor"
IAP_CONDITION = f'destination.ip == "{INSTANCE_INTERNAL_IP}" && destination.port == 22'
PROVIDER_CONDITION = " && ".join(
    [
        f"assertion.repository == '{GITHUB_REPOSITORY}'",
        "assertion.ref == 'refs/heads/main'",
        f"assertion.workflow_ref == '{WORKFLOW_REF}'",
        "assertion.event_name == 'workflow_dispatch'",
    ]
)
ATTRIBUTE_MAPPING = ",".join(
    [
        "google.subject=assertion.sub",
        "attribute.repository=assertion.repository",
        "attribute.ref=assertion.ref",
        "attribute.workflow_ref=assertion.workflow_ref",
        "attribute.event_name=assertion.event_name",
    ]
)
WIF_PRINCIPAL = (
    f"principal://iam.googleapis.com/projects/{PROJECT_NUMBER}/locations/global/"
    f"workloadIdentityPools/{POOL}/subject/repo:{GITHUB_REPOSITORY}:ref:refs/heads/main"
)

IAP_FIREWALL_RULE = "teleo-prod-allow-ssh-iap"
CHANGING_PUBLIC_RULE = "teleo-prod-allow-ssh-current-ip"
IAP_SOURCE_RANGE = "35.235.240.0/20"
DISPATCHER_VERSION = "teleo-gcp-iap-operator-v1"

FORBIDDEN_OPERATOR_ROLES = [
    "roles/compute.admin",
    "roles/secretmanager.admin",
    "roles/secretmanager.secretAccessor",
    "roles/cloudsql.admin",
    "roles/artifactregistry.admin",
    "roles/artifactregistry.writer",
]


def shell_join(parts: list[str]) -> str:
    return " ".join(shlex.quote(part) for part in parts)


def service_account_member(email: str) -> str:
    return f"serviceAccount:{email}"


def bootstrap_commands() -> list[str]:
    accounts = [
        (STATUS_SERVICE_ACCOUNT_ID, "Teleo IAP status canary"),
        (CLONE_SERVICE_ACCOUNT_ID, "Teleo IAP clone operator"),
    ]
    commands = [
        shell_join(
            [
                "gcloud",
                "services",
                "enable",
                "compute.googleapis.com",
                "iap.googleapis.com",
                "iamcredentials.googleapis.com",
                "sts.googleapis.com",
                "--project",
                PROJECT,
            ]
        ),
        shell_join(
            [
                "gcloud",
                "iam",
                "workload-identity-pools",
                "describe",
                POOL,
                "--location",
                "global",
                "--project",
                PROJECT,
            ]
        ),
        shell_join(
            [
                "gcloud",
                "iam",
                "workload-identity-pools",
                "providers",
                "create-oidc",
                PROVIDER,
                "--location",
                "global",
                "--workload-identity-pool",
                POOL,
                "--project",
                PROJECT,
                "--issuer-uri",
                "https://token.actions.githubusercontent.com",
                "--attribute-mapping",
                ATTRIBUTE_MAPPING,
                "--attribute-condition",
                PROVIDER_CONDITION,
                "--display-name",
                "Teleo fixed IAP operator",
            ]
        ),
        *[
            shell_join(
                [
                    "gcloud",
                    "iam",
                    "service-accounts",
                    "create",
                    account_id,
                    "--project",
                    PROJECT,
                    "--display-name",
                    display_name,
                ]
            )
            for account_id, display_name in accounts
        ],
        shell_join(
            [
                "gcloud",
                "iam",
                "roles",
                "create",
                CUSTOM_ROLE_ID,
                "--project",
                PROJECT,
                "--title",
                "Teleo exact-instance SSH discovery",
                "--description",
                "Only the instance and project reads required by gcloud compute ssh",
                "--permissions",
                ",".join(CUSTOM_PERMISSIONS),
                "--stage",
                "GA",
            ]
        ),
    ]

    for account in (STATUS_SERVICE_ACCOUNT, CLONE_SERVICE_ACCOUNT):
        member = service_account_member(account)
        commands.extend(
            [
                shell_join(
                    [
                        "gcloud",
                        "iam",
                        "service-accounts",
                        "add-iam-policy-binding",
                        account,
                        "--project",
                        PROJECT,
                        "--member",
                        WIF_PRINCIPAL,
                        "--role",
                        "roles/iam.workloadIdentityUser",
                    ]
                ),
                shell_join(
                    [
                        "gcloud",
                        "projects",
                        "add-iam-policy-binding",
                        PROJECT,
                        "--member",
                        member,
                        "--role",
                        CUSTOM_ROLE,
                    ]
                ),
                shell_join(
                    [
                        "gcloud",
                        "projects",
                        "add-iam-policy-binding",
                        PROJECT,
                        "--member",
                        member,
                        "--role",
                        IAP_ROLE,
                        "--condition",
                        f"expression={IAP_CONDITION},title=teleo-prod-1-ssh-only",
                    ]
                ),
                shell_join(
                    [
                        "gcloud",
                        "iam",
                        "service-accounts",
                        "add-iam-policy-binding",
                        VM_SERVICE_ACCOUNT,
                        "--project",
                        PROJECT,
                        "--member",
                        member,
                        "--role",
                        "roles/iam.serviceAccountUser",
                    ]
                ),
            ]
        )

    commands.extend(
        [
            shell_join(
                [
                    "gcloud",
                    "compute",
                    "instances",
                    "add-iam-policy-binding",
                    INSTANCE,
                    "--zone",
                    ZONE,
                    "--project",
                    PROJECT,
                    "--member",
                    service_account_member(STATUS_SERVICE_ACCOUNT),
                    "--role",
                    "roles/compute.osLogin",
                ]
            ),
            shell_join(
                [
                    "gcloud",
                    "compute",
                    "instances",
                    "add-iam-policy-binding",
                    INSTANCE,
                    "--zone",
                    ZONE,
                    "--project",
                    PROJECT,
                    "--member",
                    service_account_member(CLONE_SERVICE_ACCOUNT),
                    "--role",
                    "roles/compute.osAdminLogin",
                ]
            ),
            shell_join(
                [
                    "gcloud",
                    "compute",
                    "instances",
                    "add-metadata",
                    INSTANCE,
                    "--zone",
                    ZONE,
                    "--project",
                    PROJECT,
                    "--metadata",
                    "enable-oslogin=TRUE",
                ]
            ),
            shell_join(
                [
                    "gcloud",
                    "compute",
                    "instances",
                    "add-tags",
                    INSTANCE,
                    "--zone",
                    ZONE,
                    "--project",
                    PROJECT,
                    "--tags",
                    TARGET_TAG,
                ]
            ),
            shell_join(
                [
                    "gcloud",
                    "compute",
                    "firewall-rules",
                    "create",
                    IAP_FIREWALL_RULE,
                    "--project",
                    PROJECT,
                    "--network",
                    NETWORK,
                    "--direction",
                    "INGRESS",
                    "--action",
                    "ALLOW",
                    "--rules",
                    "tcp:22",
                    "--source-ranges",
                    IAP_SOURCE_RANGE,
                    "--target-tags",
                    TARGET_TAG,
                    "--enable-logging",
                ]
            ),
            "# Install the reviewed dispatcher through the existing admin path before removing that path.",
            'BOOTSTRAP_SSH_DIR=$(mktemp -d "${TMPDIR:-/tmp}/teleo-iap-bootstrap.XXXXXX"); '
            'chmod 0700 "${BOOTSTRAP_SSH_DIR}"; '
            'BOOTSTRAP_SSH_KEY="${BOOTSTRAP_SSH_DIR}/ssh_key"; '
            "trap 'rm -rf \"${BOOTSTRAP_SSH_DIR}\"' EXIT",
            "install -m 0600 scripts/gcp_iap_operator.sh "
            '"${BOOTSTRAP_SSH_DIR}/gcp_iap_operator.sh"; '
            '(cd "${BOOTSTRAP_SSH_DIR}" && sha256sum gcp_iap_operator.sh > dispatcher.sha256); '
            'chmod 0600 "${BOOTSTRAP_SSH_DIR}/dispatcher.sha256"',
            "gcloud compute scp "
            '"${BOOTSTRAP_SSH_DIR}/gcp_iap_operator.sh" '
            '"${BOOTSTRAP_SSH_DIR}/dispatcher.sha256" '
            f"{INSTANCE}:/tmp/ --zone {ZONE} --project {PROJECT} "
            '--tunnel-through-iap --ssh-key-file="${BOOTSTRAP_SSH_KEY}" '
            '--ssh-key-expire-after=5m',
            shell_join(
                [
                    "gcloud",
                    "compute",
                    "ssh",
                    INSTANCE,
                    "--zone",
                    ZONE,
                    "--project",
                    PROJECT,
                    "--command",
                    "cd /tmp && sha256sum -c dispatcher.sha256 && "
                    f"grep -F 'readonly DISPATCHER_VERSION=\"{DISPATCHER_VERSION}\"' gcp_iap_operator.sh && "
                    "sudo install -o root -g root -m 0755 /tmp/gcp_iap_operator.sh "
                    "/usr/local/sbin/teleo-gcp-iap-operator && "
                    "test \"$(sudo sha256sum /usr/local/sbin/teleo-gcp-iap-operator | awk '{print $1}')\" "
                    "= \"$(awk '{print $1}' /tmp/dispatcher.sha256)\" && "
                    "sudo install -d -o root -g root -m 0700 /var/lib/teleo-gcp-iap-operator/runs && "
                    "rm -f /tmp/gcp_iap_operator.sh /tmp/dispatcher.sha256",
                ]
            )
            + ' --tunnel-through-iap --ssh-key-file="${BOOTSTRAP_SSH_KEY}" '
            '--ssh-key-expire-after=5m',
        ]
    )
    return commands


def status_verification_commands() -> list[str]:
    request_id = "iap-bootstrapstatus01"
    return [
        shell_join(
            [
                "gh",
                "workflow",
                "run",
                "gcp-iap-operator.yml",
                "--repo",
                GITHUB_REPOSITORY,
                "--ref",
                "main",
                "-f",
                "operation=status",
                "-f",
                f"request_id={request_id}",
                "-f",
                "target_db=teleo_clone_status",
            ]
        ),
        shell_join(
            [
                "gh",
                "run",
                "list",
                "--repo",
                GITHUB_REPOSITORY,
                "--workflow",
                "gcp-iap-operator.yml",
                "--event",
                "workflow_dispatch",
                "--limit",
                "1",
            ]
        ),
        "IAP_RUN_ID=$(gh run list --repo "
        f"{GITHUB_REPOSITORY} --workflow gcp-iap-operator.yml --event workflow_dispatch --limit 1 "
        "--json databaseId --jq '.[0].databaseId'); test -n \"${IAP_RUN_ID}\"",
        'IAP_VERIFY_DIR=$(mktemp -d "${TMPDIR:-/tmp}/teleo-iap-verify.XXXXXX"); '
        'chmod 0700 "${IAP_VERIFY_DIR}"; trap \'rm -rf "${IAP_VERIFY_DIR}"\' EXIT',
        'gh run watch "${IAP_RUN_ID}" --repo ' + GITHUB_REPOSITORY + " --exit-status",
        'gh run download "${IAP_RUN_ID}" --repo '
        + GITHUB_REPOSITORY
        + ' --name "gcp-iap-operator-${IAP_RUN_ID}" --dir "${IAP_VERIFY_DIR}"',
        "python3 -c 'import json,sys; p=json.load(open(sys.argv[1])); "
        "import hashlib,pathlib; "
        'r=p.get("remote_result") or {}; assert p.get("status")=="pass"; '
        f'assert p.get("request_id")=="{request_id}"; '
        'assert p.get("operation")=="status"; '
        f'assert r.get("expected_internal_ip")=="{INSTANCE_INTERNAL_IP}"; '
        f'assert r.get("dispatcher_version")=="{DISPATCHER_VERSION}"; '
        'assert r.get("dispatcher_sha256")==hashlib.sha256('
        'pathlib.Path("scripts/gcp_iap_operator.sh").read_bytes()).hexdigest(); '
        'assert r.get("active_state")=="active"\' "${IAP_VERIFY_DIR}/result.json"',
        "# The old public /32 rule may be disabled only after every command above succeeds.",
        shell_join(
            [
                "gcloud",
                "compute",
                "firewall-rules",
                "update",
                CHANGING_PUBLIC_RULE,
                "--project",
                PROJECT,
                "--disabled",
            ]
        ),
        shell_join(
            [
                "gcloud",
                "compute",
                "firewall-rules",
                "describe",
                CHANGING_PUBLIC_RULE,
                "--project",
                PROJECT,
                "--format=value(disabled,sourceRanges)",
            ]
        ),
    ]


def revocation_commands() -> list[str]:
    commands = [
        "# Remove the remote dispatcher while IAP access still exists, then revoke cloud access.",
        'REVOKE_SSH_DIR=$(mktemp -d "${TMPDIR:-/tmp}/teleo-iap-revoke.XXXXXX"); '
        'chmod 0700 "${REVOKE_SSH_DIR}"; REVOKE_SSH_KEY="${REVOKE_SSH_DIR}/ssh_key"; '
        "trap 'rm -rf \"${REVOKE_SSH_DIR}\"' EXIT",
        shell_join(
            [
                "gcloud",
                "compute",
                "ssh",
                INSTANCE,
                "--zone",
                ZONE,
                "--project",
                PROJECT,
                "--tunnel-through-iap",
                "--command",
                "sudo rm -f /usr/local/sbin/teleo-gcp-iap-operator",
            ]
        )
        + ' --ssh-key-file="${REVOKE_SSH_KEY}" --ssh-key-expire-after=5m',
        shell_join(
            [
                "gcloud",
                "iam",
                "workload-identity-pools",
                "providers",
                "update-oidc",
                PROVIDER,
                "--location",
                "global",
                "--workload-identity-pool",
                POOL,
                "--project",
                PROJECT,
                "--disabled",
            ]
        ),
        shell_join(
            [
                "gcloud",
                "compute",
                "firewall-rules",
                "delete",
                IAP_FIREWALL_RULE,
                "--project",
                PROJECT,
                "--quiet",
            ]
        ),
    ]
    for account, os_role in (
        (STATUS_SERVICE_ACCOUNT, "roles/compute.osLogin"),
        (CLONE_SERVICE_ACCOUNT, "roles/compute.osAdminLogin"),
    ):
        member = service_account_member(account)
        commands.extend(
            [
                shell_join(
                    [
                        "gcloud",
                        "compute",
                        "instances",
                        "remove-iam-policy-binding",
                        INSTANCE,
                        "--zone",
                        ZONE,
                        "--project",
                        PROJECT,
                        "--member",
                        member,
                        "--role",
                        os_role,
                    ]
                ),
                shell_join(
                    [
                        "gcloud",
                        "projects",
                        "remove-iam-policy-binding",
                        PROJECT,
                        "--member",
                        member,
                        "--role",
                        CUSTOM_ROLE,
                    ]
                ),
                shell_join(
                    [
                        "gcloud",
                        "projects",
                        "remove-iam-policy-binding",
                        PROJECT,
                        "--member",
                        member,
                        "--role",
                        IAP_ROLE,
                        "--condition",
                        f"expression={IAP_CONDITION},title=teleo-prod-1-ssh-only",
                    ]
                ),
                shell_join(
                    [
                        "gcloud",
                        "iam",
                        "service-accounts",
                        "remove-iam-policy-binding",
                        VM_SERVICE_ACCOUNT,
                        "--project",
                        PROJECT,
                        "--member",
                        member,
                        "--role",
                        "roles/iam.serviceAccountUser",
                    ]
                ),
                shell_join(
                    [
                        "gcloud",
                        "iam",
                        "service-accounts",
                        "delete",
                        account,
                        "--project",
                        PROJECT,
                        "--quiet",
                    ]
                ),
            ]
        )
    commands.extend(
        [
            shell_join(
                [
                    "gcloud",
                    "iam",
                    "workload-identity-pools",
                    "providers",
                    "delete",
                    PROVIDER,
                    "--location",
                    "global",
                    "--workload-identity-pool",
                    POOL,
                    "--project",
                    PROJECT,
                    "--quiet",
                ]
            ),
            shell_join(
                [
                    "gcloud",
                    "iam",
                    "roles",
                    "delete",
                    CUSTOM_ROLE_ID,
                    "--project",
                    PROJECT,
                    "--quiet",
                ]
            ),
        ]
    )
    return commands


def build_plan() -> dict[str, object]:
    return {
        "artifact": "teleo_gcp_iap_operator_access_plan",
        "mode": "dry_run_emit_only",
        "executes_commands": False,
        "project": PROJECT,
        "target": {
            "instance": INSTANCE,
            "zone": ZONE,
            "internal_ip": INSTANCE_INTERNAL_IP,
            "network": NETWORK,
            "target_tag": TARGET_TAG,
        },
        "github": {
            "repository": GITHUB_REPOSITORY,
            "ref": "refs/heads/main",
            "workflow_ref": WORKFLOW_REF,
            "event_name": "workflow_dispatch",
        },
        "wif": {
            "provider": PROVIDER_RESOURCE,
            "provider_condition": PROVIDER_CONDITION,
            "attribute_mapping": ATTRIBUTE_MAPPING,
            "principal": WIF_PRINCIPAL,
            "long_lived_google_credentials": False,
        },
        "identities": {
            "status": {
                "service_account": STATUS_SERVICE_ACCOUNT,
                "os_login_role": "roles/compute.osLogin",
                "sudo": False,
                "operations": ["status"],
            },
            "clone_operator": {
                "service_account": CLONE_SERVICE_ACCOUNT,
                "os_login_role": "roles/compute.osAdminLogin",
                "sudo": True,
                "operations": ["direct-claim-replay", "cleanup-clone"],
                "blast_radius": (
                    "OS Admin can become root on teleo-prod-1. The identity is separated from status, "
                    "the provider is pinned to the main workflow_dispatch workflow, and the checked-in "
                    "workflow invokes only the root-owned fixed-operation dispatcher. Branch protection "
                    "and workflow review remain required controls because OIDC does not attest file content."
                ),
            },
        },
        "gcloud_ssh_custom_role": {
            "role": CUSTOM_ROLE,
            "permissions": CUSTOM_PERMISSIONS,
            "binding_scope": f"projects/{PROJECT}",
            "rationale": (
                "The official gcloud IAP SSH permission table requires instance get/list plus project lookup. "
                "This project-level role is read-only; exact destination authorization remains restricted by "
                "the IAP destination condition and instance-level OS Login binding."
            ),
        },
        "iap": {
            "role": IAP_ROLE,
            "condition": IAP_CONDITION,
            "firewall": {
                "rule": IAP_FIREWALL_RULE,
                "source_range": IAP_SOURCE_RANGE,
                "protocol_ports": "tcp:22",
                "target_tag": TARGET_TAG,
            },
        },
        "service_account_user": {
            "target_service_account": VM_SERVICE_ACCOUNT,
            "role": "roles/iam.serviceAccountUser",
            "rationale": "Required to connect to a VM that runs as the dedicated VM service account.",
        },
        "explicitly_absent": {
            "roles": FORBIDDEN_OPERATOR_ROLES,
            "permissions": [
                "compute.instances.setMetadata",
                "compute.projects.setCommonInstanceMetadata",
                "secretmanager.versions.access",
                "cloudsql.instances.update",
                "artifactregistry.repositories.uploadArtifacts",
            ],
            "credentials": [
                "Google password",
                "local user refresh token",
                "service-account JSON key",
                "API key",
            ],
        },
        "bootstrap_commands": bootstrap_commands(),
        "post_bootstrap_verification": status_verification_commands(),
        "public_ingress_cutover": {
            "old_rule": CHANGING_PUBLIC_RULE,
            "disable_only_after": "successful workflow status artifact independently inspected",
            "automatic_fallback_to_public_ssh": False,
        },
        "audit_logging": {
            "guidance": (
                "Before enabling the operator workflow, merge project IAM auditConfigs rather than replacing "
                "the policy. Enable DATA_READ for iap.googleapis.com and DATA_WRITE for iap.googleapis.com, "
                "plus ADMIN_READ for "
                "compute.googleapis.com and iam.googleapis.com; retain IAP TCP tunnel and OS Login audit "
                "entries, firewall-rule logs, "
                "GitHub run ID, request_id, target_db, and sanitized result SHA-256. Never log auth outputs."
            ),
            "readback_commands": [
                f"gcloud projects get-iam-policy {PROJECT} --format=json",
                (
                    "gcloud logging read "
                    '\'protoPayload.serviceName="iap.googleapis.com" AND '
                    f'resource.labels.project_id="{PROJECT}"\' --project {PROJECT} --limit 20 --format=json'
                ),
                (
                    "gcloud logging read "
                    f'\'resource.type="gce_firewall_rule" AND resource.labels.firewall_rule_id="{IAP_FIREWALL_RULE}"\' '
                    f"--project {PROJECT} --limit 20 --format=json"
                ),
            ],
        },
        "independent_verification": [
            "Use an auditor identity that cannot impersonate either operator service account.",
            "Confirm the provider condition, both service-account policies, instance IAM policy, IAP condition, and firewall source/tag.",
            "Dispatch status from gh on main, inspect the sanitized artifact, and confirm no public /32 rule was disabled before that pass.",
            "Confirm the workflow created no service-account key and the runner removed its five-minute SSH key and auth credential file.",
        ],
        "revocation_commands": revocation_commands(),
    }


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--format", choices=("json", "shell"), default="json")
    args = parser.parse_args()
    plan = build_plan()
    if args.format == "shell":
        print("# DRY RUN: review these commands; this program never executes them.")
        for section in ("bootstrap_commands", "post_bootstrap_verification", "revocation_commands"):
            print(f"\n# {section}")
            for command in plan[section]:
                print(command)
    else:
        print(json.dumps(plan, indent=2, sort_keys=True))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())
