FilesExpand file tree

master

/

Windows.py

Copy path

Blame

More file actions

Blame

More file actions

Latest commit

 

History

History

History

160 lines (137 loc) · 6.51 KB

master

/

Windows.py

Copy path

Top

File metadata and controls

• Code
• Blame

160 lines (137 loc) · 6.51 KB

Raw

Copy raw file

Download raw file

Open symbols panel

Edit and raw actions

1

2

3

4

5

6

7

8

9

10

11

12

13

14

15

16

17

18

19

20

21

22

23

24

25

26

27

28

29

30

31

32

33

34

35

36

37

38

39

40

41

42

43

44

45

46

47

48

49

50

51

52

53

54

55

56

57

58

59

60

61

62

63

64

65

66

67

68

69

70

71

72

73

74

75

76

77

78

79

80

81

82

83

84

85

86

87

88

89

90

91

92

93

94

95

96

97

98

99

100

101

102

103

104

105

106

107

108

109

110

111

112

113

114

115

116

117

118

119

120

121

122

123

124

125

126

127

128

129

130

131

132

133

134

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

import copy

import json

import time

from enum import Enum

try:

import wmi # pylint: disable=import-error

except Exception: # pylint: disable=broad-except

# Dummy WMI

class wmi: # pylint: disable=invalid-name

class WMI:

@staticmethod

def query(*args, **kwargs): # pylint: disable=unused-argument

return []

from base import BaseUSBMap

from Scripts import shared, usbdump

class PnpDeviceProperties(Enum):

ACPI_PATH = "DEVPKEY_Device_BiosDeviceName"

DRIVER_KEY = "DEVPKEY_Device_Driver"

LOCATION_PATHS = "DEVPKEY_Device_LocationPaths"

FRIENDLY_NAME = "DEVPKEY_Device_FriendlyName"

BUS_DEVICE_FUNCTION = "DEVPKEY_Device_LocationInfo"

BUS_REPORTED_NAME = "DEVPKEY_Device_BusReportedDeviceDesc"

BUS_RELATIONS = "DEVPKEY_Device_BusRelations"

INTERFACE = "DEVPKEY_PciDevice_ProgIf"

SERVICE = "DEVPKEY_Device_Service"

class WindowsUSBMap(BaseUSBMap):

def __init__(self):

self.usbdump = None

if shared.test_mode:

self.wmi = {}

self.wmi_cache: dict = {p["DEVPKEY_Device_InstanceId"]: p for p in json.load(shared.debug_dump_path.open())["wmitest"] if "duration" not in p}

else:

self.wmi = wmi.WMI()

self.wmi_cache = {}

self.wmi_retries = {}

super().__init__()

def update_usbdump(self):

self.usbdump = usbdump.get_controllers()

def get_property_from_wmi(self, instance_id, prop: PnpDeviceProperties):

MAX_TRIES = 2

result = None

if self.wmi_cache.get(instance_id, {}).get(prop.value):

return self.wmi_cache[instance_id][prop.value]

elif self.wmi_retries.get(instance_id, {}).get(prop.value, 0) >= MAX_TRIES:

return None

try:

result = self.wmi.query(f"SELECT * FROM Win32_PnPEntity WHERE PNPDeviceID = '{instance_id}'")[0].GetDeviceProperties([prop.value])[0][0].Data

except IndexError:

# Race condition between unplug detected in usbdump and WMI

return None

except AttributeError:

if not self.wmi_retries.get(instance_id):

self.wmi_retries[instance_id] = {prop.value: 1}

elif not self.wmi_retries[instance_id].get(prop.value):

self.wmi_retries[instance_id][prop.value] = 1

else:

self.wmi_retries[instance_id][prop.value] += 1

return None

if not self.wmi_cache.get(instance_id):

self.wmi_cache[instance_id] = {prop.value: result}

else:

self.wmi_cache[instance_id][prop.value] = result

return result

""" def get_property_from_wmi(self, instance_id, prop: PnpDeviceProperties):

value = self.wmi.setdefault(instance_id, {}).get(prop.value)

if value:

return value

else:

value = input(f"Enter value for {instance_id} {prop}: ")

self.wmi[instance_id][prop.value] = value

json.dump(self.wmi, self.wmi_path.open("w"), indent=4, sort_keys=True)

return value """

def get_name_from_wmi(self, device):

if not isinstance(device, dict):

return

if device.get("error") or not device["instance_id"]:

return

device["name"] = self.get_property_from_wmi(device["instance_id"], PnpDeviceProperties.BUS_REPORTED_NAME) or device["name"]

for i in device["devices"]:

self.get_name_from_wmi(i)

def get_controller_class(self, controller):

interface = self.get_property_from_wmi(controller["identifiers"]["instance_id"], PnpDeviceProperties.INTERFACE)

if interface:

return shared.USBControllerTypes(interface)

service = self.get_property_from_wmi(controller["identifiers"]["instance_id"], PnpDeviceProperties.SERVICE)

if not isinstance(service, str):

shared.debug(f"Unknown controller type for interface {interface} and service {service}!")

return shared.USBControllerTypes.Unknown

if service.lower() == "usbxhci":

return shared.USBControllerTypes.XHCI

elif service.lower() == "usbehci":

return shared.USBControllerTypes.EHCI

elif service.lower() == "usbohci":

return shared.USBControllerTypes.OHCI

elif service.lower() == "usbuhci":

return shared.USBControllerTypes.UHCI

else:

shared.debug(f"Unknown controller type for interface {interface} and service {service}!")

return shared.USBControllerTypes.Unknown

def get_controllers(self):

# self.update_usbdump()

for i in range(10):

try:

# time_it(self.update_usbdump, "USBdump time")

self.update_usbdump()

break

except Exception as e:

if i == 9:

raise

else:

shared.debug(e)

time.sleep(0.05 if shared.debugging else 2)

controllers = self.usbdump

for controller in controllers:

controller["name"] = self.get_property_from_wmi(controller["identifiers"]["instance_id"], PnpDeviceProperties.FRIENDLY_NAME) or controller["name"]

controller["class"] = self.get_controller_class(controller)

acpi_path = self.get_property_from_wmi(controller["identifiers"]["instance_id"], PnpDeviceProperties.ACPI_PATH)

if acpi_path:

controller["identifiers"]["acpi_path"] = acpi_path

driver_key = self.get_property_from_wmi(controller["identifiers"]["instance_id"], PnpDeviceProperties.DRIVER_KEY)

if driver_key:

controller["identifiers"]["driver_key"] = driver_key

location_paths = self.get_property_from_wmi(controller["identifiers"]["instance_id"], PnpDeviceProperties.LOCATION_PATHS)

if location_paths:

controller["identifiers"]["location_paths"] = location_paths

# controller["identifiers"]["bdf"] = self.get_property_from_wmi(controller["identifiers"]["instance_id"], PnpDeviceProperties.BUS_DEVICE_FUNCTION)

for port in controller["ports"]:

for device in port["devices"]:

self.get_name_from_wmi(device)

self.controllers = controllers

if not self.controllers_historical:

self.controllers_historical = copy.deepcopy(self.controllers)

else:

self.merge_controllers(self.controllers_historical, self.controllers)

def update_devices(self):

self.get_controllers()

e = WindowsUSBMap()