teleo-infrastructure/tests/test_apply_gcp_iap_operator_access.py

from __future__ import annotations

import copy
import hashlib
import importlib
import json
import subprocess
import sys
from pathlib import Path
from typing import Any

ROOT = Path(__file__).resolve().parents[1]
OPS = ROOT / "ops"
sys.path.insert(0, str(OPS))
apply = importlib.import_module("apply_gcp_iap_operator_access")
sys.path.pop(0)


def _binding(role: str, member: str, condition: tuple[str, str] | None = None) -> dict[str, Any]:
    result: dict[str, Any] = {"role": role, "members": [member]}
    if condition:
        result["condition"] = {"title": condition[0], "expression": condition[1]}
    return result


def _full_state() -> dict[str, Any]:
    operator_members = [f"serviceAccount:{email}" for _, email, _ in apply.SERVICE_ACCOUNTS]
    account_policies = {
        email: {"bindings": [_binding("roles/iam.workloadIdentityUser", apply.WIF_PRINCIPAL)]}
        for _, email, _ in apply.SERVICE_ACCOUNTS
    }
    project_bindings: list[dict[str, Any]] = []
    for member in operator_members:
        project_bindings.extend(
            [
                _binding(apply.CUSTOM_ROLE, member),
                _binding(apply.IAP_ROLE, member, (apply.IAP_CONDITION_TITLE, apply.IAP_CONDITION)),
            ]
        )
    return {
        "services": set(apply.REQUIRED_APIS),
        "pool": {"state": "ACTIVE", "disabled": False},
        "provider": {
            "state": "ACTIVE",
            "disabled": False,
            "displayName": apply.PROVIDER_DISPLAY_NAME,
            "oidc": {"issuerUri": apply.PROVIDER_ISSUER},
            "attributeMapping": apply.EXPECTED_ATTRIBUTE_MAPPING,
            "attributeCondition": apply.PROVIDER_CONDITION,
        },
        "accounts": {
            email: {"email": email, "displayName": display_name, "disabled": False}
            for _, email, display_name in apply.SERVICE_ACCOUNTS
        },
        "account_policies": account_policies,
        "custom_role": {
            "title": apply.CUSTOM_ROLE_TITLE,
            "description": apply.CUSTOM_ROLE_DESCRIPTION,
            "includedPermissions": list(apply.CUSTOM_PERMISSIONS),
            "stage": "GA",
        },
        "project_policy": {"bindings": project_bindings},
        "vm_policy": {"bindings": [_binding("roles/iam.serviceAccountUser", member) for member in operator_members]},
        "instance": {
            "status": "RUNNING",
            "networkInterfaces": [
                {
                    "networkIP": apply.INSTANCE_INTERNAL_IP,
                    "network": f"projects/{apply.PROJECT}/global/networks/{apply.NETWORK}",
                }
            ],
            "serviceAccounts": [{"email": apply.VM_SERVICE_ACCOUNT}],
            "metadata": {"items": [{"key": "enable-oslogin", "value": "TRUE"}]},
            "tags": {"items": [apply.TARGET_TAG]},
        },
        "instance_policy": {
            "bindings": [
                _binding("roles/compute.osLogin", operator_members[0]),
                _binding("roles/compute.osAdminLogin", operator_members[1]),
            ]
        },
        "iap_firewall": {
            "direction": "INGRESS",
            "network": f"projects/{apply.PROJECT}/global/networks/{apply.NETWORK}",
            "allowed": [{"IPProtocol": "tcp", "ports": ["22"]}],
            "sourceRanges": [apply.IAP_SOURCE_RANGE],
            "targetTags": [apply.TARGET_TAG],
            "disabled": False,
            "logConfig": {"enable": True},
        },
        "public_firewall": {"name": apply.CHANGING_PUBLIC_RULE, "disabled": False},
        "dispatcher_hash": hashlib.sha256(apply.DEFAULT_DISPATCHER.read_bytes()).hexdigest(),
        "dispatcher_file_mode": "root:root:755",
        "dispatcher_run_mode": "root:root:700",
    }


def _clean_state() -> dict[str, Any]:
    state = _full_state()
    state.update(
        {
            "services": set(),
            "pool": None,
            "provider": None,
            "accounts": {email: None for _, email, _ in apply.SERVICE_ACCOUNTS},
            "account_policies": {email: None for _, email, _ in apply.SERVICE_ACCOUNTS},
            "custom_role": None,
            "project_policy": {"bindings": []},
            "vm_policy": {"bindings": []},
            "instance_policy": {"bindings": []},
            "iap_firewall": None,
            "dispatcher_hash": None,
            "dispatcher_file_mode": None,
            "dispatcher_run_mode": None,
        }
    )
    state["instance"]["metadata"] = {"items": []}
    state["instance"]["tags"] = {"items": []}
    return state


class FakeGcloud:
    def __init__(self, state: dict[str, Any]) -> None:
        self.state = copy.deepcopy(state)
        self.commands: list[list[str]] = []
        self.mutations: list[list[str]] = []

    @staticmethod
    def _value(command: list[str], flag: str) -> str:
        return command[command.index(flag) + 1]

    @staticmethod
    def _completed(command: list[str], payload: Any = None, *, returncode: int = 0, stderr: str = ""):
        stdout = "" if payload is None else json.dumps(payload)
        return subprocess.CompletedProcess(command, returncode, stdout, stderr)

    @classmethod
    def _missing(cls, command: list[str]):
        return cls._completed(command, returncode=1, stderr="ERROR: NOT_FOUND: resource does not exist")

    @staticmethod
    def _add_binding(policy: dict[str, Any], role: str, member: str, condition: tuple[str, str] | None) -> None:
        candidate = _binding(role, member, condition)
        if candidate not in policy["bindings"]:
            policy["bindings"].append(candidate)

    def _read(self, command: list[str]):
        words = command[1:]
        if words[:2] == ["projects", "describe"]:
            return self._completed(command, {"projectId": apply.PROJECT})
        if words[:2] == ["services", "list"]:
            return self._completed(command, [{"config": {"name": name}} for name in sorted(self.state["services"])])
        if words[:3] == ["iam", "workload-identity-pools", "describe"]:
            value = self.state["pool"]
            return self._missing(command) if value is None else self._completed(command, value)
        if words[:4] == ["iam", "workload-identity-pools", "providers", "describe"]:
            value = self.state["provider"]
            return self._missing(command) if value is None else self._completed(command, value)
        if words[:3] == ["iam", "service-accounts", "describe"]:
            email = words[3]
            value = self.state["accounts"].get(email)
            return self._missing(command) if value is None else self._completed(command, value)
        if words[:3] == ["iam", "service-accounts", "get-iam-policy"]:
            email = words[3]
            value = (
                self.state["vm_policy"] if email == apply.VM_SERVICE_ACCOUNT else self.state["account_policies"][email]
            )
            return self._completed(command, value)
        if words[:3] == ["iam", "roles", "describe"]:
            value = self.state["custom_role"]
            return self._missing(command) if value is None else self._completed(command, value)
        if words[:2] == ["projects", "get-iam-policy"]:
            return self._completed(command, self.state["project_policy"])
        if words[:3] == ["compute", "instances", "describe"]:
            return self._completed(command, self.state["instance"])
        if words[:3] == ["compute", "instances", "get-iam-policy"]:
            return self._completed(command, self.state["instance_policy"])
        if words[:3] == ["compute", "firewall-rules", "describe"]:
            rule = words[3]
            value = self.state["iap_firewall"] if rule == apply.IAP_FIREWALL_RULE else self.state["public_firewall"]
            return self._missing(command) if value is None else self._completed(command, value)
        return None

    def _mutate(self, command: list[str]):
        words = command[1:]
        if words[:2] == ["services", "enable"]:
            self.state["services"].update(words[2 : words.index("--project")])
        elif words[:3] == ["iam", "workload-identity-pools", "create"]:
            self.state["pool"] = {"state": "ACTIVE", "disabled": False}
        elif words[:3] == ["iam", "workload-identity-pools", "update"]:
            self.state["pool"]["disabled"] = False
        elif words[:4] == ["iam", "workload-identity-pools", "providers", "create-oidc"]:
            self.state["provider"] = {
                "state": "ACTIVE",
                "disabled": False,
                "displayName": self._value(command, "--display-name"),
                "oidc": {"issuerUri": self._value(command, "--issuer-uri")},
                "attributeMapping": apply.EXPECTED_ATTRIBUTE_MAPPING,
                "attributeCondition": self._value(command, "--attribute-condition"),
            }
        elif words[:4] == ["iam", "workload-identity-pools", "providers", "update-oidc"]:
            self.state["provider"].update(
                {
                    "state": "ACTIVE",
                    "disabled": False,
                    "displayName": self._value(command, "--display-name"),
                    "oidc": {"issuerUri": self._value(command, "--issuer-uri")},
                    "attributeMapping": apply.EXPECTED_ATTRIBUTE_MAPPING,
                    "attributeCondition": self._value(command, "--attribute-condition"),
                }
            )
        elif words[:3] == ["iam", "service-accounts", "create"]:
            account_id = words[3]
            email = f"{account_id}@{apply.PROJECT}.iam.gserviceaccount.com"
            self.state["accounts"][email] = {
                "email": email,
                "displayName": self._value(command, "--display-name"),
                "disabled": False,
            }
            self.state["account_policies"][email] = {"bindings": []}
        elif words[:3] == ["iam", "service-accounts", "update"]:
            self.state["accounts"][words[3]]["displayName"] = self._value(command, "--display-name")
        elif words[:3] == ["iam", "roles", "create"] or words[:3] == ["iam", "roles", "update"]:
            self.state["custom_role"] = {
                "title": self._value(command, "--title"),
                "description": self._value(command, "--description"),
                "includedPermissions": self._value(command, "--permissions").split(","),
                "stage": self._value(command, "--stage"),
            }
        elif words[:3] == ["iam", "service-accounts", "add-iam-policy-binding"]:
            email = words[3]
            policy = (
                self.state["vm_policy"] if email == apply.VM_SERVICE_ACCOUNT else self.state["account_policies"][email]
            )
            self._add_binding(policy, self._value(command, "--role"), self._value(command, "--member"), None)
        elif words[:2] == ["projects", "add-iam-policy-binding"]:
            condition = None
            if "--condition" in command:
                expression, title = self._value(command, "--condition").rsplit(",title=", 1)
                condition = (title, expression.removeprefix("expression="))
            self._add_binding(
                self.state["project_policy"],
                self._value(command, "--role"),
                self._value(command, "--member"),
                condition,
            )
        elif words[:3] == ["compute", "instances", "add-iam-policy-binding"]:
            self._add_binding(
                self.state["instance_policy"],
                self._value(command, "--role"),
                self._value(command, "--member"),
                None,
            )
        elif words[:3] == ["compute", "instances", "add-metadata"]:
            self.state["instance"]["metadata"] = {"items": [{"key": "enable-oslogin", "value": "TRUE"}]}
        elif words[:3] == ["compute", "instances", "add-tags"]:
            self.state["instance"]["tags"] = {"items": [apply.TARGET_TAG]}
        elif words[:3] == ["compute", "firewall-rules", "create"]:
            self.state["iap_firewall"] = _full_state()["iap_firewall"]
        elif words[:3] == ["compute", "firewall-rules", "update"]:
            self.state["iap_firewall"].update(
                {"targetTags": [apply.TARGET_TAG], "disabled": False, "logConfig": {"enable": True}}
            )
        elif words[:2] == ["compute", "scp"]:
            pass
        elif words[:2] == ["compute", "ssh"] and "sudo install -o root" in self._value(command, "--command"):
            self.state["dispatcher_hash"] = hashlib.sha256(apply.DEFAULT_DISPATCHER.read_bytes()).hexdigest()
            self.state["dispatcher_file_mode"] = "root:root:755"
            self.state["dispatcher_run_mode"] = "root:root:700"
        else:
            return None
        self.mutations.append(command)
        return self._completed(command)

    def __call__(self, command: list[str], **_: Any):
        self.commands.append(command)
        read = self._read(command)
        if read is not None:
            return read
        if command[1:3] == ["compute", "ssh"]:
            remote = self._value(command, "--command")
            if "sudo install -o root" in remote:
                return self._mutate(command)
            file_mode = self.state["dispatcher_file_mode"] or "MISSING"
            run_mode = self.state["dispatcher_run_mode"] or "MISSING"
            lines = [f"FILE={file_mode}"]
            if self.state["dispatcher_hash"]:
                lines.append(f"SHA={self.state['dispatcher_hash']}")
            lines.append(f"RUN={run_mode}")
            return subprocess.CompletedProcess(command, 0, "\n".join(lines) + "\n", "")
        mutated = self._mutate(command)
        if mutated is not None:
            return mutated
        return subprocess.CompletedProcess(command, 2, "", "unsupported fake gcloud command")


def _run(tmp_path: Path, state: dict[str, Any]):
    fake = FakeGcloud(state)
    receipt, exitcode = apply.apply_operator_access(
        execute=True,
        output=tmp_path / "receipt.json",
        runner=fake,
    )
    return fake, receipt, exitcode


def test_clean_create_converges_and_retains_sanitized_receipt(tmp_path: Path) -> None:
    fake, receipt, exitcode = _run(tmp_path, _clean_state())

    assert exitcode == 0
    assert receipt["status"] == "applied"
    assert receipt["mutation_count"] > 0
    assert fake.state["provider"]["attributeCondition"] == apply.PROVIDER_CONDITION
    assert fake.state["dispatcher_hash"] == receipt["dispatcher"]["sha256"]
    assert receipt["credentials_logged"] is False
    assert len(apply.PROVIDER_DISPLAY_NAME) <= 32
    assert apply.PROVIDER_DISPLAY_NAME == "Teleo fixed IAP operator"
    assert json.loads((tmp_path / "receipt.json").read_text())["status"] == "applied"


def test_fully_existing_state_is_control_plane_noop(tmp_path: Path) -> None:
    fake, receipt, exitcode = _run(tmp_path, _full_state())

    assert exitcode == 0
    assert receipt["status"] == "applied"
    assert receipt["mutation_count"] == 0
    assert fake.mutations == []


def test_partial_state_recovers_only_missing_clone_identity_and_bindings(tmp_path: Path) -> None:
    state = _full_state()
    clone_member = f"serviceAccount:{apply.CLONE_SERVICE_ACCOUNT}"
    state["accounts"][apply.CLONE_SERVICE_ACCOUNT] = None
    state["account_policies"][apply.CLONE_SERVICE_ACCOUNT] = None
    for policy_name in ("project_policy", "vm_policy", "instance_policy"):
        state[policy_name]["bindings"] = [
            binding for binding in state[policy_name]["bindings"] if clone_member not in binding["members"]
        ]

    fake, receipt, exitcode = _run(tmp_path, state)
    mutation_names = [
        operation["name"] for operation in receipt["operations"] if operation["action"] not in {"probe", "verify"}
    ]

    assert exitcode == 0
    assert f"create_service_account:{apply.CLONE_SERVICE_ACCOUNT}" in mutation_names
    assert f"create_service_account:{apply.STATUS_SERVICE_ACCOUNT}" not in mutation_names
    assert all(apply.STATUS_SERVICE_ACCOUNT not in " ".join(command) for command in fake.mutations)


def test_security_sensitive_provider_drift_fails_before_mutation(tmp_path: Path) -> None:
    state = _full_state()
    state["provider"]["attributeCondition"] = "assertion.repository == 'attacker/repository'"

    fake, receipt, exitcode = _run(tmp_path, state)

    assert exitcode == 1
    assert receipt["status"] == "failed"
    assert "security-sensitive drift" in receipt["error"]
    assert fake.mutations == []


def test_disabled_pool_and_provider_are_narrowly_reenabled(tmp_path: Path) -> None:
    state = _full_state()
    state["pool"]["disabled"] = True
    state["provider"]["disabled"] = True

    fake, receipt, exitcode = _run(tmp_path, state)
    mutation_names = [
        operation["name"] for operation in receipt["operations"] if operation["action"] not in {"probe", "verify"}
    ]

    assert exitcode == 0
    assert "enable_wif_pool" in mutation_names
    assert "update_wif_provider" in mutation_names
    assert fake.state["pool"]["disabled"] is False
    assert fake.state["provider"]["disabled"] is False


def test_wrong_or_stopped_instance_fails_before_mutation(tmp_path: Path) -> None:
    state = _full_state()
    state["instance"]["status"] = "TERMINATED"

    fake, receipt, exitcode = _run(tmp_path, state)

    assert exitcode == 1
    assert receipt["status"] == "failed"
    assert "not running" in receipt["error"]
    assert fake.mutations == []


def test_bootstrap_never_disables_or_mutates_public_ssh_rule(tmp_path: Path) -> None:
    fake, receipt, exitcode = _run(tmp_path, _clean_state())

    assert exitcode == 0
    assert receipt["public_ssh_rule"] == {
        "name": apply.CHANGING_PUBLIC_RULE,
        "observed_disabled_before": False,
        "observed_disabled_after": False,
        "mutated": False,
        "disable_command_present": False,
    }
    assert not any("--disabled" in command for command in fake.commands)
    assert not any(
        apply.CHANGING_PUBLIC_RULE in command
        and command[1:3] == ["compute", "firewall-rules"]
        and command[3] != "describe"
        for command in fake.commands
    )