scan_plan.py

import json
import os
import sys
from typing import List

from google import genai
from google.genai import types
from pydantic import BaseModel


# 1. Define the exact JSON structure we want Gemini to return
class SecurityFinding(BaseModel):
    resource_name: str
    severity: str  # CRITICAL, WARNING, or INFO
    vulnerability: str
    remediation_steps: str


class SecurityReport(BaseModel):
    status: str  # MUST be "PASS" or "FAIL"
    summary: str
    findings: List[SecurityFinding]


def audit_terraform_plan():
    # Verify API key
    if not os.environ.get("GEMINI_API_KEY"):
        print("❌ Error: GEMINI_API_KEY environment variable is not set.")
        sys.exit(1)

    # Load the real plan.json generated by the user
    plan_path = "plan.json"
    if not os.path.exists(plan_path):
        print(
            f"❌ Error: {plan_path} not found. Run 'terraform show -json tfplan.binary > plan.json' first."
        )
        sys.exit(1)

    print(f"📖 Reading {plan_path}...")
    with open(plan_path, "r") as file:
        plan_data = json.load(file)

    # Clean the plan slightly to save token space (focusing on resource changes)
    resource_changes = plan_data.get("resource_changes", [])

    print(f"Securing {len(resource_changes)} planned resource modifications...")

    # 2. Design the prompt
    prompt = f"""
    You are an automated DevSecOps Cloud Security Auditor. Your job is to analyze the following 
    Terraform execution plan JSON data and evaluate it against enterprise security best practices.

    Look specifically for:
    - Publicly accessible storage (S3 buckets, Azure Blobs, etc.)
    - Overly permissive networking rules (e.g., Ingress from 0.0.0.0/0 on sensitive ports like 22, 3389, 80, 443 unless justified)
    - Unencrypted resources or databases
    - Cleartext credentials/secrets

    If you find ANY CRITICAL security violations, you MUST set the status to "FAIL".
    If the infrastructure is safe or contains only minor issues, set the status to "PASS".

    RAW TERRAFORM PLAN RESOURCE CHANGES:
    {json.dumps(resource_changes, indent=2)}
    """

    print("🤖 Analyzing plan with Gemini Security Architect...")
    client = genai.Client()

    # 3. Request structured content matching our Pydantic model
    response = client.models.generate_content(
        model="gemini-2.5-flash",
        contents=prompt,
        config=types.GenerateContentConfig(
            response_mime_type="application/json",
            response_schema=SecurityReport,
            temperature=0.1,  # Low temperature ensures strict, analytical deterministic results
        ),
    )

    # 4. Parse the structured output
    report = SecurityReport.model_validate_json(response.text)

    # 5. Print a beautiful console report
    print("\n" + "=" * 50)
    print(f"🛡️  DEVSECOPS SECURITY VERDICT: {report.status}")
    print("=" * 50)
    print(f"Summary: {report.summary}\n")

    if report.findings:
        print(f"⚠️  Found {len(report.findings)} Security Vulnerabilities:")
        for idx, finding in enumerate(report.findings, 1):
            print(f"\n  [{idx}] Resource: {finding.resource_name}")
            print(f"      Severity: {finding.severity}")
            print(f"      Vulnerability: {finding.vulnerability}")
            print(f"      Fix: {finding.remediation_steps}")
    else:
        print("✅ No critical security violations detected.")

    print("=" * 50)

    # 6. ENFORCE THE GATEKEEPER RULE (Crucial for CI/CD pipelines!)
    if report.status == "FAIL":
        print("❌ Deployment Blocked: Critical security flaws must be resolved.")
        sys.exit(1)  # Forces automated tools like GitHub Actions to halt immediately
    else:
        print("🚀 Deployment Approved: Infrastructure compliant.")
        sys.exit(0)


if __name__ == "__main__":
    audit_terraform_plan()