#!/usr/bin/env python3
"""
Validates that all GitHub Actions workflow service containers and container images
use SHA256 digests for supply chain security.

Prevents tag tampering attacks by ensuring immutable image references.
"""

import sys
from pathlib import Path
from typing import List, Tuple

try:
    import yaml
except ImportError:
    print("❌ Error: PyYAML is required. Install with: pip install pyyaml")
    sys.exit(1)


# ANSI color codes
class Colors:
    RED = "\033[0;31m"
    GREEN = "\033[0;32m"
    YELLOW = "\033[1;33m"
    BLUE = "\033[0;34m"
    NC = "\033[0m"  # No Color


def has_sha_digest(image_ref: str) -> bool:
    """Check if image reference includes SHA256 digest."""
    return "@sha256:" in image_ref


def validate_workflow(workflow_path: Path) -> List[Tuple[str, str, str]]:
    """
    Validate a workflow file for unpinned images.

    Returns:
        List of (job_name, violation_type, image) tuples
    """
    violations = []

    try:
        with open(workflow_path, encoding="utf-8") as f:
            data = yaml.safe_load(f)

        if not isinstance(data, dict) or "jobs" not in data:
            return violations

        # Check each job
        for job_name, job_def in data["jobs"].items():
            if not isinstance(job_def, dict):
                continue

            # Check container: field
            if "container" in job_def:
                container = job_def["container"]

                # Container can be a string or dict with 'image' key
                if isinstance(container, str):
                    if not has_sha_digest(container):
                        violations.append((job_name, "container", container))
                elif isinstance(container, dict) and "image" in container:
                    image = container["image"]
                    if not has_sha_digest(image):
                        violations.append((job_name, "container", image))

            # Check services: field
            if "services" in job_def and isinstance(job_def["services"], dict):
                for service_name, service_def in job_def["services"].items():
                    if isinstance(service_def, dict) and "image" in service_def:
                        image = service_def["image"]
                        if not has_sha_digest(image):
                            violations.append(
                                (job_name, f"service '{service_name}'", image)
                            )

    except yaml.YAMLError as e:
        print(f"{Colors.RED}❌ YAML parse error in {workflow_path}:{Colors.NC}")
        print(f"   {e}")
        # Return a violation to fail the check
        violations.append(("parse_error", "error", str(e)))
    except Exception as e:
        print(
            f"{Colors.RED}❌ Error processing {workflow_path}: {e}{Colors.NC}"
        )
        violations.append(("error", "error", str(e)))

    return violations


def main():
    """Main validation logic."""
    workflows_dir = Path(".github/workflows")

    if not workflows_dir.exists():
        print(
            f"{Colors.RED}❌ .github/workflows directory not found{Colors.NC}"
        )
        sys.exit(1)

    print("🔍 Validating GitHub Actions workflow images...")
    print()

    total_violations = 0
    files_checked = 0

    # Process all workflow files
    for workflow_file in sorted(workflows_dir.glob("*.yml")) + sorted(
        workflows_dir.glob("*.yaml")
    ):
        violations = validate_workflow(workflow_file)

        if violations:
            print(f"{Colors.RED}📄 {workflow_file.name}:{Colors.NC}")
            for job_name, violation_type, image in violations:
                print(
                    f"{Colors.RED}   ❌ Job '{job_name}' {violation_type}: {image}{Colors.NC}"
                )
            print()
            total_violations += len(violations)
        else:
            print(f"{Colors.GREEN}   ✓ {workflow_file.name}{Colors.NC}")

        files_checked += 1

    # Summary
    print("━" * 50)
    print("📊 Summary:")
    print(f"   Files checked: {files_checked}")
    print(f"   Violations: {total_violations}")
    print("━" * 50)

    if total_violations > 0:
        print()
        print(
            f"{Colors.RED}❌ Found {total_violations} unpinned images in workflow files{Colors.NC}"
        )
        print()
        print("Service container images must use SHA256 digests for security.")
        print()
        print("To fix:")
        print("  1. Pull the image: docker pull <image:tag>")
        print(
            "  2. Get digest: docker inspect <image:tag> | jq -r '.[0].RepoDigests[0]'"
        )
        print("  3. Update workflow:")
        print()
        print("Example:")
        print("  services:")
        print("    redis:")
        print(f"{Colors.RED}      image: redis:alpine  # Bad{Colors.NC}")
        print(
            f"{Colors.GREEN}      image: redis:alpine@sha256:...  # Good{Colors.NC}"
        )
        return 1

    print()
    print(f"{Colors.GREEN}✅ All workflow images properly pinned{Colors.NC}")
    return 0


if __name__ == "__main__":
    sys.exit(main())