PYAS_System.py

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