#!/usr/bin/python
# $Id: usb_control.py 954 2014-07-11 22:11:37Z mwall $
# Copyright 2014 Matthew Wall, all rights reserved
# turn usb port power on/off and provide details about usb devices
#
# Credit to:
#   NIIBE Yutaka for tool_hub_ctrl.py
#   Vlado Handziski for hubcontrol.py

import optparse
import usb
import inspect

VERSION = '0.6'
USB_RT_HUB = (usb.TYPE_CLASS | usb.RECIP_DEVICE)
USB_RT_PORT = (usb.TYPE_CLASS | usb.RECIP_OTHER)
USB_PORT_FEAT_POWER = 8
USB_PORT_FEAT_INDICATOR = 22
USB_DIR_IN = 0x80
CMD_POWER = 'power'
CMD_INDICATOR = 'indicator'

def scan():
    busses = usb.busses()
    if not busses:
        raise Exception, "cannot access USB"

    for bus in busses:
        for dev in bus.devices:
            mfr = 'unknown'
            prd = 'unknown'
            if dev.deviceClass == usb.CLASS_HUB:
                print "hub at %s:%03d" % (bus.dirname, dev.devnum)
                h = dev.open()
                desc = None
                try:
                    if dev.iManufacturer:
                        mfr = h.getString(dev.iManufacturer, 30)
                    if dev.iProduct:
                        prd = h.getString(dev.iProduct, 30)
                    mfr_id = dev.idVendor
                    prd_id = dev.idProduct
                    desc = h.controlMsg(requestType=USB_DIR_IN | USB_RT_HUB,
                                        request=usb.REQ_GET_DESCRIPTOR,
                                        value=usb.DT_HUB << 8,
                                        index=0, buffer=1024, timeout=1000)
                    num_ports = 'unknown'
                    indicator_support = 'unknown'
                    power_switching = 'unknown'
                    if desc:
                        num_ports = desc[2]
                        if (desc[3] & 0x80) == 0:
                            indicator_support = 'none'
                        if (desc[3] & 0x03) == 0:
                            power_switching = 'ganged'
                        elif (desc[3] & 0x03) == 1:
                            power_switching = 'individual'
                        elif (desc[3] & 0x03) == 2 or (desc[3] & 0x03) == 3:
                            power_switching = 'none'
                    print '  id: %04x:%04x' % (mfr_id, prd_id)
                    print '  manufacturer: %s' % mfr
                    print '  product: %s' % prd
                    print '  num_ports: %s' % num_ports
                    print '  power_switching: %s' % power_switching
                    print '  indicator_support: %s' % indicator_support
                    if num_ports != 'unknown':
                        for i in range(num_ports):
                            print '  ' + get_port_status(h, i+1)
                finally:
                    del h
            else:
                print "device at %s:%03d" % (bus.dirname, dev.devnum)
                h = dev.open()
#                print inspect.getmembers(dev)
#                print inspect.getmembers(dev.configurations)
                try:
                    mfr = h.getString(dev.iManufacturer, 30)
                    prd = h.getString(dev.iProduct, 30)
                    mfr_id = dev.idVendor
                    prd_id = dev.idProduct
                finally:
                    del h
                print '  id: %04x:%04x' % (mfr_id, prd_id)
                print '  manufacturer: %s' % mfr
                print '  product: %s' % prd

def get_port_status(handle, port):
    status = handle.controlMsg(requestType=USB_RT_PORT | usb.ENDPOINT_IN,
                               request=usb.REQ_GET_STATUS,
                               value=0,
                               index=port, buffer=4, timeout=1000)
    msg = "port %d: %02x%02x.%02x%02x" % (port, status[3], status[2],
                                          status[1], status[0])
    if status[1] & 0x10:
        msg += ' indicator'
    if status[1] & 0x08:
        msg += ' test'
    if status[1] & 0x04:
        msg += ' highspeed'
    if status[1] & 0x02:
        msg += ' lowspeed'
    if status[1] & 0x01:
        msg += ' power'

    if status[0] & 0x10:
        msg += ' RESET'
    if status[0] & 0x08:
        msg += ' oc'
    if status[0] & 0x04:
        msg += ' suspend'
    if status[0] & 0x02:
        msg += ' enable'
    if status[0] & 0x01:
        msg += ' connect'

    return msg

def docmd(cmd, hub, port, flag):
    if cmd == CMD_POWER:
        request = usb.REQ_SET_FEATURE if flag else usb.REQ_CLEAR_FEATURE
        index = port
        value = USB_PORT_FEAT_POWER
    elif cmd == CMD_INDICATOR:
        request = usb.REQ_SET_FEATURE
        x = 1 if not flag else 0
        index = (x << 8) | port
        value = USB_PORT_FEAT_INDICATOR

    busses = usb.busses()
    if not busses:
        raise Exception, "cannot access USB"
    device = None
    for bus in busses:
        for dev in bus.devices:
            if dev.deviceClass == usb.CLASS_HUB:
                devid = "%s:%03d" % (bus.dirname, dev.devnum)
                if hub == devid:
                    device = dev

    if device is not None:
        handle = device.open()
        try:
            handle.controlMsg(requestType=USB_RT_PORT,
                              request=request,
                              value=value,
                              index=index, buffer=None, timeout=1000)    
        finally:
            del handle

desc = """Switch USB port power/indicator on or off"""

usage = """%prog [--scan | --hub H --port P (--power (0|1) | --indicator (0|1))] [--help]"""

def main():
    parser = optparse.OptionParser(description=desc, usage=usage)
    parser.add_option("--version", dest="version", action="store_true",
                      help="display application version")
    parser.add_option("--scan", dest="scan", action="store_true",
                      help="scan for devices")
    parser.add_option("--hub", dest="hub", type=str, metavar="H",
                      help="the hub to which commands will be sent")
    parser.add_option("--port", dest="port", type=int, metavar="P",
                      help="the port for which commands are directed")
    parser.add_option("--power", dest="power", type=int, metavar="FLAG",
                      help="turn power on or off")
    parser.add_option("--indicator", dest="indicator", type=int,metavar="FLAG",
                      help="turn indicator light on or off")
    (options, args) = parser.parse_args()
    if options.version:
        print VERSION
        exit(0)
    if options.hub is not None or options.port is not None:
        if options.hub is None:
            print "No hub specified"
            exit(1)
        if options.port is None:
            print "No port specified"
            exit(1)
        if options.power is None and options.indicator is None:
            print "No action specified"
            exit(1)
        if options.power is not None:
            docmd(CMD_POWER, options.hub, options.port, options.power)
        if options.indicator is not None:
            docmd(CMD_INDICATOR, options.hub, options.port, options.indicator)
    else:
        scan()

if __name__ == '__main__':
    main()