USB / TCP Proxy

#!/usr/bin/env python

# Use the Python libusb wrapper to talk to Marvell printer Embedded Web Server
# over USB.
#
# Work in progress. A bit messy. Partial HTTP proxy implemented.
#
# https://pypi.python.org/pypi/libusb1
# https://github.com/vpelletier/python-libusb1
#
# davep 23-Apr-2014

import os
import logging
import select
import Queue
import socket
import time
from httplib import HTTPResponse
from StringIO import StringIO
import libusb1
import usb1
import pickle

import hexdump

nlog = logging.getLogger(__name__)

logfilename = "ewstraffic.dat"

# prn1's EWS
prn1 = {
    "vid" : 0x3f0,
    "pid" : 0x5617,
    "interface" : 1,
    "endpoint": 2
}

prn2 = {
    "vid": 0x1286,
    "pid": 0x0100,
    "interface":1,
    "endpoint":2
}

# TODO get the vid/pid/iface/endpoint from cmdline
#dev = prn1
dev = prn2

tcp_port = 8080

# USB timeouts are in milliseconds
SECONDS=1000

queue_timeout_seconds = 2

# big enough?
tcp_sock_recv_size = 65536

# Big enough? Depends on the firmware I guess. We risk overflows if too small.
bulk_read_buffer_size = 65535
#bulk_read_buffer_size = 64*1024*10

class TCPSocketClosed(Exception):
    def __init__(self,sock):
        Exception.__init__(self)
        self.sock = sock

class USBClosed(Exception):
    def __init__(self,interface):
        Exception.__init__(self)
        self.interface = interface

def on_transfer_complete(transfer):
    nlog.info("on_transfer_complete() status={0} endpoint={1:#x}".format(
                transfer.getStatus(), transfer.getEndpoint() ))

    # user data should be a queue to send the transfer
    queue = transfer.getUserData()
    nlog.debug("on_transfer_complete() put transfer into queue")
    queue.put(transfer,queue_timeout_seconds)

def on_hotplug(ctx,dev):
    nlog.info("on_hotplug()")

    # from the usb1.py source:
    #    Callback must return whether it must be unregistered (any true value
    #    to be unregistered, any false value to be kept registered).
    return False

def make_server_socket():
    sock = socket.socket(socket.AF_INET,socket.SOCK_STREAM)
    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    sock.bind(('',tcp_port))
    sock.listen(0)
    return sock

def read_tcp_write_to_usb(sock,usb_handle,transfer_queue):
    # Read from TCP.
    # Write to USB.
    # That's about it.
    nlog.debug("sock.recv()")
    buf = sock.recv(tcp_sock_recv_size)
    nlog.debug("conn.recv() buflen={0}".format(len(buf)))
    if len(buf)==0:
        nlog.debug("sock={0} closed".format(sock))
        raise TCPSocketClosed(sock)

    nlog.debug(hexdump.dump(buf,16))

    with open(logfilename,"ab") as outfile:
        outfile.write("@@@ read data socket={0} at {1}".format(sock.fileno(),time.ctime()))
        outfile.write(buf)

    # send the data to USB
    write_transfer = usb_handle.getTransfer()
    write_transfer.setBulk( dev["endpoint"],
                            buf,
                            on_transfer_complete,
                            transfer_queue,
                            2*SECONDS)
    retcode = write_transfer.submit()
    nlog.debug("write submit() returned {0}".format(retcode))

def clean_transfer_queue( transfer_queue ) :
    # This funcion called at interrupt and program exit.  Clean/free any
    # outstanding transfers that might be hanging around in our transfer_queue.
    while not transfer_queue.empty():
        transfer = transfer_queue.get()
        transfer_type = transfer.getType()
        transfer_endpoint = transfer.getEndpoint()
        nlog.debug("transfer={0} type={1} endpoint={2:#x}".format(transfer,transfer_type,transfer_endpoint))

        transfer.close()
        del transfer

# http://pythonwise.blogspot.com/2010/02/parse-http-response.html
class FakeSocket(StringIO):
    def makefile(self, *args, **kw):
        return self

def parse_http(buf):
    if len(buf)==0 :
        # ignore
        return

#    with open("buf.pkl","wb") as outfile:
#        pickle.dump(buf,outfile)

    # http://pythonwise.blogspot.com/2010/02/parse-http-response.html
    fakesock = FakeSocket(buf)
    response = HTTPResponse(fakesock)
    response.begin()

    nlog.debug("HTTP headers={0}".format(response.getheaders()))

    content_length = response.getheader('content-length')

    headers,body = buf.split("\r\n\r\n",1)
    nlog.debug("len(body)={0} content-length={1}".format(len(body),content_length))

def poll_transfer_queue( transfer_queue, conn ) :
    # When the USBTransfer callback is called, the finished transfer structure
    # is put into a queue. This function will read those structures from the
    # queue.
    #
    # Read transfers are written to the TCP socket.
    # Write transfers are freed.

    nlog.debug("queue status empty={0} full={1}".format(
                    transfer_queue.empty(), transfer_queue.full() ) )

    # handle USBTransfer objects from the callback
    while not transfer_queue.empty():
        transfer = transfer_queue.get()
        transfer_type = transfer.getType()
        transfer_endpoint = transfer.getEndpoint()
        nlog.debug("transfer={0} type={1} endpoint={2:#x}".format(transfer,transfer_type,transfer_endpoint))

        if transfer_type==libusb1.LIBUSB_TRANSFER_TYPE_BULK:
            nlog.debug("transfer_type=BULK")
            if transfer_endpoint & libusb1.LIBUSB_ENDPOINT_IN :
                nlog.debug("transfer_endpoint=BULK_IN")
                # read data read
                full_buffer = transfer.getBuffer()
                buflen = transfer.getActualLength()
                nlog.debug("BULK_IN buflen={0}".format(buflen))

                if buflen==0:
                    transfer.close()
                    del transfer
                    raise USBClosed(transfer_endpoint)

                buf = full_buffer[:buflen]

                nlog.debug(hexdump.dump(buf,16))

                # parse the HTTP to find the connection boundaries so we can
                # multiplex multiple TCP sockets through the same USB
                # connection
                parse_http(buf)

                # write to TCP socket
                if conn:
                    conn.send(buf)
                    nlog.debug("wrote usb data to socket {0}".format(conn.fileno()))

                # save to our logfile
                with open(logfilename,"ab") as outfile:
                    outfile.write("@@@ bulk in at {0}".format(time.ctime()))
                    outfile.write(buf)

                # resubmit the read (we use the same single read transfer
                # repeatedly)
                # TODO might need multiple read transfers outstanding at some
                # point
                retcode = transfer.submit()
                nlog.debug("read re-submit() returned {0}".format(retcode))

            else:
                # write transfer completed; release it
                # (tiny optimization would be to save this and re-use it for
                # the next write)
                nlog.debug("transfer_endpoint=BULK_OUT")
                transfer.close()
                del transfer
        else:
            assert 0, transfer_type

def main():
    sock = make_server_socket()
    nlog.info("tcp socket={0} fd={1}".format(sock,sock.fileno()))

    ctx = usb1.USBContext()
    ctx.setDebug(4)

    prn = None
    dlist = ctx.getDeviceList( skip_on_error=True)
    for d in dlist :
        vid,pid = d.getVendorID(),d.getProductID()
        nlog.debug("{0:#06x} {1:#06x}".format(vid,pid))

        if vid==dev["vid"]:
            prn = d

    nlog.info("serial_num={0}".format(prn.getSerialNumber()))

    handle = prn.open()

    kernel_active = handle.kernelDriverActive(dev["interface"])
    nlog.debug("kernelDriverActive() returned {0}".format(kernel_active))

    nlog.info("claim interface={0}".format(dev["interface"]))
    handle.claimInterface(dev["interface"])

    kernel_active = handle.kernelDriverActive(dev["interface"])
    nlog.debug("kernelDriverActive() returned {0}".format(kernel_active))

    http_get = "GET / HTTP/1.1\r\n\r\n"

    # synchronous transfer to test interface
    if 0:
        retval = handle.bulkWrite(dev["endpoint"],http_get)
        nlog.debug("bulkWrite returned {0} ".format(retval))

        buf = handle.bulkRead(dev["endpoint"],bulk_read_buffer_size,2*SECONDS)
        nlog.debug("bulkRead return buf len={0}".format(len(buf)))
        with open(logfilename,"ab") as outfile:
            outfile.write("@@@ bulk in at {0}".format(time.ctime()))
            outfile.write(buf)

    # Use async transfers. The transfer callback will send transfer
    # structures back to main loop via this queue.
    transfer_queue = Queue.Queue()

    # send an asynchronous write to test interface
    write_transfer = handle.getTransfer()
#    transfer.setBulk(libusb1.LIBUSB_ENDPOINT_OUT,http_get,on_transfer_complete,evt,2*SECONDS)
    write_transfer.setBulk(dev["endpoint"],http_get,on_transfer_complete,transfer_queue,2*SECONDS)
    retcode = write_transfer.submit()
    nlog.debug("write submit() returned {0}".format(retcode))

    # we will re-use this transfer for all reads
    # TODO might need multiple read transfers outstanding at some point
    read_transfer = handle.getTransfer()
    read_transfer.setBulk( dev["endpoint"]|libusb1.LIBUSB_ENDPOINT_IN,
                            bulk_read_buffer_size,
                            on_transfer_complete,
                            transfer_queue,
                            0 )
    read_transfer.submit()
    nlog.debug("read submit() returned {0}".format(retcode))

    # add the USB fds to the poll()
    usbfds = ctx.getPollFDList()
    poller = select.poll()
    for fd,events in usbfds:
        nlog.debug("fd={0} event={1:#x}".format(fd,events))
        poller.register( fd, events )

    # add the listening TCP socket
    poller.register(sock,select.POLLIN
                        |select.POLLOUT
                        |select.POLLPRI
                        |select.POLLERR
                        |select.POLLHUP
                        |select.POLLNVAL)

    conn = None
    connfd = -1

    # TODO catch signal, cleanly exit this loop
    quit = False
    while not quit:
        timeout = ctx.getNextTimeout()
        if not timeout:
            # If libusb doesn't want a timeout, do I need a timeout?
            # If there are any outstanding transfers, need to peek at the queue
            # again soon. There is a race condition (maybe?) between the events
            # happening and the callback pushing the data into the queue.
#            if not transfer_queue.empty() :
#                timeout = 1*SECONDS
            pass

        nlog.debug("sleeping for {0} on poll()".format(timeout))

        # Stay a while and listen.
        try :
            evt = poller.poll(timeout)
        except KeyboardInterrupt:
            nlog.info("quitting on keyboard interrupt")
            quit = True
            break

        nlog.debug("awake! evt={0}".format(evt))

        if not evt:
            # timeout
            nlog.debug("timeout; handle events")
            ctx.handleEventsTimeout(1)

        for fd,events in evt :
            nlog.debug("fd={0} event={1:#x} sockfd={2} connfd={3}".format(fd,events,sock.fileno(),connfd))
            if fd==sock.fileno():
                # new TCP connection on our listen socket
                if events & select.POLLIN:
                    if conn is None:
                        (conn,addr) = sock.accept()
                        connfd = conn.fileno()
                        nlog.debug("sock.accept() conn={0} connfd={1} addr={2}".format(conn,connfd,addr))
                        poller.register(conn,select.POLLIN
                                            |0 # select.POLLOUT
                                            |select.POLLPRI
                                            |select.POLLERR
                                            |select.POLLHUP
                                            |select.POLLNVAL)
                    else :
                        # I can't find any way to block Python from
                        # allowing a second connection. (No, listen(1)
                        # doesn't work.)  USB will probably get upset if I
                        # mix HTTP from two separate connections. For now,
                        # slam the socket closed. Webpage is missing some
                        # elements. >.<
                        #
                        # TODO find a way to allow 2+ connections but force
                        # them serialized through USB.
                        nlog.error("client attempted second socket open; slam door in face")
                        sock.accept()[0].close()
                else:
                    # not sure if this can happen but put a bucket of glass
                    # here to break its fall in case it does
                    nlog.debug("wtf?")
                    assert 0

            elif fd==connfd:
                # read data from TCP socket, write back to USB
                if events & select.POLLIN :
                    try:
                        read_tcp_write_to_usb(conn,handle,transfer_queue)
                    except TCPSocketClosed:
                        poller.unregister(conn)
                        conn.close()
                        conn = None
                        connfd = -1
                else:
                    # How can this happen? Not sure. Catch it if it does.
                    assert 0,events
            else :
                # assume a USB event
                nlog.debug("ctx.handleEvents()")
                ctx.handleEvents()

        # end for evt

        # check the USB transfers from our USB read/write callback
        try :
            poll_transfer_queue(transfer_queue,conn)
        except USBClosed:
            quit = True
            break

    # end while(1)

    # From the docs:
    # Note: cancellation happens asynchronously, so you must wait for
    #   LIBUSB_TRANSFER_CANCELLED.
    nlog.debug("read_transfer isSubmitted={0}".format(read_transfer.isSubmitted()))

    clean_transfer_queue(transfer_queue)
    nlog.info("canceling read_transfer")
    if read_transfer.isSubmitted():
        read_transfer.cancel()
        while 1:
            nlog.info("waiting for read transfer to cancel...")
            ctx.handleEventsTimeout(1)
            transfer = None
            try:
                transfer = transfer_queue.get(True,1)
            except Queue.Empty:
                pass
            if transfer:
                nlog.debug("transfer={0} type={1} endpoint={2:#x}".format(
                            transfer,
                            transfer.getType(),
                            transfer.getEndpoint()))
                transfer.close()
                del transfer
                break
    clean_transfer_queue(transfer_queue)

    if conn:
        poller.unregister(conn)
        conn.close()
    poller.unregister(sock)
    sock.close()

    nlog.debug("release interface {0}".format(dev["interface"]))
    handle.releaseInterface(dev["interface"])

    nlog.debug("close usb handle")
    handle.close()

    # will this wake up the polling thread?
    ctx.exit()

    sock.close()

if __name__=='__main__':
#    fmt = "%(message)s"
    fmt = "%(filename)s %(lineno)d %(name)s %(message)s"
#    logging.basicConfig( level=logging.INFO, format=fmt )
    logging.basicConfig( level=logging.DEBUG, format=fmt )

    # log to file and stderr/stdout
    # http://stackoverflow.com/questions/13733552/logger-configuration-to-log-to-file-and-print-to-stdout
    formatter = logging.Formatter("%(asctime)s:%(name)s:%(levelname)s:%(message)s")

    # log to file, too (useful during dev)
    filehandler = logging.FileHandler("ewsproxy.log")
    filehandler.setLevel(logging.DEBUG)
    filehandler.setFormatter(formatter)
    nlog.addHandler(filehandler)

    main()