Home | History | Annotate | Download | only in at-factory-tool 
      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 
     17 This module uses Windows APIs to get USB device information. Ctypes are used to
     18 support calling C type functions.
     19 """
     20 # pylint: disable=invalid-name
     21 import ctypes
     22 from ctypes.wintypes import BYTE
     23 from ctypes.wintypes import DWORD
     24 from ctypes.wintypes import ULONG
     25 from ctypes.wintypes import WORD
     26 
     27 NULL = None
     28 DIGCF_ALLCLASSES = 0x4
     29 DIGCF_PRESENT = 0x2
     30 ULONG_PTR = ctypes.POINTER(ULONG)
     31 SPDRP_HARDWAREID = 1
     32 SPDRP_LOCATION_INFORMATION = 0xD
     33 INVALID_HANDLE_VALUE = -1
     34 ERROR_NO_MORE_ITEMS = 0x103
     35 BUFFER_SIZE = 1024
     36 
     37 
     38 class GUID(ctypes.Structure):
     39   _fields_ = [
     40       ('Data1', DWORD),
     41       ('Data2', WORD),
     42       ('Data3', WORD),
     43       ('Data4', BYTE*8),
     44   ]
     45 
     46 
     47 class SP_DEVINFO_DATA(ctypes.Structure):
     48   _fields_ = [
     49       ('cbSize', DWORD),
     50       ('ClassGuid', GUID),
     51       ('DevInst', DWORD),
     52       ('Reserved', ULONG_PTR),
     53   ]
     54 
     55 
     56 class 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