597 lines
21 KiB
Python
597 lines
21 KiB
Python
from PYAS_Diagnostics import log_exception
|
|
import os
|
|
import time
|
|
import shutil
|
|
import winreg
|
|
import subprocess
|
|
import ctypes
|
|
import ctypes.wintypes
|
|
from PYAS_WinAPI import (
|
|
FILE_ATTRIBUTE_NORMAL,
|
|
FILE_SHARE_READ,
|
|
FILE_SHARE_WRITE,
|
|
INTERNAL_STORAGE_BUS_TYPES,
|
|
INVALID_HANDLE_VALUE,
|
|
IOCTL_STORAGE_GET_HOTPLUG_INFO,
|
|
IOCTL_STORAGE_QUERY_PROPERTY,
|
|
OPEN_EXISTING,
|
|
PROPERTY_STANDARD_QUERY,
|
|
STORAGE_DESCRIPTOR_HEADER,
|
|
STORAGE_DEVICE_DESCRIPTOR,
|
|
STORAGE_DEVICE_PROPERTY,
|
|
STORAGE_HOTPLUG_INFO,
|
|
STORAGE_PROPERTY_QUERY,
|
|
)
|
|
|
|
|
|
class SystemMixin:
|
|
def _is_internal_physical_drive(self, drive):
|
|
drive_path = rf"\\.\PhysicalDrive{drive}"
|
|
handle = self.kernel32.CreateFileW(
|
|
drive_path,
|
|
0,
|
|
FILE_SHARE_READ | FILE_SHARE_WRITE,
|
|
None,
|
|
OPEN_EXISTING,
|
|
FILE_ATTRIBUTE_NORMAL,
|
|
None,
|
|
)
|
|
|
|
if handle in (None, INVALID_HANDLE_VALUE):
|
|
return False
|
|
|
|
try:
|
|
query = STORAGE_PROPERTY_QUERY(STORAGE_DEVICE_PROPERTY, PROPERTY_STANDARD_QUERY)
|
|
header = STORAGE_DESCRIPTOR_HEADER()
|
|
bytes_returned = ctypes.wintypes.DWORD(0)
|
|
|
|
if not self.kernel32.DeviceIoControl(
|
|
handle,
|
|
IOCTL_STORAGE_QUERY_PROPERTY,
|
|
ctypes.byref(query),
|
|
ctypes.sizeof(query),
|
|
ctypes.byref(header),
|
|
ctypes.sizeof(header),
|
|
ctypes.byref(bytes_returned),
|
|
None,
|
|
):
|
|
return False
|
|
|
|
minimum_size = STORAGE_DEVICE_DESCRIPTOR.BusType.offset + ctypes.sizeof(ctypes.c_int)
|
|
|
|
if (
|
|
bytes_returned.value < ctypes.sizeof(header)
|
|
or header.Size < minimum_size
|
|
or header.Size > 65536
|
|
):
|
|
return False
|
|
|
|
descriptor_buffer = ctypes.create_string_buffer(header.Size)
|
|
bytes_returned.value = 0
|
|
|
|
if not self.kernel32.DeviceIoControl(
|
|
handle,
|
|
IOCTL_STORAGE_QUERY_PROPERTY,
|
|
ctypes.byref(query),
|
|
ctypes.sizeof(query),
|
|
descriptor_buffer,
|
|
ctypes.sizeof(descriptor_buffer),
|
|
ctypes.byref(bytes_returned),
|
|
None,
|
|
):
|
|
return False
|
|
|
|
descriptor = ctypes.cast(
|
|
descriptor_buffer, ctypes.POINTER(STORAGE_DEVICE_DESCRIPTOR)
|
|
).contents
|
|
|
|
if (
|
|
bytes_returned.value < minimum_size
|
|
or descriptor.RemovableMedia
|
|
or descriptor.BusType not in INTERNAL_STORAGE_BUS_TYPES
|
|
):
|
|
return False
|
|
|
|
hotplug = STORAGE_HOTPLUG_INFO()
|
|
hotplug.Size = ctypes.sizeof(hotplug)
|
|
bytes_returned.value = 0
|
|
|
|
if not self.kernel32.DeviceIoControl(
|
|
handle,
|
|
IOCTL_STORAGE_GET_HOTPLUG_INFO,
|
|
None,
|
|
0,
|
|
ctypes.byref(hotplug),
|
|
ctypes.sizeof(hotplug),
|
|
ctypes.byref(bytes_returned),
|
|
None,
|
|
):
|
|
return False
|
|
|
|
if bytes_returned.value < ctypes.sizeof(hotplug):
|
|
return False
|
|
|
|
return not (hotplug.MediaRemovable or hotplug.MediaHotplug or hotplug.DeviceHotplug)
|
|
except Exception:
|
|
log_exception("PYAS_System.SystemMixin._is_internal_physical_drive:80")
|
|
return False
|
|
finally:
|
|
self.kernel32.CloseHandle(handle)
|
|
|
|
def backup_mbr(self, max_drives=26):
|
|
self.mbr_backup = {}
|
|
|
|
for drive in range(max_drives):
|
|
if not self._is_internal_physical_drive(drive):
|
|
continue
|
|
|
|
try:
|
|
with open(rf"\\.\PhysicalDrive{drive}", "rb") as f:
|
|
mbr = f.read(512)
|
|
|
|
if len(mbr) == 512 and mbr[510:512] == b"\x55\xAA":
|
|
self.mbr_backup[drive] = mbr
|
|
|
|
except Exception:
|
|
log_exception("PYAS_System.SystemMixin.backup_mbr:97")
|
|
continue
|
|
|
|
def check_system_mbr(self):
|
|
for drive, mbr_value in list(self.mbr_backup.items()):
|
|
try:
|
|
with open(rf"\\.\PhysicalDrive{drive}", "rb") as f:
|
|
if f.read(512) != mbr_value:
|
|
return True
|
|
|
|
except Exception:
|
|
log_exception("PYAS_System.SystemMixin.check_system_mbr:107")
|
|
pass
|
|
|
|
return False
|
|
|
|
def repair_system_mbr(self):
|
|
for drive, mbr_value in list(self.mbr_backup.items()):
|
|
drive_path = rf"\\.\PhysicalDrive{drive}"
|
|
|
|
try:
|
|
with open(drive_path, "rb+") as f:
|
|
if f.read(512) != mbr_value:
|
|
f.seek(0)
|
|
f.write(mbr_value)
|
|
self.write_log("INFO", "MBR Repaired", source=drive_path, operate=True)
|
|
|
|
except Exception:
|
|
log_exception("PYAS_System.SystemMixin.repair_system_mbr:120")
|
|
pass
|
|
|
|
def _get_restrict_lists(self):
|
|
permissions = [
|
|
"NoControlPanel",
|
|
"NoDrives",
|
|
"NoFileMenu",
|
|
"NoFind",
|
|
"NoStartMenuPinnedList",
|
|
"NoSetFolders",
|
|
"NoSetFolderOptions",
|
|
"NoViewOnDrive",
|
|
"NoClose",
|
|
"NoDesktop",
|
|
"NoLogoff",
|
|
"NoFolderOptions",
|
|
"RestrictRun",
|
|
"NoViewContextMenu",
|
|
"HideClock",
|
|
"NoStartMenuMyGames",
|
|
"NoStartMenuMyMusic",
|
|
"DisableCMD",
|
|
"NoAddingComponents",
|
|
"NoWinKeys",
|
|
"NoStartMenuLogOff",
|
|
"NoSimpleNetIDList",
|
|
"NoLowDiskSpaceChecks",
|
|
"DisableLockWorkstation",
|
|
"Restrict_Run",
|
|
"DisableTaskMgr",
|
|
"DisableRegistryTools",
|
|
"DisableChangePassword",
|
|
"Wallpaper",
|
|
"NoComponents",
|
|
"NoStartMenuMorePrograms",
|
|
"NoActiveDesktop",
|
|
"NoSetActiveDesktop",
|
|
"NoRecentDocsMenu",
|
|
"NoWindowsUpdate",
|
|
"NoChangeStartMenu",
|
|
"NoFavoritesMenu",
|
|
"NoRecentDocsHistory",
|
|
"NoSetTaskbar",
|
|
"NoSMHelp",
|
|
"NoTrayContextMenu",
|
|
"NoManageMyComputerVerb",
|
|
"NoRealMode",
|
|
"NoRun",
|
|
"ClearRecentDocsOnExit",
|
|
"NoActiveDesktopChanges",
|
|
"NoStartMenuNetworkPlaces",
|
|
]
|
|
paths = [
|
|
(winreg.HKEY_CURRENT_USER, r"SOFTWARE\Policies\Microsoft\MMC"),
|
|
(winreg.HKEY_CURRENT_USER, r"SOFTWARE\Policies\Microsoft\Windows\System"),
|
|
(
|
|
winreg.HKEY_CURRENT_USER,
|
|
r"SOFTWARE\Microsoft\Windows\CurrentVersion\Policies\Explorer",
|
|
),
|
|
(
|
|
winreg.HKEY_CURRENT_USER,
|
|
r"SOFTWARE\Microsoft\Windows\CurrentVersion\Policies\System",
|
|
),
|
|
(
|
|
winreg.HKEY_LOCAL_MACHINE,
|
|
r"SOFTWARE\Microsoft\Windows\CurrentVersion\Policies\Explorer",
|
|
),
|
|
(
|
|
winreg.HKEY_LOCAL_MACHINE,
|
|
r"SOFTWARE\Microsoft\Windows\CurrentVersion\Policies\System",
|
|
),
|
|
(winreg.HKEY_LOCAL_MACHINE, r"SOFTWARE\Microsoft\Windows NT\CurrentVersion\Winlogon"),
|
|
(winreg.HKEY_LOCAL_MACHINE, r"SOFTWARE\Policies\Microsoft\Windows\System"),
|
|
]
|
|
return permissions, paths
|
|
|
|
def check_system_restrict(self):
|
|
permissions, paths = self._get_restrict_lists()
|
|
|
|
for hkey, path in paths:
|
|
for val in permissions:
|
|
if self._reg_read(hkey, path, val) is not None:
|
|
return True
|
|
|
|
return False
|
|
|
|
def check_system_file_type(self):
|
|
for root in [winreg.HKEY_LOCAL_MACHINE, winreg.HKEY_CURRENT_USER]:
|
|
for ext in [".exe", ".bat", ".cmd", ".com"]:
|
|
if self._reg_read(root, rf"SOFTWARE\Classes\{ext}", "") != (
|
|
"exefile" if ext == ".exe" else ext[1:] + "file"
|
|
):
|
|
return True
|
|
|
|
for cmd in ["open", "runas"]:
|
|
if (
|
|
self._reg_read(root, rf"SOFTWARE\Classes\exefile\shell\{cmd}\command", "")
|
|
!= '"%1" %*'
|
|
):
|
|
return True
|
|
|
|
return False
|
|
|
|
def check_system_file_icon(self):
|
|
return any(
|
|
self._reg_read(root, r"SOFTWARE\Classes\exefile\DefaultIcon", "") != "%1"
|
|
for root in [winreg.HKEY_LOCAL_MACHINE, winreg.HKEY_CURRENT_USER]
|
|
)
|
|
|
|
def check_system_image(self):
|
|
try:
|
|
with winreg.OpenKey(
|
|
winreg.HKEY_LOCAL_MACHINE,
|
|
r"SOFTWARE\Microsoft\Windows NT\CurrentVersion\Image File Execution Options",
|
|
0,
|
|
winreg.KEY_READ,
|
|
) as reg:
|
|
i = 0
|
|
|
|
while True:
|
|
try:
|
|
subkey = winreg.EnumKey(reg, i)
|
|
|
|
with winreg.OpenKey(reg, subkey, 0, winreg.KEY_READ) as sub_reg:
|
|
for val in ["Debugger", "UseFilter", "GlobalFlag", "MitigationOptions"]:
|
|
try:
|
|
winreg.QueryValueEx(sub_reg, val)
|
|
return True
|
|
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
i += 1
|
|
except OSError:
|
|
log_exception("PYAS_System.SystemMixin.check_system_image:185")
|
|
break
|
|
|
|
except Exception:
|
|
log_exception("PYAS_System.SystemMixin.check_system_image:188")
|
|
pass
|
|
|
|
return False
|
|
|
|
def check_system_wallpaper(self):
|
|
return self._reg_read(
|
|
winreg.HKEY_CURRENT_USER, r"Control Panel\Desktop", "Wallpaper"
|
|
) != os.path.join(self.path_system, "web", "wallpaper", "Windows", "img0.jpg")
|
|
|
|
def scan_system_repair(self):
|
|
items = []
|
|
|
|
for drive, mbr_value in list(self.mbr_backup.items()):
|
|
drive_path = rf"\\.\PhysicalDrive{drive}"
|
|
|
|
try:
|
|
with open(drive_path, "rb") as f:
|
|
if f.read(512) != mbr_value:
|
|
items.append(
|
|
{"id": f"mbr|{drive}", "display": "repair_mbr", "path": drive_path}
|
|
)
|
|
|
|
except Exception:
|
|
log_exception("PYAS_System.SystemMixin.scan_system_repair:205")
|
|
pass
|
|
|
|
permissions, paths = self._get_restrict_lists()
|
|
|
|
for hkey, path in paths:
|
|
for val in permissions:
|
|
|
|
if self._reg_read(hkey, path, val) is not None:
|
|
root = "HKLM" if hkey == winreg.HKEY_LOCAL_MACHINE else "HKCU"
|
|
items.append(
|
|
{
|
|
"id": f"restrict|{root}\\{path}\\{val}",
|
|
"display": "repair_limit",
|
|
"path": rf"{root}\{path}\{val}",
|
|
}
|
|
)
|
|
|
|
for root in [winreg.HKEY_LOCAL_MACHINE, winreg.HKEY_CURRENT_USER]:
|
|
root_str = "HKLM" if root == winreg.HKEY_LOCAL_MACHINE else "HKCU"
|
|
|
|
for ext in [".exe", ".bat", ".cmd", ".com"]:
|
|
expected = "exefile" if ext == ".exe" else ext[1:] + "file"
|
|
|
|
if self._reg_read(root, rf"SOFTWARE\Classes\{ext}", "") != expected:
|
|
items.append(
|
|
{
|
|
"id": f"file_type|{root_str}\\{ext}",
|
|
"display": "repair_assoc",
|
|
"path": rf"{root_str}\SOFTWARE\Classes\{ext}",
|
|
}
|
|
)
|
|
|
|
for cmd in ["open", "runas"]:
|
|
if (
|
|
self._reg_read(root, rf"SOFTWARE\Classes\exefile\shell\{cmd}\command", "")
|
|
!= '"%1" %*'
|
|
):
|
|
items.append(
|
|
{
|
|
"id": f"file_type|{root_str}\\exefile\\{cmd}",
|
|
"display": "repair_assoc",
|
|
"path": rf"{root_str}\SOFTWARE\Classes\exefile\shell\{cmd}\command",
|
|
}
|
|
)
|
|
|
|
if self._reg_read(root, r"SOFTWARE\Classes\exefile\DefaultIcon", "") != "%1":
|
|
items.append(
|
|
{
|
|
"id": f"file_icon|{root_str}",
|
|
"display": "repair_icon",
|
|
"path": rf"{root_str}\SOFTWARE\Classes\exefile\DefaultIcon",
|
|
}
|
|
)
|
|
|
|
try:
|
|
with winreg.OpenKey(
|
|
winreg.HKEY_LOCAL_MACHINE,
|
|
r"SOFTWARE\Microsoft\Windows NT\CurrentVersion\Image File Execution Options",
|
|
0,
|
|
winreg.KEY_READ,
|
|
) as reg:
|
|
i = 0
|
|
|
|
while True:
|
|
try:
|
|
subkey = winreg.EnumKey(reg, i)
|
|
|
|
with winreg.OpenKey(reg, subkey, 0, winreg.KEY_READ) as sub_reg:
|
|
for val in ["Debugger", "UseFilter", "GlobalFlag", "MitigationOptions"]:
|
|
try:
|
|
winreg.QueryValueEx(sub_reg, val)
|
|
items.append(
|
|
{
|
|
"id": f"image|{subkey}",
|
|
"display": "repair_hijack",
|
|
"path": rf"HKLM\SOFTWARE\Microsoft\Windows NT\CurrentVersion\Image File Execution Options\{subkey}",
|
|
}
|
|
)
|
|
break
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
i += 1
|
|
except OSError:
|
|
log_exception("PYAS_System.SystemMixin.scan_system_repair:247")
|
|
break
|
|
|
|
except Exception:
|
|
log_exception("PYAS_System.SystemMixin.scan_system_repair:250")
|
|
pass
|
|
|
|
wp = self._reg_read(winreg.HKEY_CURRENT_USER, r"Control Panel\Desktop", "Wallpaper")
|
|
expected_wp = os.path.join(self.path_system, "web", "wallpaper", "Windows", "img0.jpg")
|
|
|
|
if wp != expected_wp:
|
|
items.append(
|
|
{
|
|
"id": "wallpaper|0",
|
|
"display": "repair_wallpaper",
|
|
"path": wp if wp else r"HKCU\Control Panel\Desktop\Wallpaper",
|
|
}
|
|
)
|
|
|
|
return items
|
|
|
|
def repair_system_restrict(self):
|
|
permissions, paths = self._get_restrict_lists()
|
|
|
|
for hkey, path in paths:
|
|
for val in permissions:
|
|
self._reg_delete(hkey, path, val)
|
|
|
|
def repair_system_file_type(self):
|
|
for root in [winreg.HKEY_LOCAL_MACHINE, winreg.HKEY_CURRENT_USER]:
|
|
for ext in [".exe", ".bat", ".cmd", ".com"]:
|
|
self._reg_write(
|
|
root,
|
|
rf"SOFTWARE\Classes\{ext}",
|
|
None,
|
|
winreg.REG_SZ,
|
|
"exefile" if ext == ".exe" else ext[1:] + "file",
|
|
)
|
|
|
|
for cmd in ["open", "runas"]:
|
|
self._reg_write(
|
|
root,
|
|
rf"SOFTWARE\Classes\exefile\shell\{cmd}\command",
|
|
None,
|
|
winreg.REG_SZ,
|
|
'"%1" %*',
|
|
)
|
|
|
|
def repair_system_file_icon(self):
|
|
for root in [winreg.HKEY_LOCAL_MACHINE, winreg.HKEY_CURRENT_USER]:
|
|
self._reg_write(
|
|
root, r"SOFTWARE\Classes\exefile\DefaultIcon", None, winreg.REG_SZ, "%1"
|
|
)
|
|
|
|
def repair_system_image(self):
|
|
try:
|
|
with winreg.OpenKey(
|
|
winreg.HKEY_LOCAL_MACHINE,
|
|
r"SOFTWARE\Microsoft\Windows NT\CurrentVersion\Image File Execution Options",
|
|
0,
|
|
winreg.KEY_ALL_ACCESS,
|
|
) as reg:
|
|
i = 0
|
|
|
|
while True:
|
|
try:
|
|
subkey = winreg.EnumKey(reg, i)
|
|
|
|
with winreg.OpenKey(reg, subkey, 0, winreg.KEY_ALL_ACCESS) as sub_reg:
|
|
for val in ["Debugger", "UseFilter", "GlobalFlag", "MitigationOptions"]:
|
|
try:
|
|
winreg.DeleteValue(sub_reg, val)
|
|
except FileNotFoundError:
|
|
pass
|
|
|
|
i += 1
|
|
except OSError:
|
|
log_exception("PYAS_System.SystemMixin.repair_system_image:292")
|
|
break
|
|
|
|
except Exception as e:
|
|
log_exception("PYAS_System.SystemMixin.repair_system_image:295")
|
|
self.write_log("WARN", "repair_system_image", detail=str(e), success=False)
|
|
|
|
def repair_system_wallpaper(self):
|
|
try:
|
|
wallpaper = os.path.join(self.path_system, "web", "wallpaper", "Windows", "img0.jpg")
|
|
|
|
if not os.path.exists(wallpaper):
|
|
return
|
|
|
|
self._reg_write(
|
|
winreg.HKEY_CURRENT_USER,
|
|
r"Control Panel\Desktop",
|
|
"Wallpaper",
|
|
winreg.REG_SZ,
|
|
wallpaper,
|
|
)
|
|
theme_dir = os.path.join(self.path_appdata, "Microsoft", "Windows", "Themes")
|
|
|
|
for fname in ["TranscodedWallpaper", "TranscodedWallpaper.tmp"]:
|
|
try:
|
|
os.remove(os.path.join(theme_dir, fname))
|
|
except Exception:
|
|
log_exception("PYAS_System.SystemMixin.repair_system_wallpaper:310")
|
|
pass
|
|
|
|
shutil.rmtree(os.path.join(theme_dir, "CachedFiles"), ignore_errors=True)
|
|
self.user32.SystemParametersInfoW(20, 0, wallpaper, 3)
|
|
|
|
except Exception as e:
|
|
log_exception("PYAS_System.SystemMixin.repair_system_wallpaper:316")
|
|
self.write_log("WARN", "repair_system_wallpaper", detail=str(e), success=False)
|
|
|
|
def execute_system_repair(self, items):
|
|
try:
|
|
types = set(item.split("|")[0] for item in items)
|
|
|
|
if "mbr" in types:
|
|
self.repair_system_mbr()
|
|
|
|
if "restrict" in types:
|
|
self.repair_system_restrict()
|
|
|
|
if "file_type" in types:
|
|
self.repair_system_file_type()
|
|
|
|
if "file_icon" in types:
|
|
self.repair_system_file_icon()
|
|
|
|
if "image" in types:
|
|
self.repair_system_image()
|
|
|
|
if "wallpaper" in types:
|
|
self.repair_system_wallpaper()
|
|
|
|
self.write_log(
|
|
"INFO", "System Repair", detail=f"Repaired {len(items)} items", operate=True
|
|
)
|
|
return True
|
|
|
|
except Exception as e:
|
|
log_exception("PYAS_System.SystemMixin.execute_system_repair:339")
|
|
self.write_log(
|
|
"WARN", "execute_system_repair", detail=str(e), operate=True, success=False
|
|
)
|
|
return False
|
|
|
|
def protect_system_thread(self):
|
|
while True:
|
|
with self.lock_config:
|
|
if not self.pyas_config.get("system_switch", False):
|
|
break
|
|
|
|
try:
|
|
self.repair_system_mbr()
|
|
self.repair_system_restrict()
|
|
self.repair_system_file_type()
|
|
self.repair_system_file_icon()
|
|
self.repair_system_image()
|
|
self.check_process_survival()
|
|
time.sleep(0.5)
|
|
|
|
except Exception as e:
|
|
log_exception("PYAS_System.SystemMixin.protect_system_thread:357")
|
|
self.write_log("WARN", "protect_system_thread", detail=str(e), success=False)
|
|
|
|
def check_process_survival(self):
|
|
try:
|
|
running = any(
|
|
exe_name.lower() == "explorer.exe" for _, exe_name in self._enum_processes()
|
|
)
|
|
|
|
if not running:
|
|
explorer = self._find_windows_tool("explorer.exe")
|
|
|
|
if not explorer:
|
|
return
|
|
|
|
subprocess.Popen([explorer])
|
|
self.write_log("INFO", "System Restart", source="explorer.exe")
|
|
|
|
except Exception:
|
|
log_exception("PYAS_System.SystemMixin.check_process_survival:371")
|
|
pass
|