1765bee384462e4f64740c5aadfccbf52348b68daYu Shan# Copyright 2017 The Android Open Source Project 2765bee384462e4f64740c5aadfccbf52348b68daYu Shan# 3765bee384462e4f64740c5aadfccbf52348b68daYu Shan# Licensed under the Apache License, Version 2.0 (the "License"); 4765bee384462e4f64740c5aadfccbf52348b68daYu Shan# you may not use this file except in compliance with the License. 5765bee384462e4f64740c5aadfccbf52348b68daYu Shan# You may obtain a copy of the License at 6765bee384462e4f64740c5aadfccbf52348b68daYu Shan# 7765bee384462e4f64740c5aadfccbf52348b68daYu Shan# http://www.apache.org/licenses/LICENSE-2.0 8765bee384462e4f64740c5aadfccbf52348b68daYu Shan# 9765bee384462e4f64740c5aadfccbf52348b68daYu Shan# Unless required by applicable law or agreed to in writing, software 10765bee384462e4f64740c5aadfccbf52348b68daYu Shan# distributed under the License is distributed on an "AS IS" BASIS, 11765bee384462e4f64740c5aadfccbf52348b68daYu Shan# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12765bee384462e4f64740c5aadfccbf52348b68daYu Shan# See the License for the specific language governing permissions and 13765bee384462e4f64740c5aadfccbf52348b68daYu Shan# limitations under the License. 14765bee384462e4f64740c5aadfccbf52348b68daYu Shan 157ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan"""This module provides the USB device serial number to USB location map on Win. 167ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan 17765bee384462e4f64740c5aadfccbf52348b68daYu ShanThis module uses Windows APIs to get USB device information. Ctypes are used to 18765bee384462e4f64740c5aadfccbf52348b68daYu Shansupport calling C type functions. 197ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan""" 207ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan# pylint: disable=invalid-name 217ec8f25d62856d44690270f32119b60d8be9fd39Yu Shanimport ctypes 227ec8f25d62856d44690270f32119b60d8be9fd39Yu Shanfrom ctypes.wintypes import BYTE 237ec8f25d62856d44690270f32119b60d8be9fd39Yu Shanfrom ctypes.wintypes import DWORD 247ec8f25d62856d44690270f32119b60d8be9fd39Yu Shanfrom ctypes.wintypes import ULONG 257ec8f25d62856d44690270f32119b60d8be9fd39Yu Shanfrom ctypes.wintypes import WORD 267ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan 277ec8f25d62856d44690270f32119b60d8be9fd39Yu ShanNULL = None 287ec8f25d62856d44690270f32119b60d8be9fd39Yu ShanDIGCF_ALLCLASSES = 0x4 297ec8f25d62856d44690270f32119b60d8be9fd39Yu ShanDIGCF_PRESENT = 0x2 307ec8f25d62856d44690270f32119b60d8be9fd39Yu ShanULONG_PTR = ctypes.POINTER(ULONG) 317ec8f25d62856d44690270f32119b60d8be9fd39Yu ShanSPDRP_HARDWAREID = 1 327ec8f25d62856d44690270f32119b60d8be9fd39Yu ShanSPDRP_LOCATION_INFORMATION = 0xD 337ec8f25d62856d44690270f32119b60d8be9fd39Yu ShanINVALID_HANDLE_VALUE = -1 347ec8f25d62856d44690270f32119b60d8be9fd39Yu ShanERROR_NO_MORE_ITEMS = 0x103 357ec8f25d62856d44690270f32119b60d8be9fd39Yu ShanBUFFER_SIZE = 1024 367ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan 377ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan 387ec8f25d62856d44690270f32119b60d8be9fd39Yu Shanclass GUID(ctypes.Structure): 397ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan _fields_ = [ 407ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan ('Data1', DWORD), 417ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan ('Data2', WORD), 427ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan ('Data3', WORD), 437ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan ('Data4', BYTE*8), 447ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan ] 457ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan 467ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan 477ec8f25d62856d44690270f32119b60d8be9fd39Yu Shanclass SP_DEVINFO_DATA(ctypes.Structure): 487ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan _fields_ = [ 497ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan ('cbSize', DWORD), 507ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan ('ClassGuid', GUID), 517ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan ('DevInst', DWORD), 527ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan ('Reserved', ULONG_PTR), 537ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan ] 547ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan 557ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan 567ec8f25d62856d44690270f32119b60d8be9fd39Yu Shanclass SerialMapper(object): 577ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan """Maps serial number to its USB physical location. 587ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan 597ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan This class should run under Windows environment. It uses windows setupapi DLL 607ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan to get USB device information. This class is just a wrapper around windows C++ 617ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan library. 627ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan """ 637ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan 647ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan def __init__(self): 657ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan self.setupapi = ctypes.WinDLL('setupapi') 66cbfd415593d4fb9161dd1b932dbf912064c482a5Yu Shan self.serial_map = {} 677ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan 68cbfd415593d4fb9161dd1b932dbf912064c482a5Yu Shan def refresh_serial_map(self): 69cbfd415593d4fb9161dd1b932dbf912064c482a5Yu Shan """Refresh the serial_number -> USB location map. 707ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan """ 717ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan serial_map = {} 727ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan device_inf_set = None 737ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan SetupDiGetClassDevs = self.setupapi.SetupDiGetClassDevsA 747ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan SetupDiEnumDeviceInfo = self.setupapi.SetupDiEnumDeviceInfo 757ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan SetupDiGetDeviceRegistryProperty = ( 767ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan self.setupapi.SetupDiGetDeviceRegistryPropertyA) 777ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan SetupDiGetDeviceInstanceId = self.setupapi.SetupDiGetDeviceInstanceIdA 787ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan SetupDiDestroyDeviceInfoList = self.setupapi.SetupDiDestroyDeviceInfoList 797ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan 807ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan # Get the device information set for all the present USB devices 817ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan flags = DIGCF_ALLCLASSES | DIGCF_PRESENT 827ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan device_inf_set = SetupDiGetClassDevs(NULL, 837ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan ctypes.c_char_p('USB'), 847ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan NULL, 857ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan flags) 867ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan if device_inf_set == INVALID_HANDLE_VALUE: 877ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan raise ctypes.WinError() 887ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan 897ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan devinfo = SP_DEVINFO_DATA() 907ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan p_dev_info = ctypes.byref(devinfo) 917ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan # cbsize is the size of SP_DEVINFO_DATA, need to be set 927ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan devinfo.cbSize = ctypes.sizeof(devinfo) 937ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan i = 0 947ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan while True: 957ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan # Enumerate through the device information set until ERROR_NO_MORE_ITEMS 967ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan # i is the index 977ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan 987ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan # Fill the devinfo 997ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan result = SetupDiEnumDeviceInfo(device_inf_set, i, ctypes.byref(devinfo)) 1007ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan if not result and (ctypes.GetLastError() == ERROR_NO_MORE_ITEMS): 1017ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan # If we reach the last device. 1027ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan break 1037ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan location_buffer = ctypes.create_string_buffer(BUFFER_SIZE) 1047ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan p_location_buffer = ctypes.byref(location_buffer) 1057ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan result = SetupDiGetDeviceRegistryProperty(device_inf_set, 1067ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan p_dev_info, 1077ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan SPDRP_LOCATION_INFORMATION, 1087ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan NULL, 1097ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan p_location_buffer, 1107ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan BUFFER_SIZE, 1117ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan NULL) 1127ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan if not result: 1137ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan i += 1 1147ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan continue 1157ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan location = location_buffer.value 1167ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan device_instance_id_buffer = ctypes.create_string_buffer(BUFFER_SIZE) 1177ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan p_id_buffer = ctypes.byref(device_instance_id_buffer) 1187ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan result = SetupDiGetDeviceInstanceId(device_inf_set, 1197ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan p_dev_info, 1207ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan p_id_buffer, 1217ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan BUFFER_SIZE, 1227ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan NULL) 1237ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan if not result: 1247ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan i += 1 1257ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan continue 1267ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan 1277ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan # device instance id contains a serial number in the format of 1287ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan # [XXX]\[SERIAL] 1297ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan instance_id = device_instance_id_buffer.value 1307ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan instance_parts = instance_id.split('\\') 1317ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan if instance_parts: 132cbfd415593d4fb9161dd1b932dbf912064c482a5Yu Shan serial = instance_parts.pop().lower() 1337ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan serial_map[serial] = location 1347ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan i += 1 1357ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan 1367ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan # Destroy the device information set 1377ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan if device_inf_set is not None: 1387ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan SetupDiDestroyDeviceInfoList(device_inf_set) 139cbfd415593d4fb9161dd1b932dbf912064c482a5Yu Shan self.serial_map = serial_map 140cbfd415593d4fb9161dd1b932dbf912064c482a5Yu Shan 141cbfd415593d4fb9161dd1b932dbf912064c482a5Yu Shan def get_location(self, serial): 142cbfd415593d4fb9161dd1b932dbf912064c482a5Yu Shan """Get the USB location according to the serial number. 143cbfd415593d4fb9161dd1b932dbf912064c482a5Yu Shan 144cbfd415593d4fb9161dd1b932dbf912064c482a5Yu Shan Args: 145cbfd415593d4fb9161dd1b932dbf912064c482a5Yu Shan serial: The serial number for the device. 146cbfd415593d4fb9161dd1b932dbf912064c482a5Yu Shan Returns: 147cbfd415593d4fb9161dd1b932dbf912064c482a5Yu Shan The USB physical location for the device. 148cbfd415593d4fb9161dd1b932dbf912064c482a5Yu Shan """ 149cbfd415593d4fb9161dd1b932dbf912064c482a5Yu Shan serial_lower = serial.lower() 150cbfd415593d4fb9161dd1b932dbf912064c482a5Yu Shan if serial_lower in self.serial_map: 151cbfd415593d4fb9161dd1b932dbf912064c482a5Yu Shan return self.serial_map[serial_lower] 152cbfd415593d4fb9161dd1b932dbf912064c482a5Yu Shan return None 1537ec8f25d62856d44690270f32119b60d8be9fd39Yu Shan 154