1#!/usr/bin/python 2# Copyright 2016 The Chromium Authors. All rights reserved. 3# Use of this source code is governed by a BSD-style license that can be 4# found in the LICENSE file. 5 6import re 7import sys 8import argparse 9 10from devil.utils import cmd_helper 11from devil.utils import usb_hubs 12from devil.utils import lsusb 13 14# Note: In the documentation below, "virtual port" refers to the port number 15# as observed by the system (e.g. by usb-devices) and "physical port" refers 16# to the physical numerical label on the physical port e.g. on a USB hub. 17# The mapping between virtual and physical ports is not always the identity 18# (e.g. the port labeled "1" on a USB hub does not always show up as "port 1" 19# when you plug something into it) but, as far as we are aware, the mapping 20# between virtual and physical ports is always the same for a given 21# model of USB hub. When "port number" is referenced without specifying, it 22# means the virtual port number. 23 24 25# Wrapper functions for system commands to get output. These are in wrapper 26# functions so that they can be more easily mocked-out for tests. 27def _GetParsedLSUSBOutput(): 28 return lsusb.lsusb() 29 30 31def _GetUSBDevicesOutput(): 32 return cmd_helper.GetCmdOutput(['usb-devices']) 33 34 35def _GetTtyUSBInfo(tty_string): 36 cmd = ['udevadm', 'info', '--name=/dev/' + tty_string, '--attribute-walk'] 37 return cmd_helper.GetCmdOutput(cmd) 38 39 40def _GetCommList(): 41 return cmd_helper.GetCmdOutput('ls /dev', shell=True) 42 43 44def GetTTYList(): 45 return [x for x in _GetCommList().splitlines() if 'ttyUSB' in x] 46 47 48# Class to identify nodes in the USB topology. USB topology is organized as 49# a tree. 50class USBNode(object): 51 def __init__(self): 52 self._port_to_node = {} 53 54 @property 55 def desc(self): 56 raise NotImplementedError 57 58 @property 59 def info(self): 60 raise NotImplementedError 61 62 @property 63 def device_num(self): 64 raise NotImplementedError 65 66 @property 67 def bus_num(self): 68 raise NotImplementedError 69 70 def HasPort(self, port): 71 """Determines if this device has a device connected to the given port.""" 72 return port in self._port_to_node 73 74 def PortToDevice(self, port): 75 """Gets the device connected to the given port on this device.""" 76 return self._port_to_node[port] 77 78 def Display(self, port_chain='', info=False): 79 """Displays information about this node and its descendants. 80 81 Output format is, e.g. 1:3:3:Device 42 (ID 1234:5678 Some Device) 82 meaning that from the bus, if you look at the device connected 83 to port 1, then the device connected to port 3 of that, 84 then the device connected to port 3 of that, you get the device 85 assigned device number 42, which is Some Device. Note that device 86 numbers will be reassigned whenever a connected device is powercycled 87 or reinserted, but port numbers stay the same as long as the device 88 is reinserted back into the same physical port. 89 90 Args: 91 port_chain: [string] Chain of ports from bus to this node (e.g. '2:4:') 92 info: [bool] Whether to display detailed info as well. 93 """ 94 raise NotImplementedError 95 96 def AddChild(self, port, device): 97 """Adds child to the device tree. 98 99 Args: 100 port: [int] Port number of the device. 101 device: [USBDeviceNode] Device to add. 102 103 Raises: 104 ValueError: If device already has a child at the given port. 105 """ 106 if self.HasPort(port): 107 raise ValueError('Duplicate port number') 108 else: 109 self._port_to_node[port] = device 110 111 def AllNodes(self): 112 """Generator that yields this node and all of its descendants. 113 114 Yields: 115 [USBNode] First this node, then each of its descendants (recursively) 116 """ 117 yield self 118 for child_node in self._port_to_node.values(): 119 for descendant_node in child_node.AllNodes(): 120 yield descendant_node 121 122 def FindDeviceNumber(self, findnum): 123 """Find device with given number in tree 124 125 Searches the portion of the device tree rooted at this node for 126 a device with the given device number. 127 128 Args: 129 findnum: [int] Device number to search for. 130 131 Returns: 132 [USBDeviceNode] Node that is found. 133 """ 134 for node in self.AllNodes(): 135 if node.device_num == findnum: 136 return node 137 return None 138 139 140class USBDeviceNode(USBNode): 141 def __init__(self, bus_num=0, device_num=0, serial=None, info=None): 142 """Class that represents a device in USB tree. 143 144 Args: 145 bus_num: [int] Bus number that this node is attached to. 146 device_num: [int] Device number of this device (or 0, if this is a bus) 147 serial: [string] Serial number. 148 info: [dict] Map giving detailed device info. 149 """ 150 super(USBDeviceNode, self).__init__() 151 self._bus_num = bus_num 152 self._device_num = device_num 153 self._serial = serial 154 self._info = {} if info is None else info 155 156 #override 157 @property 158 def desc(self): 159 return self._info.get('desc') 160 161 #override 162 @property 163 def info(self): 164 return self._info 165 166 #override 167 @property 168 def device_num(self): 169 return self._device_num 170 171 #override 172 @property 173 def bus_num(self): 174 return self._bus_num 175 176 @property 177 def serial(self): 178 return self._serial 179 180 @serial.setter 181 def serial(self, serial): 182 self._serial = serial 183 184 #override 185 def Display(self, port_chain='', info=False): 186 print '%s Device %d (%s)' % (port_chain, self.device_num, self.desc) 187 if info: 188 print self.info 189 for (port, device) in self._port_to_node.iteritems(): 190 device.Display('%s%d:' % (port_chain, port), info=info) 191 192 193class USBBusNode(USBNode): 194 def __init__(self, bus_num=0): 195 """Class that represents a node (either a bus or device) in USB tree. 196 197 Args: 198 is_bus: [bool] If true, node is bus; if not, node is device. 199 bus_num: [int] Bus number that this node is attached to. 200 device_num: [int] Device number of this device (or 0, if this is a bus) 201 desc: [string] Short description of device. 202 serial: [string] Serial number. 203 info: [dict] Map giving detailed device info. 204 port_to_dev: [dict(int:USBDeviceNode)] 205 Maps port # to device connected to port. 206 """ 207 super(USBBusNode, self).__init__() 208 self._bus_num = bus_num 209 210 #override 211 @property 212 def desc(self): 213 return 'BUS %d' % self._bus_num 214 215 #override 216 @property 217 def info(self): 218 return {} 219 220 #override 221 @property 222 def device_num(self): 223 return -1 224 225 #override 226 @property 227 def bus_num(self): 228 return self._bus_num 229 230 #override 231 def Display(self, port_chain='', info=False): 232 print "=== %s ===" % self.desc 233 for (port, device) in self._port_to_node.iteritems(): 234 device.Display('%s%d:' % (port_chain, port), info=info) 235 236 237_T_LINE_REGEX = re.compile(r'T: Bus=(?P<bus>\d{2}) Lev=(?P<lev>\d{2}) ' 238 r'Prnt=(?P<prnt>\d{2,3}) Port=(?P<port>\d{2}) ' 239 r'Cnt=(?P<cnt>\d{2}) Dev#=(?P<dev>.{3}) .*') 240 241_S_LINE_REGEX = re.compile(r'S: SerialNumber=(?P<serial>.*)') 242_LSUSB_BUS_DEVICE_RE = re.compile(r'^Bus (\d{3}) Device (\d{3}): (.*)') 243 244 245def GetBusNumberToDeviceTreeMap(fast=True): 246 """Gets devices currently attached. 247 248 Args: 249 fast [bool]: whether to do it fast (only get description, not 250 the whole dictionary, from lsusb) 251 252 Returns: 253 map of {bus number: bus object} 254 where the bus object has all the devices attached to it in a tree. 255 """ 256 if fast: 257 info_map = {} 258 for line in lsusb.raw_lsusb().splitlines(): 259 match = _LSUSB_BUS_DEVICE_RE.match(line) 260 if match: 261 info_map[(int(match.group(1)), int(match.group(2)))] = ( 262 {'desc':match.group(3)}) 263 else: 264 info_map = {((int(line['bus']), int(line['device']))): line 265 for line in _GetParsedLSUSBOutput()} 266 267 268 tree = {} 269 bus_num = -1 270 for line in _GetUSBDevicesOutput().splitlines(): 271 match = _T_LINE_REGEX.match(line) 272 if match: 273 bus_num = int(match.group('bus')) 274 parent_num = int(match.group('prnt')) 275 # usb-devices starts counting ports from 0, so add 1 276 port_num = int(match.group('port')) + 1 277 device_num = int(match.group('dev')) 278 279 # create new bus if necessary 280 if bus_num not in tree: 281 tree[bus_num] = USBBusNode(bus_num=bus_num) 282 283 # create the new device 284 new_device = USBDeviceNode(bus_num=bus_num, 285 device_num=device_num, 286 info=info_map.get((bus_num, device_num), 287 {'desc': 'NOT AVAILABLE'})) 288 289 # add device to bus 290 if parent_num != 0: 291 tree[bus_num].FindDeviceNumber(parent_num).AddChild( 292 port_num, new_device) 293 else: 294 tree[bus_num].AddChild(port_num, new_device) 295 296 match = _S_LINE_REGEX.match(line) 297 if match: 298 if bus_num == -1: 299 raise ValueError('S line appears before T line in input file') 300 # put the serial number in the device 301 tree[bus_num].FindDeviceNumber(device_num).serial = match.group('serial') 302 303 return tree 304 305 306def GetHubsOnBus(bus, hub_types): 307 """Scans for all hubs on a bus of given hub types. 308 309 Args: 310 bus: [USBNode] Bus object. 311 hub_types: [iterable(usb_hubs.HubType)] Possible types of hubs. 312 313 Yields: 314 Sequence of tuples representing (hub, type of hub) 315 """ 316 for device in bus.AllNodes(): 317 for hub_type in hub_types: 318 if hub_type.IsType(device): 319 yield (device, hub_type) 320 321 322def GetPhysicalPortToNodeMap(hub, hub_type): 323 """Gets physical-port:node mapping for a given hub. 324 Args: 325 hub: [USBNode] Hub to get map for. 326 hub_type: [usb_hubs.HubType] Which type of hub it is. 327 328 Returns: 329 Dict of {physical port: node} 330 """ 331 port_device = hub_type.GetPhysicalPortToNodeTuples(hub) 332 return {port: device for (port, device) in port_device} 333 334 335def GetPhysicalPortToBusDeviceMap(hub, hub_type): 336 """Gets physical-port:(bus#, device#) mapping for a given hub. 337 Args: 338 hub: [USBNode] Hub to get map for. 339 hub_type: [usb_hubs.HubType] Which type of hub it is. 340 341 Returns: 342 Dict of {physical port: (bus number, device number)} 343 """ 344 port_device = hub_type.GetPhysicalPortToNodeTuples(hub) 345 return {port: (device.bus_num, device.device_num) 346 for (port, device) in port_device} 347 348 349def GetPhysicalPortToSerialMap(hub, hub_type): 350 """Gets physical-port:serial# mapping for a given hub. 351 352 Args: 353 hub: [USBNode] Hub to get map for. 354 hub_type: [usb_hubs.HubType] Which type of hub it is. 355 356 Returns: 357 Dict of {physical port: serial number)} 358 """ 359 port_device = hub_type.GetPhysicalPortToNodeTuples(hub) 360 return {port: device.serial 361 for (port, device) in port_device 362 if device.serial} 363 364 365def GetPhysicalPortToTTYMap(device, hub_type): 366 """Gets physical-port:tty-string mapping for a given hub. 367 Args: 368 hub: [USBNode] Hub to get map for. 369 hub_type: [usb_hubs.HubType] Which type of hub it is. 370 371 Returns: 372 Dict of {physical port: tty-string)} 373 """ 374 port_device = hub_type.GetPhysicalPortToNodeTuples(device) 375 bus_device_to_tty = GetBusDeviceToTTYMap() 376 return {port: bus_device_to_tty[(device.bus_num, device.device_num)] 377 for (port, device) in port_device 378 if (device.bus_num, device.device_num) in bus_device_to_tty} 379 380 381def CollectHubMaps(hub_types, map_func, device_tree_map=None, fast=False): 382 """Runs a function on all hubs in the system and collects their output. 383 384 Args: 385 hub_types: [usb_hubs.HubType] List of possible hub types. 386 map_func: [string] Function to run on each hub. 387 device_tree: Previously constructed device tree map, if any. 388 fast: Whether to construct device tree fast, if not already provided 389 390 Yields: 391 Sequence of dicts of {physical port: device} where the type of 392 device depends on the ident keyword. Each dict is a separate hub. 393 """ 394 if device_tree_map is None: 395 device_tree_map = GetBusNumberToDeviceTreeMap(fast=fast) 396 for bus in device_tree_map.values(): 397 for (hub, hub_type) in GetHubsOnBus(bus, hub_types): 398 yield map_func(hub, hub_type) 399 400 401def GetAllPhysicalPortToNodeMaps(hub_types, **kwargs): 402 return CollectHubMaps(hub_types, GetPhysicalPortToNodeMap, **kwargs) 403 404 405def GetAllPhysicalPortToBusDeviceMaps(hub_types, **kwargs): 406 return CollectHubMaps(hub_types, GetPhysicalPortToBusDeviceMap, **kwargs) 407 408 409def GetAllPhysicalPortToSerialMaps(hub_types, **kwargs): 410 return CollectHubMaps(hub_types, GetPhysicalPortToSerialMap, **kwargs) 411 412 413def GetAllPhysicalPortToTTYMaps(hub_types, **kwargs): 414 return CollectHubMaps(hub_types, GetPhysicalPortToTTYMap, **kwargs) 415 416 417_BUS_NUM_REGEX = re.compile(r'.*ATTRS{busnum}=="(\d*)".*') 418_DEVICE_NUM_REGEX = re.compile(r'.*ATTRS{devnum}=="(\d*)".*') 419 420 421def GetBusDeviceFromTTY(tty_string): 422 """Gets bus and device number connected to a ttyUSB port. 423 424 Args: 425 tty_string: [String] Identifier for ttyUSB (e.g. 'ttyUSB0') 426 427 Returns: 428 Tuple (bus, device) giving device connected to that ttyUSB. 429 430 Raises: 431 ValueError: If bus and device information could not be found. 432 """ 433 bus_num = None 434 device_num = None 435 # Expected output of GetCmdOutput should be something like: 436 # looking at device /devices/something/.../.../... 437 # KERNELS="ttyUSB0" 438 # SUBSYSTEMS=... 439 # DRIVERS=... 440 # ATTRS{foo}=... 441 # ATTRS{bar}=... 442 # ... 443 for line in _GetTtyUSBInfo(tty_string).splitlines(): 444 bus_match = _BUS_NUM_REGEX.match(line) 445 device_match = _DEVICE_NUM_REGEX.match(line) 446 if bus_match and bus_num == None: 447 bus_num = int(bus_match.group(1)) 448 if device_match and device_num == None: 449 device_num = int(device_match.group(1)) 450 if bus_num is None or device_num is None: 451 raise ValueError('Info not found') 452 return (bus_num, device_num) 453 454 455def GetBusDeviceToTTYMap(): 456 """Gets all mappings from (bus, device) to ttyUSB string. 457 458 Gets mapping from (bus, device) to ttyUSB string (e.g. 'ttyUSB0'), 459 for all ttyUSB strings currently active. 460 461 Returns: 462 [dict] Dict that maps (bus, device) to ttyUSB string 463 """ 464 result = {} 465 for tty in GetTTYList(): 466 result[GetBusDeviceFromTTY(tty)] = tty 467 return result 468 469 470# This dictionary described the mapping between physical and 471# virtual ports on a Plugable 7-Port Hub (model USB2-HUB7BC). 472# Keys are the virtual ports, values are the physical port. 473# The entry 4:{1:4, 2:3, 3:2, 4:1} indicates that virtual port 474# 4 connects to another 'virtual' hub that itself has the 475# virtual-to-physical port mapping {1:4, 2:3, 3:2, 4:1}. 476 477 478def TestUSBTopologyScript(): 479 """Test display and hub identification.""" 480 # Identification criteria for Plugable 7-Port Hub 481 print '==== USB TOPOLOGY SCRIPT TEST ====' 482 483 # Display devices 484 print '==== DEVICE DISPLAY ====' 485 device_trees = GetBusNumberToDeviceTreeMap() 486 for device_tree in device_trees.values(): 487 device_tree.Display() 488 print 489 490 # Display TTY information about devices plugged into hubs. 491 print '==== TTY INFORMATION ====' 492 for port_map in GetAllPhysicalPortToTTYMaps([usb_hubs.PLUGABLE_7PORT, 493 usb_hubs.PLUGABLE_7PORT_USB3_PART2, usb_hubs.PLUGABLE_7PORT_USB3_PART3], 494 device_tree_map=device_trees): 495 print port_map 496 print 497 498 # Display serial number information about devices plugged into hubs. 499 print '==== SERIAL NUMBER INFORMATION ====' 500 for port_map in GetAllPhysicalPortToSerialMaps([usb_hubs.PLUGABLE_7PORT, 501 usb_hubs.PLUGABLE_7PORT_USB3_PART2, usb_hubs.PLUGABLE_7PORT_USB3_PART3], 502 device_tree_map=device_trees): 503 print port_map 504 505 506 return 0 507 508 509def parse_options(argv): 510 """Parses and checks the command-line options. 511 512 Returns: 513 A tuple containing the options structure and a list of categories to 514 be traced. 515 """ 516 USAGE = '''./find_usb_devices [--help] 517 This script shows the mapping between USB devices and port numbers. 518 Clients are not intended to call this script from the command line. 519 Clients are intended to call the functions in this script directly. 520 For instance, GetAllPhysicalPortToSerialMaps(...) 521 Running this script with --help will display this message. 522 Running this script without --help will display information about 523 devices attached, TTY mapping, and serial number mapping, 524 for testing purposes. See design document for API documentation. 525 ''' 526 parser = argparse.ArgumentParser(usage=USAGE) 527 return parser.parse_args(argv[1:]) 528 529def main(): 530 parse_options(sys.argv) 531 TestUSBTopologyScript() 532 533if __name__ == "__main__": 534 sys.exit(main()) 535