#!/usr/bin/env python3
"""
Test suite for CVE-2026-42826 Azure DevOps Information Disclosure Exploit
Validates exploit logic, secret scanning, URL construction, and report generation.
"""

import json
import os
import re
import sys
import tempfile
from urllib.error import HTTPError
from unittest.mock import patch, MagicMock
from io import StringIO

# Add parent dir to path
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from exploit import AzureDevOpsInfoDisclosure, log

passed = 0
failed = 0
results = []

def test(name, condition, detail=""):
    global passed, failed
    status = "PASS" if condition else "FAIL"
    color = "\033[32m" if condition else "\033[31m"
    print(f"  {color}[{status}]\033[0m {name}{f' — {detail}' if detail else ''}")
    if condition:
        passed += 1
    else:
        failed += 1
    results.append({"name": name, "passed": condition, "detail": detail})


def cleanup(exploit_obj):
    """Reset extracted data between tests"""
    exploit_obj.extracted_data = {
        "projects": [], "pipelines": [], "variable_groups": [],
        "service_connections": [], "build_logs": [], "repositories": [],
        "sensitive_data": [],
    }
    exploit_obj.findings = []


# ─── Tests ───

def test_initialization():
    """Test exploit initialization"""
    e = AzureDevOpsInfoDisclosure("https://dev.azure.com/testorg")
    test("Organization URL stored correctly", e.org_url == "https://dev.azure.com/testorg")
    test("Trailing slash removed", AzureDevOpsInfoDisclosure("https://dev.azure.com/testorg/").org_url == "https://dev.azure.com/testorg")
    test("PAT defaults to None", e.pat is None)
    test("API version defaults to 7.1", e.api_version == "7.1")
    test("Timeout defaults to 15", e.timeout == 15)
    test("Debug defaults to False", e.debug is False)
    test("Extracted data initialized", "projects" in e.extracted_data)
    test("Findings list initialized", isinstance(e.findings, list))


def test_api_url_construction():
    """Test API URL construction"""
    e = AzureDevOpsInfoDisclosure("https://dev.azure.com/myorg")
    url = e._api_url("_apis/projects")
    test("API URL includes org", "dev.azure.com/myorg" in url)
    test("API URL includes version", "api-version=7.1" in url)
    test("API URL includes path", "_apis/projects" in url)


def test_secret_scanning():
    """Test sensitive pattern detection"""
    e = AzureDevOpsInfoDisclosure("https://dev.azure.com/testorg")

    # Test AWS key detection (AKIA + 16 alphanumeric chars = 20 total)
    e._scan_for_secrets("AKIA1234567890ABCDEF some text", "test:aws")
    test("AWS Access Key detected", len(e.extracted_data["sensitive_data"]) == 1)
    test("AWS key value captured", e.extracted_data["sensitive_data"][0]["value"] == "AKIA1234567890ABCDEF")
    cleanup(e)

    # Test GitHub token detection
    e._scan_for_secrets("ghp_1234567890abcdefghijklmnopqrstuvwxyz1234", "test:gh")
    test("GitHub Token detected", len(e.extracted_data["sensitive_data"]) == 1)
    cleanup(e)

    # Test JWT detection
    e._scan_for_secrets("eyJhbGciOiJIUzI1NiJ9.eyJzdWIiOiIxMjM0NTY3ODkwIn0.SflKxwRJSMeKKF2QT4fwpMeJf36POk6yJV_adQssw5c", "test:jwt")
    test("JWT Token detected", len(e.extracted_data["sensitive_data"]) == 1)
    cleanup(e)

    # Test Azure subscription reference
    e._scan_for_secrets("azureSubscription: MyProductionSub", "test:azure")
    test("Azure Subscription reference detected", len(e.extracted_data["sensitive_data"]) == 1)
    cleanup(e)

    # Test Kubernetes service connection
    e._scan_for_secrets("kubernetesServiceConnection: prod-cluster", "test:k8s")
    test("Kubernetes Service Connection detected", len(e.extracted_data["sensitive_data"]) == 1)
    cleanup(e)

    # Test email detection
    e._scan_for_secrets("Contact: admin@company.com", "test:email")
    test("Email address detected", len(e.extracted_data["sensitive_data"]) == 1)
    cleanup(e)

    # Test no false positives
    e._scan_for_secrets("This is a normal log line with no secrets", "test:clean")
    test("No false positives on clean text", len(e.extracted_data["sensitive_data"]) == 0)
    cleanup(e)

    # Test multiple patterns in one scan
    e._scan_for_secrets("""
    azureSubscription: ProdSub
    dockerRegistryServiceConnection: myRegistry
    ARM_CLIENT_ID=12345
    """, "test:multi")
    test("Multiple patterns detected in one scan", len(e.extracted_data["sensitive_data"]) >= 2)
    cleanup(e)


def test_report_generation():
    """Test report generation"""
    e = AzureDevOpsInfoDisclosure("https://dev.azure.com/testorg", pat="fake_pat")

    # Populate some data
    e.extracted_data["projects"].append({"name": "TestProj", "visibility": "public"})
    e.extracted_data["pipelines"].append({"name": "build-pipeline", "id": 1})
    e.extracted_data["variable_groups"].append({"name": "prod-vars", "variables": {}})
    e.extracted_data["service_connections"].append({"name": "azure-prod", "type": "azurerm"})
    e.extracted_data["repositories"].append({"name": "app-repo"})
    e.extracted_data["build_logs"].append({"build_id": 1, "log_id": 1})
    e.extracted_data["sensitive_data"].append({"type": "AWS Access Key", "value": "AKIA1234567890ABCDEF"})
    e.findings.append({"type": "public_project_exposure", "severity": "CRITICAL", "count": 1})

    report = e._generate_report()

    test("Report has CVE ID", report["cve"] == "CVE-2026-42826")
    test("Report has CVSS score", report["cvss"] == "10.0 Critical")
    test("Report has CWE", report["cwe"] == "CWE-200")
    test("Report has target URL", report["target"] == "https://dev.azure.com/testorg")
    test("Report has timestamp", "timestamp" in report)
    test("Report shows authenticated", report["authenticated"] is True)
    test("Report summary has project count", report["summary"]["projects"] == 1)
    test("Report summary has pipeline count", report["summary"]["pipelines"] == 1)
    test("Report summary has VG count", report["summary"]["variable_groups"] == 1)
    test("Report summary has SC count", report["summary"]["service_connections"] == 1)
    test("Report summary has repo count", report["summary"]["repositories"] == 1)
    test("Report summary has build log count", report["summary"]["build_logs"] == 1)
    test("Report summary has sensitive data count", report["summary"]["sensitive_data_points"] == 1)
    test("Report has findings", len(report["findings"]) == 1)
    test("Report has extracted_data", "extracted_data" in report)


def test_report_save():
    """Test report saving to file"""
    e = AzureDevOpsInfoDisclosure("https://dev.azure.com/testorg")
    report = e._generate_report()

    with tempfile.NamedTemporaryFile(suffix=".json", delete=False, mode='w') as f:
        tmpfile = f.name

    try:
        e.save_report(report, tmpfile)
        test("Report file created", os.path.exists(tmpfile))

        with open(tmpfile) as f:
            loaded = json.load(f)
        test("Report file is valid JSON", "cve" in loaded)
        test("Report file has correct CVE", loaded["cve"] == "CVE-2026-42826")
    finally:
        os.unlink(tmpfile)


def test_mocked_api_request():
    """Test API request with mocked HTTP"""
    e = AzureDevOpsInfoDisclosure("https://dev.azure.com/testorg")

    # Mock successful response
    mock_response = MagicMock()
    mock_response.getcode.return_value = 200
    mock_response.read.return_value = b'{"value": [{"id": "1", "name": "TestProj", "visibility": "public"}]}'
    mock_response.__enter__ = MagicMock(return_value=mock_response)
    mock_response.__exit__ = MagicMock(return_value=False)

    with patch('exploit.urlopen', return_value=mock_response):
        status, data = e._make_request("https://dev.azure.com/testorg/_apis/projects?api-version=7.1")

    test("Mocked request returns 200", status == 200)
    test("Mocked request returns parsed JSON", data is not None)
    test("Mocked request returns project list", "value" in data)
    test("Mocked project has correct name", data["value"][0]["name"] == "TestProj")


def test_mocked_401():
    """Test API request with 401 response"""
    e = AzureDevOpsInfoDisclosure("https://dev.azure.com/testorg")

    from io import BytesIO
    mock_fp = BytesIO(b'{"message": "Unauthorized"}')
    mock_error = HTTPError(
        url="https://dev.azure.com/testorg/_apis/projects?api-version=7.1",
        code=401, msg='Unauthorized', hdrs=None, fp=mock_fp,
    )
    with patch('exploit.urlopen', side_effect=mock_error):
        status, data = e._make_request("https://dev.azure.com/testorg/_apis/projects?api-version=7.1")

    test("401 response returns status 401", status == 401)
    test("401 response returns error data", data is not None)


def test_mocked_project_enumeration():
    """Test project enumeration with mocked API"""
    e = AzureDevOpsInfoDisclosure("https://dev.azure.com/testorg")

    mock_response = MagicMock()
    mock_response.getcode.return_value = 200
    mock_response.read.return_value = json.dumps({
        "value": [
            {"id": "1", "name": "PublicProj", "visibility": "public", "description": "A public project"},
            {"id": "2", "name": "PrivateProj", "visibility": "private", "description": "A private project"},
        ]
    }).encode()
    mock_response.__enter__ = MagicMock(return_value=mock_response)
    mock_response.__exit__ = MagicMock(return_value=False)

    with patch('exploit.urlopen', return_value=mock_response):
        projects = e.enumerate_projects()

    test("Enumeration returns 2 projects", len(projects) == 2)
    test("Extracted data has 2 projects", len(e.extracted_data["projects"]) == 2)
    test("First project is public", e.extracted_data["projects"][0]["visibility"] == "public")
    test("Second project is private", e.extracted_data["projects"][1]["visibility"] == "private")


def test_mocked_variable_groups():
    """Test variable group harvesting with mocked API"""
    e = AzureDevOpsInfoDisclosure("https://dev.azure.com/testorg")

    mock_response = MagicMock()
    mock_response.getcode.return_value = 200
    mock_response.read.return_value = json.dumps({
        "value": [{
            "id": 1,
            "name": "prod-secrets",
            "description": "Production variables",
            "variables": {
                "API_KEY": {"value": "sk-12345", "isSecret": False},
                "DB_PASSWORD": {"isSecret": True},
                "ENV": {"value": "production", "isSecret": False},
            }
        }]
    }).encode()
    mock_response.__enter__ = MagicMock(return_value=mock_response)
    mock_response.__exit__ = MagicMock(return_value=False)

    with patch('exploit.urlopen', return_value=mock_response):
        e.harvest_variable_groups("TestProj")

    test("Variable group extracted", len(e.extracted_data["variable_groups"]) == 1)
    vg = e.extracted_data["variable_groups"][0]
    test("VG name correct", vg["name"] == "prod-secrets")
    test("VG has 3 variables", len(vg["variables"]) == 3)
    test("API_KEY is non-secret", vg["variables"]["API_KEY"]["secret"] is False)
    test("API_KEY value captured", vg["variables"]["API_KEY"]["value"] == "sk-12345")
    test("DB_PASSWORD is secret", vg["variables"]["DB_PASSWORD"]["secret"] is True)
    test("DB_PASSWORD value masked", vg["variables"]["DB_PASSWORD"]["value"] == "***MASKED***")


def test_mocked_service_connections():
    """Test service connection extraction with mocked API"""
    e = AzureDevOpsInfoDisclosure("https://dev.azure.com/testorg")

    mock_response = MagicMock()
    mock_response.getcode.return_value = 200
    mock_response.read.return_value = json.dumps({
        "value": [{
            "id": "conn-1",
            "name": "azure-prod",
            "type": "azurerm",
            "url": "https://management.azure.com/",
            "data": {
                "subscriptionId": "12345-67890",
                "subscriptionName": "Production",
                "tenantId": "tenant-123",
            }
        }]
    }).encode()
    mock_response.__enter__ = MagicMock(return_value=mock_response)
    mock_response.__exit__ = MagicMock(return_value=False)

    with patch('exploit.urlopen', return_value=mock_response):
        e.extract_service_connections("TestProj")

    test("Service connection extracted", len(e.extracted_data["service_connections"]) == 1)
    sc = e.extracted_data["service_connections"][0]
    test("SC name correct", sc["name"] == "azure-prod")
    test("SC type correct", sc["type"] == "azurerm")
    test("SC has subscription ID", sc["data"]["subscriptionId"] == "12345-67890")
    test("SC has tenant ID", sc["data"]["tenantId"] == "tenant-123")


def test_public_exposure_detection():
    """Test public exposure check"""
    e = AzureDevOpsInfoDisclosure("https://dev.azure.com/testorg")

    mock_response = MagicMock()
    mock_response.getcode.return_value = 200
    mock_response.read.return_value = json.dumps({
        "value": [
            {"id": "1", "name": "PublicProj", "visibility": "public"},
        ]
    }).encode()
    mock_response.__enter__ = MagicMock(return_value=mock_response)
    mock_response.__exit__ = MagicMock(return_value=False)

    with patch('exploit.urlopen', return_value=mock_response):
        has_access = e.check_public_exposure()

    test("Public exposure detected", has_access is True)
    test("Finding recorded", len(e.findings) == 1)
    test("Finding is CRITICAL", e.findings[0]["severity"] == "CRITICAL")
    test("Finding has project name", "PublicProj" in e.findings[0]["projects"])


def test_no_public_exposure():
    """Test when no public projects exist"""
    e = AzureDevOpsInfoDisclosure("https://dev.azure.com/testorg")

    mock_response = MagicMock()
    mock_response.getcode.return_value = 200
    mock_response.read.return_value = json.dumps({
        "value": [
            {"id": "1", "name": "PrivateProj", "visibility": "private"},
        ]
    }).encode()
    mock_response.__enter__ = MagicMock(return_value=mock_response)
    mock_response.__exit__ = MagicMock(return_value=False)

    with patch('exploit.urlopen', return_value=mock_response):
        has_access = e.check_public_exposure()

    test("Access available but no public projects", has_access is True)
    test("No findings recorded", len(e.findings) == 0)


# ─── Main ───

def main():
    print("\n" + "=" * 70)
    print("  CVE-2026-42826 — Test Suite")
    print("  Azure DevOps Information Disclosure Exploit")
    print("=" * 70 + "\n")

    test_initialization()
    print()

    test_api_url_construction()
    print()

    test_secret_scanning()
    print()

    test_report_generation()
    print()

    test_report_save()
    print()

    test_mocked_api_request()
    print()

    test_mocked_401()
    print()

    test_mocked_project_enumeration()
    print()

    test_mocked_variable_groups()
    print()

    test_mocked_service_connections()
    print()

    test_public_exposure_detection()
    print()

    test_no_public_exposure()
    print()

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

    if failed > 0:
        print("\n  FAILED TESTS:")
        for r in results:
            if not r["passed"]:
                print(f"    ✗ {r['name']} — {r['detail']}")
        sys.exit(1)
    else:
        print("\n  ALL TESTS PASSED")
        sys.exit(0)


if __name__ == "__main__":
    main()