auto_patch.py

import os
import json
import docker
import subprocess

# -------------------------------
# Scan container image with Trivy
# -------------------------------
def scan_image(image_name):
    print(f"🔍 Scanning {image_name}...")
    result = subprocess.run(
        ["trivy", "image", "--format", "json", image_name],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True
    )
    if result.returncode != 0:
        print("❌ Trivy scan failed:", result.stderr)
        return None
    return json.loads(result.stdout)

# -------------------------------
# Parse vulnerabilities
# -------------------------------
def get_patchable_vulns(scan_result):
    patchable = []
    for res in scan_result.get("Results", []):
        for vuln in res.get("Vulnerabilities", []):
            if vuln.get("FixedVersion"):  # only if a fix exists
                patchable.append({
                    "pkg": vuln["PkgName"],
                    "installed": vuln["InstalledVersion"],
                    "fixed": vuln["FixedVersion"],
                    "cve": vuln["VulnerabilityID"],
                    "desc": vuln.get("Description", "")[:150]
                })
    return patchable

# -------------------------------
# Generate a patched Dockerfile
# -------------------------------
def generate_dockerfile(base_image, vulns):
    print("🛠 Generating patched Dockerfile...")

    lines = [f"FROM {base_image}", "RUN apk update || apt-get update -y"]

    for v in vulns:
        pkg = v["pkg"]

        # Alpine packages
        if base_image.startswith("alpine"):
            line = f"RUN apk add --no-cache --upgrade {pkg}"
        else:
            # Debian/Ubuntu
            line = f"RUN apt-get install -y --only-upgrade {pkg}"

        lines.append(line)

    lines.append("RUN rm -rf /var/lib/apt/lists/* || true")
    dockerfile = "Dockerfile.patched"

    with open(dockerfile, "w") as f:
        f.write("\n".join(lines))

    return dockerfile

# -------------------------------
# Rebuild image
# -------------------------------
def rebuild_image(dockerfile, new_image_name):
    print("🔨 Rebuilding patched image...")
    client = docker.from_env()
    image, logs = client.images.build(path=".", dockerfile=dockerfile, tag=new_image_name)
    for chunk in logs:
        if "stream" in chunk:
            print(chunk["stream"].strip())
    return image

# -------------------------------
# Main flow
# -------------------------------
if __name__ == "__main__":
    base_image = "alpine:3.12"   # 🔹 test image
    new_image = base_image.replace(":", "_") + "_patched"

    scan_result = scan_image(base_image)
    if not scan_result:
        exit(1)

    vulns = get_patchable_vulns(scan_result)
    if not vulns:
        print("✅ No patchable vulnerabilities found.")
        exit(0)

    print("⚠️ Patchable vulnerabilities found:")
    for v in vulns:
        print(f"{v['cve']} | {v['pkg']} {v['installed']} → {v['fixed']}")
        print(f"   ➡ {v['desc']}...\n")

    dockerfile = generate_dockerfile(base_image, vulns)
    patched = rebuild_image(dockerfile, new_image)

    print(f"✅ New patched image built: {new_image}")
    print("🔍 Re-scanning patched image...")
    rescan = scan_image(new_image)
    patched_vulns = get_patchable_vulns(rescan)
    if not patched_vulns:
        print("🎉 All patchable vulnerabilities fixed!")
    else:
        print("⚠️ Remaining patchable vulns:", patched_vulns)