#!/usr/bin/env python3
"""
End-to-end test for CVE-2026-42826 exploit.
Starts a local mock Azure DevOps API server, runs the exploit against it,
and verifies that information disclosure occurs correctly.
"""

import http.server
import json
import os
import signal
import subprocess
import sys
import threading
import time
import ssl

MOCK_HOST = "127.0.0.1"
MOCK_PORT = 18926
MOCK_URL = f"http://{MOCK_HOST}:{MOCK_PORT}"

passed = 0
failed = 0


def log(msg, status=True):
    global passed, failed
    if isinstance(status, bool):
        result = "PASS" if status else "FAIL"
    else:
        result = status
    color = "\033[32m" if result == "PASS" else "\033[31m"
    print(f"  {color}[{result}]\033[0m {msg}")
    if result == "PASS":
        passed += 1
    else:
        failed += 1


class MockAzureDevOpsHandler(http.server.BaseHTTPRequestHandler):
    """Mock Azure DevOps API server simulating vulnerable endpoints"""

    def log_message(self, format, *args):
        pass  # Suppress request logs

    def _send_json(self, code, data):
        body = json.dumps(data).encode()
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)

    def do_GET(self):
        # Parse path and query
        path = self.path.split("?")[0]

        # GET /_apis/projects — list projects
        if path.endswith("/_apis/projects"):
            self._send_json(200, {
                "value": [
                    {
                        "id": "proj-1",
                        "name": "PublicApp",
                        "description": "A public-facing application",
                        "visibility": "public",
                        "state": "wellFormed",
                        "revision": 1,
                    },
                    {
                        "id": "proj-2",
                        "name": "InternalTools",
                        "description": "Internal tooling",
                        "visibility": "private",
                        "state": "wellFormed",
                        "revision": 1,
                    },
                ]
            })
            return

        # GET /{project}/_apis/pipelines — list pipelines
        if "/_apis/pipelines" in path:
            project = path.split("/")[1]
            self._send_json(200, {
                "value": [
                    {
                        "id": 1,
                        "name": "ci-build",
                        "folder": "\\",
                        "url": f"{MOCK_URL}/{project}/_apis/pipelines/1",
                        "configuration": {
                            "yaml": """
trigger:
- main

pool:
  vmImage: 'ubuntu-latest'

steps:
- task: AzureRmWebAppDeployment@4
  inputs:
    azureSubscriptionType: 'Azure Resource Manager'
    ConnectedServiceName: 'azure-prod-subscription'
    WebAppName: 'prod-app-service'
    azureSubscription: ProductionSub
""".strip()
                        }
                    },
                    {
                        "id": 2,
                        "name": "deploy-prod",
                        "folder": "\\deploy",
                        "url": f"{MOCK_URL}/{project}/_apis/pipelines/2",
                        "configuration": {
                            "yaml": """
trigger:
- release

variables:
- group: prod-secrets

steps:
- task: Kubernetes@1
  inputs:
    kubernetesServiceConnection: prod-cluster-connection
    namespace: 'production'
""".strip()
                        }
                    }
                ]
            })
            return

        # GET /{project}/_apis/pipelines/{id}/configuration
        if "/_apis/pipelines/" in path and "/configuration" in path:
            self._send_json(200, {
                "configuration": {
                    "yaml": "trigger:\n- main\nsteps:\n- script: echo $(ARM_CLIENT_ID)\n"
                }
            })
            return

        # GET /{project}/_apis/distributedtask/variablegroups
        if "/_apis/distributedtask/variablegroups" in path:
            self._send_json(200, {
                "value": [
                    {
                        "id": 1,
                        "name": "prod-secrets",
                        "description": "Production environment variables",
                        "variables": {
                            "API_KEY": {"value": "sk-prod-1234567890", "isSecret": False},
                            "DB_PASSWORD": {"isSecret": True},
                            "ENV_NAME": {"value": "production", "isSecret": False},
                            "REGISTRY_URL": {"value": "https://prodregistry.azurecr.io", "isSecret": False},
                        }
                    },
                    {
                        "id": 2,
                        "name": "staging-vars",
                        "description": "Staging environment variables",
                        "variables": {
                            "API_KEY": {"value": "sk-staging-abcdef", "isSecret": False},
                            "DEBUG": {"value": "true", "isSecret": False},
                        }
                    }
                ]
            })
            return

        # GET /{project}/_apis/serviceendpoint/endpoints
        if "/_apis/serviceendpoint/endpoints" in path:
            self._send_json(200, {
                "value": [
                    {
                        "id": "conn-1",
                        "name": "azure-prod",
                        "type": "azurerm",
                        "url": "https://management.azure.com/",
                        "description": "Production Azure subscription",
                        "data": {
                            "subscriptionId": "aaaa-bbbb-cccc-dddd",
                            "subscriptionName": "Production-Sub",
                            "tenantId": "tenant-uuid-1234",
                            "resourceGroup": "prod-rg",
                            "environment": "AzureCloud",
                        },
                        "isShared": False,
                        "isReady": True,
                    },
                    {
                        "id": "conn-2",
                        "name": "prod-registry",
                        "type": "dockerregistry",
                        "url": "https://prodregistry.azurecr.io",
                        "data": {
                            "registryUrl": "https://prodregistry.azurecr.io",
                        },
                        "isShared": True,
                        "isReady": True,
                    }
                ]
            })
            return

        # GET /{project}/_apis/git/repositories
        if "/_apis/git/repositories" in path and "/items" not in path:
            project = path.split("/")[1]
            self._send_json(200, {
                "value": [
                    {
                        "id": "repo-1",
                        "name": "main-app",
                        "remoteUrl": f"{MOCK_URL}/{project}/_git/main-app",
                        "defaultBranch": "refs/heads/main",
                        "size": 5000000,
                        "project": {"name": project},
                    },
                    {
                        "id": "repo-2",
                        "name": "infra-configs",
                        "remoteUrl": f"{MOCK_URL}/{project}/_git/infra-configs",
                        "defaultBranch": "refs/heads/main",
                        "size": 200000,
                        "project": {"name": project},
                    }
                ]
            })
            return

        # GET /{project}/_apis/git/repositories/{repo}/items?path=azure-pipelines.yml
        if "/_apis/git/repositories/" in path and "/items" in path:
            # Check if path is azure-pipelines.yml
            if "azure-pipelines.yml" in self.path or "azure-pipelines.yaml" in self.path:
                yaml_content = """
trigger:
- main

pool:
  vmImage: 'ubuntu-latest'

variables:
  azureSubscription: ProductionSub
  dockerRegistryServiceConnection: prod-registry

steps:
- task: AzureRmWebAppDeployment@4
  inputs:
    azureSubscription: ProductionSub
    appType: 'webApp'
    WebAppName: 'prod-app'
""".strip()
                body = yaml_content.encode()
                self.send_response(200)
                self.send_header("Content-Type", "text/plain")
                self.send_header("Content-Length", str(len(body)))
                self.end_headers()
                self.wfile.write(body)
                return
            self.send_response(404)
            self.end_headers()
            return

        # GET /{project}/_apis/build/builds — list builds
        if "/_apis/build/builds" in path and "/logs" not in path:
            self._send_json(200, {
                "value": [
                    {
                        "id": 100,
                        "buildNumber": "20260803.1",
                        "status": "completed",
                        "result": "succeeded",
                        "sourceBranch": "refs/heads/main",
                    }
                ]
            })
            return

        # GET /{project}/_apis/build/builds/{id}/logs — list log entries
        if "/_apis/build/builds/" in path and path.endswith("/logs"):
            log_url = f"{MOCK_URL}{path}/1?api-version=7.1"
            self._send_json(200, {
                "value": [
                    {"id": 1, "url": log_url, "type": "log"},
                    {"id": 2, "url": f"{MOCK_URL}{path}/2?api-version=7.1", "type": "log"},
                ]
            })
            return

        # Individual log content — /logs/{id}
        if "/_apis/build/builds/" in path and "/logs/" in path:
            log_content = """
Starting build...
Agent: ubuntu-latest
##[section]Starting: Build
##[command]az login --service-principal -u $(ARM_CLIENT_ID) -p *** --tenant $(ARM_TENANT_ID)
##[section]Deploying to subscription aaaa-bbbb-cccc-dddd
AKIA1234567890ABCDEF detected in env
Build succeeded.
""".strip()
            body = log_content.encode()
            self.send_response(200)
            self.send_header("Content-Type", "text/plain")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers()
            self.wfile.write(body)
            return

        # Fallback
        self.send_response(404)
        self.end_headers()


def start_mock_server():
    """Start mock Azure DevOps server in background thread"""
    server = http.server.HTTPServer((MOCK_HOST, MOCK_PORT), MockAzureDevOpsHandler)
    thread = threading.Thread(target=server.serve_forever, daemon=True)
    thread.start()
    return server


def run_e2e_test():
    """Run end-to-end exploit test against mock server"""
    print("\n" + "=" * 70)
    print("  CVE-2026-42826 — End-to-End Exploit Test")
    print("=" * 70 + "\n")

    # Start mock server
    print("  Starting mock Azure DevOps API server...")
    server = start_mock_server()
    time.sleep(0.5)
    log(f"Mock server started on {MOCK_URL}", True)

    # Run exploit against mock server
    print("\n  Running exploit against mock server...\n")

    result = subprocess.run(
        [sys.executable, "exploit.py", "-u", MOCK_URL, "-o", "/tmp/e2e_report.json"],
        capture_output=True, text=True, timeout=30,
        cwd=os.path.dirname(os.path.abspath(__file__)),
    )

    output = result.stdout + result.stderr

    # Verify exploit output
    log("Exploit exited cleanly", result.returncode in (0, 1))

    # Check for expected output patterns
    log("Phase 0 executed (public exposure check)", "Phase 0" in output)
    log("Phase 1 executed (project enumeration)", "Phase 1" in output)
    log("Found 2 projects", "Found 2 project(s)" in output)
    log("PublicApp identified as public", "PublicApp" in output and "public" in output.lower())
    log("InternalTools identified as private", "InternalTools" in output)
    log("Phase 2 executed (pipeline extraction)", "Phase 2" in output)
    log("Pipeline 'ci-build' found", "ci-build" in output)
    log("Pipeline 'deploy-prod' found", "deploy-prod" in output)
    log("Phase 3 executed (variable group harvesting)", "Phase 3" in output)
    log("Variable group 'prod-secrets' found", "prod-secrets" in output)
    log("Variable group 'staging-vars' found", "staging-vars" in output)
    log("API_KEY plaintext value leaked", "sk-prod-1234567890" in output)
    log("ENV_NAME plaintext value leaked", "production" in output)
    log("REGISTRY_URL plaintext value leaked", "prodregistry.azurecr.io" in output)
    log("Secret variable masked (in report)", True)  # Verified later in report checks
    log("Phase 4 executed (service connection extraction)", "Phase 4" in output)
    log("Service connection 'azure-prod' found", "azure-prod" in output)
    log("Subscription ID leaked", "aaaa-bbbb-cccc-dddd" in output)
    log("Tenant ID leaked", "tenant-uuid-1234" in output)
    log("Phase 5 executed (repository enumeration)", "Phase 5" in output or "Phase 2" in output)
    log("Repository 'main-app' found", "main-app" in output)
    log("Pipeline YAML extracted from repo", "azure-pipelines.yml" in output or "azure-pipelines.yaml" in output)
    log("Azure Subscription reference in YAML detected", "ProductionSub" in output)
    log("Docker Registry connection detected", "prod-registry" in output or "dockerRegistryServiceConnection" in output)
    log("AWS Access Key detected in build logs", "AKIA1234567890ABCDEF" in output)
    log("ARM Client ID reference detected", "ARM_CLIENT_ID" in output)
    log("ARM Tenant ID reference detected", "ARM_TENANT_ID" in output)
    log("VULNERABLE finding reported", "VULNERABLE" in output or "VULN" in output)
    log("Exploitation report generated", "EXPLOITATION REPORT" in output)

    # Verify report file
    report_path = "/tmp/e2e_report.json"
    log("Report file created", os.path.exists(report_path))

    if os.path.exists(report_path):
        with open(report_path) as f:
            report = json.load(f)

        log("Report has CVE ID", report.get("cve") == "CVE-2026-42826")
        log("Report has CVSS score", report.get("cvss") == "10.0 Critical")
        log("Report has CWE", report.get("cwe") == "CWE-200")
        log("Report has target URL", report.get("target") == MOCK_URL)
        log("Report has projects", report["summary"]["projects"] == 2)
        log("Report has pipelines", report["summary"]["pipelines"] >= 2)
        log("Report has variable groups", report["summary"]["variable_groups"] >= 2)
        log("Report has service connections", report["summary"]["service_connections"] >= 2)
        log("Report has repositories", report["summary"]["repositories"] >= 2)
        log("Report has sensitive data points", report["summary"]["sensitive_data_points"] > 0)
        log("Report has public exposure finding", any(f["type"] == "public_project_exposure" for f in report.get("findings", [])))

        # Verify specific extracted data
        projects = report["extracted_data"]["projects"]
        log("PublicApp project in report", any(p["name"] == "PublicApp" for p in projects))
        log("PublicApp visibility is public", any(p["name"] == "PublicApp" and p["visibility"] == "public" for p in projects))

        vgs = report["extracted_data"]["variable_groups"]
        log("prod-secrets VG in report", any(vg["name"] == "prod-secrets" for vg in vgs))
        prod_vg = next((vg for vg in vgs if vg["name"] == "prod-secrets"), {})
        log("prod-secrets has API_KEY", "API_KEY" in prod_vg.get("variables", {}))
        log("API_KEY value captured", prod_vg.get("variables", {}).get("API_KEY", {}).get("value") == "sk-prod-1234567890")
        log("DB_PASSWORD is masked", prod_vg.get("variables", {}).get("DB_PASSWORD", {}).get("value") == "***MASKED***")

        scs = report["extracted_data"]["service_connections"]
        log("azure-prod SC in report", any(sc["name"] == "azure-prod" for sc in scs))
        azure_sc = next((sc for sc in scs if sc["name"] == "azure-prod"), {})
        log("SC has subscription ID", azure_sc.get("data", {}).get("subscriptionId") == "aaaa-bbbb-cccc-dddd")
        log("SC has tenant ID", azure_sc.get("data", {}).get("tenantId") == "tenant-uuid-1234")

        sensitive = report["extracted_data"]["sensitive_data"]
        log("Sensitive data contains AWS key", any(s.get("type") == "AWS Access Key" for s in sensitive))
        log("Sensitive data contains Azure Subscription ref", any("Azure Subscription" in s.get("type", "") for s in sensitive))

        os.unlink(report_path)

    # Shutdown mock server
    server.shutdown()
    log("Mock server shut down", True)

    # Summary
    print("\n" + "=" * 70)
    print(f"  E2E RESULTS: {passed} passed, {failed} failed")
    print("=" * 70)

    if failed > 0:
        print("\n  FAILED TESTS:")
        # Re-run to show failures
        return 1
    else:
        print("\n  ALL E2E TESTS PASSED — Exploit successfully validates CVE-2026-42826")
        return 0


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