add GetBattorPathFromPhoneSerial to find_usb_devices

devil/devil/utils/find_usb_devices.py

 #!/usr/bin/python
 # Copyright 2016 The Chromium Authors. All rights reserved.
 # Use of this source code is governed by a BSD-style license that can be
 # found in the LICENSE file.
 
 import re
 import sys
 import argparse
+import json
+import collections
 
+from devil import base_error
 from devil.utils import cmd_helper
+from devil.utils import usb_hubs
 from devil.utils import lsusb
 
 # Note: In the documentation below, "virtual port" refers to the port number
 # as observed by the system (e.g. by usb-devices) and "physical port" refers
 # to the physical numerical label on the physical port e.g. on a USB hub.
 # The mapping between virtual and physical ports is not always the identity
 # (e.g. the port labeled "1" on a USB hub does not always show up as "port 1"
 # when you plug something into it) but, as far as we are aware, the mapping
 # between virtual and physical ports is always the same for a given
 # model of USB hub. When "port number" is referenced without specifying, it
@@ skipping 25 matching lines @@
 
 def GetBattorList(device_tree_map):
   return [x for x in GetTTYList() if IsBattor(x, device_tree_map)]
 
 
 def IsBattor(tty_string, device_tree_map):
   (bus, device) = GetBusDeviceFromTTY(tty_string)
   node = device_tree_map[bus].FindDeviceNumber(device)
   return 'Future Technology Devices International' in node.desc
 
+def GetBattorSerialNumbers(device_tree_map):
+  for x in GetTTYList():
+    if IsBattor(x, device_tree_map):
+      (bus, device) = GetBusDeviceFromTTY(x)
+      devnode = device_tree_map[bus].FindDeviceNumber(device)
+      yield devnode.serial
+
+class BattorError(base_error.BaseError):
+  pass
 
 # Class to identify nodes in the USB topology. USB topology is organized as
 # a tree.
 class USBNode(object):
   def __init__(self):
     self._port_to_node = {}
 
   @property
   def desc(self):
     raise NotImplementedError
@@ skipping 178 matching lines @@
 
 
 _T_LINE_REGEX = re.compile(r'T:  Bus=(?P<bus>\d{2}) Lev=(?P<lev>\d{2}) '
                            r'Prnt=(?P<prnt>\d{2,3}) Port=(?P<port>\d{2}) '
                            r'Cnt=(?P<cnt>\d{2}) Dev#=(?P<dev>.{3}) .*')
 
 _S_LINE_REGEX = re.compile(r'S:  SerialNumber=(?P<serial>.*)')
 _LSUSB_BUS_DEVICE_RE = re.compile(r'^Bus (\d{3}) Device (\d{3}): (.*)')
 
 
 def GetBusNumberToDeviceTreeMap(fast=False):

[rnephew (Reviews Here), 2016/03/25 00:17:52]

Since everywhere seems to use fast=True, should that be the default?

   """Gets devices currently attached.
 
   Args:
     fast [bool]: whether to do it fast (only get description, not
     the whole dictionary, from lsusb)
 
   Returns:
     map of {bus number: bus object}
     where the bus object has all the devices attached to it in a tree.
   """
@@ skipping 39 matching lines @@
     match = _S_LINE_REGEX.match(line)
     if match:
       if bus_num == -1:
         raise ValueError('S line appears before T line in input file')
       # put the serial number in the device
       tree[bus_num].FindDeviceNumber(device_num).serial = match.group('serial')
 
   return tree
 
 
-class HubType(object):
-  def __init__(self, id_func, port_mapping):
-    """Defines a type of hub.
-
-    Args:
-      id_func: [USBNode -> bool] is a function that can be run on a node
-        to determine if the node represents this type of hub.
-      port_mapping: [dict(int:(int|dict))] maps virtual to physical port
-        numbers. For instance, {3:1, 1:2, 2:3} means that virtual port 3
-        corresponds to physical port 1, virtual port 1 corresponds to physical
-        port 2, and virtual port 2 corresponds to physical port 3. In the
-        case of hubs with "internal" topology, this is represented by nested
-        maps. For instance, {1:{1:1,2:2},2:{1:3,2:4}} means, e.g. that the
-        device plugged into physical port 3 will show up as being connected
-        to port 1, on a device which is connected to port 2 on the hub.
-    """
-    self._id_func = id_func
-    # v2p = "virtual to physical" ports
-    self._v2p_port = port_mapping
-
-  def IsType(self, node):
-    """Determines if the given Node is a hub of this type.
-
-    Args:
-      node: [USBNode] Node to check.
-    """
-    return self._id_func(node)
-
-  def GetPhysicalPortToNodeTuples(self, node):
-    """Gets devices connected to the physical ports on a hub of this type.
-
-    Args:
-      node: [USBNode] Node representing a hub of this type.
-
-    Yields:
-      A series of (int, USBNode) tuples giving a physical port
-      and the USBNode connected to it.
-
-    Raises:
-      ValueError: If the given node isn't a hub of this type.
-    """
-    if self.IsType(node):
-      for res in self._GppHelper(node, self._v2p_port):
-        yield res
-    else:
-      raise ValueError('Node must be a hub of this type')
-
-  def _GppHelper(self, node, mapping):
-    """Helper function for GetPhysicalPortToNodeMap.
-
-    Gets devices connected to physical ports, based on device tree
-    rooted at the given node and the mapping between virtual and physical
-    ports.
-
-    Args:
-      node: [USBNode] Root of tree to search for devices.
-      mapping: [dict] Mapping between virtual and physical ports.
-
-    Yields:
-      A series of (int, USBNode) tuples giving a physical port
-      and the Node connected to it.
-    """
-    for (virtual, physical) in mapping.iteritems():
-      if node.HasPort(virtual):
-        if isinstance(physical, dict):
-          for res in self._GppHelper(node.PortToDevice(virtual), physical):
-            yield res
-        else:
-          yield (physical, node.PortToDevice(virtual))
-
-
 def GetHubsOnBus(bus, hub_types):
   """Scans for all hubs on a bus of given hub types.
 
   Args:
     bus: [USBNode] Bus object.
-    hub_types: [iterable(HubType)] Possible types of hubs.
+    hub_types: [iterable(usb_hubs.HubType)] Possible types of hubs.
 
   Yields:
     Sequence of tuples representing (hub, type of hub)
   """
   for device in bus.AllNodes():
     for hub_type in hub_types:
       if hub_type.IsType(device):
         yield (device, hub_type)
 
 
 def GetPhysicalPortToNodeMap(hub, hub_type):
   """Gets physical-port:node mapping for a given hub.
   Args:
     hub: [USBNode] Hub to get map for.
-    hub_type: [HubType] Which type of hub it is.
+    hub_type: [usb_hubs.HubType] Which type of hub it is.
 
   Returns:
     Dict of {physical port: node}
   """
   port_device = hub_type.GetPhysicalPortToNodeTuples(hub)
   return {port: device for (port, device) in port_device}
 
 
 def GetPhysicalPortToBusDeviceMap(hub, hub_type):
   """Gets physical-port:(bus#, device#) mapping for a given hub.
   Args:
     hub: [USBNode] Hub to get map for.
-    hub_type: [HubType] Which type of hub it is.
+    hub_type: [usb_hubs.HubType] Which type of hub it is.
 
   Returns:
     Dict of {physical port: (bus number, device number)}
   """
   port_device = hub_type.GetPhysicalPortToNodeTuples(hub)
   return {port: (device.bus_num, device.device_num)
           for (port, device) in port_device}
 
 
 def GetPhysicalPortToSerialMap(hub, hub_type):
   """Gets physical-port:serial# mapping for a given hub.
   Args:
     hub: [USBNode] Hub to get map for.
-    hub_type: [HubType] Which type of hub it is.
+    hub_type: [usb_hubs.HubType] Which type of hub it is.
 
   Returns:
     Dict of {physical port: serial number)}
   """
   port_device = hub_type.GetPhysicalPortToNodeTuples(hub)
   return {port: device.serial
           for (port, device) in port_device
           if device.serial}
 
 
 def GetPhysicalPortToTTYMap(device, hub_type):
   """Gets physical-port:tty-string mapping for a given hub.
   Args:
     hub: [USBNode] Hub to get map for.
-    hub_type: [HubType] Which type of hub it is.
+    hub_type: [usb_hubs.HubType] Which type of hub it is.
 
   Returns:
     Dict of {physical port: tty-string)}
   """
   port_device = hub_type.GetPhysicalPortToNodeTuples(device)
   bus_device_to_tty = GetBusDeviceToTTYMap()
   return {port: bus_device_to_tty[(device.bus_num, device.device_num)]
           for (port, device) in port_device
           if (device.bus_num, device.device_num) in bus_device_to_tty}
 
 
 def CollectHubMaps(hub_types, map_func, device_tree_map=None, fast=False):
   """Runs a function on all hubs in the system and collects their output.
 
   Args:
-    hub_types: [HubType] List of possible hub types.
+    hub_types: [usb_hubs.HubType] List of possible hub types.
     map_func: [string] Function to run on each hub.
     device_tree: Previously constructed device tree map, if any.
     fast: Whether to construct device tree fast, if not already provided
 
   Yields:
     Sequence of dicts of {physical port: device} where the type of
     device depends on the ident keyword. Each dict is a separate hub.
   """
   if device_tree_map is None:
     device_tree_map = GetBusNumberToDeviceTreeMap(fast=fast)
@@ skipping 71 matching lines @@
   return result
 
 
 # This dictionary described the mapping between physical and
 # virtual ports on a Plugable 7-Port Hub (model USB2-HUB7BC).
 # Keys are the virtual ports, values are the physical port.
 # The entry 4:{1:4, 2:3, 3:2, 4:1} indicates that virtual port
 # 4 connects to another 'virtual' hub that itself has the
 # virtual-to-physical port mapping {1:4, 2:3, 3:2, 4:1}.
 
-PLUGABLE_7PORT_LAYOUT = {1:7,
-                         2:6,
-                         3:5,
-                         4:{1:4, 2:3, 3:2, 4:1}}
 
 def TestUSBTopologyScript():
   """Test display and hub identification."""
   # Identification criteria for Plugable 7-Port Hub
-  def _is_plugable_7port_hub(node):
-    """Check if a node is a Plugable 7-Port Hub
-    (Model USB2-HUB7BC)
-    The topology of this device is a 4-port hub,
-    with another 4-port hub connected on port 4.
-    """
-    if not isinstance(node, USBDeviceNode):
-      return False
-    if '4-Port HUB' not in node.desc:
-      return False
-    if not node.HasPort(4):
-      return False
-    return '4-Port HUB' in node.PortToDevice(4).desc
-
-  plugable_7port = HubType(_is_plugable_7port_hub,
-                           PLUGABLE_7PORT_LAYOUT)
   print '==== USB TOPOLOGY SCRIPT TEST ===='
 
   # Display devices
   print '==== DEVICE DISPLAY ===='
   device_trees = GetBusNumberToDeviceTreeMap(fast=True)
   for device_tree in device_trees.values():
     device_tree.Display()
   print
 
   # Display TTY information about devices plugged into hubs.
   print '==== TTY INFORMATION ===='
-  for port_map in GetAllPhysicalPortToTTYMaps([plugable_7port],
+  for port_map in GetAllPhysicalPortToTTYMaps([usb_hubs.PLUGABLE_7PORT],
                                               device_tree_map=device_trees):
     print port_map
   print
 
   # Display serial number information about devices plugged into hubs.
   print '==== SERIAL NUMBER INFORMATION ===='
-  for port_map in GetAllPhysicalPortToSerialMaps([plugable_7port],
+  for port_map in GetAllPhysicalPortToSerialMaps([usb_hubs.PLUGABLE_7PORT],
                                                  device_tree_map=device_trees):
     print port_map
+
+  print
+  # Display phone to BattOr map
+  print GenerateSerialMap()
   print ''
+
   return 0
 
+
+def ReadSerialMapFile(filename):
+  """Reads JSON file giving phone-to-battor serial number map.
+
+  Parses a JSON file consisting of a list of items of the following form:
+  [{'phone':<phone serial 1>, 'battor':<battor serial 1>},

[rnephew (Reviews Here), 2016/03/25 00:17:53]

Nit: Spaces here too. See below.

[alexandermont, 2016/03/28 18:47:29]

Done

+  {'phone':<phone serial 2>, 'battor':<battor serial 2>}, ...]
+
+  indicating which phone serial numbers should be matched with
+  which BattOr serial numbers. Returns dictionary of the form:
+
+  {<phone serial 1>: <BattOr serial 1>,
+   <phone serial 2>: <BattOr serial 2>}
+
+  Args:
+      filename: Name of file to read.
+  """
+  result = {}
+  with open(filename, 'r') as infile:
+    in_dict = json.load(infile)
+  for x in in_dict:
+    result[x['phone']] = x['battor']
+  return result
+
+def GenerateSerialMapFile(filename, hub_types=None):
+  """Writes a map of phone serial numbers to BattOr serial numbers to file.
+
+  Args:
+      filename: Name of file to write.
+  """
+  result = []
+  for (phone, battor) in GenerateSerialMap(hub_types).iteritems():
+    result.append({'phone':phone, 'battor':battor})

[rnephew (Reviews Here), 2016/03/25 00:17:53]

nit: 'phone': phone, 'battor': battor

note the spaces.

[alexandermont, 2016/03/28 18:47:29]

Done

+  with open(filename, 'w') as outfile:
+    json.dump(result, outfile)
+
+def GenerateSerialMap(hub_types=None):
+  """Generates a map of phone serial numbers to BattOr serial numbers.
+
+  Generates a JSON file consisting of a list of items of the following form:
+  [{'phone':<phone serial 1>, 'battor':<battor serial 1>},

[rnephew (Reviews Here), 2016/03/25 00:17:52]

Spaces here too.

[alexandermont, 2016/03/28 18:47:29]

Done

+  {'phone':<phone serial 2>, 'battor':<battor serial 2>}, ...]
+
+  indicating which phone serial numbers should be matched with
+  which BattOr serial numbers. Mapping is based on the physical port numbers
+  of the hubs that the BattOrs and phones are connected to.
+
+  Args:
+      hub_types: List of hub types to check for
+  """
+  hub_types = [usb_hubs.GetHubType(x)
+               for x in hub_types or ['plugable_7port']]
+  devtree = GetBusNumberToDeviceTreeMap(fast=True)
+  battor_serials = list(GetBattorSerialNumbers(devtree))
+  p2serial = GetAllPhysicalPortToSerialMaps(hub_types,
+                                            device_tree_map=devtree)
+  port_to_devices = collections.defaultdict(list)
+  result = {}
+  for hub in p2serial:
+    for (port, serial) in hub.iteritems():
+      port_to_devices[port].append(serial)
+  for (port, serial_list) in port_to_devices.iteritems():
+    battor_serial = [x for x in serial_list if x in battor_serials]
+    if len(battor_serial) >= 2:
+      raise BattorError('Duplicate BattOr serial numbers detected')
+    if len(battor_serial) == 1:
+      battor_serial = battor_serial[0]
+      phone_serial = [x for x in serial_list if x != battor_serial]
+      if len(phone_serial) >= 2:
+        raise BattorError('Multiple phones matched with same BattOr')
+      if len(phone_serial) == 0:
+        raise BattorError('BattOr has no matching phone')
+      if battor_serial in result:
+        raise BattorError('Duplicate BattOr serial numbers detected')
+      result[phone_serial[0]] = battor_serial
+  return result
+
+def _PhoneToPathMap(serial, serial_map, devtree):
+  """Maps phone serial number to TTY path, assuming serial map is provided."""
+  battor_serial = serial_map[serial]
+  for tree in devtree.values():
+    for node in tree.AllNodes():
+      if isinstance(node, USBDeviceNode):
+        if node.serial == battor_serial:
+          bus_device_to_tty = GetBusDeviceToTTYMap()
+          bus_device = (node.bus_num, node.device_num)
+          try:
+            return bus_device_to_tty[bus_device]
+          except KeyError:
+            raise BattorError('Device with given serial number not a BattOr '
+                              '(does not have TTY path)')
+
+def _PhoneToPathPorts(serial, hub_types, devtree):
+  """Maps phone serial number to TTY path, assuming hub types are provided."""
+  hub_types = [usb_hubs.GetHubType(x)
+               for x in hub_types or ['plugable_7port']]
+  p2serial = GetAllPhysicalPortToSerialMaps(hub_types,
+                                            device_tree_map=devtree)
+  p2tty = GetAllPhysicalPortToTTYMaps(hub_types,
+                                      device_tree_map=devtree)
+
+  # get the port number of this device
+  port_num = -1
+  for hub in p2serial:
+    for (port, s) in hub.iteritems():
+      if serial == s:
+        port_num = port
+  if port_num == -1:
+    raise BattorError('Device with given serial number not found.')
+
+  # get the tty for this port number
+  tty_string = None
+  for hub in p2tty:
+    x = hub.get(port_num)
+    if x is not None:
+      if tty_string:
+        raise BattorError('Two TTY devices with matching port number.')
+      else:
+        tty_string = x
+  return tty_string
+
+def GetBattorPathFromPhoneSerial(serial, serial_map=None, hub_types=None):

[nednguyen, 2016/03/25 03:05:29]

In which use case a client will not want to use serial_map but use "hub_types"
instead? I usually advocate for simpler API, so unless there is a use case which
people can't use serial_map and have to use hub_types, I would say deprecating
hub_types to simplify your design.

[rnephew (Reviews Here), 2016/03/25 05:03:19]

Thinking more about it, that would most likely happen outside of this function.

[nednguyen, 2016/03/27 16:44:32]

I see. Can we leave the hub_types param out and instead have the callsites
translate hubtypes to serial_map themselves? That's way, it's less ambiguous
about what this function will do when both serial_map & hub_types are present.

[alexandermont, 2016/03/28 18:47:29]

Done

+  """Gets the TTY path (e.g. '/dev/ttyUSB0')  to communicate with the BattOr.
+
+  (1) If serial_map is available, serial_map is treated as a dictionary mapping
+  phone serial numbers to BattOr serial numbers. This function will get the
+  TTY path for the given BattOr serial number.
+
+  (2) If serial_map is None and hub_types is available, hub_types is treated
+  as a list of strings representing the types of hubs that the phones and
+  BattOrs are plugged into (see usb_hubs.py for details). This function will
+  get the BattOr that is connected to the same physical port number,
+  on a different hub, that the phone is connected to.
+
+  (3) If serial_map is None and hub_types is None, then it will assume that the
+  type of hub the BattOrs and phones are plugged into are Plugable USB
+  7-Port Hubs, and then proceed as in (2).
+
+  Args:
+    serial: Serial number of phone connected on the same physical port that
+    the BattOr is connected to.
+    serial_map_file: Map of phone serial numbers to BattOr serial numbers.
+    hub_types: List of hub types to check for. Used only if serial_map_file
+    is None.
+
+  Returns:
+    Device string used to communicate with device.
+
+  Raises:
+    ValueError: If serial number is not given.
+    BattorError: If BattOr not found or unexpected USB topology.
+  """
+  # If there's only one BattOr connected to the system, just use that one.
+  # This allows for use on, e.g., a developer's workstation with no hubs.
+  devtree = GetBusNumberToDeviceTreeMap(fast=True)
+  all_battors = GetBattorList(devtree)
+  if len(all_battors) == 1:
+    return '/dev/' + all_battors[0]
+  if not serial:
+    raise BattorError('Two or more BattOrs connected, no serial provided')
+
+  if serial_map is not None:
+    err_string = 'with given serial number'
+    tty_string = _PhoneToPathMap(serial, serial_map, devtree)
+

[rnephew (Reviews Here), 2016/03/25 05:02:41]

nit: no space here.

[alexandermont, 2016/03/28 18:47:29]

Done

+  else:
+    err_string = 'on matching port'
+    tty_string = _PhoneToPathPorts(serial, hub_types, devtree)
+  if not tty_string:

[rnephew (Reviews Here), 2016/03/25 05:02:41]

space here instead.

[alexandermont, 2016/03/28 18:47:29]

Done

+    raise BattorError('No device %s detected.' % err_string)
+  if IsBattor(tty_string, devtree):
+    return '/dev/' + tty_string
+  else:
+    raise BattorError('Device %s is not a BattOr.' % err_string)
+
 def parse_options(argv):
   """Parses and checks the command-line options.
 
   Returns:
     A tuple containing the options structure and a list of categories to
     be traced.
   """
   USAGE = '''./find_usb_devices [--help]
     This script shows the mapping between USB devices and port numbers.
     Clients are not intended to call this script from the command line.
     Clients are intended to call the functions in this script directly.
     For instance, GetAllPhysicalPortToSerialMaps(...)
     Running this script with --help will display this message.
     Running this script without --help will display information about
     devices attached, TTY mapping, and serial number mapping,
     for testing purposes. See design document for API documentation.
   '''
   parser = argparse.ArgumentParser(usage=USAGE)
   return parser.parse_args(argv[1:])
 
 def main():
   parse_options(sys.argv)
   TestUSBTopologyScript()
 
 if __name__ == "__main__":
   sys.exit(main())