text: #!/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.
  
  
  '''
  This script provides tools to map BattOrs to phones.
  
  Phones are identified by the following string:
  
  "Phone serial number" - Serial number of the phone. This can be
  obtained via 'adb devices' or 'usb-devices', and is not expected
  to change for a given phone.
  
  BattOrs are identified by the following two strings:
  
  "BattOr serial number" - Serial number of the BattOr. This can be
  obtained via 'usb-devices', and is not expected to change for
  a given BattOr.
  
  "BattOr path" - The path of the form '/dev/ttyUSB*' that is used
  to communicate with the BattOr (the battor_agent binary takes
  this BattOr path as a parameter). The BattOr path is frequently
  reassigned by the OS, most often when the device is disconnected
  and then reconnected. Thus, the BattOr path cannot be expected
  to be stable.
  
  In a typical application, the user will require the BattOr path
  for the BattOr that is plugged into a given phone. For instance,
  the user will be running tracing on a particular phone, and will
  need to know which BattOr path to use to communicate with the BattOr
  to get the corresponding power trace.
  
  Getting this mapping requires two steps: (1) determining the
  mapping between phone serial numbers and BattOr serial numbers, and
  (2) getting the BattOr path corresponding to a given BattOr serial
  number.
  
  For step (1), we generate a JSON file giving this mapping. This
  JSON file consists of a list of items of the following form:
  [{'phone': <phone serial 1>, 'battor': <battor serial 1>},
  {'phone': <phone serial 2>, 'battor': <battor serial 2>}, ...]
  
  The default way to generate this JSON file is using the function
  GenerateSerialMapFile, which generates a mapping based on assuming
  that the system has two identical USB hubs connected to it, and
  the phone plugged into physical port number 1 on one hub corresponds
  to the BattOr plugged into physical port number 1 on the other hub,
  and similarly with physical port numbers 2, 3, etc. This generates
  the map file based on the structure at the time GenerateSerialMapFile called.
  Note that after the map file is generated, port numbers are no longer used;
  the user could move around the devices in the ports without affecting
  which phone goes with which BattOr. (Thus, if the user wanted to update the
  mapping to match the new port connections, the user would have to
  re-generate this file.)
  
  The script update_mapping.py will do this updating from the command line.
  
  If the user wanted to specify a custom mapping, the user could instead
  create the JSON file manually. (In this case, hubs would not be necessary
  and the physical ports connected would be irrelevant.)
  
  Step (2) is conducted through the function GetBattorPathFromPhoneSerial,
  which takes a serial number mapping generated via step (1) and a phone
  serial number, then gets the corresponding BattOr serial number from the
  map and determines its BattOr path (e.g. /dev/ttyUSB0). Since BattOr paths
  can change if devices are connected and disconnected (even if connected
  or disconnected via the same port) this function should be called to
  determine the BattOr path every time before connecting to the BattOr.
  
  Note that if there is only one BattOr connected to the system, then
  GetBattorPathFromPhoneSerial will always return that BattOr and will ignore
  the mapping file. Thus, if the user never has more than one BattOr connected
  to the system, the user will not need to generate mapping files.
  '''
  
  
  import json
  import collections
  
  from battor import battor_error
  from devil.utils import find_usb_devices
  from devil.utils import usb_hubs
  
  
  def GetBattorList(device_tree_map):
    return [x for x in find_usb_devices.GetTTYList()
            if IsBattor(x, device_tree_map)]
  
  
  def IsBattor(tty_string, device_tree_map):
    (bus, device) = find_usb_devices.GetBusDeviceFromTTY(tty_string)
    node = device_tree_map[bus].FindDeviceNumber(device)
    return '0403:6001' in node.desc
  
  
  def GetBattorSerialNumbers(device_tree_map):
    for x in find_usb_devices.GetTTYList():
      if IsBattor(x, device_tree_map):
        (bus, device) = find_usb_devices.GetBusDeviceFromTTY(x)
        devnode = device_tree_map[bus].FindDeviceNumber(device)
        yield devnode.serial
  
  
  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>},
    {'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 WriteSerialMapFile(filename, serial_map):
    """Writes a map of phone serial numbers to BattOr serial numbers to file.
  
    Writes a JSON file consisting of a list of items of the following form:
    [{'phone': <phone serial 1>, 'battor': <battor serial 1>},
    {'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:
        filename: Name of file to write.
        serial_map: Serial map {phone: battor}
    """
    result = []
    for (phone, battor) in serial_map.iteritems():
      result.append({'phone': phone, 'battor': battor})
    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 dict of:
    {<phone serial 1>: <battor serial 1>,
     <phone serial 2>: <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. If not specified, checks
        for all defined hub types. (see usb_hubs.py for details)
    """
    if hub_types:
      hub_types = [usb_hubs.GetHubType(x) for x in hub_types]
    else:
      hub_types = usb_hubs.ALL_HUBS
  
    devtree = find_usb_devices.GetBusNumberToDeviceTreeMap()
  
    # List of serial numbers in the system that represent BattOrs.
    battor_serials = list(GetBattorSerialNumbers(devtree))
  
    # If there's only one BattOr in the system, then a serial number ma
    # is not necessary.
    if len(battor_serials) == 1:
      return {}
  
    # List of dictionaries, one for each hub, that maps the physical
    # port number to the serial number of that hub. For instance, in a 2
    # hub system, this could return [{1:'ab', 2:'cd'}, {1:'jkl', 2:'xyz'}]
    # where 'ab' and 'cd' are the phone serial numbers and 'jkl' and 'xyz'
    # are the BattOr serial numbers.
    port_to_serial = find_usb_devices.GetAllPhysicalPortToSerialMaps(
        hub_types, device_tree_map=devtree)
  
    class serials(object):
      def __init__(self):
        self.phone = None
        self.battor = None
  
    # Map of {physical port number: [phone serial #, BattOr serial #]. This
    # map is populated by executing the code below. For instance, in the above
    # example, after the code below is executed, port_to_devices would equal
    # {1: ['ab', 'jkl'], 2: ['cd', 'xyz']}
    port_to_devices = collections.defaultdict(serials)
    for hub in port_to_serial:
      for (port, serial) in hub.iteritems():
        if serial in battor_serials:
          if port_to_devices[port].battor is not None:
            raise battor_error.BattorError('Multiple BattOrs on same port number')
          else:
            port_to_devices[port].battor = serial
        else:
          if port_to_devices[port].phone is not None:
            raise battor_error.BattorError('Multiple phones on same port number')
          else:
            port_to_devices[port].phone = serial
  
    # Turn the port_to_devices map into a map of the form
    # {phone serial number: BattOr serial number}.
    result = {}
    for pair in port_to_devices.values():
      if pair.phone is None:
        continue
      if pair.battor is None:
        raise battor_error.BattorError(
            'Phone detected with no corresponding BattOr')
      result[pair.phone] = pair.battor
    return result
  
  def GenerateSerialMapFile(filename, hub_types=None):
    """Generates a serial map file and writes it."""
    WriteSerialMapFile(filename, GenerateSerialMap(hub_types))
  
  def _PhoneToPathMap(serial, serial_map, devtree):
    """Maps phone serial number to TTY path, assuming serial map is provided."""
    try:
      battor_serial = serial_map[serial]
    except KeyError:
      raise battor_error.BattorError('Serial number not found in serial map.')
    for tree in devtree.values():
      for node in tree.AllNodes():
        if isinstance(node, find_usb_devices.USBDeviceNode):
          if node.serial == battor_serial:
            bus_device_to_tty = find_usb_devices.GetBusDeviceToTTYMap()
            bus_device = (node.bus_num, node.device_num)
            try:
              return bus_device_to_tty[bus_device]
            except KeyError:
              raise battor_error.BattorError(
                  'Device with given serial number not a BattOr '
                  '(does not have TTY path)')
  
  
  def GetBattorPathFromPhoneSerial(serial, serial_map=None,
                                   serial_map_file=None):
    """Gets the TTY path (e.g. '/dev/ttyUSB0')  to communicate with the BattOr.
  
    (1) If serial_map is given, it 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_file is given, it is treated as the name of a
    phone-to-BattOr mapping file (generated with GenerateSerialMapFile)
    and this will be loaded and used as the dict to map port numbers to
    BattOr serial numbers.
  
    You can only give one of serial_map and serial_map_file.
  
    Args:
      serial: Serial number of phone connected on the same physical port that
      the BattOr is connected to.
      serial_map: Map of phone serial numbers to BattOr serial numbers, given
      as a dictionary.
      serial_map_file: Map of phone serial numbers to BattOr serial numbers,
      given as a file.
      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 = find_usb_devices.GetBusNumberToDeviceTreeMap()
    all_battors = GetBattorList(devtree)
    if len(all_battors) == 1:
      return '/dev/' + all_battors[0]
  
    if not serial:
      raise battor_error.BattorError(
          'Two or more BattOrs connected, no serial provided')
  
    if serial_map and serial_map_file:
      raise ValueError('Cannot specify both serial_map and serial_map_file')
  
    if serial_map_file:
      serial_map = ReadSerialMapFile(serial_map_file)
  
    tty_string = _PhoneToPathMap(serial, serial_map, devtree)
  
    if not tty_string:
      raise battor_error.BattorError(
          'No device with given serial number detected.')
  
    if IsBattor(tty_string, devtree):
      return '/dev/' + tty_string
    else:
      raise battor_error.BattorError(
          'Device with given serial number is not a BattOr.')
  
  if __name__ == '__main__':
    # Main function for testing purposes
    print GenerateSerialMap()
  