1# Copyright 2017 The Chromium OS Authors. All rights reserved.
2# Use of this source code is governed by a BSD-style license that can be
3# found in the LICENSE file.
4
5"""Power cycle a usb port on DUT(device under test)."""
6
7from __future__ import print_function
8
9import logging
10import os
11import time
12
13TOKEN_NEW_BUS = '/:  '
14TOKEN_ROOT_DEVICE = '\n    |__ '
15
16# On board guado, there are three gpios that control usb port power:
17# Front left usb port:  218, port number: 2
18# Front right usb port: 219, port number: 3
19# Rear dual usb ports:  209, port number: 5,6
20PORT_NUM_DICT = {
21    'guado': {
22        'bus1': {
23            2: 'front_left',
24            3: 'front_right',
25            5: 'back_dual',
26            6: 'back_dual'
27        }
28    }
29}
30PORT_GPIO_DICT = {
31    'guado': {
32        'bus1': {
33            'front_left': 218,
34            'front_right': 219,
35            'back_dual': 209
36        }
37    }
38}
39
40
41def power_cycle_usb_gpio(dut, gpio_idx, pause=1):
42    """
43    Power cycle a usb port on dut via its gpio index.
44
45    Each usb port has corresponding gpio controling its power. If the gpio
46    index of the gpio is known, the power cycling procedure is pretty
47    straightforward.
48
49    @param dut: The handle of the device under test. Should be initialized in
50            autotest.
51    @param gpio_idx: The index of the gpio that controls power of the usb port
52            we want to reset.
53    @param pause: The waiting time before powering on usb device, unit is second.
54
55    """
56    if gpio_idx is None:
57        return
58    export_flag = False
59    if not dut.path_exists('/sys/class/gpio/gpio{}'.format(gpio_idx)):
60        export_flag = True
61        cmd = 'echo {} > /sys/class/gpio/export'.format(gpio_idx)
62        dut.run(cmd)
63    cmd = 'echo out > /sys/class/gpio/gpio{}/direction'.format(gpio_idx)
64    dut.run(cmd)
65    cmd = 'echo 0 > /sys/class/gpio/gpio{}/value'.format(gpio_idx)
66    dut.run(cmd)
67    time.sleep(pause)
68    cmd = 'echo 1 > /sys/class/gpio/gpio{}/value'.format(gpio_idx)
69    dut.run(cmd)
70    if export_flag:
71        cmd = 'echo {} > /sys/class/gpio/unexport'.format(gpio_idx)
72        dut.run(cmd)
73
74
75def power_cycle_usb_vidpid(dut, board, vid, pid):
76    """
77    Power cycle a usb port on DUT via peripharel's VID and PID.
78
79    When only the VID and PID of the peripharel is known, a search is needed
80    to decide which port it connects to by its VID and PID and look up the gpio
81    index according to the board and port number in the dictionary. Then the
82    USB port is power cycled using the gpio number.
83
84    @param dut: The handle of the device under test.
85    @param board: Board name ('guado', etc.)
86    @param vid: Vendor ID of the peripharel device.
87    @param pid: Product ID of the peripharel device.
88
89    @raise KeyError if the target device wasn't found by given VID and PID.
90
91    """
92    bus_idx, port_idx = get_port_number_from_vidpid(dut, vid, pid)
93    if port_idx is None:
94        raise KeyError('Couldn\'t find target device, {}:{}.'.format(vid, pid))
95    logging.info('found device bus {} port {}'.format(bus_idx, port_idx))
96    token_bus = 'bus{}'.format(bus_idx)
97    target_gpio_pos = (PORT_NUM_DICT.get(board, {})
98                       .get(token_bus, {}).get(port_idx, ''))
99    target_gpio = (PORT_GPIO_DICT.get(board, {})
100                   .get(token_bus, {}).get(target_gpio_pos, None))
101    logging.info('target gpio num {}'.format(target_gpio))
102    power_cycle_usb_gpio(dut, target_gpio)
103
104
105def get_port_number_from_vidpid(dut, vid, pid):
106    """
107    Get bus number and port number a device is connected to on DUT.
108
109    Get the bus number and port number of the usb port the target perpipharel
110    device is connected to.
111
112    @param dut: The handle of the device under test.
113    @param vid: Vendor ID of the peripharel device.
114    @param pid: Product ID of the peripharel device.
115
116    @returns the target bus number and port number, if device not found, returns
117          (None, None).
118
119    """
120    cmd = 'lsusb -d {}:{}'.format(vid, pid)
121    lsusb_output = dut.run(cmd, ignore_status=True).stdout
122    logging.info('lsusb output {}'.format(lsusb_output))
123    target_bus_idx, target_dev_idx = get_bus_dev_id(lsusb_output, vid, pid)
124    if target_bus_idx is None:
125        return None, None
126    cmd = 'lsusb -t'
127    lsusb_output = dut.run(cmd, ignore_status=True).stdout
128    target_port_number = get_port_number(
129        lsusb_output, target_bus_idx, target_dev_idx)
130    return target_bus_idx, target_port_number
131
132
133def get_bus_dev_id(lsusb_output, vid, pid):
134    """
135    Get bus number and device index a device is connected to on DUT.
136
137    Get the bus number and port number of the usb port the target perpipharel
138    device is connected to based on the output of command 'lsusb -d VID:PID'.
139
140    @param lsusb_output: output of command 'lsusb -d VID:PID' running on DUT.
141    @param vid: Vendor ID of the peripharel device.
142    @param pid: Product ID of the peripharel device.
143
144    @returns the target bus number and device index, if device not found,
145          returns (None, None).
146
147    """
148    if lsusb_output == '':
149        return None, None
150    lsusb_device_info = lsusb_output.strip().split('\n')
151    if len(lsusb_device_info) > 1:
152        logging.info('find more than one device with VID:PID: %s:%s', vid, pid)
153        return None, None
154    # An example of the info line is 'Bus 001 Device 006:  ID 266e:0110 ...'
155    fields = lsusb_device_info[0].split(' ')
156    assert len(fields) >= 6, 'Wrong info format: {}'.format(lsusb_device_info)
157    target_bus_idx = int(fields[1])
158    target_device_idx = int(fields[3][:-1])
159    logging.info('found target device %s:%s, bus: %d, dev: %d',
160                 vid, pid, target_bus_idx, target_device_idx)
161    return target_bus_idx, target_device_idx
162
163def get_port_number(lsusb_tree_output, bus, dev):
164    """
165    Get port number that certain device is connected to on DUT.
166
167    Get the port number of the usb port that the target peripharel device is
168    connected to based on the output of command 'lsusb -t', its bus number and
169    device index.
170    An example of lsusb_tree_output could be:
171    /:  Bus 02.Port 1: Dev 1, Class=root_hub, Driver=xhci_hcd/4p, 5000M
172        |__ Port 2: Dev 2, If 0, Class=Hub, Driver=hub/4p, 5000M
173    /:  Bus 01.Port 1: Dev 1, Class=root_hub, Driver=xhci_hcd/11p, 480M
174        |__ Port 2: Dev 52, If 0, Class=Hub, Driver=hub/4p, 480M
175            |__ Port 1: Dev 55, If 0, Class=Human Interface Device,
176                        Driver=usbhid, 12M
177            |__ Port 3: Dev 54, If 0, Class=Vendor Specific Class,
178                        Driver=udl, 480M
179        |__ Port 3: Dev 3, If 0, Class=Hub, Driver=hub/4p, 480M
180        |__ Port 4: Dev 4, If 0, Class=Wireless, Driver=btusb, 12M
181        |__ Port 4: Dev 4, If 1, Class=Wireless, Driver=btusb, 12M
182
183    @param lsusb_tree_output: The output of command 'lsusb -t' on DUT.
184    @param bus: The bus number the peripharel device is connected to.
185    @param dev: The device index of the peripharel device on DUT.
186
187    @returns the target port number, if device not found, returns None.
188
189    """
190    lsusb_device_buses = lsusb_tree_output.strip().split(TOKEN_NEW_BUS)
191    target_bus_token = 'Bus {:02d}.'.format(bus)
192    for bus_info in lsusb_device_buses:
193        if bus_info.find(target_bus_token) != 0:
194            continue
195        target_dev_token = 'Dev {}'.format(dev)
196        device_info = bus_info.strip(target_bus_token).split(TOKEN_ROOT_DEVICE)
197        for device in device_info:
198            if target_dev_token not in device:
199                continue
200            target_port_number = int(device.split(':')[0].split(' ')[1])
201            return target_port_number
202    return None
203
204
205def get_all_port_number_from_vidpid(dut, vid, pid):
206    """
207    Get the list of bus number and port number devices are connected to DUT.
208
209    Get the the list of bus number and port number of the usb ports the target
210           perpipharel devices are connected to.
211
212    @param dut: The handle of the device under test.
213    @param vid: Vendor ID of the peripharel device.
214    @param pid: Product ID of the peripharel device.
215
216    @returns the list of target bus number and port number, if device not found,
217            returns empty list.
218
219    """
220    port_number = []
221    cmd = 'lsusb -d {}:{}'.format(vid, pid)
222    lsusb_output = dut.run(cmd, ignore_status=True).stdout
223    (target_bus_idx, target_dev_idx) = get_all_bus_dev_id(lsusb_output, vid, pid)
224    if target_bus_idx is None:
225        return None, None
226    cmd = 'lsusb -t'
227    lsusb_output = dut.run(cmd, ignore_status=True).stdout
228    for bus, dev in zip(target_bus_idx, target_dev_idx):
229        port_number.append(get_port_number(
230            lsusb_output, bus, dev))
231    return (target_bus_idx, port_number)
232
233
234def get_all_bus_dev_id(lsusb_output, vid, pid):
235    """
236    Get the list of bus number and device index devices are connected to DUT.
237
238    Get the bus number and port number of the usb ports the target perpipharel
239            devices are connected to based on the output of command 'lsusb -d VID:PID'.
240
241    @param lsusb_output: output of command 'lsusb -d VID:PID' running on DUT.
242    @param vid: Vendor ID of the peripharel device.
243    @param pid: Product ID of the peripharel device.
244
245    @returns the list of target bus number and device index, if device not found,
246           returns empty list.
247
248    """
249    bus_idx = []
250    device_idx =[]
251    if lsusb_output == '':
252        return None, None
253    lsusb_device_info = lsusb_output.strip().split('\n')
254    for lsusb_device in lsusb_device_info:
255        fields = lsusb_device.split(' ')
256        assert len(fields) >= 6, 'Wrong info format: {}'.format(lsusb_device_info)
257        target_bus_idx = int(fields[1])
258        target_device_idx = int(fields[3][:-1])
259        bus_idx.append(target_bus_idx)
260        device_idx.append( target_device_idx)
261    return (bus_idx, device_idx)
262
263
264def get_target_all_gpio(dut, board, vid, pid):
265    """
266    Get GPIO for all devices with vid, pid connected to on DUT.
267
268    Get gpio of usb port the target perpipharel  devices are
269    connected to based on the output of command 'lsusb -d VID:PID'.
270
271    @param dut: The handle of the device under test.
272    @param board: Board name ('guado', etc.)
273    @param vid: Vendor ID of the peripharel device.
274    @param pid: Product ID of the peripharel device.
275
276    @returns the list of gpio, if no device found return []
277
278    """
279    gpio_list = []
280    (bus_idx, port_idx) = get_all_port_number_from_vidpid(dut, vid, pid)
281    if port_idx is None:
282        raise KeyError('Couldn\'t find target device, {}:{}.'.format(vid, pid))
283
284    for bus, port in zip(bus_idx, port_idx):
285        logging.info('found device bus {} port {}'.format(bus, port))
286        token_bus = 'bus{}'.format(bus)
287        target_gpio_pos = (PORT_NUM_DICT.get(board, {})
288                       .get(token_bus, {}).get(port, ''))
289        target_gpio = (PORT_GPIO_DICT.get(board, {})
290                   .get(token_bus, {}).get(target_gpio_pos, None))
291        logging.info('Target gpio num {}'.format(target_gpio))
292        gpio_list.append(target_gpio)
293    return gpio_list
294