1645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez#!/usr/bin/env python
2645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez# Copyright 2015 The Chromium Authors. All rights reserved.
3645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez# Use of this source code is governed by a BSD-style license that can be
4645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez# found in the LICENSE file.
5645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez
6645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavezimport argparse
7645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavezimport fcntl
8645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavezimport logging
9645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavezimport re
10645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavezimport sys
11645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez
12645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavezfrom devil.android import device_errors
13645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavezfrom devil.utils import lsusb
14645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavezfrom devil.utils import run_tests_helper
15645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez
16645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez_INDENTATION_RE = re.compile(r'^( *)')
17645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez_LSUSB_BUS_DEVICE_RE = re.compile(r'^Bus (\d{3}) Device (\d{3}):')
18645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez_LSUSB_ENTRY_RE = re.compile(r'^ *([^ ]+) +([^ ]+) *([^ ].*)?$')
19645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez_LSUSB_GROUP_RE = re.compile(r'^ *([^ ]+.*):$')
20645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez
21645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez_USBDEVFS_RESET = ord('U') << 8 | 20
22645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez
23645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez
24645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavezdef reset_usb(bus, device):
25645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez  """Reset the USB device with the given bus and device."""
26645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez  usb_file_path = '/dev/bus/usb/%03d/%03d' % (bus, device)
27645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez  with open(usb_file_path, 'w') as usb_file:
28645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez    logging.debug('fcntl.ioctl(%s, %d)', usb_file_path, _USBDEVFS_RESET)
29645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez    fcntl.ioctl(usb_file, _USBDEVFS_RESET)
30645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez
31645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez
32645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavezdef reset_android_usb(serial):
33645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez  """Reset the USB device for the given Android device."""
34645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez  lsusb_info = lsusb.lsusb()
35645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez
36645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez  bus = None
37645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez  device = None
38645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez  for device_info in lsusb_info:
39645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez    device_serial = lsusb.get_lsusb_serial(device_info)
40645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez    if device_serial == serial:
41645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez      bus = int(device_info.get('bus'))
42645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez      device = int(device_info.get('device'))
43645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez
44645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez  if bus and device:
45645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez    reset_usb(bus, device)
46645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez  else:
47645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez    raise device_errors.DeviceUnreachableError(
48645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez        'Unable to determine bus or device for device %s' % serial)
49645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez
50645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez
51645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavezdef reset_all_android_devices():
52645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez  """Reset all USB devices that look like an Android device."""
53645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez  _reset_all_matching(lambda i: bool(lsusb.get_lsusb_serial(i)))
54645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez
55645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez
56645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavezdef _reset_all_matching(condition):
57645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez  lsusb_info = lsusb.lsusb()
58645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez  for device_info in lsusb_info:
59645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez    if int(device_info.get('device')) != 1 and condition(device_info):
60645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez      bus = int(device_info.get('bus'))
61645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez      device = int(device_info.get('device'))
62645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez      try:
63645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez        reset_usb(bus, device)
64645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez        serial = lsusb.get_lsusb_serial(device_info)
65645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez        if serial:
66645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez          logging.info('Reset USB device (bus: %03d, device: %03d, serial: %s)',
67645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez              bus, device, serial)
68645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez        else:
69645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez          logging.info('Reset USB device (bus: %03d, device: %03d)',
70645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez              bus, device)
71645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez      except IOError:
72645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez        logging.error(
73645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez            'Failed to reset USB device (bus: %03d, device: %03d)',
74645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez            bus, device)
75645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez
76645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez
77645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavezdef main():
78645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez  parser = argparse.ArgumentParser()
79645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez  parser.add_argument('-v', '--verbose', action='count')
80645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez  parser.add_argument('-s', '--serial')
81645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez  parser.add_argument('--bus', type=int)
82645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez  parser.add_argument('--device', type=int)
83645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez  args = parser.parse_args()
84645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez
85645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez  run_tests_helper.SetLogLevel(args.verbose)
86645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez
87645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez  if args.serial:
88645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez    reset_android_usb(args.serial)
89645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez  elif args.bus and args.device:
90645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez    reset_usb(args.bus, args.device)
91645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez  else:
92645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez    parser.error('Unable to determine target. '
93645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez                 'Specify --serial or BOTH --bus and --device.')
94645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez
95645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez  return 0
96645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez
97645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez
98645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavezif __name__ == '__main__':
99645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez  sys.exit(main())
100645501c2ab19a559ce82a1d5a29ced159a4c30fbLuis Hector Chavez
101