#!/usr/bin/env python3
"""Fail-closed, read-only checks for the AWS production Terraform backend."""

from __future__ import annotations

import argparse
import json
import os
import re
import subprocess
import sys
from dataclasses import dataclass
from typing import Any, TextIO


PROD_STATE_KEY = "infra/aws/environments/prod/terraform.tfstate"
ADOPTION_RUNBOOK = "docs/ops/aws-terraform-state-adoption.md"
EXIT_UNSAFE_BACKEND = 2
EXIT_STATE_ABSENT = 3


class PreflightError(RuntimeError):
    """A backend property could not be verified or is unsafe."""


class StateObjectAbsent(PreflightError):
    """The exact canonical production state object does not exist."""


@dataclass(frozen=True)
class AwsCli:
    region: str
    profile: str | None = None

    def _run(self, service: str, operation: str, *arguments: str) -> str:
        command = [
            "aws",
            service,
            operation,
            "--region",
            self.region,
            "--no-cli-pager",
            *arguments,
        ]
        if self.profile:
            command.extend(["--profile", self.profile])

        environment = os.environ.copy()
        environment["AWS_PAGER"] = ""
        try:
            result = subprocess.run(
                command,
                capture_output=True,
                check=False,
                env=environment,
                text=True,
                timeout=30,
            )
        except (OSError, subprocess.TimeoutExpired) as exc:
            raise PreflightError(
                f"AWS metadata check could not run: {service} {operation}"
            ) from exc
        if result.returncode != 0:
            # Never echo AWS CLI stderr. It is not required for the decision and
            # may contain account-specific metadata supplied by a future caller.
            if (
                service == "s3api"
                and operation == "head-object"
                and _is_head_object_not_found(result.stderr)
            ):
                raise StateObjectAbsent(PROD_STATE_KEY)
            raise PreflightError(f"AWS metadata check failed: {service} {operation}")
        return result.stdout

    def bucket_versioning(self, bucket: str) -> str:
        return self._run(
            "s3api",
            "get-bucket-versioning",
            "--bucket",
            bucket,
            "--query",
            "Status",
            "--output",
            "text",
        ).strip()

    def bucket_encryption_algorithms(self, bucket: str) -> list[str]:
        output = self._run(
            "s3api",
            "get-bucket-encryption",
            "--bucket",
            bucket,
            "--query",
            (
                "ServerSideEncryptionConfiguration.Rules[]"
                ".ApplyServerSideEncryptionByDefault.SSEAlgorithm"
            ),
            "--output",
            "json",
        )
        return _json_list(output, "bucket encryption")

    def public_access_block(self, bucket: str) -> dict[str, Any]:
        output = self._run(
            "s3api",
            "get-public-access-block",
            "--bucket",
            bucket,
            "--query",
            "PublicAccessBlockConfiguration",
            "--output",
            "json",
        )
        return _json_object(output, "public access block")

    def lock_table(self, table_name: str) -> dict[str, Any]:
        output = self._run(
            "dynamodb",
            "describe-table",
            "--table-name",
            table_name,
            "--query",
            "Table.{Status:TableStatus,Keys:KeySchema}",
            "--output",
            "json",
        )
        return _json_object(output, "lock table")

    def state_object_metadata(self, bucket: str) -> dict[str, Any]:
        # HeadObject returns metadata only. Never replace this with GetObject,
        # state pull, or any command that reads the state body.
        output = self._run(
            "s3api",
            "head-object",
            "--bucket",
            bucket,
            "--key",
            PROD_STATE_KEY,
            "--query",
            "{ETag:ETag,ServerSideEncryption:ServerSideEncryption,SSEKMSKeyId:SSEKMSKeyId}",
            "--output",
            "json",
        )
        return _json_object(output, "state object metadata")


def _is_head_object_not_found(stderr: str) -> bool:
    return bool(
        re.search(
            r"An error occurred \((?:404|NoSuchKey|NotFound)\) "
            r"when calling the HeadObject operation",
            stderr,
            re.IGNORECASE,
        )
    )


def _json_list(raw: str, label: str) -> list[str]:
    try:
        value = json.loads(raw)
    except json.JSONDecodeError as exc:
        raise PreflightError(f"AWS returned invalid JSON for {label}") from exc
    if not isinstance(value, list) or not all(isinstance(item, str) for item in value):
        raise PreflightError(f"AWS returned an unexpected shape for {label}")
    return value


def _json_object(raw: str, label: str) -> dict[str, Any]:
    try:
        value = json.loads(raw)
    except json.JSONDecodeError as exc:
        raise PreflightError(f"AWS returned invalid JSON for {label}") from exc
    if not isinstance(value, dict):
        raise PreflightError(f"AWS returned an unexpected shape for {label}")
    return value


def _lock_table_is_safe(table: dict[str, Any]) -> bool:
    keys = table.get("Keys")
    return table.get("Status") == "ACTIVE" and keys == [
        {"AttributeName": "LockID", "KeyType": "HASH"}
    ]


def run_preflight(
    aws: AwsCli,
    *,
    bucket: str,
    lock_table: str,
    out: TextIO = sys.stdout,
    err: TextIO = sys.stderr,
) -> int:
    checks: tuple[tuple[str, bool], ...]
    try:
        checks = (
            ("bucket versioning is Enabled", aws.bucket_versioning(bucket) == "Enabled"),
            (
                "bucket default encryption includes aws:kms",
                "aws:kms" in aws.bucket_encryption_algorithms(bucket),
            ),
            (
                "all four S3 public-access blocks are enabled",
                aws.public_access_block(bucket)
                == {
                    "BlockPublicAcls": True,
                    "IgnorePublicAcls": True,
                    "BlockPublicPolicy": True,
                    "RestrictPublicBuckets": True,
                },
            ),
            (
                "DynamoDB lock table is ACTIVE with LockID hash key",
                _lock_table_is_safe(aws.lock_table(lock_table)),
            ),
        )
        failed = [label for label, passed in checks if not passed]
        if failed:
            for label in failed:
                print(f"FAIL: {label}", file=err)
            print("Refusing normal Terraform plan/apply: backend protections are unsafe.", file=err)
            return EXIT_UNSAFE_BACKEND

        state_metadata = aws.state_object_metadata(bucket)
        object_uses_kms = (
            state_metadata.get("ServerSideEncryption") == "aws:kms"
            and isinstance(state_metadata.get("SSEKMSKeyId"), str)
            and bool(state_metadata["SSEKMSKeyId"].strip())
        )
        if not object_uses_kms:
            print(
                "FAIL: canonical production state object is not encrypted with a KMS key",
                file=err,
            )
            print(
                "Refusing normal Terraform plan/apply: state object encryption is unsafe.",
                file=err,
            )
            return EXIT_UNSAFE_BACKEND
    except StateObjectAbsent:
        print(
            f"FAIL: canonical production state object is absent: {PROD_STATE_KEY}",
            file=err,
        )
        print(
            "Refusing normal Terraform plan/apply. This check does not create empty state. "
            f"Follow {ADOPTION_RUNBOOK}.",
            file=err,
        )
        return EXIT_STATE_ABSENT
    except PreflightError as exc:
        print(f"FAIL: {exc}", file=err)
        print(
            "Refusing normal Terraform plan/apply: backend safety could not be verified.",
            file=err,
        )
        return EXIT_UNSAFE_BACKEND

    for label, _passed in checks:
        print(f"OK: {label}", file=out)
    print(f"OK: canonical production state metadata exists at {PROD_STATE_KEY}", file=out)
    print("OK: canonical production state object uses KMS encryption", file=out)
    print("PASS: backend preflight is safe for a normal Terraform plan.", file=out)
    return 0


def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description="Read-only, fail-closed AWS production Terraform backend preflight."
    )
    parser.add_argument("--bucket", required=True, help="Terraform state S3 bucket name")
    parser.add_argument("--lock-table", required=True, help="Terraform lock DynamoDB table name")
    parser.add_argument("--region", default="ap-southeast-1", help="AWS region")
    parser.add_argument("--profile", help="Optional local AWS CLI profile")
    return parser.parse_args(argv)


def main(argv: list[str] | None = None) -> int:
    args = parse_args(argv)
    return run_preflight(
        AwsCli(region=args.region, profile=args.profile),
        bucket=args.bucket,
        lock_table=args.lock_table,
    )


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