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