atomicshop 2.21.0__py3-none-any.whl → 3.0.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of atomicshop might be problematic. Click here for more details.

Files changed (32) hide show
  1. atomicshop/__init__.py +1 -1
  2. atomicshop/basics/multiprocesses.py +228 -30
  3. atomicshop/dns.py +2 -0
  4. atomicshop/mitm/config_static.py +2 -1
  5. atomicshop/mitm/connection_thread_worker.py +124 -157
  6. atomicshop/mitm/engines/__parent/responder___parent.py +43 -43
  7. atomicshop/mitm/engines/__reference_general/responder___reference_general.py +43 -43
  8. atomicshop/mitm/engines/create_module_template.py +2 -7
  9. atomicshop/mitm/import_config.py +30 -26
  10. atomicshop/mitm/initialize_engines.py +9 -24
  11. atomicshop/mitm/mitm_main.py +187 -59
  12. atomicshop/networks.py +448 -0
  13. atomicshop/wrappers/ctyping/setup_device.py +466 -0
  14. atomicshop/wrappers/dockerw/dockerw.py +17 -21
  15. atomicshop/wrappers/mongodbw/mongodbw.py +1 -0
  16. atomicshop/wrappers/psutilw/{networks.py → psutil_networks.py} +3 -1
  17. atomicshop/wrappers/pywin32w/wmis/msft_netipaddress.py +76 -0
  18. atomicshop/wrappers/pywin32w/wmis/win32_networkadapterconfiguration.py +262 -0
  19. atomicshop/wrappers/pywin32w/wmis/win32networkadapter.py +51 -82
  20. atomicshop/wrappers/pywin32w/wmis/wmi_helpers.py +235 -0
  21. atomicshop/wrappers/socketw/accepter.py +15 -1
  22. atomicshop/wrappers/socketw/creator.py +7 -1
  23. atomicshop/wrappers/socketw/dns_server.py +33 -39
  24. atomicshop/wrappers/socketw/exception_wrapper.py +20 -11
  25. atomicshop/wrappers/socketw/socket_wrapper.py +29 -78
  26. atomicshop/wrappers/winregw/winreg_network.py +20 -0
  27. {atomicshop-2.21.0.dist-info → atomicshop-3.0.0.dist-info}/METADATA +2 -1
  28. {atomicshop-2.21.0.dist-info → atomicshop-3.0.0.dist-info}/RECORD +31 -27
  29. atomicshop/wrappers/pywin32w/wmis/helpers.py +0 -131
  30. {atomicshop-2.21.0.dist-info → atomicshop-3.0.0.dist-info}/LICENSE.txt +0 -0
  31. {atomicshop-2.21.0.dist-info → atomicshop-3.0.0.dist-info}/WHEEL +0 -0
  32. {atomicshop-2.21.0.dist-info → atomicshop-3.0.0.dist-info}/top_level.txt +0 -0
@@ -0,0 +1,262 @@
1
+ from typing import Union
2
+ import socket
3
+ import time
4
+
5
+ from win32com.client import CDispatch
6
+
7
+ from . import wmi_helpers, win32networkadapter
8
+ from ...psutilw import psutil_networks
9
+ from .... import ip_addresses
10
+
11
+
12
+ def get_network_configuration_by_adapter(
13
+ adapter: CDispatch,
14
+ wmi_instance: CDispatch = None
15
+ ) -> Union[CDispatch, None]:
16
+ """
17
+ Get the network configuration for a specific adapter index.
18
+
19
+ :param adapter: CDispatch, Win32_NetworkAdapter object.
20
+ :param wmi_instance: WMI instance. You can get it from:
21
+ wmi_helpers.get_wmi_instance()
22
+ :return: Win32_NetworkAdapterConfiguration object.
23
+ """
24
+
25
+ network_config: CDispatch = wmi_instance.ExecQuery(
26
+ f"SELECT * FROM Win32_NetworkAdapterConfiguration WHERE Index={adapter.Index}")[0]
27
+
28
+ return network_config
29
+
30
+
31
+ def get_adapter_network_configuration(
32
+ use_default_interface: bool = False,
33
+ connection_name: str = None,
34
+ mac_address: str = None,
35
+ wmi_instance: CDispatch = None
36
+ ) -> tuple:
37
+ """
38
+ Get the WMI network configuration for a network adapter.
39
+ :param use_default_interface: bool, if True, the default network interface will be used.
40
+ This is the adapter that your internet is being used from.
41
+ :param connection_name: string, adapter name as shown in the network settings.
42
+ :param mac_address: string, MAC address of the adapter. Format: '00:00:00:00:00:00'.
43
+ :param wmi_instance: WMI instance. You can get it from:
44
+ wrappers.pywin32s.wmis.wmi_helpers.get_wmi_instance()
45
+ or default will be used.
46
+ :return: tuple(Win32_NetworkAdapterConfiguration, Win32_NetworkAdapter)
47
+ """
48
+
49
+ if use_default_interface and connection_name:
50
+ raise ValueError("Only one of 'use_default_interface' or 'connection_name' must be provided.")
51
+ elif not use_default_interface and not connection_name:
52
+ raise ValueError("Either 'use_default_interface' or 'connection_name' must be provided.")
53
+
54
+ if not wmi_instance:
55
+ wmi_instance, _ = wmi_helpers.get_wmi_instance()
56
+
57
+ adapters = win32networkadapter.list_network_adapters(wmi_instance)
58
+
59
+ current_adapter = None
60
+ if use_default_interface:
61
+ default_connection_name_dict: dict = psutil_networks.get_default_connection_name()
62
+ if not default_connection_name_dict:
63
+ raise wmi_helpers.WMINetworkAdapterNotFoundError("Default network adapter not found.")
64
+ # Get the first key from the dictionary.
65
+ connection_name: str = list(default_connection_name_dict.keys())[0]
66
+
67
+ if connection_name is None and mac_address is None:
68
+ raise ValueError("Either 'connection_name' or 'mac_address' must be provided.")
69
+ elif connection_name and mac_address:
70
+ raise ValueError("Only one of 'connection_name' or 'mac_address' must be provided.")
71
+
72
+ if connection_name:
73
+ for adapter in adapters:
74
+ if connection_name == adapter.NetConnectionID:
75
+ current_adapter = adapter
76
+ break
77
+
78
+ if not current_adapter:
79
+ raise wmi_helpers.WMINetworkAdapterNotFoundError(f"Adapter with connection name '{connection_name}' not found.")
80
+ elif mac_address:
81
+ for adapter in adapters:
82
+ if mac_address == adapter.MACAddress:
83
+ current_adapter = adapter
84
+ break
85
+
86
+ if current_adapter is None:
87
+ raise wmi_helpers.WMINetworkAdapterNotFoundError(f"Adapter with MAC address '{mac_address}' not found.")
88
+
89
+ # Query the network adapter configurations
90
+ query = f"SELECT * FROM Win32_NetworkAdapterConfiguration WHERE Index='{current_adapter.DeviceID}'"
91
+ adapter_configs = wmi_instance.ExecQuery(query)
92
+
93
+ # Check if the adapter exists
94
+ if len(adapter_configs) == 0:
95
+ raise wmi_helpers.WMINetworkAdapterNotFoundError(f"Adapter with connection name '{connection_name}' not found.")
96
+
97
+ return adapter_configs[0], current_adapter
98
+
99
+
100
+ def set_static_ips(
101
+ network_config: CDispatch, # Win32_NetworkAdapterConfiguration (CDispatch)
102
+ ips: list[str], # ["192.168.157.3", ...]
103
+ masks: list[str], # ["255.255.255.0", ...]
104
+ gateways: list[str] = None,
105
+ dns_gateways: list[str] = None,
106
+ availability_wait_seconds: int = 0
107
+ ) -> None:
108
+ """
109
+ • network_config – Win32_NetworkAdapterConfiguration instance for the target NIC
110
+ (you already have it from GetObject / WMI query).
111
+ • ips – list of IPv4 strings.
112
+ • masks – matching subnet‑mask list (same length as ipv4).
113
+ • gateways – list of default gateways (optional).
114
+ • dns_gateways – list of DNS servers (optional).
115
+ • availability_wait_seconds – seconds to wait for the adapter to become available.
116
+ 0 means no wait.
117
+
118
+ Raises RuntimeError if Windows reports anything other than success (0 / 1)
119
+ or "object already exists" (22) for each operation.
120
+
121
+ ==========
122
+
123
+ Example:
124
+ cfg = wmi_instance.Get("Win32_NetworkAdapterConfiguration.Index=12") # your adapter
125
+ set_static_ips(
126
+ cfg,
127
+ ipv4=["192.168.157.129", "192.168.157.3", "192.168.157.4"],
128
+ masks=["255.255.255.0"] * 3,
129
+ # gateways=["192.168.157.2"],
130
+ # dns_gateways=["8.8.8.8", "1.1.1.1"]
131
+ )
132
+ """
133
+
134
+ initial_default_ipv4: str = socket.gethostbyname(socket.gethostname())
135
+
136
+ # -------------------- IPv4 via EnableStatic ----------------------------
137
+ if not masks or len(ips) != len(masks):
138
+ raise ValueError("ipv4 and masks must be lists of equal length")
139
+
140
+ in_params = network_config.Methods_("EnableStatic").InParameters.SpawnInstance_()
141
+ in_params.IPAddress = ips
142
+ in_params.SubnetMask = masks
143
+
144
+ rc = network_config.ExecMethod_("EnableStatic", in_params).Properties_('ReturnValue').Value
145
+ if rc not in (0, 1): # 0 = reboot required, 1 = OK
146
+ raise RuntimeError(f"EnableStatic (IPv4) failed, code {rc}")
147
+
148
+ # -------------------- Default Gateway via SetGateways -------------------
149
+ if gateways:
150
+ gateway_metrics = [1] * len(gateways)
151
+ in_params = network_config.Methods_("SetGateways").InParameters.SpawnInstance_()
152
+ in_params.DefaultIPGateway = gateways
153
+ in_params.GatewayCostMetric = [int(m) for m in gateway_metrics]
154
+
155
+ rc = network_config.ExecMethod_("SetGateways", in_params) \
156
+ .Properties_('ReturnValue').Value
157
+ if rc not in (0, 1):
158
+ raise RuntimeError(f"SetGateways failed, code {rc}")
159
+
160
+ # -------------------- DNS via SetDNSServerSearchOrder ------------------
161
+ if dns_gateways:
162
+ in_params = network_config.Methods_("SetDNSServerSearchOrder").InParameters.SpawnInstance_()
163
+ in_params.DNSServerSearchOrder = dns_gateways
164
+
165
+ rc = network_config.ExecMethod_("SetDNSServerSearchOrder", in_params).Properties_('ReturnValue').Value
166
+ if rc not in (0, 1):
167
+ raise RuntimeError(f"SetDNSServerSearchOrder failed, code {rc}")
168
+
169
+ # -------------------- Wait for the adapter to become available -----------
170
+ if availability_wait_seconds > 0:
171
+ count = 0
172
+ while count < 15:
173
+ current_default_ipv4: str = socket.gethostbyname(socket.gethostname())
174
+ if current_default_ipv4 == initial_default_ipv4:
175
+ # print(f"[+] Adapter is available: {current_default_ipv4}")
176
+ break
177
+ else:
178
+ # print(f"[!] Adapter is not available yet: [{current_default_ipv4}]")
179
+ count += 1
180
+
181
+ time.sleep(1)
182
+
183
+
184
+ def set_dynamic_ip(
185
+ nic_cfg,
186
+ reset_dns: bool = True,
187
+ reset_wins: bool = True
188
+ ) -> None:
189
+ """
190
+ Switch the adapter represented by *nic_cfg* (a Win32_NetworkAdapterConfiguration
191
+ COM object) to DHCP.
192
+
193
+ Parameters
194
+ ----------
195
+ nic_cfg : CDispatch
196
+ The adapter’s Win32_NetworkAdapterConfiguration instance (IPEnabled = TRUE).
197
+ reset_dns : bool, default True
198
+ Also clear any static DNS servers (calls SetDNSServerSearchOrder(None)).
199
+ reset_wins : bool, default True
200
+ Also clear any static WINS servers (calls SetWINSServer(None, None)).
201
+
202
+ Raises
203
+ ------
204
+ RuntimeError
205
+ If any WMI call returns a status other than 0 (“Success”) or 1 (“Restart required”).
206
+ """
207
+
208
+ # 1) Turn on DHCP for IPv4
209
+ wmi_helpers.call_method(nic_cfg, 'EnableDHCP')
210
+
211
+ # 2) Clear static gateways (otherwise Windows keeps using them)
212
+ wmi_helpers.call_method(nic_cfg, 'SetGateways', ([], [])) # empty SAFEARRAY → remove gateways
213
+
214
+ # 3) Optional: reset DNS
215
+ if reset_dns:
216
+ wmi_helpers.call_method(nic_cfg, 'SetDNSServerSearchOrder', None) # None = DHCP-provided DNS
217
+
218
+ # 4) Optional: reset WINS
219
+ if reset_wins:
220
+ wmi_helpers.call_method(nic_cfg, 'SetWINSServer', ("", ""))
221
+
222
+
223
+ def get_info_from_network_config(
224
+ network_config: CDispatch
225
+ ) -> dict:
226
+ """
227
+ Collect information about adapter that currently carries the default route.
228
+
229
+ :param network_config: CDispatch, Win32_NetworkAdapterConfiguration object.
230
+ :return: dict of the default adapter.
231
+ """
232
+
233
+ def _split_ips(config):
234
+ """Split IPAddress[] into separate v4 / v6 lists."""
235
+ current_ipv4s: list[str] = []
236
+ current_ipv4_masks: list[str] = []
237
+ current_ipv6s: list[str] = []
238
+ current_ipv6_prefixes: list[int] = []
239
+ for address_index, ip_address in enumerate(config.IPAddress):
240
+ if ip_addresses.is_ip_address(ip_address, 'ipv4'):
241
+ current_ipv4s.append(ip_address)
242
+ current_ipv4_masks.append(config.IPSubnet[address_index])
243
+ elif ip_addresses.is_ip_address(ip_address, 'ipv6'):
244
+ current_ipv6s.append(ip_address)
245
+ current_ipv6_prefixes.append(int(config.IPSubnet[address_index]))
246
+
247
+ return current_ipv4s, current_ipv6s, current_ipv4_masks, current_ipv6_prefixes
248
+
249
+ ipv4s, ipv6s, ipv4subnets, ipv6prefixes = _split_ips(network_config)
250
+ adapter = {
251
+ "description": network_config.Description,
252
+ "interface_index": network_config.InterfaceIndex,
253
+ "is_dynamic": bool(network_config.DHCPEnabled),
254
+ "ipv4s": ipv4s,
255
+ "ipv6s": ipv6s,
256
+ "ipv4_subnet_masks": ipv4subnets,
257
+ "ipv6_prefixes": ipv6prefixes,
258
+ "default_gateways": list(network_config.DefaultIPGateway or []),
259
+ "dns_gateways": list(network_config.DNSServerSearchOrder or []),
260
+ }
261
+
262
+ return adapter
@@ -1,112 +1,81 @@
1
1
  from typing import Union
2
2
 
3
- import win32com.client
3
+ from win32com.client import CDispatch
4
4
 
5
- from . import helpers
6
- from ...psutilw import networks
5
+ from . import wmi_helpers, win32_networkadapterconfiguration
6
+ from ...psutilw import psutil_networks
7
7
  from ....print_api import print_api
8
8
 
9
9
 
10
- class NetworkAdapterNotFoundError(Exception):
11
- pass
10
+ def list_network_adapters(wmi_instance: CDispatch = None) -> list[CDispatch]:
11
+ """
12
+ List all network adapters on the system, from the Win32_NetworkAdapter class.
12
13
 
14
+ :param wmi_instance: WMI instance. You can get it from:
15
+ wmi_helpers.get_wmi_instance()
16
+ :return: list of Win32_NetworkAdapter objects.
17
+ """
13
18
 
14
- def list_network_adapters() -> list:
15
- # Initialize the WMI client
16
- wmi = win32com.client.Dispatch('WbemScripting.SWbemLocator')
17
- wmi_service = wmi.ConnectServer('.', 'root\\cimv2')
19
+ if not wmi_instance:
20
+ wmi_instance, _ = wmi_helpers.get_wmi_instance()
18
21
 
19
22
  # Query all network adapters
20
- adapters = wmi_service.ExecQuery("SELECT * FROM Win32_NetworkAdapter")
23
+ adapters: list[CDispatch] = list(wmi_instance.ExecQuery("SELECT * FROM Win32_NetworkAdapter"))
21
24
 
22
25
  # Print adapter descriptions
23
26
  # for adapter in adapters:
24
27
  # print(f"Description: {adapter.Description}, IPEnabled: {adapter.IPEnabled}")
25
- return list(adapters)
28
+ return adapters
26
29
 
27
30
 
28
- def get_default_network_adapter():
29
- """
30
- Get the default network adapter.
31
- :return:
31
+ def get_network_adapter_by_device_name(
32
+ device_name: str,
33
+ wmi_instance: CDispatch = None
34
+ ) -> Union[CDispatch, None]:
32
35
  """
36
+ Get a network adapter by its name.
33
37
 
34
- default_connection_name_dict: dict = networks.get_default_connection_name()
35
- # Get the first key from the dictionary.
36
- default_connection_name: str = list(default_connection_name_dict.keys())[0]
37
- adapters = list_network_adapters()
38
-
39
- for adapter in adapters:
40
- if default_connection_name == adapter.NetConnectionID:
41
- return adapter
42
-
43
- raise NetworkAdapterNotFoundError("Default network adapter not found.")
44
-
45
-
46
- def get_wmi_network_configuration(
47
- use_default_interface: bool = False,
48
- connection_name: str = None,
49
- mac_address: str = None
50
- ) -> tuple:
51
- """
52
- Get the WMI network configuration for a network adapter.
53
- :param use_default_interface: bool, if True, the default network interface will be used.
54
- This is the adapter that your internet is being used from.
55
- :param connection_name: string, adapter name as shown in the network settings.
56
- :param mac_address: string, MAC address of the adapter. Format: '00:00:00:00:00:00'.
57
- :return: tuple(Win32_NetworkAdapterConfiguration, Win32_NetworkAdapter)
38
+ :param device_name: string, adapter name as shown in the network settings.
39
+ :param wmi_instance: WMI instance. You can get it from:
40
+ wmi_helpers.get_wmi_instance()
41
+ :return: Win32_NetworkAdapter object.
58
42
  """
59
43
 
60
- if use_default_interface and connection_name:
61
- raise ValueError("Only one of 'use_default_interface' or 'connection_name' must be provided.")
62
- elif not use_default_interface and not connection_name:
63
- raise ValueError("Either 'use_default_interface' or 'connection_name' must be provided.")
44
+ if not wmi_instance:
45
+ wmi_instance, _ = wmi_helpers.get_wmi_instance()
64
46
 
65
- adapters = list_network_adapters()
47
+ query: str = (
48
+ "SELECT * FROM Win32_NetworkAdapter "
49
+ f"WHERE Name LIKE '{device_name}'")
50
+ adapters: list[CDispatch] = list(wmi_instance.ExecQuery(query))
51
+ if not adapters:
52
+ return None
66
53
 
67
- current_adapter = None
68
- if use_default_interface:
69
- default_connection_name_dict: dict = networks.get_default_connection_name()
70
- if not default_connection_name_dict:
71
- raise NetworkAdapterNotFoundError("Default network adapter not found.")
72
- # Get the first key from the dictionary.
73
- connection_name: str = list(default_connection_name_dict.keys())[0]
54
+ return adapters[0]
74
55
 
75
- if connection_name is None and mac_address is None:
76
- raise ValueError("Either 'connection_name' or 'mac_address' must be provided.")
77
- elif connection_name and mac_address:
78
- raise ValueError("Only one of 'connection_name' or 'mac_address' must be provided.")
79
56
 
80
- if connection_name:
81
- for adapter in adapters:
82
- if connection_name == adapter.NetConnectionID:
83
- current_adapter = adapter
84
- break
85
-
86
- if not current_adapter:
87
- raise NetworkAdapterNotFoundError(f"Adapter with connection name '{connection_name}' not found.")
88
- elif mac_address:
89
- for adapter in adapters:
90
- if mac_address == adapter.MACAddress:
91
- current_adapter = adapter
92
- break
57
+ def get_default_network_adapter(wmi_instance: CDispatch = None):
58
+ """
59
+ Get the default network adapter.
93
60
 
94
- if current_adapter is None:
95
- raise NetworkAdapterNotFoundError(f"Adapter with MAC address '{mac_address}' not found.")
61
+ :param wmi_instance: WMI instance. You can get it from:
62
+ wmi_helpers.get_wmi_instance()
63
+ :return:
64
+ """
96
65
 
97
- # Initialize the WMI client
98
- wmi = win32com.client.Dispatch('WbemScripting.SWbemLocator')
99
- wmi_service = wmi.ConnectServer('.', 'root\\cimv2')
66
+ if not wmi_instance:
67
+ wmi_instance, _ = wmi_helpers.get_wmi_instance()
100
68
 
101
- # Query the network adapter configurations
102
- query = f"SELECT * FROM Win32_NetworkAdapterConfiguration WHERE Index='{current_adapter.DeviceID}'"
103
- adapter_configs = wmi_service.ExecQuery(query)
69
+ default_connection_name_dict: dict = psutil_networks.get_default_connection_name()
70
+ # Get the first key from the dictionary.
71
+ default_connection_name: str = list(default_connection_name_dict.keys())[0]
72
+ adapters: list[CDispatch] = list_network_adapters(wmi_instance)
104
73
 
105
- # Check if the adapter exists
106
- if len(adapter_configs) == 0:
107
- raise NetworkAdapterNotFoundError(f"Adapter with connection name '{connection_name}' not found.")
74
+ for adapter in adapters:
75
+ if default_connection_name == adapter.NetConnectionID:
76
+ return adapter
108
77
 
109
- return adapter_configs[0], current_adapter
78
+ raise wmi_helpers.WMINetworkAdapterNotFoundError("Default network adapter not found.")
110
79
 
111
80
 
112
81
  def set_dns_server(
@@ -127,11 +96,11 @@ def set_dns_server(
127
96
  :return:
128
97
  """
129
98
 
130
- adapter_config, current_adapter = get_wmi_network_configuration(
99
+ adapter_config, current_adapter = win32_networkadapterconfiguration.get_adapter_network_configuration(
131
100
  use_default_interface=use_default_interface, connection_name=connection_name, mac_address=mac_address)
132
101
 
133
102
  print_api(f"Adapter [{current_adapter.Description}], Connection name [{current_adapter.NetConnectionID}]\n"
134
103
  f"Setting DNS servers to {dns_servers}", color='blue')
135
104
 
136
105
  # Set DNS servers
137
- helpers.call_method(adapter_config, 'SetDNSServerSearchOrder', dns_servers)
106
+ wmi_helpers.call_method(adapter_config, 'SetDNSServerSearchOrder', dns_servers)
@@ -0,0 +1,235 @@
1
+ import os
2
+ from typing import Union
3
+
4
+ import win32com.client
5
+ from win32com.client import CDispatch
6
+
7
+
8
+ class WmiMethodExecutionError(Exception):
9
+ pass
10
+
11
+
12
+ class WmiMethodParameterError(Exception):
13
+ pass
14
+
15
+
16
+ class EmptyValue:
17
+ pass
18
+ _EMPTY_VALUE = EmptyValue()
19
+
20
+
21
+ class WMINetworkAdapterNotFoundError(Exception):
22
+ pass
23
+
24
+
25
+ LOCAL_SERVER: str = '.'
26
+
27
+
28
+ def get_wmi_instance(
29
+ server: str = LOCAL_SERVER,
30
+ namespace: str = 'root\\cimv2',
31
+ wmi_instance: CDispatch = None,
32
+ locator: CDispatch = None
33
+ ) -> tuple[CDispatch, CDispatch]:
34
+ """
35
+ Get the WMI instance.
36
+
37
+ :param server: str, The server you want to connect to. Default is '.' (local machine).
38
+ :param namespace: str, WMI namespace. Default is 'root\\cimv2'.
39
+ Other examples:
40
+ 'root\\StandardCimv2'
41
+ :param wmi_instance: WMI connected instance.
42
+ :param locator: WMI locator instance. If not provided, a new one will be created.
43
+ :return: WMI instance.
44
+
45
+ ===================
46
+
47
+ If you want to connect directly to a WMI namespace, you can use the following code:
48
+ return win32com.client.GetObject(f"winmgmts:\\\\{location}\\{namespace}")
49
+ """
50
+
51
+ if not locator:
52
+ # This is a better way to get the WMI instance, since you have more control over the WMI object.
53
+ locator: CDispatch = win32com.client.Dispatch('WbemScripting.SWbemLocator')
54
+
55
+ if wmi_instance:
56
+ server_from_instance, namespace_from_instance = get_connection_details(wmi_instance)
57
+
58
+ # If current server name of the wmi connection is not the same as was passed to the function,
59
+ # then create a new connection to the server.
60
+ if server_from_instance.lower() != server.lower():
61
+ if not (server_from_instance.lower() == os.environ["COMPUTERNAME"].lower() and server.lower() == LOCAL_SERVER):
62
+ wmi_instance = locator.ConnectServer(server, namespace)
63
+ # If the namespace is not the same as was passed to the function.
64
+ if namespace_from_instance != namespace:
65
+ wmi_instance = locator.ConnectServer(server, namespace)
66
+
67
+ else:
68
+ wmi_instance: CDispatch = locator.ConnectServer(server, namespace)
69
+
70
+ return wmi_instance, locator
71
+
72
+
73
+ def get_connection_details(wmi_instance: CDispatch) -> tuple[str, str]:
74
+ """
75
+ Get the connection details: connected server and namespace.
76
+
77
+ :param wmi_instance: WMI instance.
78
+ :return: tuple of server and namespace.
79
+ """
80
+
81
+ # Get the current connection details.
82
+ # Get the security object for the WMI instance.
83
+ security_object: CDispatch = wmi_instance.Get("__SystemSecurity=@")
84
+ # Get the Paths.
85
+ path: CDispatch = security_object.Path_
86
+
87
+ server_from_instance: str = path.Server
88
+ namespace_from_instance: str = path.Namespace
89
+
90
+ return server_from_instance, namespace_from_instance
91
+
92
+
93
+ def get_method(
94
+ wmi_object: win32com.client.CDispatch,
95
+ method_name: str
96
+ ):
97
+ """
98
+ Get the WMI method.
99
+
100
+ :param wmi_object: WMI object.
101
+ :param method_name: str, name of the method.
102
+ :return: WMI method object.
103
+ """
104
+
105
+ return wmi_object.Methods_(method_name)
106
+
107
+
108
+ def get_method_parameter_instance(
109
+ method: win32com.client.CDispatch
110
+ ):
111
+ """
112
+ Get the WMI method parameter.
113
+
114
+ :param method: WMI method object.
115
+ :return: WMI method parameter object.
116
+ """
117
+
118
+ return method.inParameters.SpawnInstance_()
119
+
120
+
121
+ def call_method(
122
+ wmi_object: win32com.client.CDispatch,
123
+ method_name: str,
124
+ value: Union[
125
+ Union[tuple, dict],
126
+ Union[bool, str, list]] = _EMPTY_VALUE
127
+ ):
128
+ """
129
+ Call the WMI method.
130
+
131
+ :param wmi_object: WMI object.
132
+ :param method_name: str, name of the method.
133
+ :param value: tuple, value to pass to the method.
134
+ tuple: If ou pass a tuple, it will be unpacked and passed as positional arguments.
135
+ Dor example if a method requires 2 parameters, you can pass a tuple with 2 values.
136
+ dict: If you pass a dictionary, it will be unpacked and passed as keyword arguments.
137
+
138
+ If you pass a single value, which is not a dict or tuple, it will be passed as a single parameter.
139
+
140
+ Methods can receive a None value if they don't require any parameters.
141
+ If the method doesn't require any parameters, leave it as 'EmptyValue' class.
142
+ :return: WMI method object.
143
+ """
144
+
145
+ if not (isinstance(value, EmptyValue) and isinstance(value, tuple) and isinstance(value, dict)):
146
+ value = (value,)
147
+
148
+ # Get the method instance out of the WMI object.
149
+ method = get_method(wmi_object, method_name)
150
+
151
+ # ── discover the method’s IN parameters up-front ─────────────────────────────
152
+ if method.InParameters:
153
+ input_properties: list = [(in_property.Name, in_property.IsArray) for in_property in method.InParameters.Properties_]
154
+ else:
155
+ input_properties: list = [] # no inputs expected
156
+
157
+ expected = len(input_properties) # how many inputs the method wants
158
+
159
+ got_tuple = isinstance(value, tuple)
160
+ got_dict = isinstance(value, dict)
161
+ got_empty = isinstance(value, EmptyValue)
162
+
163
+ # ── validate the caller’s intent ─────────────────────────────────────────────
164
+ if expected == 0 and not got_empty:
165
+ raise WmiMethodParameterError(
166
+ f"Method '{method_name}' takes no parameters, got: {value!r}"
167
+ )
168
+ if expected > 0 and got_empty:
169
+ raise WmiMethodParameterError(
170
+ f"Method '{method_name}' expects {expected} parameter(s); none given."
171
+ )
172
+ if got_tuple and len(value) != expected:
173
+ raise WmiMethodParameterError(
174
+ f"Method '{method_name}' expects {expected} parameter(s); "
175
+ f"{len(value)} positional value(s) given."
176
+ )
177
+
178
+ # ── prepare the parameter object if needed ──────────────────────────────────
179
+ if expected == 0: # simple – no inputs
180
+ result = wmi_object.ExecMethod_(method_name)
181
+
182
+ else:
183
+ param_obj = get_method_parameter_instance(method)
184
+
185
+ if got_tuple: # positional list / tuple
186
+ for (name, is_array), val in zip(input_properties, value):
187
+ setattr(param_obj, name, val)
188
+
189
+ elif got_dict: # mapping by name
190
+ for name, _ in input_properties:
191
+ if name in value:
192
+ setattr(param_obj, name, value[name])
193
+
194
+ else: # single scalar for one-input method
195
+ name, is_array = input_properties[0]
196
+ if is_array and not (isinstance(value, list) or value is None):
197
+ raise WmiMethodParameterError(
198
+ f"Parameter '{name}' must be a list.\nValue: {value!r}"
199
+ )
200
+ setattr(param_obj, name, value)
201
+
202
+ result = wmi_object.ExecMethod_(method_name, param_obj)
203
+
204
+ # ── collect OUT parameters & check return code ──────────────────────────────
205
+ out_vals = []
206
+ if method.OutParameters:
207
+ for parameter in method.OutParameters.Properties_:
208
+ out_vals.append(result.Properties_(parameter.Name).Value)
209
+
210
+ # return-code conventions: 0 = OK, 1 = OK-needs-reboot
211
+ if out_vals and out_vals[0] not in (0, 1):
212
+ result_code: int = out_vals[0]
213
+ if result_code == 91:
214
+ raise PermissionError(
215
+ f"Method '{method_name}' failed (code {result_code}) – try with admin rights."
216
+ )
217
+ if result_code == 68:
218
+ raise WmiMethodExecutionError(
219
+ f"Method '{method_name}' failed (code {result_code}) – Invalid input parameter"
220
+ )
221
+ raise WmiMethodExecutionError(
222
+ f"Method '{method_name}' failed with error code {result_code}"
223
+ )
224
+
225
+ return out_vals or None
226
+
227
+
228
+ """
229
+ # Setting SeDebugPrivilege
230
+ import win32security, ntsecuritycon, win32con, win32api
231
+ privs = ((win32security.LookupPrivilegeValue('',ntsecuritycon.SE_DEBUG_NAME), win32con.SE_PRIVILEGE_ENABLED),)
232
+ hToken = win32security.OpenProcessToken(win32api.GetCurrentProcess(), win32security.TOKEN_ALL_ACCESS)
233
+ win32security.AdjustTokenPrivileges(hToken, False, privs)
234
+ win32api.CloseHandle(hToken)
235
+ """