1# Copyright 2017 The Android Open Source Project
2#
3# Licensed under the Apache License, Version 2.0 (the "License");
4# you may not use this file except in compliance with the License.
5# You may obtain a copy of the License at
6#
7#     http://www.apache.org/licenses/LICENSE-2.0
8#
9# Unless required by applicable law or agreed to in writing, software
10# distributed under the License is distributed on an "AS IS" BASIS,
11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12# See the License for the specific language governing permissions and
13# limitations under the License.
14
15"""This module provides the USB device serial number to USB location map on Win.
16
17This module uses Windows APIs to get USB device information. Ctypes are used to
18support calling C type functions.
19"""
20# pylint: disable=invalid-name
21import ctypes
22from ctypes.wintypes import BYTE
23from ctypes.wintypes import DWORD
24from ctypes.wintypes import ULONG
25from ctypes.wintypes import WORD
26
27NULL = None
28DIGCF_ALLCLASSES = 0x4
29DIGCF_PRESENT = 0x2
30ULONG_PTR = ctypes.POINTER(ULONG)
31SPDRP_HARDWAREID = 1
32SPDRP_LOCATION_INFORMATION = 0xD
33INVALID_HANDLE_VALUE = -1
34ERROR_NO_MORE_ITEMS = 0x103
35BUFFER_SIZE = 1024
36
37
38class GUID(ctypes.Structure):
39  _fields_ = [
40      ('Data1', DWORD),
41      ('Data2', WORD),
42      ('Data3', WORD),
43      ('Data4', BYTE*8),
44  ]
45
46
47class SP_DEVINFO_DATA(ctypes.Structure):
48  _fields_ = [
49      ('cbSize', DWORD),
50      ('ClassGuid', GUID),
51      ('DevInst', DWORD),
52      ('Reserved', ULONG_PTR),
53  ]
54
55
56class SerialMapper(object):
57  """Maps serial number to its USB physical location.
58
59  This class should run under Windows environment. It uses windows setupapi DLL
60  to get USB device information. This class is just a wrapper around windows C++
61  library.
62  """
63
64  def __init__(self):
65    self.setupapi = ctypes.WinDLL('setupapi')
66    self.serial_map = {}
67
68  def refresh_serial_map(self):
69    """Refresh the serial_number -> USB location map.
70    """
71    serial_map = {}
72    device_inf_set = None
73    SetupDiGetClassDevs = self.setupapi.SetupDiGetClassDevsA
74    SetupDiEnumDeviceInfo = self.setupapi.SetupDiEnumDeviceInfo
75    SetupDiGetDeviceRegistryProperty = (
76        self.setupapi.SetupDiGetDeviceRegistryPropertyA)
77    SetupDiGetDeviceInstanceId = self.setupapi.SetupDiGetDeviceInstanceIdA
78    SetupDiDestroyDeviceInfoList = self.setupapi.SetupDiDestroyDeviceInfoList
79
80    # Get the device information set for all the present USB devices
81    flags = DIGCF_ALLCLASSES | DIGCF_PRESENT
82    device_inf_set = SetupDiGetClassDevs(NULL,
83                                         ctypes.c_char_p('USB'),
84                                         NULL,
85                                         flags)
86    if device_inf_set == INVALID_HANDLE_VALUE:
87      raise ctypes.WinError()
88
89    devinfo = SP_DEVINFO_DATA()
90    p_dev_info = ctypes.byref(devinfo)
91    # cbsize is the size of SP_DEVINFO_DATA, need to be set
92    devinfo.cbSize = ctypes.sizeof(devinfo)
93    i = 0
94    while True:
95      # Enumerate through the device information set until ERROR_NO_MORE_ITEMS
96      # i is the index
97
98      # Fill the devinfo
99      result = SetupDiEnumDeviceInfo(device_inf_set, i, ctypes.byref(devinfo))
100      if not result and (ctypes.GetLastError() == ERROR_NO_MORE_ITEMS):
101        # If we reach the last device.
102        break
103      location_buffer = ctypes.create_string_buffer(BUFFER_SIZE)
104      p_location_buffer = ctypes.byref(location_buffer)
105      result = SetupDiGetDeviceRegistryProperty(device_inf_set,
106                                                p_dev_info,
107                                                SPDRP_LOCATION_INFORMATION,
108                                                NULL,
109                                                p_location_buffer,
110                                                BUFFER_SIZE,
111                                                NULL)
112      if not result:
113        i += 1
114        continue
115      location = location_buffer.value
116      device_instance_id_buffer = ctypes.create_string_buffer(BUFFER_SIZE)
117      p_id_buffer = ctypes.byref(device_instance_id_buffer)
118      result = SetupDiGetDeviceInstanceId(device_inf_set,
119                                          p_dev_info,
120                                          p_id_buffer,
121                                          BUFFER_SIZE,
122                                          NULL)
123      if not result:
124        i += 1
125        continue
126
127      # device instance id contains a serial number in the format of
128      # [XXX]\[SERIAL]
129      instance_id = device_instance_id_buffer.value
130      instance_parts = instance_id.split('\\')
131      if instance_parts:
132        serial = instance_parts.pop().lower()
133        serial_map[serial] = location
134      i += 1
135
136    # Destroy the device information set
137    if device_inf_set is not None:
138      SetupDiDestroyDeviceInfoList(device_inf_set)
139    self.serial_map = serial_map
140
141  def get_location(self, serial):
142    """Get the USB location according to the serial number.
143
144    Args:
145      serial: The serial number for the device.
146    Returns:
147      The USB physical location for the device.
148    """
149    serial_lower = serial.lower()
150    if serial_lower in self.serial_map:
151      return self.serial_map[serial_lower]
152    return None
153
154