From 11e5fd7bdc77126debad691ce09a880f3887ca03 Mon Sep 17 00:00:00 2001 From: 87owo <85057800+87owo@users.noreply.github.com> Date: Fri, 9 Oct 2026 12:12:42 +0800 Subject: [PATCH] Add files via upload --- PYAS.py | 2909 +++++++++++++++++++++---------------------- PYAS_Autostart.py | 589 +++++++++ PYAS_Cloud.py | 557 +++++---- PYAS_Config.py | 221 ++++ PYAS_Diagnostics.py | 260 ++++ PYAS_Driver.py | 1180 ++++++++++++++++++ PYAS_Engine.py | 1337 +------------------- PYAS_Features.py | 949 ++++++++++++++ PYAS_Logs.py | 127 ++ PYAS_Maintenance.py | 278 +++++ PYAS_PE.py | 283 +++++ PYAS_Popup.py | 763 ++++++++++++ PYAS_Process.py | 342 +++++ PYAS_Protect.py | 2076 ++---------------------------- PYAS_Realtime.py | 555 +++++++++ PYAS_Rules.py | 107 ++ PYAS_Runtime.py | 255 ++++ PYAS_Scanner.py | 454 ++++--- PYAS_Scheduler.py | 160 +++ PYAS_Signature.py | 69 + PYAS_Startup.py | 201 +++ PYAS_Storage.py | 26 + PYAS_System.py | 596 +++++++++ PYAS_Threats.py | 254 ++++ PYAS_Tools.py | 1374 +++----------------- PYAS_Version.py | 33 +- PYAS_WinAPI.py | 649 ++++++++++ README.md | 207 +-- 28 files changed, 10377 insertions(+), 6434 deletions(-) create mode 100644 PYAS_Autostart.py create mode 100644 PYAS_Config.py create mode 100644 PYAS_Diagnostics.py create mode 100644 PYAS_Driver.py create mode 100644 PYAS_Features.py create mode 100644 PYAS_Logs.py create mode 100644 PYAS_Maintenance.py create mode 100644 PYAS_PE.py create mode 100644 PYAS_Popup.py create mode 100644 PYAS_Process.py create mode 100644 PYAS_Realtime.py create mode 100644 PYAS_Rules.py create mode 100644 PYAS_Runtime.py create mode 100644 PYAS_Scheduler.py create mode 100644 PYAS_Signature.py create mode 100644 PYAS_Startup.py create mode 100644 PYAS_Storage.py create mode 100644 PYAS_System.py create mode 100644 PYAS_Threats.py create mode 100644 PYAS_WinAPI.py diff --git a/PYAS.py b/PYAS.py index 62b5e16..244ddf4 100644 --- a/PYAS.py +++ b/PYAS.py @@ -1,1471 +1,1438 @@ -import os, sys, time, copy, json, uuid, queue, platform, threading, logging -import msvcrt, winreg, pystray, subprocess, webview, webbrowser -import ctypes, ctypes.wintypes - -from concurrent.futures import ThreadPoolExecutor -from http.server import SimpleHTTPRequestHandler -from PIL import Image -from socketserver import ThreadingTCPServer -from webview.dom import DOMEventHandler - -from PYAS_Engine import sign_scanner, rule_scanner, pe_scanner, cloud_scanner -from PYAS_Protect import ProtectMixin -from PYAS_Scanner import ScannerMixin -from PYAS_Tools import COPYDATASTRUCT, FILE_NOTIFY_INFORMATION, FILTER_MESSAGE_HEADER, IO_COUNTERS, LUID, LUID_AND_ATTRIBUTES, MEMORY_BASIC_INFORMATION, MIB_TCPROW_OWNER_PID, POINT, PROCESSENTRY32W, PROCESS_BASIC_INFORMATION, PYAS_FULL_MESSAGE, PYAS_MESSAGE, PYAS_USER_MESSAGE, RECT, SERVICE_STATUS_PROCESS, SHQUERYRBINFO, TOKEN_PRIVILEGES, ToolsMixin, UNICODE_STRING - -#################################################################################################### - -PYAS_WINDOW_TITLE = "PYAS Security" -WM_COPYDATA = 0x004A -SMTO_ABORTIFHUNG = 0x0002 -SMTO_ERRORONEXIT = 0x0020 -PYAS_SEND_MESSAGE_FLAGS = SMTO_ABORTIFHUNG | SMTO_ERRORONEXIT -PYAS_MESSAGE_SCAN = 1 -PYAS_MESSAGE_SHOW = 2 -PYAS_MESSAGE_CLOSE = 3 -PYAS_MESSAGE_QUIT = 4 -PYAS_MESSAGE_DRIVER_UNLOAD = 5 -PYAS_MESSAGE_DRIVER_UNINSTALL = 6 - -#################################################################################################### - -class _MainMixin: - def __init__(self): - self._window = None - self.tray_icon = None - self.logs_data = [] - self.logs_dirty = False - - self.lock_config = threading.RLock() - self.lock_logs = threading.RLock() - self.lock_virus = threading.RLock() - self.lock_file_ops = threading.RLock() - - self.init_environ() - self.init_windll() - - if not self.check_singleton("PYAS_Security_Mutex"): - forwarded_exit_code = self.forward_to_existing_instance() - if forwarded_exit_code is not None: - os._exit(forwarded_exit_code) - - self.h_recovery_mutex = self.kernel32.CreateMutexW(None, False, "PYAS_Security_Recovery_Mutex") - if ctypes.get_last_error() == 183: - os._exit(2) - - self.init_variables() - self.load_config() - self.load_logs() - - def find_existing_window(self, timeout=2.0): - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - hwnd = self.user32.FindWindowW(None, PYAS_WINDOW_TITLE) - if hwnd: - return hwnd - time.sleep(0.05) - return None - - def send_existing_window_message(self, hwnd, message_id, payload=None, timeout=1000): - if not hwnd: - return False - - cds = COPYDATASTRUCT() - cds.dwData = message_id - buffer = None - - if payload is not None: - encoded = str(payload).encode('utf-8') - buffer = ctypes.create_string_buffer(encoded + b'\x00') - cds.cbData = len(encoded) + 1 - cds.lpData = ctypes.cast(buffer, ctypes.c_void_p) - else: - cds.cbData = 0 - cds.lpData = None - - try: - return bool(self.user32.SendMessageTimeoutW(hwnd, WM_COPYDATA, 0, ctypes.byref(cds), PYAS_SEND_MESSAGE_FLAGS, timeout, None)) - except Exception: - return False - - def _is_driver_service_absent(self): - scm = self.advapi32.OpenSCManagerW(None, None, 0x0001) - if not scm: - return False - - service = None - try: - ctypes.set_last_error(0) - service = self.advapi32.OpenServiceW(scm, "PYAS_Driver", 0x0004) - if service: - return False - return ctypes.get_last_error() == 1060 - - finally: - if service: - self.advapi32.CloseServiceHandle(service) - self.advapi32.CloseServiceHandle(scm) - - def wait_for_existing_shutdown(self, timeout=30.0, require_service_removal=True): - deadline = time.monotonic() + timeout - while time.monotonic() < deadline: - window_closed = not self.user32.FindWindowW(None, PYAS_WINDOW_TITLE) - service_removed = not require_service_removal or self._is_driver_service_absent() - if window_closed and service_removed: - return True - time.sleep(0.1) - return False - - def forward_to_existing_instance(self): - maintenance_request = any( - arg in self.args_pyas - for arg in ("-quit", "-driver-uninstall", "-driver-unload") - ) - hwnd = self.find_existing_window(timeout=10.0 if maintenance_request else 2.0) - if not hwnd: - return 2 if maintenance_request else None - - if "-quit" in self.args_pyas: - if not self.send_existing_window_message(hwnd, PYAS_MESSAGE_QUIT, timeout=5000): - return 2 - return 0 if self.wait_for_existing_shutdown() else 2 - - if "-driver-uninstall" in self.args_pyas: - if not self.send_existing_window_message(hwnd, PYAS_MESSAGE_DRIVER_UNINSTALL, timeout=5000): - return 2 - return 0 if self.wait_for_existing_shutdown() else 2 - - if "-driver-unload" in self.args_pyas: - if not self.send_existing_window_message(hwnd, PYAS_MESSAGE_DRIVER_UNLOAD, timeout=5000): - return 2 - return 0 if self.wait_for_existing_shutdown(require_service_removal=False) else 2 - - if "-scan" in self.args_pyas: - try: - idx = self.args_pyas.index("-scan") - target = self.args_pyas[idx + 1] - except Exception: - return 2 - - if self.send_existing_window_message(hwnd, PYAS_MESSAGE_SCAN, target, timeout=1200): - try: - self.user32.SetForegroundWindow(hwnd) - except Exception: - pass - return 0 - return 2 - - if self.send_existing_window_message(hwnd, PYAS_MESSAGE_SHOW, timeout=1200): - try: - self.user32.SetForegroundWindow(hwnd) - except Exception: - pass - return 0 - return 2 - - def check_singleton(self, name): - try: - self.h_mutex = self.kernel32.CreateMutexW(None, False, name) - if ctypes.get_last_error() == 183: - return False - - return True - except Exception: - return False - -#################################################################################################### - - def init_environ(self): - self.python = sys.executable - if getattr(sys, 'frozen', False): - self.file_pyas = self.norm_path(sys.executable) - else: - self.file_pyas = self.norm_path(os.path.abspath(sys.argv[0])) - - self.args_pyas = sys.argv[1:] - self.path_pyas = os.path.dirname(self.file_pyas) - self.pid_pyas = int(os.getpid()) - - self.path_appdata = os.environ.get("APPDATA") - self.path_localappdata = os.environ.get("LOCALAPPDATA") or self.path_appdata - self.path_name = os.environ.get("USERNAME") - self.path_temp = os.environ.get("TEMP", f"C:\\Users\\{self.path_name}\\AppData\\Local\\Temp") - self.path_config = os.environ.get("ALLUSERSPROFILE", "C:\\ProgramData") - self.path_system = os.environ.get("SYSTEMROOT", "C:\\Windows") - self.path_user = os.environ.get("USERPROFILE", f"C:\\Users\\{self.path_name}") - - self.path_systemp = os.path.join(self.path_system, "Temp") - self.file_config = os.path.join(self.path_config, "PYAS", "Config.json") - self.file_log = os.path.join(self.path_config, "PYAS", "Report.json") - self.path_webview = os.path.join(self.path_localappdata or self.path_temp, "PYAS", "WebView2") - self.file_webview_log = os.path.join(self.path_localappdata or self.path_temp, "PYAS", "WebView2.log") - self.path_properties = os.path.join(self.path_pyas, "Engine", "Properties") - self.path_heuristic = os.path.join(self.path_pyas, "Engine", "Heuristic") - self.path_protect = os.path.join(self.path_pyas, "Plugins", "Filter") - self.path_rules = os.path.join(self.path_pyas, "Plugins", "Rules") - self.path_drivers = os.path.join(self.path_protect, "PYAS_Driver.sys") - -#################################################################################################### - - def init_windll(self): - for name in ["ntdll", "Psapi", "user32", "kernel32", "iphlpapi", "shell32", "fltlib", "advapi32"]: - try: - setattr(self, name.lower(), ctypes.WinDLL(name, use_last_error=True)) - except Exception as e: - self.write_log("WARN", "init_windll", detail=str(e), success=False) - - self.user32.FindWindowW.argtypes = [ctypes.c_wchar_p, ctypes.c_wchar_p] - self.user32.FindWindowW.restype = ctypes.wintypes.HWND - self.user32.ShowWindow.argtypes = [ctypes.wintypes.HWND, ctypes.c_int] - self.user32.ShowWindow.restype = ctypes.wintypes.BOOL - self.user32.SendMessageTimeoutW.argtypes = [ctypes.wintypes.HWND, ctypes.wintypes.UINT, ctypes.wintypes.WPARAM, ctypes.c_void_p, ctypes.wintypes.UINT, ctypes.wintypes.UINT, ctypes.c_void_p] - self.user32.SendMessageTimeoutW.restype = ctypes.wintypes.LPARAM - - self.user32.WindowFromPoint.argtypes = [POINT] - self.user32.WindowFromPoint.restype = ctypes.wintypes.HWND - self.user32.GetAncestor.argtypes = [ctypes.wintypes.HWND, ctypes.c_uint] - self.user32.GetAncestor.restype = ctypes.wintypes.HWND - self.user32.GetCursorPos.argtypes = [ctypes.POINTER(POINT)] - self.user32.GetCursorPos.restype = ctypes.wintypes.BOOL - self.user32.GetWindowRect.argtypes = [ctypes.wintypes.HWND, ctypes.POINTER(RECT)] - self.user32.GetWindowRect.restype = ctypes.wintypes.BOOL - self.user32.SetWindowPos.argtypes = [ctypes.wintypes.HWND, ctypes.wintypes.HWND, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.c_int, ctypes.wintypes.UINT] - self.user32.SetWindowPos.restype = ctypes.wintypes.BOOL - self.user32.SetLayeredWindowAttributes.argtypes = [ctypes.wintypes.HWND, ctypes.wintypes.DWORD, ctypes.c_byte, ctypes.wintypes.DWORD] - self.user32.SetLayeredWindowAttributes.restype = ctypes.wintypes.BOOL - self.user32.CreateWindowExW.restype = ctypes.wintypes.HWND - - self.ntdll.NtQueryInformationProcess.argtypes = [ctypes.wintypes.HANDLE, ctypes.wintypes.ULONG, ctypes.c_void_p, ctypes.wintypes.ULONG, ctypes.POINTER(ctypes.wintypes.ULONG)] - self.ntdll.NtQueryInformationProcess.restype = ctypes.wintypes.ULONG - self.ntdll.NtSuspendProcess.argtypes = [ctypes.wintypes.HANDLE] - self.ntdll.NtSuspendProcess.restype = ctypes.c_ulong - self.ntdll.NtResumeProcess.argtypes = [ctypes.wintypes.HANDLE] - self.ntdll.NtResumeProcess.restype = ctypes.c_ulong - - self.shell32.CommandLineToArgvW.argtypes = [ctypes.wintypes.LPCWSTR, ctypes.POINTER(ctypes.c_int)] - self.shell32.CommandLineToArgvW.restype = ctypes.POINTER(ctypes.wintypes.LPWSTR) - self.shell32.SHEmptyRecycleBinW.argtypes = [ctypes.wintypes.HWND, ctypes.c_wchar_p, ctypes.wintypes.DWORD] - self.shell32.SHEmptyRecycleBinW.restype = ctypes.c_long - self.shell32.SHQueryRecycleBinW.argtypes = [ctypes.c_wchar_p, ctypes.POINTER(SHQUERYRBINFO)] - self.shell32.SHQueryRecycleBinW.restype = ctypes.c_long - - self.fltlib.FilterConnectCommunicationPort.argtypes = [ctypes.wintypes.LPCWSTR, ctypes.wintypes.DWORD, ctypes.c_void_p, ctypes.wintypes.DWORD, ctypes.c_void_p, ctypes.POINTER(ctypes.wintypes.HANDLE)] - self.fltlib.FilterConnectCommunicationPort.restype = ctypes.c_long - self.fltlib.FilterGetMessage.argtypes = [ctypes.wintypes.HANDLE, ctypes.c_void_p, ctypes.wintypes.DWORD, ctypes.c_void_p] - self.fltlib.FilterGetMessage.restype = ctypes.c_long - self.fltlib.FilterSendMessage.argtypes = [ctypes.wintypes.HANDLE, ctypes.c_void_p, ctypes.wintypes.DWORD, ctypes.c_void_p, ctypes.wintypes.DWORD, ctypes.POINTER(ctypes.wintypes.DWORD)] - self.fltlib.FilterSendMessage.restype = ctypes.c_long - self.fltlib.FilterUnload.argtypes = [ctypes.wintypes.LPCWSTR] - self.fltlib.FilterUnload.restype = ctypes.c_long - - self.advapi32.OpenProcessToken.argtypes = [ctypes.wintypes.HANDLE, ctypes.wintypes.DWORD, ctypes.POINTER(ctypes.wintypes.HANDLE)] - self.advapi32.OpenProcessToken.restype = ctypes.wintypes.BOOL - self.advapi32.LookupPrivilegeValueW.argtypes = [ctypes.wintypes.LPCWSTR, ctypes.wintypes.LPCWSTR, ctypes.POINTER(LUID)] - self.advapi32.LookupPrivilegeValueW.restype = ctypes.wintypes.BOOL - self.advapi32.AdjustTokenPrivileges.argtypes = [ctypes.wintypes.HANDLE, ctypes.wintypes.BOOL, ctypes.POINTER(TOKEN_PRIVILEGES), ctypes.wintypes.DWORD, ctypes.POINTER(TOKEN_PRIVILEGES), ctypes.POINTER(ctypes.wintypes.DWORD)] - self.advapi32.AdjustTokenPrivileges.restype = ctypes.wintypes.BOOL - self.advapi32.OpenSCManagerW.argtypes = [ctypes.wintypes.LPCWSTR, ctypes.wintypes.LPCWSTR, ctypes.wintypes.DWORD] - self.advapi32.OpenSCManagerW.restype = ctypes.wintypes.HANDLE - self.advapi32.CreateServiceW.argtypes = [ctypes.wintypes.HANDLE, ctypes.wintypes.LPCWSTR, ctypes.wintypes.LPCWSTR, ctypes.wintypes.DWORD, ctypes.wintypes.DWORD, ctypes.wintypes.DWORD, ctypes.wintypes.DWORD, ctypes.wintypes.LPCWSTR, ctypes.wintypes.LPCWSTR, ctypes.POINTER(ctypes.wintypes.DWORD), ctypes.wintypes.LPCWSTR, ctypes.wintypes.LPCWSTR, ctypes.wintypes.LPCWSTR] - self.advapi32.CreateServiceW.restype = ctypes.wintypes.HANDLE - self.advapi32.OpenServiceW.argtypes = [ctypes.wintypes.HANDLE, ctypes.wintypes.LPCWSTR, ctypes.wintypes.DWORD] - self.advapi32.OpenServiceW.restype = ctypes.wintypes.HANDLE - self.advapi32.ChangeServiceConfigW.argtypes = [ctypes.wintypes.HANDLE, ctypes.wintypes.DWORD, ctypes.wintypes.DWORD, ctypes.wintypes.DWORD, ctypes.wintypes.LPCWSTR, ctypes.wintypes.LPCWSTR, ctypes.POINTER(ctypes.wintypes.DWORD), ctypes.wintypes.LPCWSTR, ctypes.wintypes.LPCWSTR, ctypes.wintypes.LPCWSTR, ctypes.wintypes.LPCWSTR] - self.advapi32.ChangeServiceConfigW.restype = ctypes.wintypes.BOOL - self.advapi32.StartServiceW.argtypes = [ctypes.wintypes.HANDLE, ctypes.wintypes.DWORD, ctypes.POINTER(ctypes.wintypes.LPCWSTR)] - self.advapi32.StartServiceW.restype = ctypes.wintypes.BOOL - self.advapi32.QueryServiceStatusEx.argtypes = [ctypes.wintypes.HANDLE, ctypes.c_int, ctypes.POINTER(ctypes.c_ubyte), ctypes.wintypes.DWORD, ctypes.POINTER(ctypes.wintypes.DWORD)] - self.advapi32.QueryServiceStatusEx.restype = ctypes.wintypes.BOOL - self.advapi32.DeleteService.argtypes = [ctypes.wintypes.HANDLE] - self.advapi32.DeleteService.restype = ctypes.wintypes.BOOL - self.advapi32.CloseServiceHandle.argtypes = [ctypes.wintypes.HANDLE] - self.advapi32.CloseServiceHandle.restype = ctypes.wintypes.BOOL - - self.kernel32.GetCurrentProcess.argtypes = [] - self.kernel32.GetCurrentProcess.restype = ctypes.wintypes.HANDLE - self.kernel32.CreateToolhelp32Snapshot.restype = ctypes.wintypes.HANDLE - self.kernel32.CreateToolhelp32Snapshot.argtypes = [ctypes.wintypes.DWORD, ctypes.wintypes.DWORD] - self.kernel32.Process32FirstW.restype = ctypes.wintypes.BOOL - self.kernel32.Process32FirstW.argtypes = [ctypes.wintypes.HANDLE, ctypes.c_void_p] - self.kernel32.Process32NextW.restype = ctypes.wintypes.BOOL - self.kernel32.Process32NextW.argtypes = [ctypes.wintypes.HANDLE, ctypes.c_void_p] - self.kernel32.OpenProcess.restype = ctypes.wintypes.HANDLE - self.kernel32.OpenProcess.argtypes = [ctypes.wintypes.DWORD, ctypes.wintypes.BOOL, ctypes.wintypes.DWORD] - self.kernel32.CloseHandle.restype = ctypes.wintypes.BOOL - self.kernel32.CloseHandle.argtypes = [ctypes.wintypes.HANDLE] - self.kernel32.CreateEventW.argtypes = [ctypes.c_void_p, ctypes.wintypes.BOOL, ctypes.wintypes.BOOL, ctypes.wintypes.LPCWSTR] - self.kernel32.CreateEventW.restype = ctypes.wintypes.HANDLE - self.kernel32.WaitForSingleObject.argtypes = [ctypes.wintypes.HANDLE, ctypes.wintypes.DWORD] - self.kernel32.WaitForSingleObject.restype = ctypes.wintypes.DWORD - self.kernel32.CancelIoEx.argtypes = [ctypes.wintypes.HANDLE, ctypes.c_void_p] - self.kernel32.CancelIoEx.restype = ctypes.wintypes.BOOL - self.kernel32.GetOverlappedResult.argtypes = [ctypes.wintypes.HANDLE, ctypes.c_void_p, ctypes.POINTER(ctypes.wintypes.DWORD), ctypes.wintypes.BOOL] - self.kernel32.GetOverlappedResult.restype = ctypes.wintypes.BOOL - self.kernel32.TerminateProcess.restype = ctypes.wintypes.BOOL - self.kernel32.TerminateProcess.argtypes = [ctypes.wintypes.HANDLE, ctypes.c_uint] - self.kernel32.CreateMutexW.restype = ctypes.wintypes.HANDLE - self.kernel32.CreateMutexW.argtypes = [ctypes.c_void_p, ctypes.wintypes.BOOL, ctypes.c_wchar_p] - self.kernel32.CreateFileW.restype = ctypes.wintypes.HANDLE - self.kernel32.CreateFileW.argtypes = [ctypes.c_wchar_p, ctypes.wintypes.DWORD, ctypes.wintypes.DWORD, ctypes.c_void_p, ctypes.wintypes.DWORD, ctypes.wintypes.DWORD, ctypes.wintypes.HANDLE] - self.kernel32.DeviceIoControl.restype = ctypes.wintypes.BOOL - self.kernel32.DeviceIoControl.argtypes = [ctypes.wintypes.HANDLE, ctypes.wintypes.DWORD, ctypes.c_void_p, ctypes.wintypes.DWORD, ctypes.c_void_p, ctypes.wintypes.DWORD, ctypes.POINTER(ctypes.wintypes.DWORD), ctypes.c_void_p] - self.kernel32.ReadProcessMemory.restype = ctypes.wintypes.BOOL - self.kernel32.ReadProcessMemory.argtypes = [ctypes.wintypes.HANDLE, ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, ctypes.POINTER(ctypes.c_size_t)] - self.kernel32.QueryFullProcessImageNameW.restype = ctypes.wintypes.BOOL - self.kernel32.QueryFullProcessImageNameW.argtypes = [ctypes.wintypes.HANDLE, ctypes.wintypes.DWORD, ctypes.wintypes.LPWSTR, ctypes.POINTER(ctypes.wintypes.DWORD)] - self.kernel32.QueryDosDeviceW.restype = ctypes.wintypes.DWORD - self.kernel32.QueryDosDeviceW.argtypes = [ctypes.wintypes.LPCWSTR, ctypes.wintypes.LPWSTR, ctypes.wintypes.DWORD] - self.kernel32.GetProcessIoCounters.argtypes = [ctypes.wintypes.HANDLE, ctypes.POINTER(IO_COUNTERS)] - self.kernel32.GetProcessIoCounters.restype = ctypes.wintypes.BOOL - self.kernel32.VirtualQueryEx.argtypes = [ctypes.wintypes.HANDLE, ctypes.c_void_p, ctypes.POINTER(MEMORY_BASIC_INFORMATION), ctypes.c_size_t] - self.kernel32.VirtualQueryEx.restype = ctypes.c_size_t - self.kernel32.SetProcessWorkingSetSize.argtypes = [ctypes.wintypes.HANDLE, ctypes.c_size_t, ctypes.c_size_t] - self.kernel32.SetProcessWorkingSetSize.restype = ctypes.wintypes.BOOL - - self.psapi.GetMappedFileNameW.argtypes = [ctypes.wintypes.HANDLE, ctypes.c_void_p, ctypes.wintypes.LPWSTR, ctypes.wintypes.DWORD] - self.psapi.GetMappedFileNameW.restype = ctypes.wintypes.DWORD - -#################################################################################################### - - def init_variables(self): - self.heuristic = rule_scanner() - self.properties = pe_scanner() - self.cloud = cloud_scanner() - self.cloud_queue = queue.Queue() - - self.ui_queue = queue.Queue() - self.start_daemon_thread(self.ui_dispatcher_thread) - - self.driver_port = None - self.driver_stop_event = threading.Event() - self.driver_listener_ready_event = threading.Event() - self.driver_listener_failed_event = threading.Event() - self.driver_listener_thread = None - self.ui_ready_event = threading.Event() - self.engine_initialized = False - self.scan_running = False - self.scan_preparing = False - self.scan_stop_requested = False - self.scan_finished = False - self.virus_lock = {} - self.virus_results = [] - self.scan_count = 0 - self.scan_events = {} - self.hash_cache = {} - self.file_task_timers = {} - self.file_task_generations = {} - self.mbr_backup = {} - self.cloud_pending = set() - self.last_io_counters = {} - self.last_io_time = time.time() - self.suspended_procs = set() - - self.lock_driver = threading.RLock() - self.lock_driver_unload = threading.Lock() - self.driver_unload_worker = None - self.driver_unload_result = None - self.closing = False - self.lock_update = threading.RLock() - self.lock_proc = threading.RLock() - self.lock_net = threading.RLock() - self.lock_io = threading.RLock() - - self.pyas_default = { - "version": "3.7.0", - "api_host": "https://pyas-security.com/", - "api_key": "fBRZxYS1UxykM-qzNOlKOEl63WILzlvgNMn6QfsG6FXCAAIktCrOPTAfY5_hEyuZ", - "suffix": [ - ".exe", ".dll", ".sys", ".ocx", ".scr", ".efi", ".acm", ".ax", ".cpl", ".drv", ".com", ".mui", ".pyd", ".wfx", ".api", ".awx", ".rll", ".winmd", - ".bat", ".cmd", ".ps1", ".vbs", ".vbe", ".wsf", ".reg", ".html", ".js", ".jse", ".jsp", ".php", ".hta", ".lnk", ".py", ".sh", ".url", ".rtf", ".ini", - ], - "size": 256 * 1024 * 1024, - "language": "english_switch", - "theme": "system_switch", - "first_launch": True, - "process_switch": False, - "suspend_switch": True, - "load_switch": True, - "document_switch": False, - "system_switch": False, - "driver_switch": False, - "network_switch": False, - "extension_switch": False, - "sensitive_switch": False, - "cloud_switch": False, - "suffix_switch": True, - "autostart_switch": True, - "context_switch": True, - "custom_rule": [], - "white_list": [], - "quarantine": [], - "block_list": [] - } - - self.pass_windows = [ - {"exe": "System Idle Process", "class": "", "title": ""}, - {"exe": "", "class": "Windows.UI.Core.CoreWindow", "title": ""}, - {"exe": "explorer.exe", "class": "", "title": ""} - ] - - self.scan_pool = ThreadPoolExecutor(max_workers=2) - self.protect_pool = ThreadPoolExecutor(max_workers=8) - self.proc_pool = ThreadPoolExecutor(max_workers=16) - self.start_daemon_thread(self.log_flush_thread) - - for _ in range(2): - self.start_daemon_thread(self.cloud_worker) - -#################################################################################################### - - def ui_dispatcher_thread(self): - batch = [] - while True: - try: - task = self.ui_queue.get(timeout=0.1) - batch.append(task) - - while not self.ui_queue.empty() and len(batch) < 50: - try: - batch.append(self.ui_queue.get_nowait()) - except queue.Empty: - break - - if self._window and batch: - js_script = "".join(batch) - try: - self._window.evaluate_js(js_script) - except Exception: - pass - - for _ in batch: - self.ui_queue.task_done() - batch.clear() - - except queue.Empty: - continue - except Exception: - batch.clear() - -#################################################################################################### - - def _loc(self, text_dict): - with self.lock_config: - lang = self.pyas_config.get("language", "traditional_switch") if hasattr(self, 'pyas_config') else "traditional_switch" - return text_dict.get(lang, text_dict.get("traditional_switch", "")) - - def load_config(self): - with self.lock_config: - - if not os.path.exists(self.file_config): - self.pyas_config = copy.deepcopy(self.pyas_default) - self.write_log("INFO", "Config Update", detail="Create default config") - self.save_config() - - else: - try: - with open(self.file_config, "r", encoding="utf-8") as f: - data = json.load(f) - self.pyas_config = copy.deepcopy(self.pyas_default) - self.pyas_config.update(data) - - self.pyas_config["version"] = self.pyas_default["version"] - whitelist_metadata_updated = False - for item in self.pyas_config.get("white_list", []): - if not isinstance(item, dict) or not item.get("file") or "is_dir" in item: - continue - - item_path = self.norm_path(item["file"], must_exist=False) - if item_path and os.path.exists(item_path): - item["is_dir"] = os.path.isdir(item_path) - whitelist_metadata_updated = True - - if whitelist_metadata_updated: - self.save_config() - except Exception as e: - self.pyas_config = copy.deepcopy(self.pyas_default) - self.write_log("WARN", "load_config", detail=str(e), success=False) - - def save_config(self): - with self.lock_config: - try: - os.makedirs(os.path.dirname(self.file_config), exist_ok=True) - with open(self.file_config, "w", encoding="utf-8") as f: - json.dump(self.pyas_config, f, indent=4, ensure_ascii=False) - - except Exception as e: - self.write_log("WARN", "save_config", detail=str(e), success=False) - - def update_config(self, key, value): - with self.lock_update: - defer_driver_disable = key == "driver_switch" and not value - - with self.lock_config: - old_value = self.pyas_config.get(key) - if old_value == value: - return value - - if not defer_driver_disable: - self.pyas_config[key] = value - - if key in ["extension_switch", "sensitive_switch"]: - with self.lock_file_ops: - if hasattr(self, 'hash_cache'): - self.hash_cache.clear() - - success = True - if value: - try: - if key == "process_switch": - self.start_daemon_thread(self.protect_proc_thread) - elif key == "document_switch": - self.start_daemon_thread(self.protect_file_thread) - elif key == "system_switch": - self.start_daemon_thread(self.protect_system_thread) - elif key == "network_switch": - self.start_daemon_thread(self.protect_net_thread) - elif key == "driver_switch": - if self.install_system_driver() and self.start_driver_listener(wait_ready=True): - success = True - else: - self.stop_system_driver() - success = False - - elif key == "context_switch": - self.register_context_menu(True) - elif key == "autostart_switch": - success = self.manage_autostart(True) - - except Exception as e: - self.write_log("WARN", "Feature Start", detail=str(e), success=False) - success = False - else: - if key == "driver_switch": - success = self.stop_system_driver() - - elif key == "context_switch": - self.register_context_menu(False) - elif key == "autostart_switch": - success = self.manage_autostart(False) - - elif key == "document_switch": - self._cancel_pending_file_scans() - with self.lock_file_ops: - if getattr(self, 'h_dir_file', None): - try: - self.kernel32.CloseHandle(self.h_dir_file) - except Exception: - pass - self.h_dir_file = None - - elif key == "suspend_switch": - with self.lock_proc: - if hasattr(self, 'suspended_procs'): - for h in list(self.suspended_procs): - try: - self.ntdll.NtResumeProcess(h) - except Exception: - pass - self.suspended_procs.clear() - - if success: - with self.lock_config: - if defer_driver_disable: - self.pyas_config[key] = value - - self.write_log("INFO", "Config Update", detail=f"[{key}] {old_value} -> {value}") - self.save_config() - - if key == "language" and self.tray_icon: - try: - self.tray_icon.update_menu() - except Exception: - pass - - return value - - with self.lock_config: - self.pyas_config[key] = old_value - - if self._window: - self._window.evaluate_js(f"if(window.revertSwitch) window.revertSwitch('{key}');") - - return old_value - - def get_config(self): - with self.lock_config: - cfg = self.pyas_config.copy() - - rules = [] - if os.path.exists(self.path_rules): - for f in os.listdir(self.path_rules): - if f.lower().endswith(".json"): - fp = os.path.join(self.path_rules, f) - rules.append({"file": fp, "time": os.path.getmtime(fp)}) - - cfg["custom_rule"] = rules - return cfg - - def reset_config(self): - with self.lock_config: - self.pyas_config = copy.deepcopy(self.pyas_default) - self.write_log("INFO", "Config Update", detail="Reset to default") - self.save_config() - - return True - - def load_logs(self): - with self.lock_logs: - pending_logs = list(self.logs_data) - if os.path.exists(self.file_log): - try: - with open(self.file_log, "r", encoding="utf-8") as f: - loaded_logs = json.load(f) - - self.logs_data = (loaded_logs + pending_logs)[-1000:] - except Exception: - self.logs_data = pending_logs - - def log_flush_thread(self): - while True: - time.sleep(5) - with self.lock_logs: - if getattr(self, 'logs_dirty', False): - try: - os.makedirs(os.path.dirname(self.file_log), exist_ok=True) - with open(self.file_log, "w", encoding="utf-8") as f: - json.dump(self.logs_data, f, indent=4, ensure_ascii=False) - - self.logs_dirty = False - except Exception: - pass - -#################################################################################################### - - def write_log(self, level, action, detail=None, code=None, pid=None, file_hash=None, source=None, target=None, operate=None, success=True): - with self.lock_logs: - entry = { - "id": str(uuid.uuid4()), - "timestamp": time.time(), - "time_str": time.strftime("%Y-%m-%d %H:%M:%S"), - "level": level, - "action": action, - "detail": detail, - "code": code, - "pid": pid, - "hash": file_hash, - "source": source, - "target": target, - "operate": operate, - "success": success - } - self.logs_data.append(entry) - - if len(self.logs_data) > 1000: - self.logs_data = self.logs_data[-1000:] - self.logs_dirty = True - - if self._window: - js_cmd = f"if(window.updateLogs) window.updateLogs({json.dumps(entry)});" - self.ui_queue.put(js_cmd) - - if level == "BLOCK" and self.tray_icon: - self.trigger_block_notification(action, source, target, code) - - def get_logs(self): - with self.lock_logs: - return self.logs_data.copy() - - def clear_logs(self, log_ids=None): - with self.lock_logs: - if log_ids is None: - self.logs_data = [] - else: - self.logs_data = [log for log in self.logs_data if log['id'] not in log_ids] - - try: - if not self.logs_data: - if os.path.exists(self.file_log): - os.remove(self.file_log) - else: - os.makedirs(os.path.dirname(self.file_log), exist_ok=True) - with open(self.file_log, "w", encoding="utf-8") as f: - json.dump(self.logs_data, f, indent=4, ensure_ascii=False) - - except Exception: - pass - self.logs_dirty = False - - def export_logs(self, log_ids=None): - if self._window: - path = self._window.create_file_dialog(webview.FileDialog.SAVE, directory='', save_filename='PYAS_Logs.json') - if path: - target_path = path[0] if isinstance(path, (tuple, list)) else path - - with self.lock_logs: - export_data = self.logs_data - if log_ids is not None: - export_data = [log for log in self.logs_data if log['id'] in log_ids] - - try: - with open(target_path, 'w', encoding='utf-8') as f: - json.dump(export_data, f, indent=4, ensure_ascii=False) - - return True - except Exception: - pass - return False - -#################################################################################################### - - def get_tray_text(self, key): - texts = { - "open_ui": { - "traditional_switch": "開啟介面", "simplified_switch": "打开界面", "english_switch": "Open PYAS", - "japanese_switch": "PYAS を開く", "korean_switch": "PYAS 열기", "french_switch": "Ouvrir PYAS", - "spanish_switch": "Abrir PYAS", "hindi_switch": "PYAS खोलें", "arabic_switch": "فتح PYAS", - "russian_switch": "Открыть PYAS", "slovenian_switch": "Odpri PYAS" - }, - "optimize_mem": { - "traditional_switch": "一鍵加速", "simplified_switch": "一键加速", "english_switch": "Memory Boost", - "japanese_switch": "メモリ最適化", "korean_switch": "메모리 최적화", "french_switch": "Optimiser", - "spanish_switch": "Optimizar", "hindi_switch": "मेमोरी बूस्ट", "arabic_switch": "تسريع", - "russian_switch": "Ускорение", "slovenian_switch": "Optimizacija" - }, - "check_update": { - "traditional_switch": "檢查更新", "simplified_switch": "检查更新", "english_switch": "Check Update", - "japanese_switch": "更新を確認", "korean_switch": "업데이트 확인", "french_switch": "Vérifier la mise à jour", - "spanish_switch": "Buscar actualizaciones", "hindi_switch": "अद्यतन जाँचे", "arabic_switch": "التحقق من التحديثات", - "russian_switch": "Проверить обновления", "slovenian_switch": "Preveri posodobitve" - }, - "exit_app": { - "traditional_switch": "退出防護", "simplified_switch": "退出防护", "english_switch": "Exit Security", - "japanese_switch": "保護を終了", "korean_switch": "보호 종료", "french_switch": "Quitter la sécurité", - "spanish_switch": "Salir de la seguridad", "hindi_switch": "सुरक्षा से बाहर निकलें", "arabic_switch": "خروج من الحماية", - "russian_switch": "Выйти из защиты", "slovenian_switch": "Izhod iz zaščite" - } - } - return self._loc(texts.get(key, {})) - - def get_app_icon(self): - icon_path = os.path.join(self.path_pyas, "Interface", "static", "img", "icon.ico") - if os.path.exists(icon_path): - try: - return Image.open(icon_path) - except Exception: - pass - - def show_tray(self): - if self.tray_icon is not None: - return - - menu = pystray.Menu( - pystray.MenuItem(lambda item: self.get_tray_text("open_ui"), self.restore_from_tray, default=True), - pystray.MenuItem(lambda item: self.get_tray_text("optimize_mem"), self.optimize_memory), - pystray.MenuItem(lambda item: self.get_tray_text("check_update"), self.tray_check_update), - pystray.MenuItem(lambda item: self.get_tray_text("exit_app"), self.close) - ) - self.tray_icon = pystray.Icon("PYAS", self.get_app_icon(), "PYAS Security", menu) - self.tray_icon.run_detached() - - def tray_check_update(self, icon=None, item=None): - def _check(): - res = self.check_update() - - title_error = self._loc({ - "traditional_switch": "錯誤", "simplified_switch": "错误", "english_switch": "Error", - "japanese_switch": "エラー", "korean_switch": "오류", "french_switch": "Erreur", - "spanish_switch": "Error", "hindi_switch": "त्रुटि", "arabic_switch": "خطأ", - "russian_switch": "Ошибка", "slovenian_switch": "Napaka"}) - - title_prompt = self._loc({ - "traditional_switch": "提示", "simplified_switch": "提示", "english_switch": "Prompt", - "japanese_switch": "プロンプト", "korean_switch": "프롬프트", "french_switch": "Indication", - "spanish_switch": "Aviso", "hindi_switch": "सुझाव", "arabic_switch": "تلميح", - "russian_switch": "Подсказка", "slovenian_switch": "Namig"}) - - msg_fail = self._loc({ - "traditional_switch": "檢查更新失敗", "simplified_switch": "检查更新失败", "english_switch": "Update check failed", - "japanese_switch": "アップデートの確認に失敗しました", "korean_switch": "업데이트 확인 실패", "french_switch": "Échec de la vérification des mises à jour", - "spanish_switch": "Fallo al buscar actualizaciones", "hindi_switch": "अपडेट की जाँच विफल रही", "arabic_switch": "فشل التحقق من التحديثات", - "russian_switch": "Ошибка проверки обновлений", "slovenian_switch": "Preverjanje posodobitev ni uspelo"}) - - msg_new = self._loc({ - "traditional_switch": "發現新版本", "simplified_switch": "发现新版本", "english_switch": "New version found", - "japanese_switch": "新しいバージョンが見つかりました", "korean_switch": "새 버전을 찾았습니다", "french_switch": "Nouvelle version trouvée", - "spanish_switch": "Nueva versión encontrada", "hindi_switch": "नया संस्करण मिला", "arabic_switch": "تم العثور على إصدار جديد", - "russian_switch": "Найдена новая версия", "slovenian_switch": "Najdena nova različica"}) - - msg_latest = self._loc({ - "traditional_switch": "當前已是最新版本", "simplified_switch": "当前已是最新版本", "english_switch": "Currently at latest version", - "japanese_switch": "現在は最新バージョンです", "korean_switch": "현재 최신 버전입니다", "french_switch": "Actuellement à la dernière version", - "spanish_switch": "Actualmente en la última versión", "hindi_switch": "वर्तमान में नवीनतम संस्करण है", "arabic_switch": "أنت تستخدم أحدث إصدار حاليًا", - "russian_switch": "Установлена последняя версия", "slovenian_switch": "Trenutno imate najnovejšo različico"}) - - if res.get("error"): - self.show_alert(title_error, msg_fail, "error") - - elif res.get("has_update"): - msg = f"{msg_new} {res.get('latest')}\n({res.get('current')} -> {res.get('latest')})" - if self.show_confirm(title_prompt, msg): - self.open_url(res.get("url")) - - else: - msg = f"{msg_latest} {res.get('current')}" - self.show_alert(title_prompt, msg, "info") - - threading.Thread(target=_check, daemon=True).start() - - def restore_from_tray(self, icon=None, item=None): - if self._window: - self._window.restore() - self._window.show() - - hwnd = self.user32.FindWindowW(None, PYAS_WINDOW_TITLE) - if hwnd: - self.user32.SetForegroundWindow(hwnd) - -#################################################################################################### - - def set_window(self, window): - self._window = window - - def minimize(self): - if self._window: - self._window.minimize() - else: - hwnd = self.user32.FindWindowW(None, PYAS_WINDOW_TITLE) - if hwnd: - self.user32.ShowWindow(hwnd, 6) - - def hide_window(self): - if self._window: - self._window.hide() - - def destroy_window_with_timeout(self, window, timeout=1.0): - if not window: - return - - done = threading.Event() - - def destroy_window(): - try: - window.destroy() - except Exception: - pass - finally: - done.set() - - threading.Thread(target=destroy_window, daemon=True).start() - done.wait(timeout) - - def flush_logs_now(self): - with self.lock_logs: - if not getattr(self, 'logs_dirty', False): - return - - try: - os.makedirs(os.path.dirname(self.file_log), exist_ok=True) - with open(self.file_log, "w", encoding="utf-8") as f: - json.dump(self.logs_data, f, indent=4, ensure_ascii=False) - self.logs_dirty = False - except Exception: - pass - - def close(self, *args, uninstall_driver=False, **kwargs): - if getattr(self, 'closing', False): - return True - - self.closing = True - - with self.lock_config: - driver_enabled = self.pyas_config.get("driver_switch", False) - - driver_loaded = driver_enabled or self.check_system_driver() - if uninstall_driver: - driver_stopped, driver_error = self.uninstall_system_driver() - elif driver_loaded: - driver_stopped = self.stop_system_driver() - driver_error = 0 if driver_stopped else 1051 - else: - driver_stopped, driver_error = True, 0 - - if not driver_stopped: - self.closing = False - action = "uninstall" if uninstall_driver else "unload" - self.write_log("WARN", "Driver Protection", detail=f"Controlled {action} failed: 0x{driver_error & 0xFFFFFFFF:08X}", success=False) - if self._window: - try: - self._window.show() - except Exception: - pass - return False - - if self.tray_icon: - try: - self.tray_icon.stop() - except Exception: - pass - self.tray_icon = None - - window = self._window - self._window = None - if window: - try: - window.hide() - except Exception: - pass - - with self.lock_config: - self.pyas_config["process_switch"] = False - self.pyas_config["document_switch"] = False - self.pyas_config["system_switch"] = False - self.pyas_config["driver_switch"] = False - self.pyas_config["network_switch"] = False - - self._cancel_pending_file_tasks() - - with self.lock_file_ops: - if getattr(self, 'h_dir_file', None): - try: - self.kernel32.CloseHandle(self.h_dir_file) - except Exception: - pass - self.h_dir_file = None - - if hasattr(self, 'virus_lock'): - for file_path, (fd, lock_size) in list(self.virus_lock.items()): - try: - msvcrt.locking(fd, msvcrt.LK_UNLCK, lock_size) - except Exception: - pass - try: - os.close(fd) - except Exception: - pass - self.virus_lock.clear() - - with self.lock_proc: - if hasattr(self, 'suspended_procs'): - for h in list(self.suspended_procs): - try: - self.ntdll.NtResumeProcess(h) - self.kernel32.CloseHandle(h) - except Exception: - pass - self.suspended_procs.clear() - - for mutex_name in ('h_mutex', 'h_recovery_mutex'): - handle = getattr(self, mutex_name, None) - if handle: - try: - self.kernel32.CloseHandle(handle) - except Exception: - pass - setattr(self, mutex_name, None) - - self.flush_logs_now() - self.destroy_window_with_timeout(window) - os._exit(0) - - def init_ui_ready(self): - self.ui_ready_event.set() - with self.lock_config: - if self.engine_initialized: - return - - self.engine_initialized = True - self.start_daemon_thread(self.init_engine_thread) - -#################################################################################################### - - def trigger_block_notification(self, action, source, target, code): - if action not in ["Process Block", "Process DLL Block", "File Block", "Network Block", "Driver Block"]: - return - - titles = { - "Process Block": { - "traditional_switch": "進程防護", "simplified_switch": "进程防护", "english_switch": "Process Protection", - "japanese_switch": "プロセス保護", "korean_switch": "프로세스 보호", "french_switch": "Protection des Processus", - "spanish_switch": "Protección de Procesos", "hindi_switch": "प्रक्रिया सुरक्षा", "arabic_switch": "حماية العمليات", - "russian_switch": "Защита процессов", "slovenian_switch": "Zaščita procesov" - }, - "Process DLL Block": { - "traditional_switch": "記憶體防護", "simplified_switch": "内存防护", "english_switch": "Memory Protection", - "japanese_switch": "メモリ保護", "korean_switch": "메모리 보호", "french_switch": "Protection de la mémoire", - "spanish_switch": "Protección de memoria", "hindi_switch": "मेमोरी सुरक्षा", "arabic_switch": "حماية الذاكرة", - "russian_switch": "Защита памяти", "slovenian_switch": "Zaščita pomnilnika" - }, - "File Block": { - "traditional_switch": "檔案防護", "simplified_switch": "文件防护", "english_switch": "File Protection", - "japanese_switch": "ファイル保護", "korean_switch": "파일 보호", "french_switch": "Protection des Fichiers", - "spanish_switch": "Protección de Archivos", "hindi_switch": "फ़ाइल सुरक्षा", "arabic_switch": "حماية الملفات", - "russian_switch": "Защита файлов", "slovenian_switch": "Zaščita datotek" - }, - "Network Block": { - "traditional_switch": "網路防護", "simplified_switch": "网络防护", "english_switch": "Network Protection", - "japanese_switch": "ネットワーク保護", "korean_switch": "네트워크 보호", "french_switch": "Protection Réseau", - "spanish_switch": "Protección de Red", "hindi_switch": "नेटवर्क सुरक्षा", "arabic_switch": "حماية الشبكة", - "russian_switch": "Сетевая защита", "slovenian_switch": "Omrežna zaščita" - }, - "Driver Block": { - "traditional_switch": "驅動防護", "simplified_switch": "驱动防护", "english_switch": "Driver Protection", - "japanese_switch": "ドライバー保護", "korean_switch": "드라이버 보호", "french_switch": "Protection des Pilotes", - "spanish_switch": "Protección de Controladores", "hindi_switch": "ड्राइवर सुरक्षा", "arabic_switch": "حماية برامج التشغيل", - "russian_switch": "Защита драйверов", "slovenian_switch": "Zaščita gonilnikov" - } - } - - path = source or "" - messages = { - "traditional_switch": f"威脅已終止: {path}", "simplified_switch": f"威胁已终止: {path}", "english_switch": f"Threat terminated: {path}", - "japanese_switch": f"脅威が終了しました: {path}", "korean_switch": f"위협이 종료되었습니다: {path}", "french_switch": f"Menace terminée : {path}", - "spanish_switch": f"Amenaza terminada: {path}", "hindi_switch": f"खतरा समाप्त: {path}", "arabic_switch": f"تم إنهاء التهديد: {path}", - "russian_switch": f"Угроза устранена: {path}", "slovenian_switch": f"Grožnja odpravljena: {path}" - } - - title = self._loc(titles.get(action, {})) - message = self._loc(messages) - - try: - self.tray_icon.notify(message, title) - except Exception: - pass - -#################################################################################################### - - def show_notification(self, title, message): - try: - if self.tray_icon: - self.tray_icon.notify(message, title) - - except Exception as e: - self.write_log("WARN", "show_notification", detail=str(e), success=False) - - def show_alert(self, title, message, style="info"): - flags = 0x00000000 | (0x00000010 if style == "error" else 0x00000030 if style == "warning" else 0x00000040) - self.user32.MessageBoxW(0, message, title, flags) - return True - - def show_confirm(self, title, message): - return self.user32.MessageBoxW(0, message, title, 0x00000004 | 0x00000020) == 6 - -#################################################################################################### - - def register_context_menu(self, enable): - paths = [r"Software\Classes\*\shell\PYAS_Scan", r"Software\Classes\Directory\shell\PYAS_Scan"] - cmd_path = f'"{self.file_pyas}"' if getattr(sys, 'frozen', False) else f'"{self.python}" "{self.file_pyas}"' - - try: - for path in paths: - if enable: - self._reg_write(winreg.HKEY_CURRENT_USER, path, None, winreg.REG_SZ, "PYAS Security Scan") - self._reg_write(winreg.HKEY_CURRENT_USER, path, "Icon", winreg.REG_SZ, f'{cmd_path},0') - self._reg_write(winreg.HKEY_CURRENT_USER, rf"{path}\command", None, winreg.REG_SZ, f'{cmd_path} -scan "%1"') - - else: - self._reg_delete(winreg.HKEY_CURRENT_USER, rf"{path}\command") - self._reg_delete(winreg.HKEY_CURRENT_USER, path) - - except Exception as e: - self.write_log("WARN", "register_context_menu", detail=str(e), success=False) - - def trigger_context_scan(self, target): - if self._window: - self._window.evaluate_js(f"if(window.triggerContextScan) window.triggerContextScan({json.dumps(target.replace(os.sep, '/'))});") - -#################################################################################################### - - def on_drop(self, e): - def _process_drop(): - try: - files = e.get('dataTransfer', {}).get('files', []) - paths = [f.get('pywebviewFullPath') for f in files if f.get('pywebviewFullPath')] - - if paths and self._window: - self._window.evaluate_js(f"if(window.triggerContextScan) window.triggerContextScan({json.dumps(paths)});") - - except Exception as ex: - self.write_log("WARN", "on_drop", detail=str(ex), success=False) - - threading.Thread(target=_process_drop, daemon=True).start() - - def select_files(self, file_types=None): - if self._window: - kwargs = {"allow_multiple": True} - if file_types: - kwargs["file_types"] = tuple(file_types) - - return self._window.create_file_dialog(getattr(webview, 'OPEN_DIALOG', 10), **kwargs) or [] - return [] - - def select_folder(self): - if self._window: - return self._window.create_file_dialog(getattr(webview, 'FOLDER_DIALOG', 20)) or [] - return [] - -#################################################################################################### - - def open_file_location(self, file_path): - if not file_path: - return False - - expanded_path = os.path.expandvars(file_path).strip('"').strip("'") - reg_prefixes = ("HKLM", "HKCU", "HKCR", "HKU", "HKCC", "HKEY_") - - if expanded_path.upper().startswith(reg_prefixes): - try: - full_path = expanded_path - if full_path.startswith("HKLM"): - full_path = full_path.replace("HKLM", "HKEY_LOCAL_MACHINE", 1) - - elif full_path.startswith("HKCU"): - full_path = full_path.replace("HKCU", "HKEY_CURRENT_USER", 1) - - elif full_path.startswith("HKCR"): - full_path = full_path.replace("HKCR", "HKEY_CLASSES_ROOT", 1) - - elif full_path.startswith("HKU"): - full_path = full_path.replace("HKU", "HKEY_USERS", 1) - - elif full_path.startswith("HKCC"): - full_path = full_path.replace("HKCC", "HKEY_CURRENT_CONFIG", 1) - - self._run_windows_tool( - ("taskkill.exe",), - ["/F", "/IM", "regedit.exe"], - stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL - ) - self._reg_write(winreg.HKEY_CURRENT_USER, r"Software\Microsoft\Windows\CurrentVersion\Applets\Regedit", "LastKey", winreg.REG_SZ, full_path) - regedit = self._find_windows_tool("regedit.exe") - if not regedit: - return False - subprocess.Popen([regedit]) - - return True - except Exception: - pass - return False - - if os.path.exists(expanded_path): - try: - clean_path = os.path.normpath(expanded_path) - explorer = self._find_windows_tool("explorer.exe") - if not explorer: - return False - subprocess.Popen([explorer, "/select,", clean_path]) - return True - - except Exception: - pass - return False - - def open_website(self): - try: - return webbrowser.open(self.pyas_config.get("api_host")) - except Exception: - return False - - def open_url(self, url): - try: - return webbrowser.open(url) - except Exception: - return False - -#################################################################################################### - -class WindowAPI(_MainMixin, ScannerMixin, ToolsMixin, ProtectMixin): - pass - - -def get_base_path(): - if getattr(sys, 'frozen', False): - return os.path.dirname(sys.executable) - - return os.path.dirname(os.path.abspath(__file__)) - - -def get_frontend_asset_errors(): - base_path = get_base_path() - required_files = [ - os.path.join("Interface", "templates", "index.html"), - os.path.join("Interface", "static", "css", "style.css"), - os.path.join("Interface", "static", "js", "i18n.js"), - os.path.join("Interface", "static", "js", "main.js"), - ] - errors = [] - - for rel_path in required_files: - abs_path = os.path.join(base_path, rel_path) - if not os.path.isfile(abs_path): - errors.append(f"Missing UI asset: {rel_path}") - elif os.path.getsize(abs_path) <= 0: - errors.append(f"Empty UI asset: {rel_path}") - - return errors - - -def show_startup_error(message): - try: - ctypes.windll.user32.MessageBoxW(None, message, PYAS_WINDOW_TITLE, 0x00000010) - except Exception: - pass - -#################################################################################################### - -class NoCacheRequestHandler(SimpleHTTPRequestHandler): - def __init__(self, *args, **kwargs): - super().__init__(*args, directory=os.path.join(get_base_path(), "Interface"), **kwargs) - - def do_GET(self): - if self.path == '/': - self.path = '/templates/index.html' - return super().do_GET() - - def end_headers(self): - self.send_header("Cache-Control", "no-cache, no-store, must-revalidate") - self.send_header("Expires", "0") - self.send_header("Pragma", "no-cache") - super().end_headers() - -#################################################################################################### - -class WindowHook: - def __init__(self, title, api_ref=None): - self.title, self.api_ref = title, api_ref - self.old_wndproc = None - self.WM_DPICHANGED = 0x02E0 - self.WM_NCHITTEST = 0x0084 - self.WM_CLOSE = 0x0010 - self.WM_COPYDATA = WM_COPYDATA - self.WM_SYSCOMMAND = 0x0112 - self.SC_CLOSE = 0xF060 - self.HTCAPTION = 2 - self.GWLP_WNDPROC = -4 - - class RECT(ctypes.Structure): - _fields_ = [("left", ctypes.c_long), ("top", ctypes.c_long), ("right", ctypes.c_long), ("bottom", ctypes.c_long)] - - self.RECT = RECT - self.WNDPROC = ctypes.WINFUNCTYPE(ctypes.c_void_p, ctypes.c_void_p, ctypes.c_uint, ctypes.c_void_p, ctypes.c_void_p) - - self.user32 = ctypes.windll.user32 - self.user32.FindWindowW.argtypes, self.user32.FindWindowW.restype = [ctypes.c_wchar_p, ctypes.c_wchar_p], ctypes.wintypes.HWND - self.user32.CallWindowProcW.argtypes, self.user32.CallWindowProcW.restype = [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_uint, ctypes.c_void_p, ctypes.c_void_p], ctypes.c_void_p - self.user32.DefWindowProcW.argtypes, self.user32.DefWindowProcW.restype = [ctypes.c_void_p, ctypes.c_uint, ctypes.c_void_p, ctypes.c_void_p], ctypes.c_void_p - - if ctypes.sizeof(ctypes.c_void_p) == 8: - self.SetWindowLong, self.GetWindowLong = self.user32.SetWindowLongPtrW, self.user32.GetWindowLongPtrW - else: - self.SetWindowLong, self.GetWindowLong = self.user32.SetWindowLongW, self.user32.GetWindowLongW - - self.SetWindowLong.argtypes, self.SetWindowLong.restype = [ctypes.c_void_p, ctypes.c_int, self.WNDPROC], ctypes.c_void_p - self.GetWindowLong.argtypes, self.GetWindowLong.restype = [ctypes.c_void_p, ctypes.c_int], ctypes.c_void_p - self.new_wndproc_cb = self.WNDPROC(self.wndproc) - - def hook(self): - if self.old_wndproc: - return - - hwnd = self.user32.FindWindowW(None, self.title) - if hwnd: - self.old_wndproc = self.GetWindowLong(hwnd, self.GWLP_WNDPROC) - self.SetWindowLong(hwnd, self.GWLP_WNDPROC, self.new_wndproc_cb) - - try: - self.user32.ChangeWindowMessageFilterEx.argtypes = [ctypes.c_void_p, ctypes.c_uint, ctypes.c_uint, ctypes.c_void_p] - self.user32.ChangeWindowMessageFilterEx(hwnd, self.WM_COPYDATA, 1, None) - except Exception: - pass - - def call_default(self, hwnd, msg, wparam, lparam): - if self.old_wndproc: - return self.user32.CallWindowProcW(self.old_wndproc, hwnd, msg, wparam, lparam) - return self.user32.DefWindowProcW(hwnd, msg, wparam, lparam) - - def wndproc(self, hwnd, msg, wparam, lparam): - if msg == self.WM_COPYDATA: - try: - cds = COPYDATASTRUCT.from_address(lparam) - if cds.dwData == PYAS_MESSAGE_SCAN: - path = ctypes.string_at(cds.lpData, cds.cbData).decode('utf-8').strip('\x00') - if self.api_ref: - threading.Thread(target=self.api_ref.restore_from_tray, daemon=True).start() - threading.Thread(target=self.api_ref.trigger_context_scan, args=(path,), daemon=True).start() - - elif cds.dwData == PYAS_MESSAGE_SHOW and self.api_ref: - threading.Thread(target=self.api_ref.restore_from_tray, daemon=True).start() - - elif cds.dwData == PYAS_MESSAGE_CLOSE and self.api_ref: - threading.Thread(target=self.api_ref.close, daemon=True).start() - - elif cds.dwData in (PYAS_MESSAGE_QUIT, PYAS_MESSAGE_DRIVER_UNINSTALL) and self.api_ref: - threading.Thread(target=self.api_ref.close, kwargs={"uninstall_driver": True}, daemon=True).start() - - elif cds.dwData == PYAS_MESSAGE_DRIVER_UNLOAD and self.api_ref: - threading.Thread(target=self.api_ref.close, daemon=True).start() - - except Exception: - pass - return 1 - - if msg == self.WM_CLOSE or (msg == self.WM_SYSCOMMAND and (wparam & 0xFFF0) == self.SC_CLOSE): - if self.api_ref and not getattr(self.api_ref, 'closing', False): - threading.Thread(target=self.api_ref.hide_window, daemon=True).start() - return 0 - return self.call_default(hwnd, msg, wparam, lparam) - - if msg == self.WM_NCHITTEST: - x, y = lparam & 0xFFFF, (lparam >> 16) & 0xFFFF - if x >= 32768: - x -= 65536 - if y >= 32768: - y -= 65536 - - rect = self.RECT() - self.user32.GetWindowRect(hwnd, ctypes.byref(rect)) - if rect.top <= y <= rect.top + 44 and rect.left <= x <= rect.right - 150: - return self.HTCAPTION - - if msg == self.WM_DPICHANGED: - try: - rect = self.RECT.from_address(lparam) - self.user32.SetWindowPos(hwnd, None, rect.left, rect.top, rect.right - rect.left, rect.bottom - rect.top, 0x0004 | 0x0010 | 0x0020) - except Exception: - pass - - return self.call_default(hwnd, msg, wparam, lparam) - -#################################################################################################### - -def configure_webview_logging(log_path): - try: - os.makedirs(os.path.dirname(log_path), exist_ok=True) - logger = logging.getLogger("pywebview") - logger.setLevel(logging.DEBUG) - normalized = os.path.normcase(os.path.abspath(log_path)) - for handler in logger.handlers: - if isinstance(handler, logging.FileHandler) and os.path.normcase(os.path.abspath(handler.baseFilename)) == normalized: - return - handler = logging.FileHandler(log_path, encoding="utf-8") - handler.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(message)s")) - logger.addHandler(handler) - except Exception: - pass - - -def start_api(port_container, error_container, ready_event): - ThreadingTCPServer.allow_reuse_address = True - ThreadingTCPServer.daemon_threads = True - try: - with ThreadingTCPServer(("127.0.0.1", 0), NoCacheRequestHandler) as httpd: - port_container.append(httpd.server_address[1]) - ready_event.set() - httpd.serve_forever() - - except Exception as e: - error_container.append(str(e)) - ready_event.set() - -if __name__ == "__main__": - if "-quit" in sys.argv or "-driver-unload" in sys.argv or "-driver-uninstall" in sys.argv: - controller = WindowAPI() - if "-quit" in sys.argv or "-driver-uninstall" in sys.argv: - success, delete_error = controller.uninstall_system_driver() - if not success: - controller.write_log("WARN", "Driver Service", detail=f"DeleteService failed: 0x{delete_error & 0xFFFFFFFF:08X}", success=False) - else: - success = controller.stop_system_driver() - controller.flush_logs_now() - os._exit(0 if success else 2) - - hide_on_start = "-h" in sys.argv or "-hide" in sys.argv - init_width, init_height = 980, 670 - - os.environ["WEBVIEW2_ADDITIONAL_BROWSER_ARGUMENTS"] = "--proxy-bypass-list=<-loopback>;localhost;127.0.0.1" - - frontend_errors = get_frontend_asset_errors() - if frontend_errors: - show_startup_error("PYAS UI files are incomplete.\n\n" + "\n".join(frontend_errors)) - os._exit(4) - - port_container = [] - server_errors = [] - server_ready = threading.Event() - api_thread = threading.Thread(target=start_api, args=(port_container, server_errors, server_ready), daemon=True) - api_thread.start() - server_ready.wait(timeout=5.0) - - if not port_container: - show_startup_error("PYAS could not start its local UI server.\n\n" + (server_errors[0] if server_errors else "Unknown error")) - os._exit(1) - - js_api = WindowAPI() - configure_webview_logging(js_api.file_webview_log) - try: - os.makedirs(js_api.path_webview, exist_ok=True) - except Exception as e: - js_api.write_log("WARN", "WebView2", detail=f"User data folder creation failed: {e}", success=False) - - user32 = ctypes.windll.user32 - pos_x, pos_y = (user32.GetSystemMetrics(0) - init_width) // 2, (user32.GetSystemMetrics(1) - init_height) // 2 - - startup_url = f"http://127.0.0.1:{port_container[0]}/" - window = webview.create_window( - title=PYAS_WINDOW_TITLE, url=startup_url, width=init_width, height=init_height, x=pos_x, y=pos_y, - frameless=True, easy_drag=False, js_api=js_api, background_color='#e0e0e0', hidden=hide_on_start) - - if platform.system() == "Windows": - window_hook = WindowHook(PYAS_WINDOW_TITLE, js_api) - window.events.shown += window_hook.hook - - js_api.set_window(window) - js_api.show_tray() - - webview_loaded_event = threading.Event() - dnd_state = {"bound": False} - - def bind_dnd(): - webview_loaded_event.set() - if dnd_state["bound"]: - return - try: - window.dom.document.events.drop += DOMEventHandler(js_api.on_drop, True, True) - dnd_state["bound"] = True - except Exception: - pass - - def webview_startup_watchdog(): - if js_api.ui_ready_event.wait(30.0): - return - - if webview_loaded_event.is_set(): - try: - js_api.write_log("WARN", "WebView2", detail="UI bridge was not ready after 30 seconds; reloading once", success=False) - window.reload() - except Exception as e: - js_api.write_log("WARN", "WebView2", detail=f"Reload failed: {e}", success=False) - - if js_api.ui_ready_event.wait(30.0): - return - - js_api.write_log("WARN", "WebView2", detail="UI initialization timed out after recovery attempt", success=False) - js_api.flush_logs_now() - if not hide_on_start: - show_startup_error( - "PYAS WebView2 failed to initialize.\n\n" - "Please repair or update Microsoft Edge WebView2 Runtime, then start PYAS again.\n\n" - f"Diagnostic log: {js_api.file_webview_log}" - ) - os._exit(3) - - window.events.loaded += bind_dnd - threading.Thread(target=webview_startup_watchdog, daemon=True).start() - try: - webview.start(gui="edgechromium", private_mode=False, storage_path=js_api.path_webview) - except Exception as e: - js_api.write_log("WARN", "WebView2", detail=f"Startup failed: {e}", success=False) - js_api.flush_logs_now() - if not hide_on_start: - show_startup_error(f"PYAS WebView2 failed to start.\n\n{e}\n\nDiagnostic log: {js_api.file_webview_log}") - os._exit(3) +from PYAS_Diagnostics import log_exception, bind_report_logging +import os +import sys +import time +import copy +import json +import uuid +import queue +import platform +import threading +import logging +from PYAS_Startup import ( + bootstrap_startup, + check_webview_dependencies, + prepare_webview_profile, + restart_with_new_webview_profile, + wait_for_recovery_parent, +) + +STARTUP_LOG_PATH = bootstrap_startup() if __name__ == "__main__" else None +import msvcrt +import winreg +import pystray +import subprocess +import webview +import webbrowser +import ctypes +import ctypes.wintypes +from concurrent.futures import ThreadPoolExecutor +from http.server import SimpleHTTPRequestHandler +from PIL import Image +from socketserver import ThreadingTCPServer +from webview.dom import DOMEventHandler +from PYAS_Engine import sign_scanner, rule_scanner, pe_scanner, cloud_scanner +from PYAS_Protect import ProtectMixin +from PYAS_Scanner import ScannerMixin +from PYAS_Tools import ( + COPYDATASTRUCT, + FILE_NOTIFY_INFORMATION, + FILTER_MESSAGE_HEADER, + IO_COUNTERS, + LUID, + LUID_AND_ATTRIBUTES, + MEMORY_BASIC_INFORMATION, + MIB_TCPROW_OWNER_PID, + POINT, + PROCESSENTRY32W, + PROCESS_BASIC_INFORMATION, + PYAS_FULL_MESSAGE, + PYAS_MESSAGE, + PYAS_USER_MESSAGE, + RECT, + SERVICE_STATUS_PROCESS, + SHQUERYRBINFO, + TOKEN_PRIVILEGES, + ToolsMixin, + UNICODE_STRING, +) + + +PYAS_WINDOW_TITLE = "PYAS Security" +WM_COPYDATA = 0x004A +SMTO_ABORTIFHUNG = 0x0002 +SMTO_ERRORONEXIT = 0x0020 +PYAS_SEND_MESSAGE_FLAGS = SMTO_ABORTIFHUNG | SMTO_ERRORONEXIT +PYAS_MESSAGE_SCAN = 1 +PYAS_MESSAGE_SHOW = 2 +PYAS_MESSAGE_CLOSE = 3 +PYAS_MESSAGE_QUIT = 4 +PYAS_MESSAGE_DRIVER_UNLOAD = 5 +PYAS_MESSAGE_DRIVER_UNINSTALL = 6 +PYAS_MAINTENANCE_MESSAGE = 0x8000 + 0x501 +PYAS_MAINTENANCE_COMMANDS = frozenset( + { + PYAS_MESSAGE_CLOSE, + PYAS_MESSAGE_QUIT, + PYAS_MESSAGE_DRIVER_UNLOAD, + PYAS_MESSAGE_DRIVER_UNINSTALL, + } +) + +from PYAS_WinAPI import WindowsMixin +from PYAS_Config import ConfigMixin +from PYAS_Logs import LogMixin +from PYAS_Runtime import RuntimeMixin + + +class _MainMixin(WindowsMixin, ConfigMixin, LogMixin, RuntimeMixin): + def __init__(self): + self._window = None + self.tray_icon = None + self.logs_data = [] + self.logs_dirty = False + + self.lock_config = threading.RLock() + self.lock_logs = threading.RLock() + self.lock_virus = threading.RLock() + self.lock_file_ops = threading.RLock() + + self.init_environ() + self.init_windll() + + if not self.check_singleton("PYAS_Security_Mutex"): + forwarded_exit_code = self.forward_to_existing_instance() + + if forwarded_exit_code is not None: + os._exit(forwarded_exit_code) + + self.h_recovery_mutex = self.kernel32.CreateMutexW( + None, False, "PYAS_Security_Recovery_Mutex" + ) + + if ctypes.get_last_error() == 183: + os._exit(2) + + self.init_variables() + self.load_config() + self.load_logs() + bind_report_logging(self) + + def find_existing_window(self, timeout=2.0): + deadline = time.monotonic() + timeout + + while time.monotonic() < deadline: + hwnd = self.user32.FindWindowW(None, PYAS_WINDOW_TITLE) + + if hwnd: + return hwnd + + time.sleep(0.05) + + return None + + def send_existing_window_message(self, hwnd, message_id, payload=None, timeout=1000): + if not hwnd: + return False + + if message_id in PYAS_MAINTENANCE_COMMANDS: + try: + return bool( + self.user32.SendMessageTimeoutW( + hwnd, + PYAS_MAINTENANCE_MESSAGE, + message_id, + None, + PYAS_SEND_MESSAGE_FLAGS, + timeout, + None, + ) + ) + except Exception: + log_exception("MainMixin.send_maintenance_message") + return False + + cds = COPYDATASTRUCT() + cds.dwData = message_id + buffer = None + + if payload is not None: + encoded = str(payload).encode("utf-8") + buffer = ctypes.create_string_buffer(encoded + b"\x00") + cds.cbData = len(encoded) + 1 + cds.lpData = ctypes.cast(buffer, ctypes.c_void_p) + else: + cds.cbData = 0 + cds.lpData = None + + try: + return bool( + self.user32.SendMessageTimeoutW( + hwnd, WM_COPYDATA, 0, ctypes.byref(cds), PYAS_SEND_MESSAGE_FLAGS, timeout, None + ) + ) + except Exception: + log_exception("PYAS._MainMixin.send_existing_window_message:96") + return False + + def _is_driver_service_absent(self): + scm = self.advapi32.OpenSCManagerW(None, None, 0x0001) + + if not scm: + return False + + service = None + + try: + ctypes.set_last_error(0) + service = self.advapi32.OpenServiceW(scm, "PYAS_Driver", 0x0004) + + if service: + return False + + return ctypes.get_last_error() == 1060 + + finally: + if service: + self.advapi32.CloseServiceHandle(service) + + self.advapi32.CloseServiceHandle(scm) + + def wait_for_existing_shutdown(self, timeout=30.0, require_service_removal=True): + deadline = time.monotonic() + timeout + + while time.monotonic() < deadline: + window_closed = not self.user32.FindWindowW(None, PYAS_WINDOW_TITLE) + service_removed = not require_service_removal or self._is_driver_service_absent() + + if window_closed and service_removed: + return True + + time.sleep(0.1) + + return False + + def forward_to_existing_instance(self): + maintenance_request = any( + arg in self.args_pyas for arg in ("-quit", "-driver-uninstall", "-driver-unload") + ) + hwnd = self.find_existing_window(timeout=10.0 if maintenance_request else 2.0) + + if not hwnd: + return 2 if maintenance_request else None + + if "-quit" in self.args_pyas: + if not self.send_existing_window_message(hwnd, PYAS_MESSAGE_QUIT, timeout=5000): + return 2 + + return 0 if self.wait_for_existing_shutdown() else 2 + + if "-driver-uninstall" in self.args_pyas: + if not self.send_existing_window_message( + hwnd, PYAS_MESSAGE_DRIVER_UNINSTALL, timeout=5000 + ): + return 2 + + return 0 if self.wait_for_existing_shutdown() else 2 + + if "-driver-unload" in self.args_pyas: + if not self.send_existing_window_message( + hwnd, PYAS_MESSAGE_DRIVER_UNLOAD, timeout=5000 + ): + return 2 + + return 0 if self.wait_for_existing_shutdown(require_service_removal=False) else 2 + + if "-scan" in self.args_pyas: + try: + idx = self.args_pyas.index("-scan") + target = self.args_pyas[idx + 1] + except Exception: + log_exception("PYAS._MainMixin.forward_to_existing_instance:155") + return 2 + + if self.send_existing_window_message(hwnd, PYAS_MESSAGE_SCAN, target, timeout=1200): + try: + self.user32.SetForegroundWindow(hwnd) + except Exception: + log_exception("PYAS._MainMixin.forward_to_existing_instance:161") + pass + + return 0 + + return 2 + + if self.send_existing_window_message(hwnd, PYAS_MESSAGE_SHOW, timeout=1200): + try: + self.user32.SetForegroundWindow(hwnd) + except Exception: + log_exception("PYAS._MainMixin.forward_to_existing_instance:169") + pass + + return 0 + + return 2 + + def check_singleton(self, name): + try: + self.h_mutex = self.kernel32.CreateMutexW(None, False, name) + + if ctypes.get_last_error() == 183: + return False + + return True + except Exception: + log_exception("PYAS._MainMixin.check_singleton:181") + return False + + def get_tray_text(self, key): + texts = { + "open_ui": { + "traditional_switch": "開啟介面", + "simplified_switch": "打开界面", + "english_switch": "Open PYAS", + "japanese_switch": "PYAS を開く", + "korean_switch": "PYAS 열기", + "french_switch": "Ouvrir PYAS", + "spanish_switch": "Abrir PYAS", + "hindi_switch": "PYAS खोलें", + "arabic_switch": "فتح PYAS", + "russian_switch": "Открыть PYAS", + "slovenian_switch": "Odpri PYAS", + }, + "optimize_mem": { + "traditional_switch": "一鍵加速", + "simplified_switch": "一键加速", + "english_switch": "Memory Boost", + "japanese_switch": "メモリ最適化", + "korean_switch": "메모리 최적화", + "french_switch": "Optimiser", + "spanish_switch": "Optimizar", + "hindi_switch": "मेमोरी बूस्ट", + "arabic_switch": "تسريع", + "russian_switch": "Ускорение", + "slovenian_switch": "Optimizacija", + }, + "check_update": { + "traditional_switch": "檢查更新", + "simplified_switch": "检查更新", + "english_switch": "Check Update", + "japanese_switch": "更新を確認", + "korean_switch": "업데이트 확인", + "french_switch": "Vérifier la mise à jour", + "spanish_switch": "Buscar actualizaciones", + "hindi_switch": "अद्यतन जाँचे", + "arabic_switch": "التحقق من التحديثات", + "russian_switch": "Проверить обновления", + "slovenian_switch": "Preveri posodobitve", + }, + "exit_app": { + "traditional_switch": "退出防護", + "simplified_switch": "退出防护", + "english_switch": "Exit Security", + "japanese_switch": "保護を終了", + "korean_switch": "보호 종료", + "french_switch": "Quitter la sécurité", + "spanish_switch": "Salir de la seguridad", + "hindi_switch": "सुरक्षा से बाहर निकलें", + "arabic_switch": "خروج من الحماية", + "russian_switch": "Выйти из защиты", + "slovenian_switch": "Izhod iz zaščite", + }, + } + return self._loc(texts.get(key, {})) + + def get_app_icon(self): + icon_path = os.path.join(self.path_pyas, "Interface", "static", "img", "icon.ico") + + if os.path.exists(icon_path): + try: + return Image.open(icon_path) + except Exception: + log_exception("PYAS._MainMixin.get_app_icon:218") + pass + + def show_tray(self): + if self.tray_icon is not None: + return + + menu = pystray.Menu( + pystray.MenuItem( + lambda item: self.get_tray_text("open_ui"), self.restore_from_tray, default=True + ), + pystray.MenuItem(lambda item: self.get_tray_text("optimize_mem"), self.optimize_memory), + pystray.MenuItem( + lambda item: self.get_tray_text("check_update"), self.tray_check_update + ), + pystray.MenuItem(lambda item: self.get_tray_text("exit_app"), self.close), + ) + self.tray_icon = pystray.Icon("PYAS", self.get_app_icon(), "PYAS Security", menu) + self.tray_icon.run_detached() + + def tray_check_update(self, icon=None, item=None): + def _check(): + res = self.check_update() + + title_error = self._loc( + { + "traditional_switch": "錯誤", + "simplified_switch": "错误", + "english_switch": "Error", + "japanese_switch": "エラー", + "korean_switch": "오류", + "french_switch": "Erreur", + "spanish_switch": "Error", + "hindi_switch": "त्रुटि", + "arabic_switch": "خطأ", + "russian_switch": "Ошибка", + "slovenian_switch": "Napaka", + } + ) + + title_prompt = self._loc( + { + "traditional_switch": "提示", + "simplified_switch": "提示", + "english_switch": "Prompt", + "japanese_switch": "プロンプト", + "korean_switch": "프롬프트", + "french_switch": "Indication", + "spanish_switch": "Aviso", + "hindi_switch": "सुझाव", + "arabic_switch": "تلميح", + "russian_switch": "Подсказка", + "slovenian_switch": "Namig", + } + ) + + msg_fail = self._loc( + { + "traditional_switch": "檢查更新失敗", + "simplified_switch": "检查更新失败", + "english_switch": "Update check failed", + "japanese_switch": "アップデートの確認に失敗しました", + "korean_switch": "업데이트 확인 실패", + "french_switch": "Échec de la vérification des mises à jour", + "spanish_switch": "Fallo al buscar actualizaciones", + "hindi_switch": "अपडेट की जाँच विफल रही", + "arabic_switch": "فشل التحقق من التحديثات", + "russian_switch": "Ошибка проверки обновлений", + "slovenian_switch": "Preverjanje posodobitev ni uspelo", + } + ) + + msg_new = self._loc( + { + "traditional_switch": "發現新版本", + "simplified_switch": "发现新版本", + "english_switch": "New version found", + "japanese_switch": "新しいバージョンが見つかりました", + "korean_switch": "새 버전을 찾았습니다", + "french_switch": "Nouvelle version trouvée", + "spanish_switch": "Nueva versión encontrada", + "hindi_switch": "नया संस्करण मिला", + "arabic_switch": "تم العثور على إصدار جديد", + "russian_switch": "Найдена новая версия", + "slovenian_switch": "Najdena nova različica", + } + ) + + msg_latest = self._loc( + { + "traditional_switch": "當前已是最新版本", + "simplified_switch": "当前已是最新版本", + "english_switch": "Currently at latest version", + "japanese_switch": "現在は最新バージョンです", + "korean_switch": "현재 최신 버전입니다", + "french_switch": "Actuellement à la dernière version", + "spanish_switch": "Actualmente en la última versión", + "hindi_switch": "वर्तमान में नवीनतम संस्करण है", + "arabic_switch": "أنت تستخدم أحدث إصدار حاليًا", + "russian_switch": "Установлена последняя версия", + "slovenian_switch": "Trenutno imate najnovejšo različico", + } + ) + + if res.get("error"): + self.show_alert(title_error, msg_fail, "error") + + elif res.get("has_update"): + msg = ( + f"{msg_new} {res.get('latest')}\n({res.get('current')} -> {res.get('latest')})" + ) + + if self.show_confirm(title_prompt, msg): + self.open_url(res.get("url")) + + else: + msg = f"{msg_latest} {res.get('current')}" + self.show_alert(title_prompt, msg, "info") + + threading.Thread(target=_check, daemon=True).start() + + def restore_from_tray(self, icon=None, item=None): + if self._window: + self._window.restore() + self._window.show() + + hwnd = self.user32.FindWindowW(None, PYAS_WINDOW_TITLE) + + if hwnd: + self.user32.SetForegroundWindow(hwnd) + + def set_window(self, window): + self._window = window + + def minimize(self): + if self._window: + self._window.minimize() + else: + hwnd = self.user32.FindWindowW(None, PYAS_WINDOW_TITLE) + + if hwnd: + self.user32.ShowWindow(hwnd, 6) + + def hide_window(self): + if self._window: + self._window.hide() + + def destroy_window_with_timeout(self, window, timeout=1.0): + if not window: + return + + done = threading.Event() + + def destroy_window(): + try: + window.destroy() + except Exception: + log_exception("PYAS._MainMixin.destroy_window_with_timeout.destroy_window:315") + pass + finally: + done.set() + + threading.Thread(target=destroy_window, daemon=True).start() + done.wait(timeout) + + def close(self, *args, uninstall_driver=False, **kwargs): + if getattr(self, "closing", False): + return True + + self.closing = True + + with self.lock_config: + driver_enabled = self.pyas_config.get("driver_switch", False) + + driver_loaded = driver_enabled or self.check_system_driver() + + if uninstall_driver: + driver_stopped, driver_error = self.uninstall_system_driver() + elif driver_loaded: + driver_stopped = self.stop_system_driver() + driver_error = 0 if driver_stopped else 1051 + else: + driver_stopped, driver_error = True, 0 + + if not driver_stopped: + self.closing = False + action = "uninstall" if uninstall_driver else "unload" + self.write_log( + "WARN", + "Driver Protection", + detail=f"Controlled {action} failed: 0x{driver_error & 0xFFFFFFFF:08X}", + success=False, + ) + + if self._window: + try: + self._window.show() + except Exception: + log_exception("PYAS._MainMixin.close:348") + pass + + return False + + if self.tray_icon: + try: + self.tray_icon.stop() + except Exception: + log_exception("PYAS._MainMixin.close:355") + pass + + self.tray_icon = None + + window = self._window + self._window = None + + if window: + try: + window.hide() + except Exception: + log_exception("PYAS._MainMixin.close:364") + pass + + with self.lock_config: + self.pyas_config["process_switch"] = False + self.pyas_config["document_switch"] = False + self.pyas_config["system_switch"] = False + self.pyas_config["driver_switch"] = False + self.pyas_config["network_switch"] = False + + self._cancel_pending_file_tasks() + scheduler = getattr(self, "file_scheduler", None) + + if scheduler: + scheduler.close() + + self.cloud_cancel_event.set() + + with self.lock_file_ops: + if getattr(self, "h_dir_file", None): + try: + self.kernel32.CloseHandle(self.h_dir_file) + except Exception: + log_exception("PYAS._MainMixin.close:380") + pass + + self.h_dir_file = None + + if hasattr(self, "virus_lock"): + for file_path, (fd, lock_size) in list(self.virus_lock.items()): + try: + msvcrt.locking(fd, msvcrt.LK_UNLCK, lock_size) + except Exception: + log_exception("PYAS._MainMixin.close:388") + pass + + try: + os.close(fd) + except Exception: + log_exception("PYAS._MainMixin.close:392") + pass + + self.virus_lock.clear() + + with self.lock_proc: + if hasattr(self, "suspended_procs"): + for h in list(self.suspended_procs): + try: + self.ntdll.NtResumeProcess(h) + self.kernel32.CloseHandle(h) + except Exception: + log_exception("PYAS._MainMixin.close:402") + pass + + self.suspended_procs.clear() + + for mutex_name in ("h_mutex", "h_recovery_mutex"): + handle = getattr(self, mutex_name, None) + + if handle: + try: + self.kernel32.CloseHandle(handle) + except Exception: + log_exception("PYAS._MainMixin.close:411") + pass + + setattr(self, mutex_name, None) + + self.flush_logs_now() + self.destroy_window_with_timeout(window) + os._exit(0) + + def report_ui_error(self, stage, message): + logging.getLogger("PYAS.Startup").warning( + "UI stage=%s: %s", str(stage)[:80], str(message)[:2000] + ) + return True + + def init_ui_ready(self): + logging.getLogger("PYAS.Startup").info("Python UI bridge ready") + self.ui_ready_event.set() + + with self.lock_config: + if self.engine_initialized: + return + + self.engine_initialized = True + + self.start_daemon_thread(self.init_engine_thread) + + def trigger_block_notification(self, action, source, target, code): + if action not in [ + "Process Block", + "Process DLL Block", + "File Block", + "Network Block", + "Driver Block", + ]: + return + + titles = { + "Process Block": { + "traditional_switch": "進程防護", + "simplified_switch": "进程防护", + "english_switch": "Process Protection", + "japanese_switch": "プロセス保護", + "korean_switch": "프로세스 보호", + "french_switch": "Protection des Processus", + "spanish_switch": "Protección de Procesos", + "hindi_switch": "प्रक्रिया सुरक्षा", + "arabic_switch": "حماية العمليات", + "russian_switch": "Защита процессов", + "slovenian_switch": "Zaščita procesov", + }, + "Process DLL Block": { + "traditional_switch": "記憶體防護", + "simplified_switch": "内存防护", + "english_switch": "Memory Protection", + "japanese_switch": "メモリ保護", + "korean_switch": "메모리 보호", + "french_switch": "Protection de la mémoire", + "spanish_switch": "Protección de memoria", + "hindi_switch": "मेमोरी सुरक्षा", + "arabic_switch": "حماية الذاكرة", + "russian_switch": "Защита памяти", + "slovenian_switch": "Zaščita pomnilnika", + }, + "File Block": { + "traditional_switch": "檔案防護", + "simplified_switch": "文件防护", + "english_switch": "File Protection", + "japanese_switch": "ファイル保護", + "korean_switch": "파일 보호", + "french_switch": "Protection des Fichiers", + "spanish_switch": "Protección de Archivos", + "hindi_switch": "फ़ाइल सुरक्षा", + "arabic_switch": "حماية الملفات", + "russian_switch": "Защита файлов", + "slovenian_switch": "Zaščita datotek", + }, + "Network Block": { + "traditional_switch": "網路防護", + "simplified_switch": "网络防护", + "english_switch": "Network Protection", + "japanese_switch": "ネットワーク保護", + "korean_switch": "네트워크 보호", + "french_switch": "Protection Réseau", + "spanish_switch": "Protección de Red", + "hindi_switch": "नेटवर्क सुरक्षा", + "arabic_switch": "حماية الشبكة", + "russian_switch": "Сетевая защита", + "slovenian_switch": "Omrežna zaščita", + }, + "Driver Block": { + "traditional_switch": "驅動防護", + "simplified_switch": "驱动防护", + "english_switch": "Driver Protection", + "japanese_switch": "ドライバー保護", + "korean_switch": "드라이버 보호", + "french_switch": "Protection des Pilotes", + "spanish_switch": "Protección de Controladores", + "hindi_switch": "ड्राइवर सुरक्षा", + "arabic_switch": "حماية برامج التشغيل", + "russian_switch": "Защита драйверов", + "slovenian_switch": "Zaščita gonilnikov", + }, + } + + path = source or "" + messages = { + "traditional_switch": f"威脅已終止: {path}", + "simplified_switch": f"威胁已终止: {path}", + "english_switch": f"Threat terminated: {path}", + "japanese_switch": f"脅威が終了しました: {path}", + "korean_switch": f"위협이 종료되었습니다: {path}", + "french_switch": f"Menace terminée : {path}", + "spanish_switch": f"Amenaza terminada: {path}", + "hindi_switch": f"खतरा समाप्त: {path}", + "arabic_switch": f"تم إنهاء التهديد: {path}", + "russian_switch": f"Угроза устранена: {path}", + "slovenian_switch": f"Grožnja odpravljena: {path}", + } + + title = self._loc(titles.get(action, {})) + message = self._loc(messages) + + try: + self.tray_icon.notify(message, title) + except Exception: + log_exception("PYAS._MainMixin.trigger_block_notification:483") + pass + + def show_notification(self, title, message): + try: + if self.tray_icon: + self.tray_icon.notify(message, title) + + except Exception as e: + log_exception("PYAS._MainMixin.show_notification:491") + self.write_log("WARN", "show_notification", detail=str(e), success=False) + + def show_alert(self, title, message, style="info"): + flags = 0x00000000 | ( + 0x00000010 if style == "error" else 0x00000030 if style == "warning" else 0x00000040 + ) + self.user32.MessageBoxW(0, message, title, flags) + return True + + def show_confirm(self, title, message): + return self.user32.MessageBoxW(0, message, title, 0x00000004 | 0x00000020) == 6 + + def register_context_menu(self, enable): + paths = [ + r"Software\Classes\*\shell\PYAS_Scan", + r"Software\Classes\Directory\shell\PYAS_Scan", + ] + cmd_path = ( + f'"{self.file_pyas}"' + if getattr(sys, "frozen", False) + else f'"{self.python}" "{self.file_pyas}"' + ) + + try: + success = True + + for path in paths: + if enable: + success = ( + self._reg_write( + winreg.HKEY_CURRENT_USER, + path, + None, + winreg.REG_SZ, + "PYAS Security Scan", + ) + and success + ) + success = ( + self._reg_write( + winreg.HKEY_CURRENT_USER, path, "Icon", winreg.REG_SZ, f"{cmd_path},0" + ) + and success + ) + success = ( + self._reg_write( + winreg.HKEY_CURRENT_USER, + rf"{path}\command", + None, + winreg.REG_SZ, + f'{cmd_path} -scan "%1"', + ) + and success + ) + else: + for entry in (rf"{path}\command", path): + try: + with winreg.OpenKey( + winreg.HKEY_CURRENT_USER, entry, 0, winreg.KEY_READ + ): + pass + except FileNotFoundError: + continue + + success = self._reg_delete(winreg.HKEY_CURRENT_USER, entry) and success + + if not success: + self.write_log( + "WARN", + "register_context_menu", + detail="Registry operation failed", + success=False, + ) + + return success + except Exception as error: + log_exception("PYAS._MainMixin.register_context_menu:523") + self.write_log("WARN", "register_context_menu", detail=str(error), success=False) + return False + + def trigger_context_scan(self, target): + if self._window: + self._window.evaluate_js( + f"if(window.triggerContextScan) window.triggerContextScan({json.dumps(target.replace(os.sep, '/'))});" + ) + + def on_drop(self, e): + def _process_drop(): + try: + files = e.get("dataTransfer", {}).get("files", []) + paths = [f.get("pywebviewFullPath") for f in files if f.get("pywebviewFullPath")] + + if paths and self._window: + self._window.evaluate_js( + f"if(window.triggerContextScan) window.triggerContextScan({json.dumps(paths)});" + ) + + except Exception as ex: + log_exception("PYAS._MainMixin.on_drop._process_drop:540") + self.write_log("WARN", "on_drop", detail=str(ex), success=False) + + threading.Thread(target=_process_drop, daemon=True).start() + + def select_files(self, file_types=None): + if self._window: + kwargs = {"allow_multiple": True} + + if file_types: + kwargs["file_types"] = tuple(file_types) + + return ( + self._window.create_file_dialog(getattr(webview, "OPEN_DIALOG", 10), **kwargs) or [] + ) + + return [] + + def select_folder(self): + if self._window: + return self._window.create_file_dialog(getattr(webview, "FOLDER_DIALOG", 20)) or [] + + return [] + + def open_file_location(self, file_path): + if not file_path: + return False + + expanded_path = os.path.expandvars(file_path).strip('"').strip("'") + reg_prefixes = ("HKLM", "HKCU", "HKCR", "HKU", "HKCC", "HKEY_") + + if expanded_path.upper().startswith(reg_prefixes): + try: + full_path = expanded_path + + if full_path.startswith("HKLM"): + full_path = full_path.replace("HKLM", "HKEY_LOCAL_MACHINE", 1) + + elif full_path.startswith("HKCU"): + full_path = full_path.replace("HKCU", "HKEY_CURRENT_USER", 1) + + elif full_path.startswith("HKCR"): + full_path = full_path.replace("HKCR", "HKEY_CLASSES_ROOT", 1) + + elif full_path.startswith("HKU"): + full_path = full_path.replace("HKU", "HKEY_USERS", 1) + + elif full_path.startswith("HKCC"): + full_path = full_path.replace("HKCC", "HKEY_CURRENT_CONFIG", 1) + + self._run_windows_tool( + ("taskkill.exe",), + ["/F", "/IM", "regedit.exe"], + stdout=subprocess.DEVNULL, + stderr=subprocess.DEVNULL, + ) + self._reg_write( + winreg.HKEY_CURRENT_USER, + r"Software\Microsoft\Windows\CurrentVersion\Applets\Regedit", + "LastKey", + winreg.REG_SZ, + full_path, + ) + regedit = self._find_windows_tool("regedit.exe") + + if not regedit: + return False + + subprocess.Popen([regedit]) + + return True + except Exception: + log_exception("PYAS._MainMixin.open_file_location:597") + pass + + return False + + if os.path.exists(expanded_path): + try: + clean_path = os.path.normpath(expanded_path) + explorer = self._find_windows_tool("explorer.exe") + + if not explorer: + return False + + subprocess.Popen([explorer, "/select,", clean_path]) + return True + + except Exception: + log_exception("PYAS._MainMixin.open_file_location:610") + pass + + return False + + def open_website(self): + try: + return webbrowser.open(self.pyas_config.get("api_host")) + except Exception: + log_exception("PYAS._MainMixin.open_website:617") + return False + + def open_url(self, url): + try: + return webbrowser.open(url) + except Exception: + log_exception("PYAS._MainMixin.open_url:623") + return False + + +class WindowAPI(_MainMixin, ScannerMixin, ToolsMixin, ProtectMixin): + pass + + +def get_base_path(): + if getattr(sys, "frozen", False): + return os.path.dirname(sys.executable) + + return os.path.dirname(os.path.abspath(__file__)) + + +def get_frontend_asset_errors(): + base_path = get_base_path() + required_files = [ + os.path.join("Interface", "templates", "index.html"), + os.path.join("Interface", "static", "css", "style.css"), + os.path.join("Interface", "static", "js", "i18n.js"), + os.path.join("Interface", "static", "js", "main.js"), + ] + errors = [] + + for rel_path in required_files: + abs_path = os.path.join(base_path, rel_path) + + if not os.path.isfile(abs_path): + errors.append(f"Missing UI asset: {rel_path}") + elif os.path.getsize(abs_path) <= 0: + errors.append(f"Empty UI asset: {rel_path}") + + return errors + + +def show_startup_error(message): + try: + ctypes.windll.user32.MessageBoxW(None, message, PYAS_WINDOW_TITLE, 0x00000010) + except Exception: + log_exception("PYAS.show_startup_error:662") + pass + + +class NoCacheRequestHandler(SimpleHTTPRequestHandler): + def __init__(self, *args, **kwargs): + super().__init__(*args, directory=os.path.join(get_base_path(), "Interface"), **kwargs) + + def log_message(self, format, *args): + logging.getLogger("PYAS.Startup").info("HTTP %s: " + format, self.client_address[0], *args) + + def do_GET(self): + if self.path == "/": + self.path = "/templates/index.html" + + return super().do_GET() + + def end_headers(self): + self.send_header("Cache-Control", "no-cache, no-store, must-revalidate") + self.send_header("Expires", "0") + self.send_header("Pragma", "no-cache") + super().end_headers() + + +class WindowHook: + def __init__(self, title, api_ref=None): + self.title, self.api_ref = title, api_ref + self.old_wndproc = None + self.WM_DPICHANGED = 0x02E0 + self.WM_NCHITTEST = 0x0084 + self.WM_CLOSE = 0x0010 + self.WM_COPYDATA = WM_COPYDATA + self.WM_SYSCOMMAND = 0x0112 + self.SC_CLOSE = 0xF060 + self.HTCAPTION = 2 + self.GWLP_WNDPROC = -4 + + self.RECT = RECT + self.WNDPROC = ctypes.WINFUNCTYPE( + ctypes.c_void_p, ctypes.c_void_p, ctypes.c_uint, ctypes.c_void_p, ctypes.c_void_p + ) + + self.user32 = ctypes.windll.user32 + self.user32.FindWindowW.argtypes, self.user32.FindWindowW.restype = [ + ctypes.c_wchar_p, + ctypes.c_wchar_p, + ], ctypes.wintypes.HWND + self.user32.CallWindowProcW.argtypes, self.user32.CallWindowProcW.restype = [ + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_uint, + ctypes.c_void_p, + ctypes.c_void_p, + ], ctypes.c_void_p + self.user32.DefWindowProcW.argtypes, self.user32.DefWindowProcW.restype = [ + ctypes.c_void_p, + ctypes.c_uint, + ctypes.c_void_p, + ctypes.c_void_p, + ], ctypes.c_void_p + + if ctypes.sizeof(ctypes.c_void_p) == 8: + self.SetWindowLong, self.GetWindowLong = ( + self.user32.SetWindowLongPtrW, + self.user32.GetWindowLongPtrW, + ) + else: + self.SetWindowLong, self.GetWindowLong = ( + self.user32.SetWindowLongW, + self.user32.GetWindowLongW, + ) + + self.SetWindowLong.argtypes, self.SetWindowLong.restype = [ + ctypes.c_void_p, + ctypes.c_int, + self.WNDPROC, + ], ctypes.c_void_p + self.GetWindowLong.argtypes, self.GetWindowLong.restype = [ + ctypes.c_void_p, + ctypes.c_int, + ], ctypes.c_void_p + self.new_wndproc_cb = self.WNDPROC(self.wndproc) + + def hook(self): + if self.old_wndproc: + return + + hwnd = self.user32.FindWindowW(None, self.title) + + if hwnd: + self.old_wndproc = self.GetWindowLong(hwnd, self.GWLP_WNDPROC) + self.SetWindowLong(hwnd, self.GWLP_WNDPROC, self.new_wndproc_cb) + + try: + self.user32.ChangeWindowMessageFilterEx.argtypes = [ + ctypes.c_void_p, + ctypes.c_uint, + ctypes.c_uint, + ctypes.c_void_p, + ] + self.user32.ChangeWindowMessageFilter.argtypes = [ctypes.c_uint, ctypes.c_uint] + self.user32.ChangeWindowMessageFilter.restype = ctypes.wintypes.BOOL + + if not self.user32.ChangeWindowMessageFilter(PYAS_MAINTENANCE_MESSAGE, 0): + raise OSError("Could not remove process-wide maintenance message allowance") + + if not self.user32.ChangeWindowMessageFilterEx( + hwnd, PYAS_MAINTENANCE_MESSAGE, 2, None + ): + raise OSError("Could not protect maintenance message channel") + + self.maintenance_channel_ready = True + self.user32.ChangeWindowMessageFilterEx(hwnd, self.WM_COPYDATA, 1, None) + except Exception: + log_exception("PYAS.WindowHook.hook:732") + pass + + def call_default(self, hwnd, msg, wparam, lparam): + if self.old_wndproc: + return self.user32.CallWindowProcW(self.old_wndproc, hwnd, msg, wparam, lparam) + + return self.user32.DefWindowProcW(hwnd, msg, wparam, lparam) + + def wndproc(self, hwnd, msg, wparam, lparam): + if msg == PYAS_MAINTENANCE_MESSAGE: + if ( + not getattr(self, "maintenance_channel_ready", False) + or wparam not in PYAS_MAINTENANCE_COMMANDS + or not self.api_ref + ): + return 0 + + uninstall = wparam in (PYAS_MESSAGE_QUIT, PYAS_MESSAGE_DRIVER_UNINSTALL) + threading.Thread( + target=self.api_ref.close, kwargs={"uninstall_driver": uninstall}, daemon=True + ).start() + return 1 + + if msg == self.WM_COPYDATA: + try: + if not lparam: + return 0 + + cds = COPYDATASTRUCT.from_address(lparam) + + if cds.dwData not in (PYAS_MESSAGE_SCAN, PYAS_MESSAGE_SHOW): + return 0 + + if cds.dwData == PYAS_MESSAGE_SCAN and ( + not cds.lpData or not 0 < cds.cbData <= 131072 + ): + return 0 + + if cds.dwData == PYAS_MESSAGE_SCAN: + path = ctypes.string_at(cds.lpData, cds.cbData).decode("utf-8").strip("\x00") + + if self.api_ref: + threading.Thread(target=self.api_ref.restore_from_tray, daemon=True).start() + threading.Thread( + target=self.api_ref.trigger_context_scan, args=(path,), daemon=True + ).start() + + elif cds.dwData == PYAS_MESSAGE_SHOW and self.api_ref: + threading.Thread(target=self.api_ref.restore_from_tray, daemon=True).start() + + except Exception: + log_exception("PYAS.WindowHook.wndproc:762") + pass + + return 1 + + if msg == self.WM_CLOSE or ( + msg == self.WM_SYSCOMMAND and (wparam & 0xFFF0) == self.SC_CLOSE + ): + if self.api_ref and not getattr(self.api_ref, "closing", False): + threading.Thread(target=self.api_ref.hide_window, daemon=True).start() + return 0 + + return self.call_default(hwnd, msg, wparam, lparam) + + if msg == self.WM_NCHITTEST: + x, y = lparam & 0xFFFF, (lparam >> 16) & 0xFFFF + + if x >= 32768: + x -= 65536 + + if y >= 32768: + y -= 65536 + + rect = self.RECT() + self.user32.GetWindowRect(hwnd, ctypes.byref(rect)) + + if rect.top <= y <= rect.top + 44 and rect.left <= x <= rect.right - 150: + return self.HTCAPTION + + if msg == self.WM_DPICHANGED: + try: + rect = self.RECT.from_address(lparam) + self.user32.SetWindowPos( + hwnd, + None, + rect.left, + rect.top, + rect.right - rect.left, + rect.bottom - rect.top, + 0x0004 | 0x0010 | 0x0020, + ) + except Exception: + log_exception("PYAS.WindowHook.wndproc:788") + pass + + return self.call_default(hwnd, msg, wparam, lparam) + + +class StartupHTTPServer(ThreadingTCPServer): + allow_reuse_address = True + daemon_threads = True + + def handle_error(self, request, client_address): + logging.getLogger("PYAS.Startup").exception("HTTP request failed: %s", client_address) + + +def start_api(port_container, error_container, ready_event): + logger = logging.getLogger("PYAS.Startup") + + try: + with StartupHTTPServer(("127.0.0.1", 0), NoCacheRequestHandler) as httpd: + port_container.append(httpd.server_address[1]) + logger.info("HTTP listener ready: 127.0.0.1:%s", port_container[0]) + ready_event.set() + httpd.serve_forever() + except Exception as e: + logger.exception("HTTP listener failed") + error_container.append(str(e)) + ready_event.set() + + +def verify_ui_server(port): + from urllib.request import build_opener, ProxyHandler + + opener = build_opener(ProxyHandler({})) + + for path in ("/", "/static/css/style.css", "/static/js/i18n.js", "/static/js/main.js"): + with opener.open(f"http://127.0.0.1:{port}{path}", timeout=3) as response: + if response.status != 200 or not response.read(1): + raise RuntimeError(f"UI HTTP health check failed: {path}") + + logging.getLogger("PYAS.Startup").info("HTTP UI health checks passed") + + +def start_ui(): + logger = logging.getLogger("PYAS.Startup") + hide_on_start = "-h" in sys.argv or "-hide" in sys.argv + init_width, init_height = 980, 670 + stage = "assets" + js_api = None + webview_completed = threading.Event() + recovery_lock = threading.Lock() + + def fail_startup(reason, recover=False): + with recovery_lock: + if webview_completed.is_set(): + return + + webview_completed.set() + logger.error("Startup failed: %s", reason) + + if js_api: + js_api.write_log("WARN", "WebView2", detail=reason, success=False) + js_api.flush_logs_now() + + if ( + recover + and js_api + and restart_with_new_webview_profile(js_api.path_webview, __file__) + ): + os._exit(3) + + if not hide_on_start: + show_startup_error( + f"PYAS could not initialize its interface.\n\n{reason}\n\n" + f"Diagnostic log: {STARTUP_LOG_PATH or 'unavailable'}" + ) + + os._exit(3) + + try: + wait_for_recovery_parent() + arguments = os.environ.get("WEBVIEW2_ADDITIONAL_BROWSER_ARGUMENTS", "") + os.environ["WEBVIEW2_ADDITIONAL_BROWSER_ARGUMENTS"] = ( + arguments + " --proxy-bypass-list=localhost;127.0.0.1" + ).strip() + frontend_errors = get_frontend_asset_errors() + + if frontend_errors: + raise RuntimeError("UI files are incomplete: " + "; ".join(frontend_errors)) + + logger.info("UI assets verified") + + stage = "http" + port_container, server_errors = [], [] + server_ready = threading.Event() + threading.Thread( + target=start_api, args=(port_container, server_errors, server_ready), daemon=True + ).start() + + if not server_ready.wait(5.0) or not port_container or server_errors: + raise RuntimeError("Local UI server did not start: " + "; ".join(server_errors)) + + verify_ui_server(port_container[0]) + + stage = "api" + logger.info("Initializing Python API") + js_api = WindowAPI() + js_api.file_webview_log = STARTUP_LOG_PATH or js_api.file_webview_log + + stage = "profile" + preferred_profile = os.environ.get("PYAS_WEBVIEW_PROFILE", js_api.path_webview) + js_api.path_webview = prepare_webview_profile(preferred_profile) + + stage = "runtime" + check_webview_dependencies(webview) + + stage = "window" + user32 = ctypes.windll.user32 + pos_x = (user32.GetSystemMetrics(0) - init_width) // 2 + pos_y = (user32.GetSystemMetrics(1) - init_height) // 2 + startup_url = f"http://127.0.0.1:{port_container[0]}/" + window = webview.create_window( + title=PYAS_WINDOW_TITLE, + url=startup_url, + width=init_width, + height=init_height, + x=pos_x, + y=pos_y, + frameless=True, + easy_drag=False, + js_api=js_api, + background_color="#e0e0e0", + hidden=hide_on_start, + ) + + if platform.system() == "Windows": + window_hook = WindowHook(PYAS_WINDOW_TITLE, js_api) + window.events.shown += window_hook.hook + + js_api.set_window(window) + logger.info("Native window configured; hidden=%s", hide_on_start) + + stage = "tray" + js_api.show_tray() + logger.info("Tray initialized") + webview_loaded_event = threading.Event() + dnd_state = {"bound": False} + + def bind_dnd(): + webview_loaded_event.set() + logger.info("WebView document loaded") + + if dnd_state["bound"]: + return + + try: + window.dom.document.events.drop += DOMEventHandler(js_api.on_drop, True, True) + dnd_state["bound"] = True + except Exception: + logger.exception("Drag and drop binding failed") + + def webview_startup_watchdog(): + if js_api.ui_ready_event.wait(30.0) or webview_completed.is_set(): + return + + if webview_loaded_event.is_set(): + logger.warning("UI bridge timeout; reloading once") + + try: + window.load_url(startup_url) + except Exception: + logger.exception("UI reload failed") + + if js_api.ui_ready_event.wait(30.0) or webview_completed.is_set(): + return + + reason = ( + "Python bridge or configuration initialization timed out" + if webview_loaded_event.is_set() + else ("WebView2 did not load the local UI document") + ) + fail_startup(reason, recover=True) + + def on_closed(): + webview_completed.set() + + window.events.loaded += bind_dnd + window.events.closed += on_closed + threading.Thread(target=webview_startup_watchdog, daemon=True).start() + stage = "webview" + logger.info("Starting edgechromium") + webview.start(gui="edgechromium", private_mode=False, storage_path=js_api.path_webview) + webview_completed.set() + except Exception as e: + logger.exception("Startup exception at stage=%s", stage) + fail_startup(f"{stage}: {e}", recover=stage == "webview") + + +if __name__ == "__main__": + if "-quit" in sys.argv or "-driver-unload" in sys.argv or "-driver-uninstall" in sys.argv: + controller = WindowAPI() + + if "-quit" in sys.argv or "-driver-uninstall" in sys.argv: + success, delete_error = controller.uninstall_system_driver() + + if not success: + controller.write_log( + "WARN", + "Driver Service", + detail=f"DeleteService failed: 0x{delete_error & 0xFFFFFFFF:08X}", + success=False, + ) + else: + success = controller.stop_system_driver() + + controller.flush_logs_now() + os._exit(0 if success else 2) + + start_ui() diff --git a/PYAS_Autostart.py b/PYAS_Autostart.py new file mode 100644 index 0000000..caeb570 --- /dev/null +++ b/PYAS_Autostart.py @@ -0,0 +1,589 @@ +from PYAS_Diagnostics import log_exception +import os +import io +import csv +import json +import sys +import winreg +import subprocess +import ctypes +import ctypes.wintypes +from concurrent.futures import ThreadPoolExecutor + + +class AutostartMixin: + def _set_registry_autostart(self, enable): + run_key = r"Software\Microsoft\Windows\CurrentVersion\Run" + value_name = "PYAS_Security" + + if enable: + command = f'"{self.file_pyas}" -hide' + return self._reg_write( + winreg.HKEY_CURRENT_USER, run_key, value_name, winreg.REG_SZ, command + ) + + try: + with winreg.OpenKey(winreg.HKEY_CURRENT_USER, run_key, 0, winreg.KEY_SET_VALUE) as reg: + try: + winreg.DeleteValue(reg, value_name) + except FileNotFoundError: + pass + + return True + except FileNotFoundError: + return True + except Exception: + log_exception("PYAS_Autostart.AutostartMixin._set_registry_autostart:25") + return False + + def _get_autostart_user_sid(self): + user32 = ctypes.WinDLL("user32", use_last_error=True) + kernel32 = ctypes.WinDLL("kernel32", use_last_error=True) + advapi32 = ctypes.WinDLL("advapi32", use_last_error=True) + user32.GetShellWindow.restype = ctypes.wintypes.HWND + user32.GetWindowThreadProcessId.argtypes = [ + ctypes.wintypes.HWND, + ctypes.POINTER(ctypes.wintypes.DWORD), + ] + user32.GetWindowThreadProcessId.restype = ctypes.wintypes.DWORD + kernel32.OpenProcess.argtypes = [ + ctypes.wintypes.DWORD, + ctypes.wintypes.BOOL, + ctypes.wintypes.DWORD, + ] + kernel32.OpenProcess.restype = ctypes.wintypes.HANDLE + kernel32.CloseHandle.argtypes = [ctypes.wintypes.HANDLE] + kernel32.LocalFree.argtypes = [ctypes.c_void_p] + kernel32.LocalFree.restype = ctypes.c_void_p + advapi32.OpenProcessToken.argtypes = [ + ctypes.wintypes.HANDLE, + ctypes.wintypes.DWORD, + ctypes.POINTER(ctypes.wintypes.HANDLE), + ] + advapi32.GetTokenInformation.argtypes = [ + ctypes.wintypes.HANDLE, + ctypes.c_int, + ctypes.c_void_p, + ctypes.wintypes.DWORD, + ctypes.POINTER(ctypes.wintypes.DWORD), + ] + advapi32.ConvertSidToStringSidW.argtypes = [ + ctypes.c_void_p, + ctypes.POINTER(ctypes.c_void_p), + ] + hwnd = user32.GetShellWindow() + + if not hwnd: + raise OSError("No interactive Windows desktop is available for autostart") + + pid = ctypes.wintypes.DWORD() + user32.GetWindowThreadProcessId(hwnd, ctypes.byref(pid)) + process = kernel32.OpenProcess(0x1000, False, pid.value) + + if not process: + raise ctypes.WinError(ctypes.get_last_error()) + + token = ctypes.wintypes.HANDLE() + sid_string = ctypes.c_void_p() + + try: + if not advapi32.OpenProcessToken(process, 0x0008, ctypes.byref(token)): + raise ctypes.WinError(ctypes.get_last_error()) + + size = ctypes.wintypes.DWORD() + advapi32.GetTokenInformation(token, 1, None, 0, ctypes.byref(size)) + + if not size.value: + raise ctypes.WinError(ctypes.get_last_error()) + + buffer = ctypes.create_string_buffer(size.value) + + if not advapi32.GetTokenInformation(token, 1, buffer, size, ctypes.byref(size)): + raise ctypes.WinError(ctypes.get_last_error()) + + sid = ctypes.cast(buffer, ctypes.POINTER(ctypes.c_void_p)).contents + + if not advapi32.ConvertSidToStringSidW(sid, ctypes.byref(sid_string)): + raise ctypes.WinError(ctypes.get_last_error()) + + return ctypes.wstring_at(sid_string) + finally: + if sid_string.value: + kernel32.LocalFree(sid_string) + + if token.value: + kernel32.CloseHandle(token) + + kernel32.CloseHandle(process) + + def _autostart_task_script(self, enable, sid): + executable = self.file_pyas if getattr(sys, "frozen", False) else sys.executable + arguments = ( + "-hide" + if getattr(sys, "frozen", False) + else subprocess.list2cmdline([self.file_pyas, "-hide"]) + ) + quote = lambda value: "'" + value.replace("'", "''") + "'" + values = ( + "$ErrorActionPreference = 'Stop'; " + "[Console]::OutputEncoding = [Text.UTF8Encoding]::new(); " + f"$Sid = {quote(sid)}; $Executable = {quote(os.path.normpath(executable))}; " + f"$Arguments = {quote(arguments)}; $WorkingDirectory = {quote(os.path.dirname(self.file_pyas))}; " + f"$Enable = {'$true' if enable else '$false'}; " + ) + return ( + values + + r""" +function Resolve-Sid($UserId) { + if ($UserId -like 'S-1-*') { return $UserId } + return ([Security.Principal.NTAccount]::new($UserId)).Translate([Security.Principal.SecurityIdentifier]).Value +} +$Identity = [Security.Principal.WindowsIdentity]::GetCurrent() +if ($Identity.User.Value -ne $Sid) { + throw 'The elevated account differs from the signed-in desktop user. Configure autostart while signed in as the target administrator.' +} +if (-not ([Security.Principal.WindowsPrincipal]::new($Identity)).IsInRole([Security.Principal.WindowsBuiltInRole]::Administrator)) { + throw 'Administrator privileges are required to configure elevated autostart.' +} +$Service = New-Object -ComObject 'Schedule.Service' +$Service.Connect() +$Folder = $Service.GetFolder('\') +$TaskName = 'PYAS_Security_ATS_' + $Sid +function Find-Task($Name) { + foreach ($Task in $Folder.GetTasks(1)) { + if ($Task.Name -eq $Name) { return $Task } + } + return $null +} +function Save-Task($Task) { + $BackupDirectory = Join-Path ([Environment]::GetFolderPath('CommonApplicationData')) 'PYAS\StartupTasks' + [IO.Directory]::CreateDirectory($BackupDirectory) | Out-Null + $BackupPath = Join-Path $BackupDirectory ($Task.Name + '-' + [Guid]::NewGuid().ToString('N') + '.xml') + [IO.File]::WriteAllText($BackupPath, $Task.Xml, [Text.Encoding]::UTF8) +} +function Test-Task($Task) { + if (-not $Task -or -not $Task.Enabled) { return $false } + $Definition = $Task.Definition + if ($Definition.Actions.Count -ne 1 -or $Definition.Triggers.Count -ne 1) { return $false } + $Action = $Definition.Actions.Item(1) + $Trigger = $Definition.Triggers.Item(1) + $Settings = $Definition.Settings + return ((Resolve-Sid $Definition.Principal.UserId) -eq $Sid -and + $Definition.Principal.LogonType -eq 3 -and $Definition.Principal.RunLevel -eq 1 -and + $Action.Type -eq 0 -and $Action.Path -eq $Executable -and + $Action.Arguments -eq $Arguments -and $Action.WorkingDirectory -eq $WorkingDirectory -and + $Trigger.Type -eq 9 -and $Trigger.Enabled -and (Resolve-Sid $Trigger.UserId) -eq $Sid -and + $Trigger.Delay -eq 'PT10S' -and -not $Settings.DisallowStartIfOnBatteries -and + -not $Settings.StopIfGoingOnBatteries -and $Settings.ExecutionTimeLimit -eq 'PT0S' -and + $Settings.MultipleInstances -eq 2 -and $Settings.RestartCount -eq 3 -and + $Settings.RestartInterval -eq 'PT1M') +} +$Existing = Find-Task $TaskName +$ExistingXml = if ($Existing) { $Existing.Xml } else { $null } +if ($Enable) { + if (-not (Test-Task $Existing)) { + if ($Existing) { Save-Task $Existing } + $Definition = $Service.NewTask(0) + $Definition.RegistrationInfo.Description = 'PYAS elevated autostart for the signed-in user' + $Definition.Principal.UserId = $Sid + $Definition.Principal.LogonType = 3 + $Definition.Principal.RunLevel = 1 + $Trigger = $Definition.Triggers.Create(9) + $Trigger.UserId = $Sid + $Trigger.Delay = 'PT10S' + $Trigger.Enabled = $true + $Action = $Definition.Actions.Create(0) + $Action.Path = $Executable + $Action.Arguments = $Arguments + $Action.WorkingDirectory = $WorkingDirectory + $Definition.Settings.Enabled = $true + $Definition.Settings.DisallowStartIfOnBatteries = $false + $Definition.Settings.StopIfGoingOnBatteries = $false + $Definition.Settings.ExecutionTimeLimit = 'PT0S' + $Definition.Settings.MultipleInstances = 2 + $Definition.Settings.RestartCount = 3 + $Definition.Settings.RestartInterval = 'PT1M' + try { + $Folder.RegisterTaskDefinition($TaskName, $Definition, 6, $Sid, $null, 3) | Out-Null + if (-not (Test-Task (Find-Task $TaskName))) { throw 'Autostart task verification failed' } + } catch { + $Failure = $_ + if ($Existing) { + $Folder.RegisterTask($TaskName, $ExistingXml, 6, $Sid, $null, 3) | Out-Null + } else { + $Created = Find-Task $TaskName + if ($Created) { $Folder.DeleteTask($TaskName, 0) } + } + throw $Failure + } + } +} elseif ($Existing) { + Save-Task $Existing + $Folder.DeleteTask($TaskName, 0) + if (Find-Task $TaskName) { throw 'Autostart task removal verification failed' } +} +$Legacy = Find-Task 'PYAS_Security_ATS' +if ($Legacy -and (Resolve-Sid $Legacy.Definition.Principal.UserId) -eq $Sid -and + $Legacy.Definition.Actions.Count -eq 1 -and $Legacy.Definition.Actions.Item(1).Path -eq $Executable) { + Save-Task $Legacy + $Folder.DeleteTask($Legacy.Name, 0) +} +@{verified=$true; taskName=$TaskName; userSid=$Sid; enabled=$Enable} | ConvertTo-Json -Compress +""" + ) + + def manage_autostart(self, enable): + try: + sid = self._get_autostart_user_sid() + command = self._autostart_task_script(enable, sid) + result = self._run_powershell( + command, + capture_output=True, + text=True, + encoding="utf-8", + errors="replace", + timeout=30, + ) + + if result is None: + raise OSError("PowerShell is unavailable; elevated autostart cannot be configured") + + if result.returncode != 0: + raise OSError( + ( + result.stderr or result.stdout or f"PowerShell exit {result.returncode}" + ).strip() + ) + + state = json.loads(result.stdout) + + if ( + state.get("verified") is not True + or state.get("userSid") != sid + or state.get("taskName") != "PYAS_Security_ATS_" + sid + or state.get("enabled") is not bool(enable) + ): + raise OSError("Autostart task readback did not match the requested state") + + if not self._set_registry_autostart(False): + raise OSError("The legacy HKCU Run entry could not be removed") + + self.autostart_mode = "task" if enable else "disabled" + self.write_log( + "INFO", + "manage_autostart", + detail=f"Verified {self.autostart_mode}: {state['taskName']}", + success=True, + ) + return True + except Exception as e: + log_exception("PYAS_Autostart.AutostartMixin.manage_autostart:201") + self.autostart_mode = "unavailable" + self.write_log("WARN", "manage_autostart", detail=str(e), success=False) + return False + + def get_startup_list(self): + items = [] + + def get_reg_startup(): + res = [] + paths = [ + ( + winreg.HKEY_LOCAL_MACHINE, + r"SOFTWARE\Microsoft\Windows\CurrentVersion\Run", + "enabled", + ), + ( + winreg.HKEY_CURRENT_USER, + r"SOFTWARE\Microsoft\Windows\CurrentVersion\Run", + "enabled", + ), + ( + winreg.HKEY_LOCAL_MACHINE, + r"SOFTWARE\Microsoft\Windows\CurrentVersion\Run_Disabled", + "disabled", + ), + ( + winreg.HKEY_CURRENT_USER, + r"SOFTWARE\Microsoft\Windows\CurrentVersion\Run_Disabled", + "disabled", + ), + ] + + for root, path, status in paths: + try: + with winreg.OpenKey(root, path, 0, winreg.KEY_READ) as reg: + i = 0 + + while True: + try: + name, val, _ = winreg.EnumValue(reg, i) + file_path = self.extract_paths_from_cmdline(val) + fp = file_path[0] if file_path else val + res.append( + { + "id": f"reg|{root}|{path}|{name}", + "type": "type_reg", + "name": name, + "status": status, + "path": fp, + } + ) + i += 1 + + except OSError: + log_exception( + "PYAS_Autostart.AutostartMixin.get_startup_list.get_reg_startup:230" + ) + break + except Exception: + log_exception( + "PYAS_Autostart.AutostartMixin.get_startup_list.get_reg_startup:232" + ) + pass + + return res + + def get_services(): + res = [] + + try: + with winreg.OpenKey( + winreg.HKEY_LOCAL_MACHINE, + r"SYSTEM\CurrentControlSet\Services", + 0, + winreg.KEY_READ, + ) as reg: + i = 0 + + while True: + try: + name = winreg.EnumKey(reg, i) + i += 1 + + try: + with winreg.OpenKey(reg, name, 0, winreg.KEY_READ) as sub: + stype, _ = winreg.QueryValueEx(sub, "Type") + + if stype not in (16, 32, 256): + continue + + start, _ = winreg.QueryValueEx(sub, "Start") + + if start not in (2, 3, 4): + continue + + path, _ = winreg.QueryValueEx(sub, "ImagePath") + + if ( + path + and "svchost" not in path.lower() + and "windows" not in path.lower() + ): + fp = self.extract_paths_from_cmdline(path) + fpp = fp[0] if fp else path + status = "enabled" if start == 2 else "disabled" + res.append( + { + "id": f"srv|{name}", + "type": "type_srv", + "name": name, + "status": status, + "path": fpp, + } + ) + + except Exception: + log_exception( + "PYAS_Autostart.AutostartMixin.get_startup_list.get_services:262" + ) + pass + except OSError: + log_exception( + "PYAS_Autostart.AutostartMixin.get_startup_list.get_services:264" + ) + break + + except Exception: + log_exception("PYAS_Autostart.AutostartMixin.get_startup_list.get_services:267") + pass + + return res + + def get_tasks(): + res = [] + + try: + proc = self._run_windows_tool( + ("schtasks.exe",), + ["/query", "/fo", "csv", "/v"], + capture_output=True, + text=True, + ) + + if proc is not None and proc.returncode == 0: + reader = csv.reader(io.StringIO(proc.stdout)) + next(reader, None) + + for row in reader: + if len(row) > 8 and row[1].strip() and row[8].strip() and row[8] != "N/A": + name = row[1].split("\\")[-1] + path = row[8] + status = row[3] + + if path and "windows" not in path.lower(): + fp = self.extract_paths_from_cmdline(path) + fpp = fp[0] if fp else path + res.append( + { + "id": f"tsk|{row[1]}", + "type": "type_tsk", + "name": name, + "status": ( + "disabled" + if status.lower() == "disabled" + else "enabled" + ), + "path": fpp, + } + ) + + return res + + command = ( + "[Console]::OutputEncoding = [Text.UTF8Encoding]::new(); " + "Get-ScheduledTask -ErrorAction SilentlyContinue | ForEach-Object { " + "$Action = @($_.Actions)[0]; " + "[pscustomobject]@{Name=$_.TaskName; FullName=($_.TaskPath + $_.TaskName); " + "State=[string]$_.State; Execute=$Action.Execute; Arguments=$Action.Arguments} " + "} | ConvertTo-Json -Compress" + ) + proc = self._run_powershell( + command, capture_output=True, text=True, encoding="utf-8", errors="replace" + ) + + if not proc or proc.returncode != 0 or not proc.stdout.strip(): + return res + + task_data = json.loads(proc.stdout) + + if isinstance(task_data, dict): + task_data = [task_data] + + for task in task_data: + name = str(task.get("Name") or "").strip() + full_name = str(task.get("FullName") or name).strip() + execute = str(task.get("Execute") or "").strip() + arguments = str(task.get("Arguments") or "").strip() + command_line = " ".join(part for part in (execute, arguments) if part) + + if not name or not command_line or "windows" in command_line.lower(): + continue + + paths = self.extract_paths_from_cmdline(command_line) + path = paths[0] if paths else execute or command_line + status = ( + "disabled" + if str(task.get("State") or "").lower() == "disabled" + else "enabled" + ) + res.append( + { + "id": f"tsk|{full_name}", + "type": "type_tsk", + "name": name, + "status": status, + "path": path, + } + ) + + except Exception: + log_exception("PYAS_Autostart.AutostartMixin.get_startup_list.get_tasks:326") + pass + + return res + + with ThreadPoolExecutor(max_workers=3) as executor: + f1 = executor.submit(get_reg_startup) + f2 = executor.submit(get_services) + f3 = executor.submit(get_tasks) + + items.extend(f1.result()) + items.extend(f2.result()) + items.extend(f3.result()) + + return items + + def manage_startup(self, items, action): + def process_item(item_id): + try: + parts = item_id.split("|", 3) + stype = parts[0] + + if stype == "reg": + root, old_path, name = int(parts[1]), parts[2], parts[3] + + if action == "delete": + return self._reg_delete(root, old_path, name) + + else: + new_path = ( + old_path.replace("_Disabled", "") + if action == "enable" + else (old_path if "_Disabled" in old_path else old_path + "_Disabled") + ) + + if old_path != new_path: + with winreg.OpenKey(root, old_path, 0, winreg.KEY_READ) as reg: + val, value_type = winreg.QueryValueEx(reg, name) + + return self._reg_write( + root, new_path, name, value_type, val + ) and self._reg_delete(root, old_path, name) + + return True + + elif stype == "srv": + name = parts[1] + + if action == "delete": + result = self._run_windows_tool( + ("sc.exe",), ["delete", name], capture_output=True + ) + + else: + mode = "auto" if action == "enable" else "disabled" + result = self._run_windows_tool( + ("sc.exe",), ["config", name, "start=", mode], capture_output=True + ) + + if result is not None and result.returncode == 0: + return True + + return self._manage_service_native(name, action) + + elif stype == "tsk": + name = item_id.partition("|")[2] + return self._manage_scheduled_task(name, action) + + return False + + except Exception as e: + log_exception("PYAS_Autostart.AutostartMixin.manage_startup.process_item:379") + self.write_log("WARN", "manage_startup", detail=str(e), success=False) + return False + + with ThreadPoolExecutor(max_workers=5) as executor: + results = list(executor.map(process_item, items)) + + success = all(results) if results else True + + if not success: + self.write_log( + "WARN", + "manage_startup", + detail="One or more startup items could not be updated", + success=False, + ) + + return success diff --git a/PYAS_Cloud.py b/PYAS_Cloud.py index c6d04a8..61368d3 100644 --- a/PYAS_Cloud.py +++ b/PYAS_Cloud.py @@ -1,248 +1,309 @@ -import os, time, uuid, logging, requests, hashlib - -from pathlib import Path -from typing import Optional, List, Dict, Any, Union -from concurrent.futures import ThreadPoolExecutor - -#################################################################################################### - -class PYAS_Client: - def __init__(self, api_key: str, hosts: Union[str, List[str]], timeout: int = 30): - self.timeout = timeout - self.logger = logging.getLogger("PYAS_API") - self.session = requests.Session() - self.session.headers.update({"X-API-Key": api_key, "User-Agent": "PYAS-API/1.1"}) - - if isinstance(hosts, str): - self.base_url = hosts.rstrip('/') - else: - self.base_url = self._get_fastest_host(hosts) - - self.logger.info(f"Initialized PYAS API Client targeting {self.base_url}") - - def _get_fastest_host(self, hosts: List[str]) -> str: - def ping(host: str): - url = host.rstrip('/') - try: - start = time.perf_counter() - r = self.session.head(url, timeout=3) - if r.status_code < 500: - return url, time.perf_counter() - start - except Exception: - pass - return url, float('inf') - - self.logger.info("Testing host latencies...") - with ThreadPoolExecutor(max_workers=min(len(hosts), 10)) as executor: - results = list(executor.map(ping, hosts)) - - valid_results = [] - for url, latency in results: - if latency != float('inf'): - self.logger.info(f"Host: {url} -> {latency:.4f}s") - valid_results.append((url, latency)) - else: - self.logger.warning(f"Host: {url} -> Timeout/Error") - - if not valid_results: - self.logger.warning("All host latency tests failed. Falling back to the first host.") - return hosts[0].rstrip('/') - - best_host, best_time = min(valid_results, key=lambda x: x[1]) - self.logger.info(f"Selected fastest host: {best_host} ({best_time:.4f}s)") - return best_host - - def _request(self, method: str, endpoint: str, **kwargs) -> Optional[requests.Response]: - self.logger.debug(f"Request: {method} {endpoint}") - try: - r = self.session.request(method, f"{self.base_url}{endpoint}", timeout=self.timeout, **kwargs) - if r.status_code == 200: - return r - - error_map = {401: "API Key invalid", 402: "Insufficient points", 404: "Not found", 413: "Payload too large", 500: "Server Error"} - self.logger.error(f"[{endpoint}] Failed: {error_map.get(r.status_code, f'HTTP {r.status_code}')}") - - if 400 <= r.status_code < 500: - setattr(r, 'is_fatal', True) - return r - - except Exception as e: - self.logger.error(f"Connection error: {e}") - return None - -#################################################################################################### - - def _calculate_sha256(self, path: Path) -> str: - self.logger.debug(f"Calculating SHA256 for {path}") - sha256_hash = hashlib.sha256() - - with open(path, "rb") as f: - for byte_block in iter(lambda: f.read(1048576), b""): - sha256_hash.update(byte_block) - - return sha256_hash.hexdigest() - - def upload_file(self, file_path: str, chunk_size: int = 4194304, max_retries: int = 3) -> Optional[str]: - path = Path(file_path) - if not path.exists(): - self.logger.error(f"File not found: {file_path}") - return None - - file_size = path.stat().st_size - self.logger.info(f"Initiating unified chunked upload for {path.name} ({file_size} bytes)") - - if file_size > 268435456: - self.logger.error(f"File size exceeds 256MB limit: {file_size} bytes") - return None - - local_sha256 = self._calculate_sha256(path) - status_req = self._request("GET", f"/api/processing_status/{local_sha256}") - - if status_req and status_req.status_code == 200: - if status_req.json().get('status') in ['done', 'processing', 'queued']: - self.logger.info(f"File already exists or processing on server: {local_sha256}") - return local_sha256 - - return self._upload_chunked(path, chunk_size, file_size, max_retries) - -#################################################################################################### - - def _upload_chunked(self, path: Path, chunk_size: int, file_size: int, max_retries: int = 3) -> Optional[str]: - total_chunks = max(1, (file_size + chunk_size - 1) // chunk_size) - upload_id = uuid.uuid4().hex - self.logger.info(f"Starting chunked upload for {path.name} (Chunks: {total_chunks}, ID: {upload_id})") - - with open(path, 'rb') as f: - for i in range(total_chunks): - chunk_data = f.read(chunk_size) - headers = { - "X-Chunk-Index": str(i), - "X-Total-Chunks": str(total_chunks), - "X-Upload-ID": upload_id - } - - chunk_success = False - for attempt in range(max_retries): - r = self._request("POST", "/api/upload", files={'file': (path.name, chunk_data)}, headers=headers) - - if r and r.status_code == 200: - chunk_success = True - self.logger.debug(f"Uploaded chunk {i + 1}/{total_chunks}") - - if i == total_chunks - 1: - try: - return r.json().get('url', '').split('/')[-1] or None - except Exception as e: - self.logger.error(f"Failed to parse final chunk response: {e}") - return None - break - - if hasattr(r, 'is_fatal'): - self.logger.error(f"Fatal upload error encountered at chunk {i}. Aborting.") - return None - - wait_time = 2 ** attempt - self.logger.warning(f"Chunk upload retry {attempt + 1}/{max_retries} for chunk {i}. Waiting {wait_time}s") - time.sleep(wait_time) - - if not chunk_success: - self.logger.error(f"Failed to upload chunk {i + 1}/{total_chunks} after {max_retries} attempts") - return None - return None - -#################################################################################################### - - def wait_for_analysis(self, sha256: str, interval: int = 5, max_retries: int = 60) -> bool: - self.logger.info(f"Waiting for analysis completion: {sha256}") - for attempt in range(max_retries): - r = self._request("GET", f"/api/processing_status/{sha256}") - if r and r.status_code == 200: - status = r.json().get('status') - self.logger.debug(f"Analysis status: {status}") - - if status == 'done': - self.logger.info(f"Analysis completed: {sha256}") - return True - - if status in ['missing', 'error', 'failed']: - self.logger.error(f"Analysis failed with status: {status}") - return False - - time.sleep(interval) - - self.logger.error(f"Analysis polling timed out for {sha256}") - return False - - def get_report(self, sha256: str) -> Optional[Dict[str, Any]]: - self.logger.info(f"Fetching report for {sha256}") - r = self._request("GET", f"/api/report/{sha256}") - if r: - return r.json() - - self.logger.error(f"Failed to fetch report for {sha256}") - return None - -#################################################################################################### - - def rescan(self, sha256: str) -> bool: - self.logger.info(f"Triggering rescan for {sha256}") - r = self._request("POST", f"/api/rescan/{sha256}") - - success = bool(r and r.json().get('status') == 'success') - if success: - self.logger.info(f"Rescan triggered successfully for {sha256}") - else: - self.logger.error(f"Rescan failed for {sha256}") - return success - - def search_files(self, query: str, limit: int = 50, page: int = 0, snapshot: Optional[str] = None) -> List[Dict[str, Any]]: - self.logger.info(f"Searching files (Query: '{query}', Limit: {limit}, Page: {page})") - payload = {'query': query, 'page': page} - if snapshot: - payload['snapshot'] = snapshot - - r = self._request("POST", "/api/search_more", data=payload) - results = r.json().get('results', [])[:limit] if r else [] - self.logger.info(f"Search returned {len(results)} results") - return results - - def download_sample(self, sha256: str, save_dir: str = "downloads") -> Optional[str]: - self.logger.info(f"Downloading sample {sha256}") - os.makedirs(save_dir, exist_ok=True) - path = os.path.join(save_dir, f"{sha256}.zip") - - r = self._request("GET", f"/api/download/{sha256}", stream=True) - if r: - with open(path, 'wb') as f: - for chunk in r.iter_content(8192): - if chunk: - f.write(chunk) - self.logger.info(f"Sample downloaded successfully to {path}") - return path - - self.logger.error(f"Failed to download sample {sha256}") - return None - -#################################################################################################### - -if __name__ == "__main__": - logging.basicConfig(level=logging.DEBUG, format='[%(asctime)s] %(levelname)s [%(name)s] - %(message)s', datefmt='%Y-%m-%d %H:%M:%S') - - API_KEY = "" # https://pyas-security.com/setting?tab=application - HOSTS = [ - "https://pyas-security.com/", - "https://cn.pyas-security.com/" - ] - - client = PYAS_Client(API_KEY, HOSTS) - target_file = r"C:/Windows/explorer.exe" - - sha = client.upload_file(target_file) - if sha and client.wait_for_analysis(sha): - report = client.get_report(sha) - if report: - meta = report.get('data', {}).get('metadata', {}) - logging.info(f"Result: {meta.get('label')} ({meta.get('score')}%)") - - for item in client.search_files('date>=2020; ext=exe | dll', limit=3): - client.download_sample(item['sha256']) +from PYAS_Diagnostics import log_exception +import os +import time +import threading +import requests + + +class CloudScanner: + def __init__(self): + self.api_host = None + self.api_key = None + self.timeout = 30 + + self.lock = threading.RLock() + self.local = threading.local() + + def _get_session(self, api_host, api_key): + if ( + not hasattr(self.local, "session") + or getattr(self.local, "host", None) != api_host + or getattr(self.local, "key", None) != api_key + ): + previous = getattr(self.local, "session", None) + + if previous is not None: + previous.close() + + session = requests.Session() + session.headers.update({"X-API-Key": api_key, "User-Agent": "PYAS-Engine/1.1"}) + self.local.session = session + self.local.host = api_host + self.local.key = api_key + + return self.local.session + + def _request(self, method, endpoint, api_host, api_key, **kwargs): + session = self._get_session(api_host, api_key) + + try: + r = session.request(method, f"{api_host}{endpoint}", timeout=self.timeout, **kwargs) + + if r.status_code == 200: + return r + + except Exception: + log_exception("PYAS_Cloud.CloudScanner._request:31") + pass + + return None + + def rescan(self, sha256, api_host, api_key): + r = self._request("POST", f"/api/rescan/{sha256}", api_host, api_key) + return r is not None and r.json().get("status") == "success" + + def upload_file( + self, + file_path, + api_host, + api_key, + chunk_size=4194304, + need_rescan=False, + max_retries=3, + file_hash=None, + cancel_event=None, + ): + try: + sha256 = file_hash + cancel_event = cancel_event or threading.Event() + + if cancel_event.is_set() or not sha256: + return False, None + + status_req = self._request("GET", f"/api/processing_status/{sha256}", api_host, api_key) + + if cancel_event.is_set(): + return False, sha256 + + if status_req: + current_status = status_req.json().get("status") + + if current_status == "done": + if need_rescan: + self.rescan(sha256, api_host, api_key) + + return True, sha256 + + elif current_status in ("queued", "processing"): + return True, sha256 + + file_size = os.path.getsize(file_path) + + if file_size > 104857600: + return False, sha256 + + total_chunks = max(1, (file_size + chunk_size - 1) // chunk_size) + upload_id = os.urandom(16).hex() + + with open(file_path, "rb") as f: + for i in range(total_chunks): + if cancel_event.is_set(): + return False, sha256 + + chunk_data = f.read(chunk_size) + headers = { + "X-Chunk-Index": str(i), + "X-Total-Chunks": str(total_chunks), + "X-Upload-ID": upload_id, + } + + chunk_success = False + + for attempt in range(max_retries): + if cancel_event.is_set(): + return False, sha256 + + r = self._request( + "POST", + "/api/upload", + api_host, + api_key, + files={"file": (os.path.basename(file_path), chunk_data)}, + headers=headers, + ) + + if cancel_event.is_set(): + return False, sha256 + + if r: + if i == total_chunks - 1: + try: + resp = r.json() + + if "url" in resp: + sha256 = resp.get("url", "").split("/")[-1] + except Exception: + log_exception("PYAS_Cloud.CloudScanner.upload_file:89") + pass + + chunk_success = True + break + + if cancel_event.wait(2**attempt): + return False, sha256 + + if not chunk_success: + return False, sha256 + + return True, sha256 + except Exception: + log_exception("PYAS_Cloud.CloudScanner.upload_file:101") + return False, None + + def get_result(self, sha256, api_host, api_key, max_retries=6, interval=10): + try: + if not sha256: + return False + + is_done = False + + for _ in range(max_retries): + r = self._request("GET", f"/api/processing_status/{sha256}", api_host, api_key) + + if r: + st = r.json().get("status", "error") + + if st == "done": + is_done = True + break + + if st in ["error", "failed"]: + return False + + time.sleep(interval) + + if not is_done: + return False + + r = self._request("GET", f"/api/report/{sha256}", api_host, api_key) + + if r: + data = r.json().get("data", {}) + metadata = data.get("metadata", {}) + label = metadata.get("label", "Unsupport") + score = metadata.get("score", 0) + sims = data.get("similar", []) + + is_malicious = "General" in label + sim_malicious_count = 0 + valid_sim_count = 0 + + for s in sims: + if s.get("similarity", 0) > 80: + valid_sim_count += 1 + + if "General" in s.get("label", ""): + sim_malicious_count += 1 + + if is_malicious and ( + valid_sim_count == 0 or sim_malicious_count == valid_sim_count + ): + return f"Malware:WinPE/General.{score}!cl" + + except Exception: + log_exception("PYAS_Cloud.CloudScanner.get_result:149") + pass + + return False + + +class CloudQueueMixin: + def cloud_check(self, file_path): + with self.lock_config: + if not self.pyas_config.get("cloud_switch", False): + return + + cancel_event = self.cloud_cancel_event + + norm_path = self.norm_path(file_path) + + if not norm_path or cancel_event.is_set(): + return + + cache_key = (os.path.normcase(norm_path), cancel_event) + + with self.lock_file_ops: + if cache_key in self.cloud_pending: + return + + self.cloud_pending.add(cache_key) + + self.cloud_queue.put((norm_path, cancel_event)) + + def cloud_worker(self): + while True: + try: + file_path, cancel_event = self.cloud_queue.get() + cache_key = (os.path.normcase(file_path), cancel_event) + + try: + self.perform_cloud_scan(file_path, cancel_event) + finally: + with self.lock_file_ops: + self.cloud_pending.discard(cache_key) + + self.cloud_queue.task_done() + except Exception: + log_exception("PYAS_Cloud.CloudQueueMixin.cloud_worker:31") + pass + + def perform_cloud_scan(self, file_path, cancel_event=None): + was_locked = False + + try: + with self.lock_config: + if not self.pyas_config.get("cloud_switch", False): + return False + + cancel_event = cancel_event or self.cloud_cancel_event + api_host, api_key, max_size = ( + self.pyas_config.get("api_host"), + self.pyas_config.get("api_key"), + self.pyas_config.get("size", 256 * 1024 * 1024), + ) + + if ( + cancel_event.is_set() + or not os.path.exists(file_path) + or not os.path.isfile(file_path) + ): + return False + + with self.lock_file_ops: + if file_path in self.virus_lock: + self.lock_file(file_path, False) + was_locked = True + + if os.path.getsize(file_path) > max_size: + if was_locked: + self.lock_file(file_path, True) + + return False + + file_hash = self.calc_file_hash(file_path) + success, sha256 = self.cloud.upload_file( + file_path, api_host, api_key, file_hash=file_hash, cancel_event=cancel_event + ) + + if was_locked: + self.lock_file(file_path, True) + was_locked = False + + if not success and not cancel_event.is_set(): + self.write_log( + "WARN", "Cloud API", source=file_path, detail="Failed", success=False + ) + + except Exception as e: + log_exception("PYAS_Cloud.CloudQueueMixin.perform_cloud_scan:66") + self.write_log("WARN", "perform_cloud_scan", detail=str(e), success=False) + + finally: + if was_locked: + try: + self.lock_file(file_path, True) + except Exception: + log_exception("PYAS_Cloud.CloudQueueMixin.perform_cloud_scan:73") + pass + + return False diff --git a/PYAS_Config.py b/PYAS_Config.py new file mode 100644 index 0000000..afe8edd --- /dev/null +++ b/PYAS_Config.py @@ -0,0 +1,221 @@ +from PYAS_Diagnostics import log_exception +import os +import copy +import json +import threading +from PYAS_Storage import atomic_write_json + + +class ConfigMixin: + def _loc(self, text_dict): + with self.lock_config: + lang = ( + self.pyas_config.get("language", "traditional_switch") + if hasattr(self, "pyas_config") + else "traditional_switch" + ) + + return text_dict.get(lang, text_dict.get("traditional_switch", "")) + + def load_config(self): + with self.lock_config: + + if not os.path.exists(self.file_config): + self.pyas_config = copy.deepcopy(self.pyas_default) + self.write_log("INFO", "Config Update", detail="Create default config") + self.save_config() + + else: + try: + with open(self.file_config, "r", encoding="utf-8") as f: + data = json.load(f) + self.pyas_config = copy.deepcopy(self.pyas_default) + self.pyas_config.update(data) + + self.pyas_config["version"] = self.pyas_default["version"] + whitelist_metadata_updated = False + + for item in self.pyas_config.get("white_list", []): + if not isinstance(item, dict) or not item.get("file") or "is_dir" in item: + continue + + item_path = self.norm_path(item["file"], must_exist=False) + + if item_path and os.path.exists(item_path): + item["is_dir"] = os.path.isdir(item_path) + whitelist_metadata_updated = True + + if whitelist_metadata_updated: + self.save_config() + except Exception as e: + log_exception("PYAS_Config.ConfigMixin.load_config:40") + self.pyas_config = copy.deepcopy(self.pyas_default) + self.write_log("WARN", "load_config", detail=str(e), success=False) + + def save_config(self): + with self.lock_config: + try: + atomic_write_json(self.file_config, self.pyas_config) + return True + except Exception as error: + log_exception("PYAS_Config.ConfigMixin.save_config:49") + self.write_log("WARN", "save_config", detail=str(error), success=False) + return False + + def update_config(self, key, value): + with self.lock_update: + with self.lock_config: + old_value = self.pyas_config.get(key) + + if old_value == value: + return value + + defer_driver_disable = key == "driver_switch" and not value + + if not defer_driver_disable: + self.pyas_config[key] = value + + def apply_switch(state): + if key == "cloud_switch": + with self.lock_config: + if state: + self.cloud_cancel_event = threading.Event() + else: + self.cloud_cancel_event.set() + elif ( + key in ("process_switch", "document_switch", "system_switch", "network_switch") + and state + ): + target = { + "process_switch": self.protect_proc_thread, + "document_switch": self.protect_file_thread, + "system_switch": self.protect_system_thread, + "network_switch": self.protect_net_thread, + }[key] + self.start_feature_thread(target, key) + elif key == "driver_switch": + if state: + if self.install_system_driver() and self.start_driver_listener( + wait_ready=True + ): + return True + + self.stop_system_driver() + return False + + return bool(self.stop_system_driver()) + elif key == "context_switch": + return self.register_context_menu(state) is not False + elif key == "autostart_switch": + return bool(self.manage_autostart(state)) + elif key == "document_switch" and not state: + self._cancel_pending_file_scans() + + with self.lock_file_ops: + if getattr(self, "h_dir_file", None): + if not self.kernel32.CloseHandle(self.h_dir_file): + raise OSError("Could not close directory watcher") + + self.h_dir_file = None + elif key == "suspend_switch" and not state: + with self.lock_proc: + for handle in list(getattr(self, "suspended_procs", ())): + if self.ntdll.NtResumeProcess(handle) != 0: + raise OSError("Could not resume suspended process") + + self.suspended_procs.discard(handle) + + return True + + action_succeeded = False + + try: + if key in ("extension_switch", "sensitive_switch"): + with self.lock_file_ops: + self.hash_cache.clear() + + action_succeeded = apply_switch(value) + + if not action_succeeded: + raise RuntimeError(f"Could not apply {key}") + + with self.lock_config: + self.pyas_config[key] = value + + if self.save_config() is False: + raise OSError("Could not persist configuration") + + self.write_log("INFO", "Config Update", detail=f"[{key}] {old_value} -> {value}") + + if key == "language" and self.tray_icon: + try: + self.tray_icon.update_menu() + except Exception: + log_exception("PYAS_Config.ConfigMixin.update_config:116") + pass + + return value + except Exception as error: + log_exception("PYAS_Config.ConfigMixin.update_config:119") + self.write_log("WARN", "Config Update", detail=f"[{key}] {error}", success=False) + + with self.lock_config: + self.pyas_config[key] = old_value + + if action_succeeded: + try: + if apply_switch(old_value) is False: + raise RuntimeError(f"Could not roll back {key}") + except Exception as rollback_error: + log_exception("PYAS_Config.ConfigMixin.update_config:127") + self.write_log( + "WARN", + "Config Rollback", + detail=f"[{key}] {rollback_error}", + success=False, + ) + + if key == "driver_switch": + with self.lock_config: + self.pyas_config[key] = bool(self.check_system_driver()) + + if self._window: + try: + self._window.evaluate_js( + f"if(window.revertSwitch) window.revertSwitch({json.dumps(key)});" + ) + except Exception: + log_exception("PYAS_Config.ConfigMixin.update_config:135") + pass + + return self.pyas_config.get(key, old_value) + + def get_config(self): + with self.lock_config: + cfg = self.pyas_config.copy() + cfg["autostart_mode"] = self.autostart_mode + + rules = [] + + if os.path.exists(self.path_rules): + for f in os.listdir(self.path_rules): + if f.lower().endswith(".json"): + fp = os.path.join(self.path_rules, f) + rules.append({"file": fp, "time": os.path.getmtime(fp)}) + + cfg["custom_rule"] = rules + return cfg + + def reset_config(self): + with self.lock_update: + with self.lock_config: + previous = self.pyas_config + self.pyas_config = copy.deepcopy(self.pyas_default) + + if self.save_config() is False: + self.pyas_config = previous + return False + + self.write_log("INFO", "Config Update", detail="Reset to default") + + return True diff --git a/PYAS_Diagnostics.py b/PYAS_Diagnostics.py new file mode 100644 index 0000000..9759714 --- /dev/null +++ b/PYAS_Diagnostics.py @@ -0,0 +1,260 @@ +import logging +import atexit +import json +import os +import sys +import tempfile +import threading +import time +import uuid +import weakref +from collections import OrderedDict +from PYAS_Storage import atomic_write_json +from PYAS_Version import VERSION + +LOG_FORMAT = "%(asctime)s %(levelname)s %(name)s %(message)s" +MAX_LOG_ENTRIES = 10000 + +_lock = threading.Lock() +_recent = OrderedDict() +_local = threading.local() + + +def create_log_entry( + level, + action, + detail=None, + code=None, + pid=None, + file_hash=None, + source=None, + target=None, + operate=None, + success=True, + timestamp=None, +): + timestamp = time.time() if timestamp is None else timestamp + + return { + "id": str(uuid.uuid4()), + "timestamp": timestamp, + "time_str": time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(timestamp)), + "version": VERSION, + "level": level, + "action": action, + "detail": detail, + "code": code, + "pid": pid, + "hash": file_hash, + "source": source, + "target": target, + "operate": operate, + "success": success, + } + + +def report_log_path(): + for handler in logging.getLogger("PYAS").handlers: + if isinstance(handler, ReportHandler): + return handler.baseFilename + + return os.path.join(os.environ.get("ALLUSERSPROFILE", "C:\\ProgramData"), "PYAS", "Report.json") + + +class ReportHandler(logging.Handler): + def __init__(self, path): + super().__init__() + self.baseFilename = os.path.abspath(path) + self.pending = [] + self.pending_lock = threading.RLock() + self.application = None + self.local = threading.local() + self.setFormatter(logging.Formatter("%(message)s")) + + def handle(self, record): + accepted = self.filter(record) + + if accepted: + self.emit(record) + + return accepted + + def emit(self, record): + if self._closed or getattr(self.local, "busy", False): + return + + self.local.busy = True + + try: + entry = create_log_entry( + "WARN" if record.levelno == logging.WARNING else record.levelname, + "Exception" if record.exc_info else "Diagnostic", + detail=self.format(record), + pid=record.process, + source=record.name, + success=record.levelno < logging.WARNING, + timestamp=record.created, + ) + application = self.application() if self.application else None + + if application is not None: + application._append_log_entry(entry) + + if record.levelno >= logging.ERROR: + application.flush_logs_now() + else: + with self.pending_lock: + application = self.application() if self.application else None + + if application is None: + self.pending.append(entry) + del self.pending[:-MAX_LOG_ENTRIES] + + if record.levelno >= logging.ERROR: + self._flush_pending() + + if application is not None: + application._append_log_entry(entry) + except Exception: + pass + finally: + self.local.busy = False + + def _flush_pending(self): + if not self.pending: + return True + + try: + existing = [] + + if os.path.exists(self.baseFilename): + with open(self.baseFilename, "r", encoding="utf-8") as stream: + existing = json.load(stream) + + if not isinstance(existing, list): + return False + + identifiers = {entry.get("id") for entry in existing if isinstance(entry, dict)} + merged = existing + [entry for entry in self.pending if entry["id"] not in identifiers] + atomic_write_json(self.baseFilename, merged[-MAX_LOG_ENTRIES:]) + return True + except Exception: + return False + + def flush(self): + if self._closed or getattr(self.local, "busy", False): + return + + self.local.busy = True + + try: + application = self.application() if self.application else None + + if application is not None: + application.flush_logs_now() + else: + with self.pending_lock: + self._flush_pending() + except Exception: + pass + finally: + self.local.busy = False + + def bind(self, application): + with application.lock_logs: + with self.pending_lock: + identifiers = {entry.get("id") for entry in application.logs_data} + + for entry in self.pending: + if entry["id"] not in identifiers: + application._append_log_entry(entry) + + self.pending.clear() + self.application = weakref.ref(application) + + +def configure_report_logging(): + logger = logging.getLogger("PYAS") + + for handler in logger.handlers: + if isinstance(handler, ReportHandler): + return handler.baseFilename + + directories = [os.environ.get("ALLUSERSPROFILE", "C:\\ProgramData"), tempfile.gettempdir()] + + for directory in directories: + try: + path = os.path.join(directory, "PYAS", "Report.json") + os.makedirs(os.path.dirname(path), exist_ok=True) + descriptor, probe = tempfile.mkstemp(prefix=".pyas-log-", dir=os.path.dirname(path)) + os.close(descriptor) + os.unlink(probe) + handler = ReportHandler(path) + + for name in ("PYAS", "pywebview"): + target = logging.getLogger(name) + target.setLevel(logging.DEBUG) + target.addHandler(handler) + target.propagate = False + + atexit.register(handler.flush) + return handler.baseFilename + except OSError: + continue + + logger.addHandler(logging.NullHandler()) + return None + + +def bind_report_logging(application): + for handler in logging.getLogger("PYAS").handlers: + if isinstance(handler, ReportHandler): + handler.bind(application) + + +def log_exception(context, level=logging.WARNING): + exc_info = sys.exc_info() + + if exc_info[0] is None: + return + + _local.exceptions = getattr(_local, "exceptions", 0) + 1 + key = (context, exc_info[0]) + now = time.monotonic() + + try: + with _lock: + previous, suppressed = _recent.pop(key, (float("-inf"), 0)) + + if now - previous < 30: + _recent[key] = (previous, suppressed + 1) + return + + _recent[key] = (now, 0) + + while len(_recent) > 128: + _recent.popitem(last=False) + + logging.getLogger("PYAS.Exception").log( + level, "context=%s; suppressed=%d", context, suppressed, exc_info=exc_info + ) + except Exception: + return + + +def observe_future(future, context): + def completed(result): + if result.cancelled(): + return + + try: + result.result() + except Exception: + log_exception(context) + + future.add_done_callback(completed) + return future + + +def exception_sequence(): + return getattr(_local, "exceptions", 0) diff --git a/PYAS_Driver.py b/PYAS_Driver.py new file mode 100644 index 0000000..5d45a76 --- /dev/null +++ b/PYAS_Driver.py @@ -0,0 +1,1180 @@ +from PYAS_Diagnostics import log_exception +import os +import time +import winreg +import threading +import ctypes +import ctypes.wintypes +from PYAS_WinAPI import ( + LUID, + PYAS_FULL_MESSAGE, + PYAS_USER_MESSAGE, + SERVICE_STATUS_PROCESS, + TOKEN_PRIVILEGES, +) +from PYAS_WinAPI import ( + ERROR_NOT_FOUND, + ERROR_OPERATION_ABORTED, + FLT_PORT_FLAG_SYNC_HANDLE, + HRESULT_IO_PENDING, + OVERLAPPED, + PYAS_CONNECTION_CONTEXT, + PYAS_CONNECTION_MAGIC, + PYAS_CONNECTION_VERSION, + WAIT_OBJECT_0, + WAIT_TIMEOUT, +) + + +class DriverMixin: + def _query_service_state_handle(self, service): + status = SERVICE_STATUS_PROCESS() + needed = ctypes.wintypes.DWORD(0) + buffer = ctypes.cast(ctypes.byref(status), ctypes.POINTER(ctypes.c_ubyte)) + + if not self.advapi32.QueryServiceStatusEx( + service, 0, buffer, ctypes.sizeof(status), ctypes.byref(needed) + ): + return None, ctypes.get_last_error() + + return int(status.dwCurrentState), 0 + + def _query_driver_service_state(self): + scm = self.advapi32.OpenSCManagerW(None, None, 0x0001) + + if not scm: + return None + + service = None + + try: + service = self.advapi32.OpenServiceW(scm, "PYAS_Driver", 0x0004) + + if not service: + return None + + state, _ = self._query_service_state_handle(service) + return state + finally: + if service: + self.advapi32.CloseServiceHandle(service) + + self.advapi32.CloseServiceHandle(scm) + + def _get_driver_client_image_path(self): + buffer = ctypes.create_unicode_buffer(32768) + size = ctypes.wintypes.DWORD(len(buffer)) + + if not self.kernel32.QueryFullProcessImageNameW( + self.kernel32.GetCurrentProcess(), 0, buffer, ctypes.byref(size) + ): + return None + + image_path = os.path.abspath(buffer.value) + drive, tail = os.path.splitdrive(image_path) + + if not drive or not tail: + return None + + device_buffer = ctypes.create_unicode_buffer(32768) + + if not self.kernel32.QueryDosDeviceW(drive, device_buffer, len(device_buffer)): + return None + + return device_buffer.value.rstrip("\\") + tail + + def _configure_driver_identity(self): + image_path = self._get_driver_client_image_path() + + if not image_path: + return False, ctypes.get_last_error() or 1 + + try: + access = winreg.KEY_SET_VALUE | getattr(winreg, "KEY_WOW64_64KEY", 0) + + with winreg.CreateKeyEx( + winreg.HKEY_LOCAL_MACHINE, + r"SYSTEM\CurrentControlSet\Services\PYAS_Driver", + 0, + access, + ) as key: + winreg.SetValueEx(key, "ClientImagePath", 0, winreg.REG_SZ, image_path) + + return True, 0 + except OSError as e: + log_exception("PYAS_Driver.DriverMixin._configure_driver_identity:74") + return False, int(getattr(e, "winerror", 0) or 1) + + def _ensure_driver_service(self): + desired_service_access = 0x0002 | 0x0004 | 0x0010 | 0x00010000 + dependencies = ctypes.create_unicode_buffer("FltMgr\0\0") + dependencies_ptr = ctypes.cast(dependencies, ctypes.wintypes.LPCWSTR) + last_error = 0 + + for attempt in range(20): + scm = self.advapi32.OpenSCManagerW(None, None, 0x0001 | 0x0002) + + if not scm: + return False, ctypes.get_last_error() + + service = None + + try: + ctypes.set_last_error(0) + service = self.advapi32.CreateServiceW( + scm, + "PYAS_Driver", + "PYAS_Driver", + desired_service_access, + 0x00000001, + 0x00000003, + 0x00000001, + self.path_drivers, + "FSFilter Activity Monitor", + None, + dependencies_ptr, + None, + None, + ) + + if not service: + last_error = ctypes.get_last_error() + + if last_error in (1073, 1078): + service = self.advapi32.OpenServiceW( + scm, "PYAS_Driver", desired_service_access + ) + + if not service: + last_error = ctypes.get_last_error() + + elif last_error == 1072: + pass + + if service: + ctypes.set_last_error(0) + changed = self.advapi32.ChangeServiceConfigW( + service, + 0x00000001, + 0x00000003, + 0x00000001, + self.path_drivers, + "FSFilter Activity Monitor", + None, + dependencies_ptr, + None, + None, + None, + ) + + if not changed: + last_error = ctypes.get_last_error() + return False, last_error + + return True, 0 + finally: + if service: + self.advapi32.CloseServiceHandle(service) + + self.advapi32.CloseServiceHandle(scm) + + if last_error != 1072: + break + + if attempt < 19: + time.sleep(0.05) + + return False, last_error or 1 + + def _start_driver_service(self): + scm = self.advapi32.OpenSCManagerW(None, None, 0x0001) + + if not scm: + return False, ctypes.get_last_error() + + service = None + + try: + service = self.advapi32.OpenServiceW(scm, "PYAS_Driver", 0x0004 | 0x0010) + + if not service: + return False, ctypes.get_last_error() + + state, error = self._query_service_state_handle(service) + + if state in (2, 4): + return True, 0 + + if error: + return False, error + + ctypes.set_last_error(0) + + if self.advapi32.StartServiceW(service, 0, None): + return True, 0 + + error = ctypes.get_last_error() + + if error == 1056: + return True, 0 + + return False, error + finally: + if service: + self.advapi32.CloseServiceHandle(service) + + self.advapi32.CloseServiceHandle(scm) + + def uninstall_system_driver(self): + if not self.stop_system_driver(): + return False, 1051 + + return self._delete_driver_service() + + def _delete_driver_service(self, timeout=15.0): + deadline = time.monotonic() + timeout + last_error = 0 + delete_requested = False + + while time.monotonic() < deadline: + scm = self.advapi32.OpenSCManagerW(None, None, 0x0001) + + if not scm: + return False, ctypes.get_last_error() + + service = None + + try: + ctypes.set_last_error(0) + service = self.advapi32.OpenServiceW(scm, "PYAS_Driver", 0x0004 | 0x00010000) + + if not service: + last_error = ctypes.get_last_error() + + if last_error == 1060: + return True, 0 + + if last_error != 1072: + return False, last_error + else: + state, status_error = self._query_service_state_handle(service) + + if status_error: + return False, status_error + + if state not in (2, 3, 4) and not delete_requested: + ctypes.set_last_error(0) + + if self.advapi32.DeleteService(service): + delete_requested = True + last_error = 1072 + else: + last_error = ctypes.get_last_error() + + if last_error == 1060: + return True, 0 + + if last_error == 1072: + delete_requested = True + else: + return False, last_error + elif state in (2, 3, 4): + last_error = 1051 + finally: + if service: + self.advapi32.CloseServiceHandle(service) + + self.advapi32.CloseServiceHandle(scm) + + time.sleep(0.1) + + return False, last_error or 1072 + + def install_system_driver(self): + try: + ensured, service_error = self._ensure_driver_service() + + if not ensured: + self.write_log( + "WARN", + "Driver Service", + detail=f"CreateServiceW/ChangeServiceConfigW failed: 0x{service_error & 0xFFFFFFFF:08X}", + success=False, + ) + return False + + configured, identity_error = self._configure_driver_identity() + + if not configured: + self.write_log( + "WARN", + "Driver Identity", + detail=f"ClientImagePath configuration failed: 0x{identity_error & 0xFFFFFFFF:08X}", + success=False, + ) + return False + + self._reg_write( + winreg.HKEY_LOCAL_MACHINE, + r"SYSTEM\CurrentControlSet\Services\PYAS_Driver\Instances", + "DefaultInstance", + winreg.REG_SZ, + "PYAS Instance", + ) + self._reg_write( + winreg.HKEY_LOCAL_MACHINE, + r"SYSTEM\CurrentControlSet\Services\PYAS_Driver\Instances\PYAS Instance", + "Altitude", + winreg.REG_SZ, + "320000", + ) + self._reg_write( + winreg.HKEY_LOCAL_MACHINE, + r"SYSTEM\CurrentControlSet\Services\PYAS_Driver\Instances\PYAS Instance", + "Flags", + winreg.REG_DWORD, + 0, + ) + + started, start_error = self._start_driver_service() + + if not started: + self.write_log( + "WARN", + "Driver Service", + detail=f"StartServiceW failed: 0x{start_error & 0xFFFFFFFF:08X}", + success=False, + ) + return False + + final_state = None + last_start_attempt = time.monotonic() + deadline = time.monotonic() + 3.0 + + while time.monotonic() < deadline: + service_state = self._query_driver_service_state() + + if service_state == 4: + final_state = self._query_driver_state() + + if final_state == 2: + return True + + if final_state in (4, 5): + break + elif service_state == 1: + now = time.monotonic() + + if now - last_start_attempt >= 0.1: + started, start_error = self._start_driver_service() + last_start_attempt = now + + if not started: + break + elif service_state is None: + break + + time.sleep(0.02) + + if self.check_system_driver(): + self.stop_system_driver() + + detail = ( + "Driver entered unload-retry state during startup" + if final_state == 4 + else "Driver did not reach running state" + ) + self.write_log("WARN", "Driver Protection", detail=detail, success=False) + return False + + except Exception as e: + log_exception("PYAS_Driver.DriverMixin.install_system_driver:267") + self.write_log("WARN", "install_system_driver", detail=str(e), success=False) + return False + + def _driver_handle_value(self, handle): + if handle is None: + return None + + value = getattr(handle, "value", handle) + + if value is None: + return None + + try: + return int(value) + except Exception: + log_exception("PYAS_Driver.DriverMixin._driver_handle_value:281") + return None + + def _connect_driver_port(self, asynchronous=False): + context = PYAS_CONNECTION_CONTEXT() + context.Size = ctypes.sizeof(context) + context.Version = PYAS_CONNECTION_VERSION + context.Magic = PYAS_CONNECTION_MAGIC + context.ProcessId = os.getpid() + + temp_port = ctypes.wintypes.HANDLE() + status = self.fltlib.FilterConnectCommunicationPort( + "\\PYAS_Output_Pipe", + 0 if asynchronous else FLT_PORT_FLAG_SYNC_HANDLE, + ctypes.byref(context), + ctypes.sizeof(context), + None, + ctypes.byref(temp_port), + ) + + if status != 0: + return None + + return temp_port + + def _detach_driver_port(self, expected_port=None): + expected_value = self._driver_handle_value(expected_port) + + with self.lock_driver: + current_port = self.driver_port + current_value = self._driver_handle_value(current_port) + + if current_value is None: + return None + + if expected_port is not None and current_value != expected_value: + return None + + self.driver_port = None + return current_port + + def _close_driver_port(self, expected_port=None): + port = self._detach_driver_port(expected_port) + + if not port: + return False + + try: + self.kernel32.CloseHandle(port) + except Exception: + log_exception("PYAS_Driver.DriverMixin._close_driver_port:326") + pass + + return True + + def _driver_listener_should_run(self): + if self.driver_stop_event.is_set(): + return False + + with self.lock_config: + return bool(self.pyas_config.get("driver_switch", False)) + + def start_driver_listener(self, wait_ready=True): + with self.lock_driver: + current = self.driver_listener_thread + + if current and current.is_alive(): + return ( + self.driver_listener_ready_event.is_set() + and not self.driver_stop_event.is_set() + ) + + self.driver_stop_event.clear() + self.driver_listener_ready_event.clear() + self.driver_listener_failed_event.clear() + thread = threading.Thread(target=self.pipe_server_thread, daemon=True) + self.driver_listener_thread = thread + thread.start() + + if not wait_ready: + return True + + deadline = time.monotonic() + 5.0 + + while time.monotonic() < deadline: + if self.driver_listener_ready_event.wait(0.02): + return True + + if self.driver_listener_failed_event.is_set() or not thread.is_alive(): + break + + self._stop_driver_listener() + return False + + def _stop_driver_listener(self): + self.driver_stop_event.set() + + with self.lock_driver: + thread = self.driver_listener_thread + port = self.driver_port + + if port: + try: + self.kernel32.CancelIoEx(port, None) + except Exception: + log_exception("PYAS_Driver.DriverMixin._stop_driver_listener:373") + pass + + if thread and thread is not threading.current_thread(): + thread.join(timeout=2.0) + + alive = bool(thread and thread.is_alive()) + + if not alive: + with self.lock_driver: + if self.driver_listener_thread is thread: + self.driver_listener_thread = None + + self.driver_listener_ready_event.clear() + self._close_driver_port() + + return not alive + + def _resume_driver_listener(self): + if not self.check_system_driver(): + return False + + with self.lock_config: + enabled = bool(self.pyas_config.get("driver_switch", False)) + + if not enabled: + return False + + self._close_driver_port() + return self.start_driver_listener(wait_ready=True) + + def _set_driver_unload_authorization(self, enabled): + msg = PYAS_USER_MESSAGE() + msg.Command = 5 if enabled else 6 + msg.Path = "" + + for _ in range(3): + if not self._ensure_driver_port(): + time.sleep(0.1) + continue + + with self.lock_driver: + current_port = self.driver_port + + if not current_port: + continue + + bytes_returned = ctypes.wintypes.DWORD(0) + + try: + status = self.fltlib.FilterSendMessage( + current_port, + ctypes.byref(msg), + ctypes.sizeof(msg), + None, + 0, + ctypes.byref(bytes_returned), + ) + except Exception as e: + log_exception("PYAS_Driver.DriverMixin._set_driver_unload_authorization:427") + self.write_log( + "WARN", "Driver Unload Authorization", detail=str(e), success=False + ) + status = -1 + + if status == 0: + return True + + self._close_driver_port(current_port) + time.sleep(0.1) + + return False + + def _ensure_driver_port(self): + for _ in range(20): + with self.lock_driver: + if self.driver_port: + return True + + temp_port = self._connect_driver_port() + + if temp_port: + with self.lock_driver: + if not self.driver_port: + self.driver_port = temp_port + return True + + try: + self.kernel32.CloseHandle(temp_port) + except Exception: + log_exception("PYAS_Driver.DriverMixin._ensure_driver_port:454") + pass + + return True + + time.sleep(0.05) + + return False + + def _query_driver_state(self): + for use_existing in (True, False): + local_port = None + port = None + + if use_existing: + with self.lock_driver: + if self.driver_port: + port = self.driver_port + + if not port: + continue + else: + local_port = self._connect_driver_port() + + if not local_port: + return None + + port = local_port + + msg = PYAS_USER_MESSAGE() + msg.Command = 7 + msg.Path = "" + state = ctypes.wintypes.DWORD(0) + bytes_returned = ctypes.wintypes.DWORD(0) + + try: + status = self.fltlib.FilterSendMessage( + port, + ctypes.byref(msg), + ctypes.sizeof(msg), + ctypes.byref(state), + ctypes.sizeof(state), + ctypes.byref(bytes_returned), + ) + + if status == 0 and bytes_returned.value == ctypes.sizeof(state): + return int(state.value) + except Exception: + log_exception("PYAS_Driver.DriverMixin._query_driver_state:496") + pass + finally: + if local_port: + try: + self.kernel32.CloseHandle(local_port) + except Exception: + log_exception("PYAS_Driver.DriverMixin._query_driver_state:502") + pass + + if use_existing: + self._close_driver_port(port) + + return None + + def _enable_token_privilege(self, privilege_name): + token = ctypes.wintypes.HANDLE() + desired_access = 0x0020 | 0x0008 + + if not self.advapi32.OpenProcessToken( + self.kernel32.GetCurrentProcess(), desired_access, ctypes.byref(token) + ): + return None, ctypes.get_last_error() + + luid = LUID() + + if not self.advapi32.LookupPrivilegeValueW(None, privilege_name, ctypes.byref(luid)): + error = ctypes.get_last_error() + self.kernel32.CloseHandle(token) + return None, error + + new_state = TOKEN_PRIVILEGES() + new_state.PrivilegeCount = 1 + new_state.Privileges[0].Luid = luid + new_state.Privileges[0].Attributes = 0x00000002 + + previous_state = TOKEN_PRIVILEGES() + return_length = ctypes.wintypes.DWORD(0) + ctypes.set_last_error(0) + + adjusted = self.advapi32.AdjustTokenPrivileges( + token, + False, + ctypes.byref(new_state), + ctypes.sizeof(previous_state), + ctypes.byref(previous_state), + ctypes.byref(return_length), + ) + error = ctypes.get_last_error() + + if not adjusted or error == 1300: + self.kernel32.CloseHandle(token) + return None, error or 1 + + return (token, previous_state, return_length.value), 0 + + def _restore_token_privilege(self, privilege_state): + if not privilege_state: + return + + token, previous_state, previous_length = privilege_state + + try: + if previous_length: + self.advapi32.AdjustTokenPrivileges( + token, False, ctypes.byref(previous_state), 0, None, None + ) + finally: + self.kernel32.CloseHandle(token) + + def _filter_unload_with_privilege(self, filter_name): + result = { + "status": None, + "privilege_error": 0, + "authorization_failed": False, + "exception": None, + "attempts": 0, + } + privilege_state = None + unload_authorized = False + + self.driver_unload_worker = None + self.driver_unload_result = result + + try: + privilege_state, privilege_error = self._enable_token_privilege("SeLoadDriverPrivilege") + + if not privilege_state: + result["privilege_error"] = privilege_error + return result + + retryable_statuses = {0x80070522, 0x80070005, 0x800700AA} + + for delay in (0.0, 0.05, 0.15): + if delay: + time.sleep(delay) + + if not self.check_system_driver(): + result["status"] = 0 + unload_authorized = False + return result + + if not self._set_driver_unload_authorization(True): + continue + + unload_authorized = True + result["attempts"] += 1 + unload_status = self.fltlib.FilterUnload(filter_name) + result["status"] = unload_status + + if unload_status == 0: + unload_authorized = False + return result + + revoked = False + + if self.check_system_driver(): + try: + revoked = self._set_driver_unload_authorization(False) + except Exception: + log_exception("PYAS_Driver.DriverMixin._filter_unload_with_privilege:606") + revoked = False + + unload_authorized = not revoked + + if (unload_status & 0xFFFFFFFF) not in retryable_statuses: + return result + + if result["status"] is None: + result["authorization_failed"] = True + + except Exception as e: + log_exception("PYAS_Driver.DriverMixin._filter_unload_with_privilege:616") + result["exception"] = str(e) + finally: + if unload_authorized and self.check_system_driver(): + try: + self._set_driver_unload_authorization(False) + except Exception: + log_exception("PYAS_Driver.DriverMixin._filter_unload_with_privilege:622") + pass + + self._restore_token_privilege(privilege_state) + + return result + + def stop_system_driver(self): + with self.lock_driver_unload: + runtime_unloaded = False + unload_confirmed = False + + try: + if not self._stop_driver_listener(): + self.write_log( + "WARN", + "Driver Protection", + detail="Driver listener cancellation did not complete", + success=False, + ) + return False + + if not self.check_system_driver(): + runtime_unloaded = True + unload_confirmed = True + else: + if not self._ensure_driver_port(): + if not self.check_system_driver(): + runtime_unloaded = True + unload_confirmed = True + else: + self.write_log( + "WARN", + "Driver Protection", + detail="Control port unavailable", + success=False, + ) + return False + + if not runtime_unloaded: + unload_result = self._filter_unload_with_privilege("PYAS_Driver") + + if unload_result["privilege_error"]: + error = unload_result["privilege_error"] + self.write_log( + "WARN", + "Driver Protection", + detail=f"SeLoadDriverPrivilege unavailable: 0x{error & 0xFFFFFFFF:08X}", + success=False, + ) + return False + + if unload_result["authorization_failed"]: + self.write_log( + "WARN", + "Driver Protection", + detail="Unload authorization rejected", + success=False, + ) + return False + + if unload_result["exception"]: + self.write_log( + "WARN", + "stop_system_driver", + detail=unload_result["exception"], + success=False, + ) + return False + + unload_status = unload_result["status"] + + if unload_status == 0: + runtime_unloaded = True + unload_confirmed = True + elif not self.check_system_driver(): + runtime_unloaded = True + unload_confirmed = True + else: + attempts = unload_result["attempts"] + self.write_log( + "WARN", + "Driver Protection", + detail=f"FilterUnload failed after {attempts} attempt(s): 0x{unload_status & 0xFFFFFFFF:08X}", + success=False, + ) + return False + + self._close_driver_port() + + if unload_confirmed: + deadline = time.monotonic() + 0.5 + + while time.monotonic() < deadline and self.check_system_driver(): + time.sleep(0.01) + + if self.check_system_driver(): + self.write_log( + "INFO", + "Driver Protection", + detail="Filter unloaded; SCM state is still settling", + ) + + return runtime_unloaded + except Exception as e: + log_exception("PYAS_Driver.DriverMixin.stop_system_driver:689") + self.write_log("WARN", "stop_system_driver", detail=str(e), success=False) + return False + finally: + if not runtime_unloaded and self.check_system_driver(): + self._resume_driver_listener() + + def check_system_driver(self): + runtime_state = self._query_driver_state() + + if runtime_state in (1, 2, 3, 4): + return True + + service_state = self._query_driver_service_state() + return service_state in (2, 3, 4) + + def clear_driver_rules(self): + with self.lock_driver: + if not self.driver_port: + return False + + msg = PYAS_USER_MESSAGE() + msg.Command = 4 + msg.Path = "" + bytes_returned = ctypes.wintypes.DWORD(0) + + try: + return ( + self.fltlib.FilterSendMessage( + self.driver_port, + ctypes.byref(msg), + ctypes.sizeof(msg), + None, + 0, + ctypes.byref(bytes_returned), + ) + == 0 + ) + except Exception: + log_exception("PYAS_Driver.DriverMixin.clear_driver_rules:723") + return False + + def load_driver_rule_file(self, json_path): + with self.lock_driver: + if not self.driver_port: + return False + + norm_path = os.path.abspath(json_path) + + if not os.path.exists(norm_path): + return False + + nt_path = f"\\??\\{norm_path}" + + msg = PYAS_USER_MESSAGE() + msg.Command = 3 + msg.Path = nt_path + bytes_returned = ctypes.wintypes.DWORD(0) + + try: + return ( + self.fltlib.FilterSendMessage( + self.driver_port, + ctypes.byref(msg), + ctypes.sizeof(msg), + None, + 0, + ctypes.byref(bytes_returned), + ) + == 0 + ) + except Exception: + log_exception("PYAS_Driver.DriverMixin.load_driver_rule_file:751") + return False + + def _cancel_driver_receive(self, port, overlapped): + try: + ctypes.set_last_error(0) + cancelled = self.kernel32.CancelIoEx(port, ctypes.byref(overlapped)) + error = ctypes.get_last_error() + + if not cancelled and error not in (0, ERROR_NOT_FOUND): + return False + + wait_status = self.kernel32.WaitForSingleObject(overlapped.hEvent, 1000) + + if wait_status != WAIT_OBJECT_0: + return False + + transferred = ctypes.wintypes.DWORD(0) + self.kernel32.GetOverlappedResult( + port, ctypes.byref(overlapped), ctypes.byref(transferred), False + ) + return True + except Exception: + log_exception("PYAS_Driver.DriverMixin._cancel_driver_receive:774") + return False + + def _receive_driver_message(self, port): + message = PYAS_FULL_MESSAGE() + overlapped = OVERLAPPED() + event_handle = self.kernel32.CreateEventW(None, True, False, None) + + if not event_handle: + return "error", None + + overlapped.hEvent = event_handle + + try: + status = self.fltlib.FilterGetMessage( + port, + ctypes.byref(message), + ctypes.sizeof(PYAS_FULL_MESSAGE), + ctypes.byref(overlapped), + ) + + if status == 0: + return "message", message + + if (status & 0xFFFFFFFF) != HRESULT_IO_PENDING: + return "disconnect", None + + while True: + if self.driver_stop_event.is_set(): + self._cancel_driver_receive(port, overlapped) + return "stop", None + + wait_status = self.kernel32.WaitForSingleObject(event_handle, 20) + + if wait_status == WAIT_TIMEOUT: + continue + + if wait_status != WAIT_OBJECT_0: + self._cancel_driver_receive(port, overlapped) + return "error", None + + transferred = ctypes.wintypes.DWORD(0) + + if self.kernel32.GetOverlappedResult( + port, ctypes.byref(overlapped), ctypes.byref(transferred), False + ): + return "message", message + + error = ctypes.get_last_error() + + if error == ERROR_OPERATION_ABORTED and self.driver_stop_event.is_set(): + return "stop", None + + return "disconnect", None + finally: + self.kernel32.CloseHandle(event_handle) + + def pipe_server_thread(self): + owned_port = None + ready = False + + try: + while self._driver_listener_should_run(): + temp_port = self._connect_driver_port(asynchronous=True) + + if not temp_port: + self.driver_stop_event.wait(0.02) + continue + + if not self._driver_listener_should_run(): + self.kernel32.CloseHandle(temp_port) + break + + with self.lock_driver: + if self.driver_port: + accepted = False + else: + self.driver_port = temp_port + owned_port = temp_port + accepted = True + + if not accepted: + self.kernel32.CloseHandle(temp_port) + self.driver_stop_event.wait(0.02) + continue + + rule_files = [] + + if os.path.exists(self.path_rules): + rule_files = sorted( + os.path.join(self.path_rules, name) + for name in os.listdir(self.path_rules) + if name.lower().endswith(".json") and name.lower() != "rules_driver_p1.json" + ) + + load_failed = False + + for rule_file in rule_files: + if self.driver_stop_event.is_set() or not self.load_driver_rule_file(rule_file): + load_failed = True + break + + if load_failed: + self.driver_listener_failed_event.set() + break + + with self.lock_config: + whitelist = list(self.pyas_config.get("white_list", [])) + + for item in whitelist: + if self.driver_stop_event.is_set(): + break + + if isinstance(item, dict) and item.get("file"): + self.sync_driver_whitelist(item["file"], True, item.get("is_dir")) + + if self.driver_stop_event.is_set(): + break + + ready = True + self.driver_listener_ready_event.set() + + while self._driver_listener_should_run(): + with self.lock_driver: + current_port = self.driver_port + owns_current = self._driver_handle_value( + current_port + ) == self._driver_handle_value(owned_port) + + if not owns_current: + break + + receive_state, message = self._receive_driver_message(owned_port) + + if receive_state != "message": + break + + code = message.Data.MessageCode + pid = message.Data.ProcessId + target = message.Data.Path + exe_info = self.get_exe_info(pid) + safe_source = exe_info[1] if exe_info and exe_info[1] else "Unknown" + + if not self.is_in_whitelist(safe_source): + self.write_log( + "BLOCK", + "Driver Block", + detail="Driver protection event", + pid=pid, + source=safe_source, + target=target, + code=code, + file_hash=self.calc_file_hash(safe_source), + operate=None, + success=True, + ) + + self._close_driver_port(owned_port) + owned_port = None + ready = False + self.driver_listener_ready_event.clear() + + if self.driver_stop_event.is_set(): + break + + self.driver_stop_event.wait(0.02) + + except Exception as e: + log_exception("PYAS_Driver.DriverMixin.pipe_server_thread:928") + self.driver_listener_failed_event.set() + + if not self.driver_stop_event.is_set(): + self.write_log("WARN", "pipe_server_thread", detail=str(e), success=False) + + finally: + if owned_port: + try: + self.kernel32.CancelIoEx(owned_port, None) + except Exception: + log_exception("PYAS_Driver.DriverMixin.pipe_server_thread:937") + pass + + self._close_driver_port(owned_port) + + if not ready and not self.driver_stop_event.is_set(): + self.driver_listener_failed_event.set() + + self.driver_listener_ready_event.clear() + + with self.lock_driver: + if self.driver_listener_thread is threading.current_thread(): + self.driver_listener_thread = None diff --git a/PYAS_Engine.py b/PYAS_Engine.py index 75c0a53..17c41c7 100644 --- a/PYAS_Engine.py +++ b/PYAS_Engine.py @@ -1,1322 +1,15 @@ -import os, re, yara, time, math, json, zlib, numpy -import pefile, threading, onnxruntime, datetime, requests -import ctypes, ctypes.wintypes - -#################################################################################################### - -class GUID(ctypes.Structure): - _fields_ = [ - ("Data1", ctypes.wintypes.DWORD), - ("Data2", ctypes.wintypes.WORD), - ("Data3", ctypes.wintypes.WORD), - ("Data4", ctypes.c_ubyte * 8) - ] - -class WINTRUST_FILE_INFO(ctypes.Structure): - _fields_ = [ - ("cbStruct", ctypes.wintypes.DWORD), - ("pcwszFilePath", ctypes.wintypes.LPCWSTR), - ("hFile", ctypes.wintypes.HANDLE), - ("pgKnownSubject", ctypes.wintypes.LPVOID) - ] - -class WINTRUST_DATA_UNION(ctypes.Union): - _fields_ = [ - ("pFile", ctypes.POINTER(WINTRUST_FILE_INFO)), - ("pCatalog", ctypes.wintypes.LPVOID), - ("pBlob", ctypes.wintypes.LPVOID), - ("pSgnr", ctypes.wintypes.LPVOID), - ("pCert", ctypes.wintypes.LPVOID) - ] - -class WINTRUST_DATA(ctypes.Structure): - _fields_ = [ - ("cbStruct", ctypes.wintypes.DWORD), - ("pPolicyCallbackData", ctypes.wintypes.LPVOID), - ("pSIPClientData", ctypes.wintypes.LPVOID), - ("dwUIChoice", ctypes.wintypes.DWORD), - ("fdwRevocationChecks", ctypes.wintypes.DWORD), - ("dwUnionChoice", ctypes.wintypes.DWORD), - ("u", WINTRUST_DATA_UNION), - ("dwStateAction", ctypes.wintypes.DWORD), - ("hWVTStateData", ctypes.wintypes.HANDLE), - ("pwszURLReference", ctypes.wintypes.LPCWSTR), - ("dwProvFlags", ctypes.wintypes.DWORD), - ("dwUIContext", ctypes.wintypes.DWORD), - ("pSignatureSettings", ctypes.wintypes.LPVOID) - ] - -#################################################################################################### - -class sign_scanner: - def __init__(self): - self.verify = GUID(0x00AAC56B, 0xCD44, 0x11D0, (ctypes.c_ubyte * 8)(0x8C, 0xC2, 0x00, 0xC0, 0x4F, 0xC2, 0x95, 0xEE)) - - def init_windll(self, path): - for name in path: - try: - setattr(self, name.lower(), ctypes.WinDLL(name, use_last_error=True)) - except Exception: - pass - - try: - self.WinVerifyTrust = self.wintrust.WinVerifyTrust - self.WinVerifyTrust.restype = ctypes.wintypes.LONG - self.WinVerifyTrust.argtypes = [ctypes.wintypes.HWND, ctypes.POINTER(GUID), ctypes.c_void_p] - except Exception: - pass - -#################################################################################################### - - def sign_verify(self, file_path): - if os.name != 'nt': - return False - try: - fi = WINTRUST_FILE_INFO() - fi.cbStruct = ctypes.sizeof(WINTRUST_FILE_INFO) - fi.pcwszFilePath = os.path.abspath(file_path) - fi.hFile = None - fi.pgKnownSubject = None - - wt_union = WINTRUST_DATA_UNION() - wt_union.pFile = ctypes.pointer(fi) - - td = WINTRUST_DATA() - td.cbStruct = ctypes.sizeof(WINTRUST_DATA) - td.pPolicyCallbackData = None - td.pSIPClientData = None - td.dwUIChoice = 2 - td.fdwRevocationChecks = 0 - td.dwUnionChoice = 1 - td.u = wt_union - td.dwStateAction = 0 - td.hWVTStateData = None - td.pwszURLReference = None - td.dwProvFlags = 0 - td.dwUIContext = 0 - td.pSignatureSettings = None - - return self.WinVerifyTrust(None, ctypes.byref(self.verify), ctypes.byref(td)) == 0 - except Exception: - return False - -#################################################################################################### - -class rule_scanner: - def __init__(self): - self.rules = None - self.network = [] - - def load_path(self, path, callback=None): - yara_files = {} - for root, _, files in os.walk(path): - for file in files: - - full_path = os.path.join(root, file) - if callback: - callback(full_path) - - ext = os.path.splitext(file)[1].lower() - if ext in ('.yara', '.yar'): - namespace = os.path.relpath(full_path, path).replace(os.sep, '_') - yara_files[namespace] = full_path - - elif ext in ('.yc', '.yrc'): - self.load_compiled_rule(full_path) - - elif ext in ('.ip', '.txt'): - self.load_network_list(full_path) - - if yara_files: - self.compile_all_rules(yara_files) - -#################################################################################################### - - def load_compiled_rule(self, file): - try: - self.rules = yara.load(file) - except Exception: - pass - - def load_network_list(self, file): - try: - with open(file, "r", encoding="utf-8", errors="ignore") as f: - self.network.extend(line.strip() for line in f if line.strip()) - except Exception: - pass - - def compile_all_rules(self, file_map): - try: - self.rules = yara.compile(filepaths=file_map) - except Exception: - pass - -#################################################################################################### - - def yara_scan(self, file_path): - try: - if not self.rules: - return False, False - - matches = self.rules.match(filepath=file_path) - if matches: - rule_name = str(matches[0]) - try: - label = rule_name.split("_")[0] - platform = rule_name.split("_")[1] - family = rule_name.split("_")[2] - - return f"{label}:{platform}/{family}.{len(matches)}!yr", 100 - except Exception: - return rule_name, 100 - - return False, False - except Exception: - return False, False - - def yara_mem_scan(self, pid): - try: - if not self.rules: - return False, False - - matches = self.rules.match(pid=pid) - if matches: - rule_name = str(matches[0]) - try: - label = rule_name.split("_")[0] - platform = rule_name.split("_")[1] - family = rule_name.split("_")[2] - - return f"{label}:{platform}/{family}.{len(matches)}!ym", 100 - except Exception: - return rule_name, 100 - - return False, False - except Exception: - return False, False - -#################################################################################################### - -def _entropy_from_counts(counts, size): - if size <= 0: - return 0.0 - nonzero = counts[counts > 0] - return float(numpy.log2(size) - numpy.sum(nonzero * numpy.log2(nonzero)) / size) - - -def _bounded_entropy(data, batch_size=1048576): - size = len(data) - if size == 0: - return 0.0 - counts = numpy.zeros(256, dtype=numpy.int64) - view = memoryview(data) - try: - for start in range(0, size, batch_size): - chunk = numpy.frombuffer( - view[start:min(start + batch_size, size)], dtype=numpy.uint8 - ) - counts += numpy.bincount(chunk, minlength=256) - finally: - view.release() - return _entropy_from_counts(counts, size) - - -def _accumulate_string_chunk(chunk, array, state, pattern): - separators = (array < 0x20) | (array > 0x7E) - if not separators.any(): - state["pending_counts"] += numpy.bincount(array, minlength=256) - state["pending_length"] += len(array) - return - - first_separator = int(numpy.argmax(separators)) - last_separator = len(separators) - 1 - int(numpy.argmax(separators[::-1])) - if first_separator: - state["pending_counts"] += numpy.bincount( - array[:first_separator], minlength=256 - ) - state["pending_length"] += first_separator - if state["pending_length"] >= 5: - state["count"] += 1 - state["total_length"] += state["pending_length"] - state["counts"] += state["pending_counts"] - state["pending_length"] = 0 - state["pending_counts"].fill(0) - - middle_start = first_separator + 1 - middle_end = last_separator - if middle_end > middle_start: - matches = pattern.findall(chunk[middle_start:middle_end]) - if matches: - combined = b"".join(matches) - state["count"] += len(matches) - state["total_length"] += len(combined) - state["counts"] += numpy.bincount( - numpy.frombuffer(combined, dtype=numpy.uint8), minlength=256 - ) - - trailing_start = last_separator + 1 - if trailing_start < len(array): - trailing = array[trailing_start:] - state["pending_counts"] += numpy.bincount(trailing, minlength=256) - state["pending_length"] = len(trailing) - - -def _accumulate_entropy_windows(data, entropy_histogram, c_log_c_table): - block_count = len(data) // 1024 - if block_count < 2: - return data - - complete_size = block_count * 1024 - blocks = numpy.frombuffer( - data, dtype=numpy.uint8, count=complete_size - ).reshape(block_count, 1024) - encoded = blocks.astype(numpy.int32) - encoded += numpy.arange(block_count, dtype=numpy.int32)[:, None] * 256 - block_counts = numpy.bincount( - encoded.ravel(), minlength=block_count * 256 - ).reshape(block_count, 256) - window_counts = block_counts[:-1] + block_counts[1:] - sum_c_log_c = numpy.sum(c_log_c_table[window_counts], axis=1) - entropies = 11.0 - sum_c_log_c / 2048.0 - entropy_bins = (entropies * 2.0).astype(numpy.int32) - numpy.clip(entropy_bins, 0, 15, out=entropy_bins) - byte_bins = window_counts.reshape(block_count - 1, 16, 16).sum(axis=2) - flat_indices = (entropy_bins[:, None] * 16 + numpy.arange(16)).ravel() - numpy.add.at(entropy_histogram, flat_indices, byte_bins.ravel()) - return data[(block_count - 1) * 1024:] - - -def _collect_resource_ranges(pe, file_size): - ranges = [] - if not hasattr(pe, "DIRECTORY_ENTRY_RESOURCE"): - return ranges - for entry_type in pe.DIRECTORY_ENTRY_RESOURCE.entries: - for entry_id in getattr( - getattr(entry_type, "directory", None), "entries", [] - ): - for entry_lang in getattr( - getattr(entry_id, "directory", None), "entries", [] - ): - if not hasattr(entry_lang, "data"): - continue - try: - struct = entry_lang.data.struct - offset = int(pe.get_offset_from_rva(int(struct.OffsetToData))) - size = int(struct.Size) - start = max(0, offset) - end = min(file_size, start + max(0, size)) - ranges.append((start, max(start, end))) - except Exception: - ranges.append((0, 0)) - return ranges - - -def _extract_stream_features(file_path, file_size, ranges, extractor): - normalized = [] - for start, end in ranges: - start = max(0, min(file_size, int(start))) - end = max(start, min(file_size, int(end))) - normalized.append((start, end)) - - byte_counts = numpy.zeros(256, dtype=numpy.int64) - range_counts = [numpy.zeros(256, dtype=numpy.int64) for _ in normalized] - entropy_histogram = numpy.zeros(256, dtype=numpy.float64) - string_state = { - "count": 0, - "total_length": 0, - "counts": numpy.zeros(256, dtype=numpy.int64), - "pending_length": 0, - "pending_counts": numpy.zeros(256, dtype=numpy.int64), - } - histogram_tail = b"" - position = 0 - - with open(file_path, "rb") as source: - while True: - chunk = source.read(4194304) - if not chunk: - break - array = numpy.frombuffer(chunk, dtype=numpy.uint8) - _accumulate_string_chunk( - chunk, array, string_state, extractor._STRING_PATTERN - ) - histogram_tail = _accumulate_entropy_windows( - histogram_tail + chunk, - entropy_histogram, - extractor._C_LOG_C_TABLE, - ) - - chunk_end = position + len(chunk) - boundaries = {0, len(chunk)} - for start, end in normalized: - if position < start < chunk_end: - boundaries.add(start - position) - if position < end < chunk_end: - boundaries.add(end - position) - - ordered = sorted(boundaries) - for index in range(len(ordered) - 1): - local_start = ordered[index] - local_end = ordered[index + 1] - if local_end <= local_start: - continue - counts = numpy.bincount( - array[local_start:local_end], minlength=256 - ) - byte_counts += counts - absolute_start = position + local_start - absolute_end = position + local_end - for range_index, (start, end) in enumerate(normalized): - if start <= absolute_start and absolute_end <= end: - range_counts[range_index] += counts - position = chunk_end - - if string_state["pending_length"] >= 5: - string_state["count"] += 1 - string_state["total_length"] += string_state["pending_length"] - string_state["counts"] += string_state["pending_counts"] - - string_count = string_state["count"] - total_string_length = string_state["total_length"] - result = { - "StringCount": float(string_count), - "StringMeanLength": ( - float(total_string_length) / string_count if string_count else 0.0 - ), - "StringEntropy": _entropy_from_counts( - string_state["counts"], total_string_length - ), - } - result.update(dict.fromkeys(extractor._STRING_KEYS, 0.0)) - if total_string_length: - valid_counts = string_state["counts"][0x20:0x7F] - result.update(dict(zip( - extractor._STRING_KEYS, - valid_counts.astype(numpy.float64) / total_string_length, - ))) - - result["FileEntropy"] = _entropy_from_counts(byte_counts, file_size) - result.update(dict(zip( - extractor._BYTE_HIST_KEYS, - byte_counts.astype(numpy.float64) / file_size, - ))) - entropy_sum = numpy.sum(entropy_histogram) - if entropy_sum: - entropy_histogram /= entropy_sum - result.update(dict(zip(extractor._BYTE_ENT_KEYS, entropy_histogram))) - - range_entropies = [ - _entropy_from_counts(counts, end - start) - for counts, (start, end) in zip(range_counts, normalized) - ] - return result, range_entropies - -#################################################################################################### - -class pe_scanner: - _STRING_PATTERN = re.compile(b'[\x20-\x7E]{5,}') - _STRING_FAST_PATH_BYTES = 4194304 - _STRING_BATCH_BYTES = 4194304 - _HISTOGRAM_BATCH_WINDOWS = 1024 - - _BYTE_HIST_KEYS = [f'ByteHist_{i:02X}' for i in range(256)] - _BYTE_ENT_KEYS = [f'ByteEnt_{i:02X}' for i in range(256)] - _STRING_KEYS = [f'StringHist_{i:02d}' for i in range(95)] - _SECTION_HASH_KEYS = [f'SectionHash_{i:02d}' for i in range(50)] - - _CHAR_FLAGS = [0x00000020, 0x00000040, 0x00000080, 0x02000000, 0x20000000, 0x40000000, 0x80000000] - _CHAR_COUNT_KEYS = {flag: f'Char_{flag:08X}_Count' for flag in _CHAR_FLAGS} - _CHAR_ENT_KEYS = {flag: f'Char_{flag:08X}_MeanEntropy' for flag in _CHAR_FLAGS} - - _C_LOG_C_TABLE = numpy.zeros(2049, dtype=numpy.float64) - _C_LOG_C_TABLE[1:] = numpy.arange(1, 2049) * numpy.log2(numpy.arange(1, 2049)) - - _OPTIONAL_HEADER_FIELDS = [ - 'Magic', 'MajorLinkerVersion', 'MinorLinkerVersion', 'SizeOfCode', - 'SizeOfInitializedData', 'SizeOfUninitializedData', 'AddressOfEntryPoint', - 'BaseOfCode', 'ImageBase', 'SectionAlignment', 'FileAlignment', - 'MajorOperatingSystemVersion', 'MinorOperatingSystemVersion', - 'MajorImageVersion', 'MinorImageVersion', 'MajorSubsystemVersion', - 'MinorSubsystemVersion', 'SizeOfImage', 'SizeOfHeaders', 'CheckSum', - 'Subsystem', 'DllCharacteristics', 'SizeOfStackReserve', 'SizeOfStackCommit', - 'SizeOfHeapReserve', 'SizeOfHeapCommit', 'LoaderFlags', 'NumberOfRvaAndSizes' - ] - - _SECTION_DEFAULTS = { - 'SectionCount': 0.0, - 'SectionMaxEntropy': 0.0, - 'SectionMinEntropy': 0.0, - 'SectionMeanEntropy': 0.0, - 'SectionMaxRawSize': 0.0, - 'SectionMinRawSize': 0.0, - 'SectionMeanRawSize': 0.0, - 'SectionMaxVSize': 0.0, - 'SectionMinVSize': 0.0, - 'SectionMeanVSize': 0.0, - 'ExecutableSections': 0.0, - 'WritableSections': 0.0, - 'ReadableSections': 0.0, - 'SectionException': 0.0 - } - - _CONDITIONAL_DEFAULTS = { - 'IsDriver': 0.0, - 'HasTlsCallbacks': 0.0, - 'IsDebug': 0.0, - 'IsPreRelease': 0.0, - 'IsPatched': 0.0, - 'IsPrivateBuild': 0.0, - 'IsSpecialBuild': 0.0, - 'ImportCount': 0.0, - 'ImportFunctionCount': 0.0, - 'ExportCount': 0.0, - 'IconCount': 0.0, - 'DebugCount': 0.0 - } - -#################################################################################################### - - def __init__(self): - self.model = None - self.input_name = None - self.feature_order = [] - self.dll_hash_dim = 512 - self.api_hash_dim = 4096 - self.dll_hash_pad = 3 - self.api_hash_pad = 4 - self.signer = sign_scanner() - self.signer.init_windll(["wintrust"]) - - def load_path(self, path, callback=None): - for root, _, files in os.walk(path): - for file in files: - full_path = os.path.join(root, file) - - if callback: - callback(full_path) - - if full_path.endswith('.onnx'): - self.load_model(full_path) - - def load_model(self, model_path): - if not os.path.exists(model_path): - return - - try: - self.model = onnxruntime.InferenceSession(model_path, providers=['CPUExecutionProvider']) - self.input_name = self.model.get_inputs()[0].name - - feat_path = os.path.join(os.path.dirname(model_path), "features.json") - if os.path.exists(feat_path): - with open(feat_path, 'r', encoding='utf-8') as f: - self.feature_order = json.load(f) - - self._parse_hash_dims() - except Exception: - pass - -#################################################################################################### - - def _parse_hash_dims(self): - max_dll = -1 - max_api = -1 - - for feat in self.feature_order: - if feat.startswith("DllHash_"): - try: - val_str = feat.split("_")[1] - max_dll = max(max_dll, int(val_str)) - self.dll_hash_pad = len(val_str) - except Exception: - pass - elif feat.startswith("ApiHash_"): - try: - val_str = feat.split("_")[1] - max_api = max(max_api, int(val_str)) - self.api_hash_pad = len(val_str) - except Exception: - pass - - if max_dll >= 0: - self.dll_hash_dim = max_dll + 1 - if max_api >= 0: - self.api_hash_dim = max_api + 1 - - def _safe_float(self, val): - try: - f = float(val) - if math.isinf(f) or math.isnan(f): - return 0.0 - - return f - except Exception: - return 0.0 - - def _base_defaults(self): - base = dict.fromkeys(self._OPTIONAL_HEADER_FIELDS, 0.0) - base.update(self._SECTION_DEFAULTS) - base.update(self._CONDITIONAL_DEFAULTS) - - for i in range(16): - base[f'DataDirectory_{i}_Size'] = 0.0 - base[f'DataDirectory_{i}_VA'] = 0.0 - - return base - - def _calc_entropy(self, data): - return _bounded_entropy(data) - - def _extract_strings(self, file_bytes): - res = {"StringCount": 0.0, "StringMeanLength": 0.0, "StringEntropy": 0.0} - res.update(dict.fromkeys(self._STRING_KEYS, 0.0)) - - if len(file_bytes) <= self._STRING_FAST_PATH_BYTES: - matches = self._STRING_PATTERN.findall(file_bytes) - if not matches: - return res - - combined_data = b''.join(matches) - count = len(matches) - total_len = len(combined_data) - global_counts = numpy.bincount(numpy.frombuffer(combined_data, dtype=numpy.uint8), minlength=256) - else: - count = 0 - total_len = 0 - global_counts = numpy.zeros(256, dtype=numpy.int64) - mv = memoryview(file_bytes) - pos = 0 - size = len(file_bytes) - - while pos < size: - chunk_end = min(pos + self._STRING_BATCH_BYTES, size) - chunk_arr = numpy.frombuffer(mv[pos:chunk_end], dtype=numpy.uint8) - separators = (chunk_arr < 0x20) | (chunk_arr > 0x7E) - - if separators.any(): - last_separator = len(separators) - 1 - int(numpy.argmax(separators[::-1])) - segment_end = pos + last_separator + 1 - matches = self._STRING_PATTERN.findall(mv[pos:segment_end]) - - if matches: - combined_data = b''.join(matches) - count += len(matches) - total_len += len(combined_data) - global_counts += numpy.bincount(numpy.frombuffer(combined_data, dtype=numpy.uint8), minlength=256) - - pos = segment_end - continue - - pending_len = 0 - pending_counts = numpy.zeros(256, dtype=numpy.int64) - - while pos < size: - chunk_end = min(pos + self._STRING_BATCH_BYTES, size) - chunk_arr = numpy.frombuffer(mv[pos:chunk_end], dtype=numpy.uint8) - separators = (chunk_arr < 0x20) | (chunk_arr > 0x7E) - - if separators.any(): - first_separator = int(numpy.argmax(separators)) - if first_separator: - pending_counts += numpy.bincount(chunk_arr[:first_separator], minlength=256) - pending_len += first_separator - pos += first_separator + 1 - break - - pending_counts += numpy.bincount(chunk_arr, minlength=256) - pending_len += len(chunk_arr) - pos = chunk_end - - if pending_len >= 5: - count += 1 - total_len += pending_len - global_counts += pending_counts - - if count == 0 or total_len == 0: - return res - - res["StringCount"] = float(count) - res["StringMeanLength"] = float(total_len) / count - - counts_nonzero = global_counts[global_counts > 0] - if len(counts_nonzero) > 0: - res["StringEntropy"] = float(numpy.log2(total_len) - numpy.sum(counts_nonzero * numpy.log2(counts_nonzero)) / total_len) - - valid_chars = global_counts[0x20:0x7F] - res.update(dict(zip(self._STRING_KEYS, valid_chars.astype(numpy.float64) / total_len))) - - return res - -#################################################################################################### - - def _extract_histograms(self, file_bytes): - arr = numpy.frombuffer(file_bytes, dtype=numpy.uint8) - sz = len(arr) - res = {} - - if sz == 0: - res['FileEntropy'] = 0.0 - res.update(dict.fromkeys(self._BYTE_HIST_KEYS, 0.0)) - res.update(dict.fromkeys(self._BYTE_ENT_KEYS, 0.0)) - return res - - byte_counts = numpy.zeros(256, dtype=numpy.int64) - ent_hist = numpy.zeros(256, dtype=numpy.float64) - window = 2048 - step = 1024 - - if sz < window: - byte_counts += numpy.bincount(arr, minlength=256) - p = byte_counts[byte_counts > 0] / float(sz) - entropy = float(-numpy.sum(p * numpy.log2(p))) - ent_bin = min(int(entropy * 2.0), 15) - byte_bins = byte_counts.reshape(16, 16).sum(axis=1) - idx_start = ent_bin * 16 - ent_hist[idx_start : idx_start+16] += byte_bins - else: - num_windows = (sz - window) // step + 1 - - for w_start in range(0, num_windows, self._HISTOGRAM_BATCH_WINDOWS): - w_end = min(w_start + self._HISTOGRAM_BATCH_WINDOWS, num_windows) - b_num = w_end - w_start - - start_byte = w_start * step - end_byte = start_byte + b_num * step + step - encoded = arr[start_byte:end_byte].astype(numpy.int32).reshape(b_num + 1, step) - encoded += numpy.arange(b_num + 1, dtype=numpy.int32)[:, None] * 256 - - chunk_counts = numpy.bincount(encoded.ravel(), minlength=(b_num + 1) * 256).reshape(b_num + 1, 256) - del encoded - - byte_counts += numpy.sum(chunk_counts[:-1], axis=0) - if w_end == num_windows: - byte_counts += chunk_counts[-1] - - window_counts = chunk_counts[:-1] + chunk_counts[1:] - sum_c_log_c = numpy.sum(self._C_LOG_C_TABLE[window_counts], axis=1) - entropies = 11.0 - sum_c_log_c / 2048.0 - - ent_bins = (entropies * 2.0).astype(numpy.int32) - numpy.clip(ent_bins, 0, 15, out=ent_bins) - - byte_bins = window_counts.reshape(b_num, 16, 16).sum(axis=2) - flat_ent_bins = (ent_bins[:, None] * 16 + numpy.arange(16)).ravel() - numpy.add.at(ent_hist, flat_ent_bins, byte_bins.ravel()) - - covered_size = (num_windows + 1) * step - if covered_size < sz: - byte_counts += numpy.bincount(arr[covered_size:], minlength=256) - - counts_nonzero = byte_counts[byte_counts > 0] - res['FileEntropy'] = float(numpy.log2(sz) - numpy.sum(counts_nonzero * numpy.log2(counts_nonzero)) / sz) - - byte_hist = byte_counts.astype(numpy.float64) / sz - res.update(dict(zip(self._BYTE_HIST_KEYS, byte_hist))) - - ent_sum = numpy.sum(ent_hist) - if ent_sum > 0: - ent_hist /= ent_sum - - res.update(dict(zip(self._BYTE_ENT_KEYS, ent_hist))) - return res - - def _extract_overlay_features(self, pe, fsize): - overlay_offset = pe.get_overlay_data_start_offset() - if not overlay_offset or overlay_offset >= fsize: - return {"HasOverlay": 0.0, "OverlaySize": 0.0, "OverlayRatio": 0.0, "OverlayEntropy": 0.0} - - overlay_data = pe.get_overlay() - if not overlay_data: - return {"HasOverlay": 0.0, "OverlaySize": 0.0, "OverlayRatio": 0.0, "OverlayEntropy": 0.0} - - sz = len(overlay_data) - return { - "HasOverlay": 1.0, - "OverlaySize": float(sz), - "OverlayRatio": float(sz) / float(fsize), - "OverlayEntropy": self._calc_entropy(overlay_data) - } - - def _extract_rich_header(self, pe): - if not hasattr(pe, 'RICH_HEADER') or not pe.RICH_HEADER: - return {"HasRichHeader": 0.0, "RichHeaderCount": 0.0} - - return { - "HasRichHeader": 1.0, - "RichHeaderCount": float(len(pe.RICH_HEADER.values) // 2) if hasattr(pe.RICH_HEADER, 'values') else 0.0 - } - - def _extract_ep_anomalies(self, pe): - if not hasattr(pe, 'OPTIONAL_HEADER') or not hasattr(pe, 'sections') or not pe.sections: - return {"EntryPointSectionIndex": -1.0, "EntryPointInExecutable": 0.0, "EntryPointInLastSection": 0.0} - - ep = pe.OPTIONAL_HEADER.AddressOfEntryPoint - ep_section_idx = -1 - ep_in_exec = 0.0 - - for idx, sec in enumerate(pe.sections): - if sec.VirtualAddress <= ep < (sec.VirtualAddress + sec.Misc_VirtualSize): - ep_section_idx = idx - if sec.Characteristics & 0x20000000: - ep_in_exec = 1.0 - break - - return { - "EntryPointSectionIndex": float(ep_section_idx), - "EntryPointInExecutable": ep_in_exec, - "EntryPointInLastSection": 1.0 if ep_section_idx == len(pe.sections) - 1 else 0.0 - } - - def _extract_advanced_resources(self, pe, range_entropies=None): - res_data = { - "ResourceMaxEntropy": 0.0, - "ResourceMinEntropy": 0.0, - "ResourceMeanEntropy": 0.0, - "ResourceLangCount": 0.0, - "ResourceRCDataCount": 0.0 - } - - if not hasattr(pe, 'DIRECTORY_ENTRY_RESOURCE'): - return res_data - - entropies = [] - entropy_iter = iter(range_entropies or ()) - langs = set() - rcdata_count = 0 - - for entry_type in pe.DIRECTORY_ENTRY_RESOURCE.entries: - if hasattr(entry_type, 'id') and entry_type.id == 10: - rcdata_count += len(getattr(entry_type.directory, 'entries', [])) - - if not hasattr(entry_type, 'directory'): - continue - - for entry_id in entry_type.directory.entries: - if not hasattr(entry_id, 'directory'): - continue - - for entry_lang in entry_id.directory.entries: - if hasattr(entry_lang, 'id'): - langs.add(entry_lang.id) - - if hasattr(entry_lang, 'data'): - try: - entropies.append(float(next(entropy_iter))) - except (StopIteration, TypeError, ValueError): - entropies.append(0.0) - - res_data["ResourceLangCount"] = float(len(langs)) - res_data["ResourceRCDataCount"] = float(rcdata_count) - - if entropies: - res_data["ResourceMaxEntropy"] = max(entropies) - res_data["ResourceMinEntropy"] = min(entropies) - res_data["ResourceMeanEntropy"] = sum(entropies) / len(entropies) - - return res_data - - def _extract_load_config(self, pe): - cfg = {"HasLoadConfig": 0.0, "HasCFG": 0.0, "HasSEHTable": 0.0} - - if hasattr(pe, 'DIRECTORY_ENTRY_LOAD_CONFIG'): - cfg["HasLoadConfig"] = 1.0 - struct = pe.DIRECTORY_ENTRY_LOAD_CONFIG.struct - - if getattr(struct, 'GuardCFFunctionTable', 0) != 0: - cfg["HasCFG"] = 1.0 - - if getattr(struct, 'SEHandlerTable', 0) != 0: - cfg["HasSEHTable"] = 1.0 - - return cfg - - def _extract_security_directory(self, pe, file_path, fsize): - res = {"HasSignature": 0.0, "SignatureCount": 0.0} - if not hasattr(pe, 'OPTIONAL_HEADER') or not hasattr(pe.OPTIONAL_HEADER, 'DATA_DIRECTORY'): - return res - directories = pe.OPTIONAL_HEADER.DATA_DIRECTORY - if len(directories) <= 4: - return res - sec_dir = directories[4] - offset = int(sec_dir.VirtualAddress) - size = int(sec_dir.Size) - if offset <= 0 or size <= 0 or offset + size > fsize: - return res - - res["HasSignature"] = 1.0 - count = 0 - curr = offset - end = offset + size - with open(file_path, "rb") as source: - while curr + 8 <= end: - source.seek(curr) - header = source.read(4) - if len(header) != 4: - break - length = int.from_bytes(header, 'little') - if length < 8 or curr + length > end: - break - count += 1 - curr += (length + 7) & ~7 - res["SignatureCount"] = float(count) - return res - - def extract_features(self, file_path): - pe = None - - try: - fsize = os.path.getsize(file_path) - if fsize == 0 or fsize > 4294967296: - return None - - base = self._base_defaults() - dlls = set() - apis = set() - pe = pefile.PE(name=file_path, fast_load=True) - - try: - pe.parse_rich_header() - pe.parse_data_directories(directories=[ - pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_IMPORT'], - pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_EXPORT'], - pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_RESOURCE'], - pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_DEBUG'], - pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_TLS'], - pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG'] - ]) - except Exception: - pass - - sections = list(getattr(pe, 'sections', [])) - section_ranges = [] - for section in sections: - start = max(0, int(section.get_PointerToRawData_adj())) - end = min(fsize, start + max(0, int(section.SizeOfRawData))) - section_ranges.append((start, max(start, end))) - resource_ranges = _collect_resource_ranges(pe, fsize) - overlay_offset = pe.get_overlay_data_start_offset() - if not overlay_offset or overlay_offset >= fsize: - overlay_offset = fsize - all_ranges = section_ranges + resource_ranges + [ - (int(overlay_offset), fsize) - ] - stream_features, range_entropies = _extract_stream_features( - file_path, fsize, all_ranges, self - ) - section_entropies = range_entropies[:len(section_ranges)] - resource_start = len(section_ranges) - resource_entropies = range_entropies[ - resource_start:resource_start + len(resource_ranges) - ] - overlay_entropy = range_entropies[-1] - - base['TrustSigned'] = 1.0 if self.signer.sign_verify(file_path) else 0.0 - base['FileSize'] = float(fsize) - - fh = pe.FILE_HEADER - base['Machine'] = getattr(fh, 'Machine', 0) - base['NumberOfSections'] = getattr(fh, 'NumberOfSections', 0) - base['TimeDateStamp'] = getattr(fh, 'TimeDateStamp', 0) - base['PointerToSymbolTable'] = getattr(fh, 'PointerToSymbolTable', 0) - base['NumberOfSymbols'] = getattr(fh, 'NumberOfSymbols', 0) - base['SizeOfOptionalHeader'] = getattr(fh, 'SizeOfOptionalHeader', 0) - base['Characteristics'] = getattr(fh, 'Characteristics', 0) - - curr_ts = datetime.datetime.utcnow().timestamp() - base['HasInvalidTimestamp'] = 1.0 if (base['TimeDateStamp'] < 631152000 or base['TimeDateStamp'] > curr_ts + 2592000) else 0.0 - base['FileTimeException'] = 1.0 if base['TimeDateStamp'] == 0 else 0.0 - base['Is64Bit'] = 1.0 if base['Machine'] in (0x8664, 0xAA64, 0x0200) else 0.0 - base['IsExe'] = 1.0 if pe.is_exe() else 0.0 - base['IsDll'] = 1.0 if pe.is_dll() else 0.0 - base['ExceptionCount'] = 0.0 - - if hasattr(pe, 'OPTIONAL_HEADER'): - op = pe.OPTIONAL_HEADER - fields = [ - 'Magic', 'MajorLinkerVersion', 'MinorLinkerVersion', 'SizeOfCode', - 'SizeOfInitializedData', 'SizeOfUninitializedData', 'AddressOfEntryPoint', - 'BaseOfCode', 'ImageBase', 'SectionAlignment', 'FileAlignment', - 'MajorOperatingSystemVersion', 'MinorOperatingSystemVersion', - 'MajorImageVersion', 'MinorImageVersion', 'MajorSubsystemVersion', - 'MinorSubsystemVersion', 'SizeOfImage', 'SizeOfHeaders', 'CheckSum', - 'Subsystem', 'DllCharacteristics', 'SizeOfStackReserve', 'SizeOfStackCommit', - 'SizeOfHeapReserve', 'SizeOfHeapCommit', 'LoaderFlags', 'NumberOfRvaAndSizes' - ] - for f_field in fields: - base[f_field] = getattr(op, f_field, 0) - - base['IsDriver'] = 1.0 if base.get('Subsystem') == 1 else 0.0 - - if hasattr(op, 'DATA_DIRECTORY'): - for i, directory in enumerate(op.DATA_DIRECTORY): - base[f'DataDirectory_{i}_Size'] = directory.Size - base[f'DataDirectory_{i}_VA'] = directory.VirtualAddress - - exc_idx = pefile.DIRECTORY_ENTRY.get('IMAGE_DIRECTORY_ENTRY_EXCEPTION', 3) - if len(op.DATA_DIRECTORY) > exc_idx and op.DATA_DIRECTORY[exc_idx].Size > 0: - base['ExceptionCount'] = op.DATA_DIRECTORY[exc_idx].Size // 12 if base['Machine'] in (0x8664, 0xAA64) else 0.0 - - res_flags_count = dict.fromkeys(self._CHAR_COUNT_KEYS.values(), 0.0) - res_flags_ent = dict.fromkeys(self._CHAR_ENT_KEYS.values(), 0.0) - res_sec_hash = dict.fromkeys(self._SECTION_HASH_KEYS, 0.0) - - if hasattr(pe, 'sections'): - base['SectionCount'] = len(pe.sections) - entropies = [] - raw_sizes = [] - v_sizes = [] - exec_sec = write_sec = read_sec = sec_exc = 0 - - for section_index, section in enumerate(pe.sections): - sec_name = section.Name.rstrip(b'\x00') - hash_val = zlib.crc32(sec_name) % 50 - res_sec_hash[self._SECTION_HASH_KEYS[hash_val]] += 1.0 - - s_entropy = ( - section_entropies[section_index] - if section_index < len(section_entropies) else 0.0 - ) - - entropies.append(s_entropy) - raw_sizes.append(section.SizeOfRawData) - v_sizes.append(section.Misc_VirtualSize) - - if section.Characteristics & 0x20000000: - exec_sec += 1 - if section.Characteristics & 0x80000000: - write_sec += 1 - if section.Characteristics & 0x40000000: - read_sec += 1 - if section.SizeOfRawData + section.PointerToRawData > fsize: - sec_exc = 1 - - for flag in self._CHAR_FLAGS: - if section.Characteristics & flag: - res_flags_count[self._CHAR_COUNT_KEYS[flag]] += 1.0 - res_flags_ent[self._CHAR_ENT_KEYS[flag]] += s_entropy - - base['SectionMaxEntropy'] = max(entropies) if entropies else 0.0 - base['SectionMinEntropy'] = min(entropies) if entropies else 0.0 - base['SectionMeanEntropy'] = sum(entropies) / len(entropies) if entropies else 0.0 - base['SectionMaxRawSize'] = max(raw_sizes) if raw_sizes else 0.0 - base['SectionMinRawSize'] = min(raw_sizes) if raw_sizes else 0.0 - base['SectionMeanRawSize'] = sum(raw_sizes) / len(raw_sizes) if raw_sizes else 0.0 - base['SectionMaxVSize'] = max(v_sizes) if v_sizes else 0.0 - base['SectionMinVSize'] = min(v_sizes) if v_sizes else 0.0 - base['SectionMeanVSize'] = sum(v_sizes) / len(v_sizes) if v_sizes else 0.0 - base['ExecutableSections'] = float(exec_sec) - base['WritableSections'] = float(write_sec) - base['ReadableSections'] = float(read_sec) - base['SectionException'] = float(sec_exc) - - for flag in self._CHAR_FLAGS: - cnt_key = self._CHAR_COUNT_KEYS[flag] - if res_flags_count[cnt_key] > 0: - res_flags_ent[self._CHAR_ENT_KEYS[flag]] /= res_flags_count[cnt_key] - - base.update(res_flags_count) - base.update(res_flags_ent) - base.update(res_sec_hash) - - string_keys = ['FileDescription', 'FileVersion', 'ProductName', 'ProductVersion', 'CompanyName', 'LegalCopyright', 'Comments', 'InternalName', 'LegalTrademarks', 'SpecialBuild', 'PrivateBuild'] - for key in string_keys: - base[f'{key}Length'] = 0.0 - - if hasattr(pe, 'FileInfo'): - for fileinfo_list in pe.FileInfo: - for fileinfo in fileinfo_list: - - if getattr(fileinfo, 'name', '') in ('StringFileInfo', b'StringFileInfo'): - for st in getattr(fileinfo, 'StringTable', []): - for key, val in st.entries.items(): - try: - k = key.decode('utf-8', 'ignore') if isinstance(key, bytes) else str(key) - if k in string_keys: - v = val.decode('utf-8', 'ignore') if isinstance(val, bytes) else str(val) - base[f'{k}Length'] = float(len(v)) - - except Exception: - continue - - if hasattr(pe, 'DIRECTORY_ENTRY_TLS') and hasattr(pe.DIRECTORY_ENTRY_TLS, 'struct'): - if getattr(pe.DIRECTORY_ENTRY_TLS.struct, 'AddressOfCallBacks', 0) != 0: - base['HasTlsCallbacks'] = 1.0 - - if hasattr(pe, 'VS_FIXEDFILEINFO') and len(pe.VS_FIXEDFILEINFO) > 0: - flags = getattr(pe.VS_FIXEDFILEINFO[0], 'FileFlags', 0) - base['IsDebug'] = 1.0 if flags & 0x1 else 0.0 - base['IsPreRelease'] = 1.0 if flags & 0x2 else 0.0 - base['IsPatched'] = 1.0 if flags & 0x4 else 0.0 - base['IsPrivateBuild'] = 1.0 if flags & 0x8 else 0.0 - base['IsSpecialBuild'] = 1.0 if flags & 0x20 else 0.0 - - if hasattr(pe, 'DIRECTORY_ENTRY_IMPORT'): - base['ImportCount'] = float(len(pe.DIRECTORY_ENTRY_IMPORT)) - func_count = 0 - - for entry in pe.DIRECTORY_ENTRY_IMPORT: - if getattr(entry, 'dll', None): - try: - dlls.add(entry.dll.decode('ascii', 'ignore').lower()) - except Exception: - pass - - for imp in getattr(entry, 'imports', []): - func_count += 1 - if getattr(imp, 'name', None): - try: - apis.add(imp.name.decode('ascii', 'ignore')) - except Exception: - pass - - base['ImportFunctionCount'] = float(func_count) - else: - base['ImportCount'] = 0.0 - base['ImportFunctionCount'] = 0.0 - - if hasattr(pe, 'DIRECTORY_ENTRY_EXPORT') and hasattr(pe.DIRECTORY_ENTRY_EXPORT, 'symbols'): - base['ExportCount'] = float(len(pe.DIRECTORY_ENTRY_EXPORT.symbols)) - else: - base['ExportCount'] = 0.0 - - if hasattr(pe, 'DIRECTORY_ENTRY_RESOURCE'): - icon_count = 0 - for entry in pe.DIRECTORY_ENTRY_RESOURCE.entries: - if getattr(entry, 'id', None) == 3 and hasattr(entry, 'directory'): - icon_count += len(getattr(entry.directory, 'entries', [])) - - base['IconCount'] = float(icon_count) - else: - base['IconCount'] = 0.0 - - if hasattr(pe, 'DIRECTORY_ENTRY_DEBUG'): - base['DebugCount'] = float(len(pe.DIRECTORY_ENTRY_DEBUG)) - else: - base['DebugCount'] = 0.0 - - base.update(stream_features) - overlay_size = fsize - int(overlay_offset) - base.update({ - "HasOverlay": 1.0 if overlay_size else 0.0, - "OverlaySize": float(overlay_size), - "OverlayRatio": float(overlay_size) / float(fsize) if overlay_size else 0.0, - "OverlayEntropy": overlay_entropy if overlay_size else 0.0 - }) - base.update(self._extract_rich_header(pe)) - base.update(self._extract_ep_anomalies(pe)) - base.update(self._extract_advanced_resources(pe, resource_entropies)) - base.update(self._extract_load_config(pe)) - base.update(self._extract_security_directory(pe, file_path, fsize)) - - for k in base: - base[k] = self._safe_float(base[k]) - - return {"Base": base, "DLLs": list(dlls), "APIs": list(apis)} - - except Exception: - return None - finally: - if pe: - pe.close() - -#################################################################################################### - - def pe_scan(self, file_path, enhanced_mode=False): - if not self.model or not self.feature_order: - return False, False - - try: - raw_data = self.extract_features(file_path) - if not raw_data: - return False, False - - vec = numpy.zeros((1, len(self.feature_order)), dtype=numpy.float32) - base = raw_data.get('Base', {}) - dlls = set(raw_data.get('DLLs', [])) - apis = set(raw_data.get('APIs', [])) - feat_map = {feat: i for i, feat in enumerate(self.feature_order)} - - for k, v in base.items(): - if k in feat_map: - vec[0, feat_map[k]] = v - - for d in dlls: - old_fn = f"Dll_{d}" - if old_fn in feat_map: - vec[0, feat_map[old_fn]] = 1.0 - - h = zlib.crc32(d.encode('utf-8', 'ignore')) % self.dll_hash_dim - new_fn = f"DllHash_{h:0{self.dll_hash_pad}d}" - if new_fn in feat_map: - vec[0, feat_map[new_fn]] += 1.0 - - for a in apis: - old_fn = f"Api_{a}" - if old_fn in feat_map: - vec[0, feat_map[old_fn]] = 1.0 - - h = zlib.crc32(a.encode('utf-8', 'ignore')) % self.api_hash_dim - new_fn = f"ApiHash_{h:0{self.api_hash_pad}d}" - if new_fn in feat_map: - vec[0, feat_map[new_fn]] += 1.0 - - outputs = self.model.run(None, {self.input_name: vec}) - - del raw_data, vec - result = outputs[1] - prob = 0.0 - - if isinstance(result, list) and len(result) > 0: - prob_dict = result[0] - if hasattr(prob_dict, 'get'): - prob = float(prob_dict.get(1, prob_dict.get('1', 0.0))) - - elif isinstance(result, numpy.ndarray): - if result.ndim == 2 and result.shape[1] > 1: - prob = float(result[0][1]) - - score = int(prob * 100) - if score >= 80: - return f"Malware:WinPE/General.{score}!ml", score - elif enhanced_mode and score >= 50: - return f"Suspicious:WinPE/General.{score}!ml", score - - return False, False - except Exception: - return False, False - -#################################################################################################### - -class cloud_scanner: - def __init__(self): - self.api_host = None - self.api_key = None - self.timeout = 30 - self.lock = threading.RLock() - self.local = threading.local() - - def _get_session(self, api_host, api_key): - if not hasattr(self.local, 'session') or getattr(self.local, 'host', None) != api_host: - self.local.session = requests.Session() - self.local.session.headers.update({"X-API-Key": api_key, "User-Agent": "PYAS-Engine/1.1"}) - self.local.host = api_host - - return self.local.session - - def _request(self, method, endpoint, api_host, api_key, **kwargs): - session = self._get_session(api_host, api_key) - try: - r = session.request(method, f"{api_host}{endpoint}", timeout=self.timeout, **kwargs) - if r.status_code == 200: - return r - - except Exception: - pass - return None - -#################################################################################################### - - def rescan(self, sha256, api_host, api_key): - r = self._request("POST", f"/api/rescan/{sha256}", api_host, api_key) - return r is not None and r.json().get('status') == 'success' - - def upload_file(self, file_path, api_host, api_key, chunk_size=4194304, need_rescan=False, max_retries=3, file_hash=None): - try: - sha256 = file_hash - if not sha256: - return False, None - - status_req = self._request("GET", f"/api/processing_status/{sha256}", api_host, api_key) - if status_req: - current_status = status_req.json().get('status') - - if current_status == 'done': - if need_rescan: - self.rescan(sha256, api_host, api_key) - return True, sha256 - - elif current_status in ('queued', 'processing'): - return True, sha256 - - file_size = os.path.getsize(file_path) - if file_size > 104857600: - return False, sha256 - - total_chunks = max(1, (file_size + chunk_size - 1) // chunk_size) - upload_id = os.urandom(16).hex() - - with open(file_path, 'rb') as f: - for i in range(total_chunks): - chunk_data = f.read(chunk_size) - headers = {"X-Chunk-Index": str(i), "X-Total-Chunks": str(total_chunks), "X-Upload-ID": upload_id} - - chunk_success = False - for attempt in range(max_retries): - r = self._request("POST", "/api/upload", api_host, api_key, files={'file': (os.path.basename(file_path), chunk_data)}, headers=headers) - if r: - if i == total_chunks - 1: - try: - resp = r.json() - if 'url' in resp: - sha256 = resp.get('url', '').split('/')[-1] - except Exception: - pass - - chunk_success = True - break - time.sleep(2 ** attempt) - - if not chunk_success: - return False, sha256 - - return True, sha256 - except Exception: - return False, None - -#################################################################################################### - - def get_result(self, sha256, api_host, api_key, max_retries=6, interval=10): - try: - if not sha256: - return False - - is_done = False - for _ in range(max_retries): - r = self._request("GET", f"/api/processing_status/{sha256}", api_host, api_key) - if r: - st = r.json().get('status', 'error') - if st == 'done': - is_done = True - break - - if st in ['error', 'failed']: - return False - - time.sleep(interval) - - if not is_done: - return False - - r = self._request("GET", f"/api/report/{sha256}", api_host, api_key) - if r: - data = r.json().get('data', {}) - metadata = data.get('metadata', {}) - label = metadata.get('label', 'Unsupport') - score = metadata.get('score', 0) - sims = data.get('similar', []) - - is_malicious = 'General' in label - sim_malicious_count = 0 - valid_sim_count = 0 - - for s in sims: - if s.get('similarity', 0) > 80: - valid_sim_count += 1 - if "General" in s.get('label', ''): - sim_malicious_count += 1 - - if is_malicious and (valid_sim_count == 0 or sim_malicious_count == valid_sim_count): - return f"Malware:WinPE/General.{score}!cl" - - except Exception: - pass - return False +from PYAS_Signature import SignatureScanner as sign_scanner +from PYAS_Rules import RuleScanner as rule_scanner +from PYAS_PE import PEScanner as pe_scanner +from PYAS_Cloud import CloudScanner as cloud_scanner + +__all__ = ["sign_scanner", "rule_scanner", "pe_scanner", "cloud_scanner"] +from PYAS_WinAPI import GUID, WINTRUST_FILE_INFO, WINTRUST_DATA_UNION, WINTRUST_DATA +from PYAS_Features import ( + _entropy_from_counts, + _bounded_entropy, + _accumulate_string_chunk, + _accumulate_entropy_windows, + _collect_resource_ranges, + _extract_stream_features, +) diff --git a/PYAS_Features.py b/PYAS_Features.py new file mode 100644 index 0000000..5a53adc --- /dev/null +++ b/PYAS_Features.py @@ -0,0 +1,949 @@ +from PYAS_Diagnostics import log_exception +import os +import math +import zlib +import numpy +import pefile +import datetime + + +def _entropy_from_counts(counts, size): + if size <= 0: + return 0.0 + + nonzero = counts[counts > 0] + return float(numpy.log2(size) - numpy.sum(nonzero * numpy.log2(nonzero)) / size) + + +def _bounded_entropy(data, batch_size=1048576): + size = len(data) + + if size == 0: + return 0.0 + + counts = numpy.zeros(256, dtype=numpy.int64) + view = memoryview(data) + + try: + for start in range(0, size, batch_size): + chunk = numpy.frombuffer(view[start : min(start + batch_size, size)], dtype=numpy.uint8) + counts += numpy.bincount(chunk, minlength=256) + finally: + view.release() + + return _entropy_from_counts(counts, size) + + +def _accumulate_string_chunk(chunk, array, state, pattern): + separators = (array < 0x20) | (array > 0x7E) + + if not separators.any(): + state["pending_counts"] += numpy.bincount(array, minlength=256) + state["pending_length"] += len(array) + return + + first_separator = int(numpy.argmax(separators)) + last_separator = len(separators) - 1 - int(numpy.argmax(separators[::-1])) + + if first_separator: + state["pending_counts"] += numpy.bincount(array[:first_separator], minlength=256) + state["pending_length"] += first_separator + + if state["pending_length"] >= 5: + state["count"] += 1 + state["total_length"] += state["pending_length"] + state["counts"] += state["pending_counts"] + + state["pending_length"] = 0 + state["pending_counts"].fill(0) + + middle_start = first_separator + 1 + middle_end = last_separator + + if middle_end > middle_start: + matches = pattern.findall(chunk[middle_start:middle_end]) + + if matches: + combined = b"".join(matches) + state["count"] += len(matches) + state["total_length"] += len(combined) + state["counts"] += numpy.bincount( + numpy.frombuffer(combined, dtype=numpy.uint8), minlength=256 + ) + + trailing_start = last_separator + 1 + + if trailing_start < len(array): + trailing = array[trailing_start:] + state["pending_counts"] += numpy.bincount(trailing, minlength=256) + state["pending_length"] = len(trailing) + + +def _accumulate_entropy_windows(data, entropy_histogram, c_log_c_table): + block_count = len(data) // 1024 + + if block_count < 2: + return data + + complete_size = block_count * 1024 + blocks = numpy.frombuffer(data, dtype=numpy.uint8, count=complete_size).reshape( + block_count, 1024 + ) + encoded = blocks.astype(numpy.int32) + encoded += numpy.arange(block_count, dtype=numpy.int32)[:, None] * 256 + block_counts = numpy.bincount(encoded.ravel(), minlength=block_count * 256).reshape( + block_count, 256 + ) + window_counts = block_counts[:-1] + block_counts[1:] + sum_c_log_c = numpy.sum(c_log_c_table[window_counts], axis=1) + entropies = 11.0 - sum_c_log_c / 2048.0 + entropy_bins = (entropies * 2.0).astype(numpy.int32) + numpy.clip(entropy_bins, 0, 15, out=entropy_bins) + byte_bins = window_counts.reshape(block_count - 1, 16, 16).sum(axis=2) + flat_indices = (entropy_bins[:, None] * 16 + numpy.arange(16)).ravel() + numpy.add.at(entropy_histogram, flat_indices, byte_bins.ravel()) + return data[(block_count - 1) * 1024 :] + + +def _collect_resource_ranges(pe, file_size): + ranges = [] + + if not hasattr(pe, "DIRECTORY_ENTRY_RESOURCE"): + return ranges + + for entry_type in pe.DIRECTORY_ENTRY_RESOURCE.entries: + for entry_id in getattr(getattr(entry_type, "directory", None), "entries", []): + for entry_lang in getattr(getattr(entry_id, "directory", None), "entries", []): + if not hasattr(entry_lang, "data"): + continue + + try: + struct = entry_lang.data.struct + offset = int(pe.get_offset_from_rva(int(struct.OffsetToData))) + size = int(struct.Size) + start = max(0, offset) + end = min(file_size, start + max(0, size)) + ranges.append((start, max(start, end))) + except Exception: + log_exception("PYAS_Features._collect_resource_ranges:109") + ranges.append((0, 0)) + + return ranges + + +def _extract_stream_features(file_path, file_size, ranges, extractor): + normalized = [] + + for start, end in ranges: + start = max(0, min(file_size, int(start))) + end = max(start, min(file_size, int(end))) + normalized.append((start, end)) + + byte_counts = numpy.zeros(256, dtype=numpy.int64) + range_counts = [numpy.zeros(256, dtype=numpy.int64) for _ in normalized] + entropy_histogram = numpy.zeros(256, dtype=numpy.float64) + string_state = { + "count": 0, + "total_length": 0, + "counts": numpy.zeros(256, dtype=numpy.int64), + "pending_length": 0, + "pending_counts": numpy.zeros(256, dtype=numpy.int64), + } + histogram_tail = b"" + position = 0 + + with open(file_path, "rb") as source: + while True: + chunk = source.read(4194304) + + if not chunk: + break + + array = numpy.frombuffer(chunk, dtype=numpy.uint8) + _accumulate_string_chunk(chunk, array, string_state, extractor._STRING_PATTERN) + histogram_tail = _accumulate_entropy_windows( + histogram_tail + chunk, + entropy_histogram, + extractor._C_LOG_C_TABLE, + ) + + chunk_end = position + len(chunk) + boundaries = {0, len(chunk)} + + for start, end in normalized: + if position < start < chunk_end: + boundaries.add(start - position) + + if position < end < chunk_end: + boundaries.add(end - position) + + ordered = sorted(boundaries) + + for index in range(len(ordered) - 1): + local_start = ordered[index] + local_end = ordered[index + 1] + + if local_end <= local_start: + continue + + counts = numpy.bincount(array[local_start:local_end], minlength=256) + byte_counts += counts + absolute_start = position + local_start + absolute_end = position + local_end + + for range_index, (start, end) in enumerate(normalized): + if start <= absolute_start and absolute_end <= end: + range_counts[range_index] += counts + + position = chunk_end + + if string_state["pending_length"] >= 5: + string_state["count"] += 1 + string_state["total_length"] += string_state["pending_length"] + string_state["counts"] += string_state["pending_counts"] + + string_count = string_state["count"] + total_string_length = string_state["total_length"] + result = { + "StringCount": float(string_count), + "StringMeanLength": (float(total_string_length) / string_count if string_count else 0.0), + "StringEntropy": _entropy_from_counts(string_state["counts"], total_string_length), + } + result.update(dict.fromkeys(extractor._STRING_KEYS, 0.0)) + + if total_string_length: + valid_counts = string_state["counts"][0x20:0x7F] + result.update( + dict( + zip( + extractor._STRING_KEYS, + valid_counts.astype(numpy.float64) / total_string_length, + ) + ) + ) + + result["FileEntropy"] = _entropy_from_counts(byte_counts, file_size) + result.update( + dict( + zip( + extractor._BYTE_HIST_KEYS, + byte_counts.astype(numpy.float64) / file_size, + ) + ) + ) + entropy_sum = numpy.sum(entropy_histogram) + + if entropy_sum: + entropy_histogram /= entropy_sum + + result.update(dict(zip(extractor._BYTE_ENT_KEYS, entropy_histogram))) + + range_entropies = [ + _entropy_from_counts(counts, end - start) + for counts, (start, end) in zip(range_counts, normalized) + ] + return result, range_entropies + + +class PEFeatureMixin: + def _safe_float(self, val): + try: + f = float(val) + + if math.isinf(f) or math.isnan(f): + return 0.0 + + return f + except Exception: + log_exception("PYAS_Features.PEFeatureMixin._safe_float:221") + return 0.0 + + def _base_defaults(self): + base = dict.fromkeys(self._OPTIONAL_HEADER_FIELDS, 0.0) + base.update(self._SECTION_DEFAULTS) + base.update(self._CONDITIONAL_DEFAULTS) + + for i in range(16): + base[f"DataDirectory_{i}_Size"] = 0.0 + base[f"DataDirectory_{i}_VA"] = 0.0 + + return base + + def _calc_entropy(self, data): + return _bounded_entropy(data) + + def _extract_strings(self, file_bytes): + res = {"StringCount": 0.0, "StringMeanLength": 0.0, "StringEntropy": 0.0} + res.update(dict.fromkeys(self._STRING_KEYS, 0.0)) + + if len(file_bytes) <= self._STRING_FAST_PATH_BYTES: + matches = self._STRING_PATTERN.findall(file_bytes) + + if not matches: + return res + + combined_data = b"".join(matches) + count = len(matches) + total_len = len(combined_data) + global_counts = numpy.bincount( + numpy.frombuffer(combined_data, dtype=numpy.uint8), minlength=256 + ) + else: + count = 0 + total_len = 0 + global_counts = numpy.zeros(256, dtype=numpy.int64) + mv = memoryview(file_bytes) + pos = 0 + size = len(file_bytes) + + while pos < size: + chunk_end = min(pos + self._STRING_BATCH_BYTES, size) + chunk_arr = numpy.frombuffer(mv[pos:chunk_end], dtype=numpy.uint8) + separators = (chunk_arr < 0x20) | (chunk_arr > 0x7E) + + if separators.any(): + last_separator = len(separators) - 1 - int(numpy.argmax(separators[::-1])) + segment_end = pos + last_separator + 1 + matches = self._STRING_PATTERN.findall(mv[pos:segment_end]) + + if matches: + combined_data = b"".join(matches) + count += len(matches) + total_len += len(combined_data) + global_counts += numpy.bincount( + numpy.frombuffer(combined_data, dtype=numpy.uint8), minlength=256 + ) + + pos = segment_end + continue + + pending_len = 0 + pending_counts = numpy.zeros(256, dtype=numpy.int64) + + while pos < size: + chunk_end = min(pos + self._STRING_BATCH_BYTES, size) + chunk_arr = numpy.frombuffer(mv[pos:chunk_end], dtype=numpy.uint8) + separators = (chunk_arr < 0x20) | (chunk_arr > 0x7E) + + if separators.any(): + first_separator = int(numpy.argmax(separators)) + + if first_separator: + pending_counts += numpy.bincount( + chunk_arr[:first_separator], minlength=256 + ) + pending_len += first_separator + + pos += first_separator + 1 + break + + pending_counts += numpy.bincount(chunk_arr, minlength=256) + pending_len += len(chunk_arr) + pos = chunk_end + + if pending_len >= 5: + count += 1 + total_len += pending_len + global_counts += pending_counts + + if count == 0 or total_len == 0: + return res + + res["StringCount"] = float(count) + res["StringMeanLength"] = float(total_len) / count + + counts_nonzero = global_counts[global_counts > 0] + + if len(counts_nonzero) > 0: + res["StringEntropy"] = float( + numpy.log2(total_len) + - numpy.sum(counts_nonzero * numpy.log2(counts_nonzero)) / total_len + ) + + valid_chars = global_counts[0x20:0x7F] + res.update(dict(zip(self._STRING_KEYS, valid_chars.astype(numpy.float64) / total_len))) + + return res + + def _extract_histograms(self, file_bytes): + arr = numpy.frombuffer(file_bytes, dtype=numpy.uint8) + sz = len(arr) + res = {} + + if sz == 0: + res["FileEntropy"] = 0.0 + res.update(dict.fromkeys(self._BYTE_HIST_KEYS, 0.0)) + res.update(dict.fromkeys(self._BYTE_ENT_KEYS, 0.0)) + return res + + byte_counts = numpy.zeros(256, dtype=numpy.int64) + ent_hist = numpy.zeros(256, dtype=numpy.float64) + window = 2048 + step = 1024 + + if sz < window: + byte_counts += numpy.bincount(arr, minlength=256) + p = byte_counts[byte_counts > 0] / float(sz) + entropy = float(-numpy.sum(p * numpy.log2(p))) + ent_bin = min(int(entropy * 2.0), 15) + byte_bins = byte_counts.reshape(16, 16).sum(axis=1) + idx_start = ent_bin * 16 + ent_hist[idx_start : idx_start + 16] += byte_bins + else: + num_windows = (sz - window) // step + 1 + + for w_start in range(0, num_windows, self._HISTOGRAM_BATCH_WINDOWS): + w_end = min(w_start + self._HISTOGRAM_BATCH_WINDOWS, num_windows) + b_num = w_end - w_start + + start_byte = w_start * step + end_byte = start_byte + b_num * step + step + encoded = arr[start_byte:end_byte].astype(numpy.int32).reshape(b_num + 1, step) + encoded += numpy.arange(b_num + 1, dtype=numpy.int32)[:, None] * 256 + + chunk_counts = numpy.bincount(encoded.ravel(), minlength=(b_num + 1) * 256).reshape( + b_num + 1, 256 + ) + del encoded + + byte_counts += numpy.sum(chunk_counts[:-1], axis=0) + + if w_end == num_windows: + byte_counts += chunk_counts[-1] + + window_counts = chunk_counts[:-1] + chunk_counts[1:] + sum_c_log_c = numpy.sum(self._C_LOG_C_TABLE[window_counts], axis=1) + entropies = 11.0 - sum_c_log_c / 2048.0 + + ent_bins = (entropies * 2.0).astype(numpy.int32) + numpy.clip(ent_bins, 0, 15, out=ent_bins) + + byte_bins = window_counts.reshape(b_num, 16, 16).sum(axis=2) + flat_ent_bins = (ent_bins[:, None] * 16 + numpy.arange(16)).ravel() + numpy.add.at(ent_hist, flat_ent_bins, byte_bins.ravel()) + + covered_size = (num_windows + 1) * step + + if covered_size < sz: + byte_counts += numpy.bincount(arr[covered_size:], minlength=256) + + counts_nonzero = byte_counts[byte_counts > 0] + res["FileEntropy"] = float( + numpy.log2(sz) - numpy.sum(counts_nonzero * numpy.log2(counts_nonzero)) / sz + ) + + byte_hist = byte_counts.astype(numpy.float64) / sz + res.update(dict(zip(self._BYTE_HIST_KEYS, byte_hist))) + + ent_sum = numpy.sum(ent_hist) + + if ent_sum > 0: + ent_hist /= ent_sum + + res.update(dict(zip(self._BYTE_ENT_KEYS, ent_hist))) + return res + + def _extract_overlay_features(self, pe, fsize): + overlay_offset = pe.get_overlay_data_start_offset() + + if not overlay_offset or overlay_offset >= fsize: + return { + "HasOverlay": 0.0, + "OverlaySize": 0.0, + "OverlayRatio": 0.0, + "OverlayEntropy": 0.0, + } + + overlay_data = pe.get_overlay() + + if not overlay_data: + return { + "HasOverlay": 0.0, + "OverlaySize": 0.0, + "OverlayRatio": 0.0, + "OverlayEntropy": 0.0, + } + + sz = len(overlay_data) + return { + "HasOverlay": 1.0, + "OverlaySize": float(sz), + "OverlayRatio": float(sz) / float(fsize), + "OverlayEntropy": self._calc_entropy(overlay_data), + } + + def _extract_rich_header(self, pe): + if not hasattr(pe, "RICH_HEADER") or not pe.RICH_HEADER: + return {"HasRichHeader": 0.0, "RichHeaderCount": 0.0} + + return { + "HasRichHeader": 1.0, + "RichHeaderCount": ( + float(len(pe.RICH_HEADER.values) // 2) if hasattr(pe.RICH_HEADER, "values") else 0.0 + ), + } + + def _extract_ep_anomalies(self, pe): + if not hasattr(pe, "OPTIONAL_HEADER") or not hasattr(pe, "sections") or not pe.sections: + return { + "EntryPointSectionIndex": -1.0, + "EntryPointInExecutable": 0.0, + "EntryPointInLastSection": 0.0, + } + + ep = pe.OPTIONAL_HEADER.AddressOfEntryPoint + ep_section_idx = -1 + ep_in_exec = 0.0 + + for idx, sec in enumerate(pe.sections): + if sec.VirtualAddress <= ep < (sec.VirtualAddress + sec.Misc_VirtualSize): + ep_section_idx = idx + + if sec.Characteristics & 0x20000000: + ep_in_exec = 1.0 + + break + + return { + "EntryPointSectionIndex": float(ep_section_idx), + "EntryPointInExecutable": ep_in_exec, + "EntryPointInLastSection": 1.0 if ep_section_idx == len(pe.sections) - 1 else 0.0, + } + + def _extract_advanced_resources(self, pe, range_entropies=None): + res_data = { + "ResourceMaxEntropy": 0.0, + "ResourceMinEntropy": 0.0, + "ResourceMeanEntropy": 0.0, + "ResourceLangCount": 0.0, + "ResourceRCDataCount": 0.0, + } + + if not hasattr(pe, "DIRECTORY_ENTRY_RESOURCE"): + return res_data + + entropies = [] + entropy_iter = iter(range_entropies or ()) + langs = set() + rcdata_count = 0 + + for entry_type in pe.DIRECTORY_ENTRY_RESOURCE.entries: + if hasattr(entry_type, "id") and entry_type.id == 10: + rcdata_count += len(getattr(entry_type.directory, "entries", [])) + + if not hasattr(entry_type, "directory"): + continue + + for entry_id in entry_type.directory.entries: + if not hasattr(entry_id, "directory"): + continue + + for entry_lang in entry_id.directory.entries: + if hasattr(entry_lang, "id"): + langs.add(entry_lang.id) + + if hasattr(entry_lang, "data"): + try: + entropies.append(float(next(entropy_iter))) + except (StopIteration, TypeError, ValueError): + entropies.append(0.0) + + res_data["ResourceLangCount"] = float(len(langs)) + res_data["ResourceRCDataCount"] = float(rcdata_count) + + if entropies: + res_data["ResourceMaxEntropy"] = max(entropies) + res_data["ResourceMinEntropy"] = min(entropies) + res_data["ResourceMeanEntropy"] = sum(entropies) / len(entropies) + + return res_data + + def _extract_load_config(self, pe): + cfg = {"HasLoadConfig": 0.0, "HasCFG": 0.0, "HasSEHTable": 0.0} + + if hasattr(pe, "DIRECTORY_ENTRY_LOAD_CONFIG"): + cfg["HasLoadConfig"] = 1.0 + struct = pe.DIRECTORY_ENTRY_LOAD_CONFIG.struct + + if getattr(struct, "GuardCFFunctionTable", 0) != 0: + cfg["HasCFG"] = 1.0 + + if getattr(struct, "SEHandlerTable", 0) != 0: + cfg["HasSEHTable"] = 1.0 + + return cfg + + def _extract_security_directory(self, pe, file_path, fsize): + res = {"HasSignature": 0.0, "SignatureCount": 0.0} + + if not hasattr(pe, "OPTIONAL_HEADER") or not hasattr(pe.OPTIONAL_HEADER, "DATA_DIRECTORY"): + return res + + directories = pe.OPTIONAL_HEADER.DATA_DIRECTORY + + if len(directories) <= 4: + return res + + sec_dir = directories[4] + offset = int(sec_dir.VirtualAddress) + size = int(sec_dir.Size) + + if offset <= 0 or size <= 0 or offset + size > fsize: + return res + + res["HasSignature"] = 1.0 + count = 0 + curr = offset + end = offset + size + + with open(file_path, "rb") as source: + while curr + 8 <= end: + source.seek(curr) + header = source.read(4) + + if len(header) != 4: + break + + length = int.from_bytes(header, "little") + + if length < 8 or curr + length > end: + break + + count += 1 + curr += (length + 7) & ~7 + + res["SignatureCount"] = float(count) + return res + + def extract_features(self, file_path): + pe = None + + try: + fsize = os.path.getsize(file_path) + + if fsize == 0 or fsize > 4294967296: + return None + + base = self._base_defaults() + dlls = set() + apis = set() + pe = pefile.PE(name=file_path, fast_load=True) + + try: + pe.parse_rich_header() + pe.parse_data_directories( + directories=[ + pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_IMPORT"], + pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_EXPORT"], + pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_RESOURCE"], + pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_DEBUG"], + pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_TLS"], + pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG"], + ] + ) + except Exception: + log_exception("PYAS_Features.PEFeatureMixin.extract_features:553") + pass + + sections = list(getattr(pe, "sections", [])) + section_ranges = [] + + for section in sections: + start = max(0, int(section.get_PointerToRawData_adj())) + end = min(fsize, start + max(0, int(section.SizeOfRawData))) + section_ranges.append((start, max(start, end))) + + resource_ranges = _collect_resource_ranges(pe, fsize) + overlay_offset = pe.get_overlay_data_start_offset() + + if not overlay_offset or overlay_offset >= fsize: + overlay_offset = fsize + + all_ranges = section_ranges + resource_ranges + [(int(overlay_offset), fsize)] + stream_features, range_entropies = _extract_stream_features( + file_path, fsize, all_ranges, self + ) + section_entropies = range_entropies[: len(section_ranges)] + resource_start = len(section_ranges) + resource_entropies = range_entropies[ + resource_start : resource_start + len(resource_ranges) + ] + overlay_entropy = range_entropies[-1] + + base["TrustSigned"] = 1.0 if self.signer.sign_verify(file_path) else 0.0 + base["FileSize"] = float(fsize) + + fh = pe.FILE_HEADER + base["Machine"] = getattr(fh, "Machine", 0) + base["NumberOfSections"] = getattr(fh, "NumberOfSections", 0) + base["TimeDateStamp"] = getattr(fh, "TimeDateStamp", 0) + base["PointerToSymbolTable"] = getattr(fh, "PointerToSymbolTable", 0) + base["NumberOfSymbols"] = getattr(fh, "NumberOfSymbols", 0) + base["SizeOfOptionalHeader"] = getattr(fh, "SizeOfOptionalHeader", 0) + base["Characteristics"] = getattr(fh, "Characteristics", 0) + + curr_ts = datetime.datetime.utcnow().timestamp() + base["HasInvalidTimestamp"] = ( + 1.0 + if (base["TimeDateStamp"] < 631152000 or base["TimeDateStamp"] > curr_ts + 2592000) + else 0.0 + ) + base["FileTimeException"] = 1.0 if base["TimeDateStamp"] == 0 else 0.0 + base["Is64Bit"] = 1.0 if base["Machine"] in (0x8664, 0xAA64, 0x0200) else 0.0 + base["IsExe"] = 1.0 if pe.is_exe() else 0.0 + base["IsDll"] = 1.0 if pe.is_dll() else 0.0 + base["ExceptionCount"] = 0.0 + + if hasattr(pe, "OPTIONAL_HEADER"): + op = pe.OPTIONAL_HEADER + fields = [ + "Magic", + "MajorLinkerVersion", + "MinorLinkerVersion", + "SizeOfCode", + "SizeOfInitializedData", + "SizeOfUninitializedData", + "AddressOfEntryPoint", + "BaseOfCode", + "ImageBase", + "SectionAlignment", + "FileAlignment", + "MajorOperatingSystemVersion", + "MinorOperatingSystemVersion", + "MajorImageVersion", + "MinorImageVersion", + "MajorSubsystemVersion", + "MinorSubsystemVersion", + "SizeOfImage", + "SizeOfHeaders", + "CheckSum", + "Subsystem", + "DllCharacteristics", + "SizeOfStackReserve", + "SizeOfStackCommit", + "SizeOfHeapReserve", + "SizeOfHeapCommit", + "LoaderFlags", + "NumberOfRvaAndSizes", + ] + + for f_field in fields: + base[f_field] = getattr(op, f_field, 0) + + base["IsDriver"] = 1.0 if base.get("Subsystem") == 1 else 0.0 + + if hasattr(op, "DATA_DIRECTORY"): + for i, directory in enumerate(op.DATA_DIRECTORY): + base[f"DataDirectory_{i}_Size"] = directory.Size + base[f"DataDirectory_{i}_VA"] = directory.VirtualAddress + + exc_idx = pefile.DIRECTORY_ENTRY.get("IMAGE_DIRECTORY_ENTRY_EXCEPTION", 3) + + if len(op.DATA_DIRECTORY) > exc_idx and op.DATA_DIRECTORY[exc_idx].Size > 0: + base["ExceptionCount"] = ( + op.DATA_DIRECTORY[exc_idx].Size // 12 + if base["Machine"] in (0x8664, 0xAA64) + else 0.0 + ) + + res_flags_count = dict.fromkeys(self._CHAR_COUNT_KEYS.values(), 0.0) + res_flags_ent = dict.fromkeys(self._CHAR_ENT_KEYS.values(), 0.0) + res_sec_hash = dict.fromkeys(self._SECTION_HASH_KEYS, 0.0) + + if hasattr(pe, "sections"): + base["SectionCount"] = len(pe.sections) + entropies = [] + raw_sizes = [] + v_sizes = [] + exec_sec = write_sec = read_sec = sec_exc = 0 + + for section_index, section in enumerate(pe.sections): + sec_name = section.Name.rstrip(b"\x00") + hash_val = zlib.crc32(sec_name) % 50 + res_sec_hash[self._SECTION_HASH_KEYS[hash_val]] += 1.0 + + s_entropy = ( + section_entropies[section_index] + if section_index < len(section_entropies) + else 0.0 + ) + + entropies.append(s_entropy) + raw_sizes.append(section.SizeOfRawData) + v_sizes.append(section.Misc_VirtualSize) + + if section.Characteristics & 0x20000000: + exec_sec += 1 + + if section.Characteristics & 0x80000000: + write_sec += 1 + + if section.Characteristics & 0x40000000: + read_sec += 1 + + if section.SizeOfRawData + section.PointerToRawData > fsize: + sec_exc = 1 + + for flag in self._CHAR_FLAGS: + if section.Characteristics & flag: + res_flags_count[self._CHAR_COUNT_KEYS[flag]] += 1.0 + res_flags_ent[self._CHAR_ENT_KEYS[flag]] += s_entropy + + base["SectionMaxEntropy"] = max(entropies) if entropies else 0.0 + base["SectionMinEntropy"] = min(entropies) if entropies else 0.0 + base["SectionMeanEntropy"] = sum(entropies) / len(entropies) if entropies else 0.0 + base["SectionMaxRawSize"] = max(raw_sizes) if raw_sizes else 0.0 + base["SectionMinRawSize"] = min(raw_sizes) if raw_sizes else 0.0 + base["SectionMeanRawSize"] = sum(raw_sizes) / len(raw_sizes) if raw_sizes else 0.0 + base["SectionMaxVSize"] = max(v_sizes) if v_sizes else 0.0 + base["SectionMinVSize"] = min(v_sizes) if v_sizes else 0.0 + base["SectionMeanVSize"] = sum(v_sizes) / len(v_sizes) if v_sizes else 0.0 + base["ExecutableSections"] = float(exec_sec) + base["WritableSections"] = float(write_sec) + base["ReadableSections"] = float(read_sec) + base["SectionException"] = float(sec_exc) + + for flag in self._CHAR_FLAGS: + cnt_key = self._CHAR_COUNT_KEYS[flag] + + if res_flags_count[cnt_key] > 0: + res_flags_ent[self._CHAR_ENT_KEYS[flag]] /= res_flags_count[cnt_key] + + base.update(res_flags_count) + base.update(res_flags_ent) + base.update(res_sec_hash) + + string_keys = [ + "FileDescription", + "FileVersion", + "ProductName", + "ProductVersion", + "CompanyName", + "LegalCopyright", + "Comments", + "InternalName", + "LegalTrademarks", + "SpecialBuild", + "PrivateBuild", + ] + + for key in string_keys: + base[f"{key}Length"] = 0.0 + + if hasattr(pe, "FileInfo"): + for fileinfo_list in pe.FileInfo: + for fileinfo in fileinfo_list: + + if getattr(fileinfo, "name", "") in ("StringFileInfo", b"StringFileInfo"): + for st in getattr(fileinfo, "StringTable", []): + for key, val in st.entries.items(): + try: + k = ( + key.decode("utf-8", "ignore") + if isinstance(key, bytes) + else str(key) + ) + + if k in string_keys: + v = ( + val.decode("utf-8", "ignore") + if isinstance(val, bytes) + else str(val) + ) + base[f"{k}Length"] = float(len(v)) + + except Exception: + log_exception( + "PYAS_Features.PEFeatureMixin.extract_features:704" + ) + continue + + if hasattr(pe, "DIRECTORY_ENTRY_TLS") and hasattr(pe.DIRECTORY_ENTRY_TLS, "struct"): + if getattr(pe.DIRECTORY_ENTRY_TLS.struct, "AddressOfCallBacks", 0) != 0: + base["HasTlsCallbacks"] = 1.0 + + if hasattr(pe, "VS_FIXEDFILEINFO") and len(pe.VS_FIXEDFILEINFO) > 0: + flags = getattr(pe.VS_FIXEDFILEINFO[0], "FileFlags", 0) + base["IsDebug"] = 1.0 if flags & 0x1 else 0.0 + base["IsPreRelease"] = 1.0 if flags & 0x2 else 0.0 + base["IsPatched"] = 1.0 if flags & 0x4 else 0.0 + base["IsPrivateBuild"] = 1.0 if flags & 0x8 else 0.0 + base["IsSpecialBuild"] = 1.0 if flags & 0x20 else 0.0 + + if hasattr(pe, "DIRECTORY_ENTRY_IMPORT"): + base["ImportCount"] = float(len(pe.DIRECTORY_ENTRY_IMPORT)) + func_count = 0 + + for entry in pe.DIRECTORY_ENTRY_IMPORT: + if getattr(entry, "dll", None): + try: + dlls.add(entry.dll.decode("ascii", "ignore").lower()) + except Exception: + log_exception("PYAS_Features.PEFeatureMixin.extract_features:727") + pass + + for imp in getattr(entry, "imports", []): + func_count += 1 + + if getattr(imp, "name", None): + try: + apis.add(imp.name.decode("ascii", "ignore")) + except Exception: + log_exception("PYAS_Features.PEFeatureMixin.extract_features:735") + pass + + base["ImportFunctionCount"] = float(func_count) + else: + base["ImportCount"] = 0.0 + base["ImportFunctionCount"] = 0.0 + + if hasattr(pe, "DIRECTORY_ENTRY_EXPORT") and hasattr( + pe.DIRECTORY_ENTRY_EXPORT, "symbols" + ): + base["ExportCount"] = float(len(pe.DIRECTORY_ENTRY_EXPORT.symbols)) + else: + base["ExportCount"] = 0.0 + + if hasattr(pe, "DIRECTORY_ENTRY_RESOURCE"): + icon_count = 0 + + for entry in pe.DIRECTORY_ENTRY_RESOURCE.entries: + if getattr(entry, "id", None) == 3 and hasattr(entry, "directory"): + icon_count += len(getattr(entry.directory, "entries", [])) + + base["IconCount"] = float(icon_count) + else: + base["IconCount"] = 0.0 + + if hasattr(pe, "DIRECTORY_ENTRY_DEBUG"): + base["DebugCount"] = float(len(pe.DIRECTORY_ENTRY_DEBUG)) + else: + base["DebugCount"] = 0.0 + + base.update(stream_features) + overlay_size = fsize - int(overlay_offset) + base.update( + { + "HasOverlay": 1.0 if overlay_size else 0.0, + "OverlaySize": float(overlay_size), + "OverlayRatio": float(overlay_size) / float(fsize) if overlay_size else 0.0, + "OverlayEntropy": overlay_entropy if overlay_size else 0.0, + } + ) + base.update(self._extract_rich_header(pe)) + base.update(self._extract_ep_anomalies(pe)) + base.update(self._extract_advanced_resources(pe, resource_entropies)) + base.update(self._extract_load_config(pe)) + base.update(self._extract_security_directory(pe, file_path, fsize)) + + for k in base: + base[k] = self._safe_float(base[k]) + + return {"Base": base, "DLLs": list(dlls), "APIs": list(apis)} + + except Exception: + log_exception("PYAS_Features.PEFeatureMixin.extract_features:782") + return None + finally: + if pe: + pe.close() diff --git a/PYAS_Logs.py b/PYAS_Logs.py new file mode 100644 index 0000000..176cbbc --- /dev/null +++ b/PYAS_Logs.py @@ -0,0 +1,127 @@ +from PYAS_Diagnostics import log_exception, create_log_entry, MAX_LOG_ENTRIES +import os +import time +import json +import uuid +import webview +from PYAS_Storage import atomic_write_json + + +class LogMixin: + def load_logs(self): + with self.lock_logs: + pending_logs = list(self.logs_data) + + if os.path.exists(self.file_log): + try: + with open(self.file_log, "r", encoding="utf-8") as f: + loaded_logs = json.load(f) + + self.logs_data = (loaded_logs + pending_logs)[-MAX_LOG_ENTRIES:] + except Exception: + log_exception("PYAS_Logs.LogMixin.load_logs:17") + self.logs_data = pending_logs + + def log_flush_thread(self): + while not getattr(self, "closing", False): + time.sleep(5) + self.flush_logs_now() + + def write_log( + self, + level, + action, + detail=None, + code=None, + pid=None, + file_hash=None, + source=None, + target=None, + operate=None, + success=True, + ): + entry = create_log_entry( + level, action, detail, code, pid, file_hash, source, target, operate, success + ) + self._append_log_entry(entry) + + if level == "BLOCK" and self.tray_icon: + self.trigger_block_notification(action, source, target, code) + + def _append_log_entry(self, entry): + with self.lock_logs: + self.logs_data.append(entry) + + if len(self.logs_data) > MAX_LOG_ENTRIES: + del self.logs_data[:-MAX_LOG_ENTRIES] + + self.logs_dirty = True + + if self._window: + js_cmd = f"if(window.updateLogs) window.updateLogs({json.dumps(entry)});" + self.ui_queue.put(js_cmd) + + def get_logs(self): + with self.lock_logs: + return self.logs_data.copy() + + def clear_logs(self, log_ids=None): + with self.lock_logs: + if log_ids is None: + self.logs_data = [] + else: + self.logs_data = [log for log in self.logs_data if log["id"] not in log_ids] + + try: + if not self.logs_data: + if os.path.exists(self.file_log): + os.remove(self.file_log) + else: + atomic_write_json(self.file_log, self.logs_data) + + except Exception: + log_exception("PYAS_Logs.LogMixin.clear_logs:73") + self.logs_dirty = True + return False + + self.logs_dirty = False + return True + + def export_logs(self, log_ids=None): + if self._window: + path = self._window.create_file_dialog( + webview.FileDialog.SAVE, directory="", save_filename="PYAS_Logs.json" + ) + + if path: + target_path = path[0] if isinstance(path, (tuple, list)) else path + + with self.lock_logs: + export_data = self.logs_data + + if log_ids is not None: + export_data = [log for log in self.logs_data if log["id"] in log_ids] + + try: + with open(target_path, "w", encoding="utf-8") as f: + json.dump(export_data, f, indent=4, ensure_ascii=False) + + return True + except Exception: + log_exception("PYAS_Logs.LogMixin.export_logs:95") + pass + + return False + + def flush_logs_now(self): + with self.lock_logs: + if not getattr(self, "logs_dirty", False): + return True + + try: + atomic_write_json(self.file_log, self.logs_data) + self.logs_dirty = False + return True + except Exception: + log_exception("PYAS_Logs.LogMixin.flush_logs_now:107") + return False diff --git a/PYAS_Maintenance.py b/PYAS_Maintenance.py new file mode 100644 index 0000000..e1aaa28 --- /dev/null +++ b/PYAS_Maintenance.py @@ -0,0 +1,278 @@ +from PYAS_Diagnostics import log_exception +import os +import threading +import ctypes +import ctypes.wintypes +from PYAS_WinAPI import SHQUERYRBINFO + + +class MaintenanceMixin: + def _traverse_delete(self, path): + deleted = 0 + + if not path or not os.path.exists(path): + return deleted + + try: + items = os.listdir(path) + except Exception: + log_exception("PYAS_Maintenance.MaintenanceMixin._traverse_delete:16") + return deleted + + for fd in items: + file = os.path.join(path, fd) + + try: + if os.path.isdir(file): + if self._is_reparse_point(file): + continue + + deleted += self._traverse_delete(file) + + try: + os.rmdir(file) + except Exception: + log_exception("PYAS_Maintenance.MaintenanceMixin._traverse_delete:29") + pass + + else: + if file.lower().endswith((".sys", ".dll", ".exe", ".ini", ".dat")): + continue + + size = os.path.getsize(file) + os.remove(file) + deleted += size + + except Exception: + log_exception("PYAS_Maintenance.MaintenanceMixin._traverse_delete:40") + continue + + return deleted + + def _get_junk_dirs(self): + return [ + self.path_temp, + self.path_systemp, + os.path.join(self.path_system, "SoftwareDistribution", "Download"), + ] + + def _yield_log_files(self): + log_dir = os.path.join(self.path_system, "Logs") + + if os.path.exists(log_dir): + for root, _, files in os.walk(log_dir): + for file in files: + if file.lower().endswith((".log", ".etl", ".evtx")): + yield os.path.join(root, file) + + def _get_recycle_bin_size(self): + try: + info = SHQUERYRBINFO() + info.cbSize = ctypes.sizeof(SHQUERYRBINFO) + + if self.shell32.SHQueryRecycleBinW(None, ctypes.byref(info)) == 0 and info.i64Size > 0: + return info.i64Size + except Exception: + log_exception("PYAS_Maintenance.MaintenanceMixin._get_recycle_bin_size:67") + pass + + return 0 + + def scan_system_junk(self): + junk_list = [] + + try: + for path in self._get_junk_dirs(): + if os.path.exists(path): + for root, _, files in os.walk(path): + for file in files: + if not file.lower().endswith((".sys", ".dll", ".exe", ".ini", ".dat")): + try: + fp = os.path.join(root, file) + junk_list.append({"path": fp, "size": os.path.getsize(fp)}) + except Exception: + log_exception( + "PYAS_Maintenance.MaintenanceMixin.scan_system_junk:83" + ) + pass + + for fp in self._yield_log_files(): + try: + junk_list.append({"path": fp, "size": os.path.getsize(fp)}) + except Exception: + log_exception("PYAS_Maintenance.MaintenanceMixin.scan_system_junk:89") + pass + + rb_size = self._get_recycle_bin_size() + + if rb_size > 0: + junk_list.append({"path": "Recycle Bin", "size": rb_size}) + + return junk_list + + except Exception as e: + log_exception("PYAS_Maintenance.MaintenanceMixin.scan_system_junk:98") + self.write_log("WARN", "scan_system_junk", detail=str(e), success=False) + return [] + + def clean_system_junk(self, paths_to_delete=None): + total_deleted = 0 + + try: + if paths_to_delete is not None: + for path in paths_to_delete: + if path == "Recycle Bin": + rb_size = self._get_recycle_bin_size() + + if rb_size > 0 and self.shell32.SHEmptyRecycleBinW(None, None, 7) == 0: + total_deleted += rb_size + + continue + + if path.lower().endswith((".sys", ".dll", ".exe", ".ini", ".dat")): + continue + + try: + size = os.path.getsize(path) + os.remove(path) + total_deleted += size + except Exception: + log_exception("PYAS_Maintenance.MaintenanceMixin.clean_system_junk:120") + pass + else: + for path in self._get_junk_dirs(): + total_deleted += self._traverse_delete(path) + + for fp in self._yield_log_files(): + try: + size = os.path.getsize(fp) + os.remove(fp) + total_deleted += size + except Exception: + log_exception("PYAS_Maintenance.MaintenanceMixin.clean_system_junk:131") + pass + + rb_size = self._get_recycle_bin_size() + + if rb_size > 0 and self.shell32.SHEmptyRecycleBinW(None, None, 7) == 0: + total_deleted += rb_size + + try: + for log_type in ["Application", "Security", "Setup", "System"]: + log_path = os.path.join( + self.path_system, "System32", "winevt", "Logs", f"{log_type}.evtx" + ) + + size = os.path.getsize(log_path) if os.path.exists(log_path) else 0 + + if self._clear_event_log(log_type): + total_deleted += size + + except Exception: + log_exception("PYAS_Maintenance.MaintenanceMixin.clean_system_junk:146") + pass + + self.write_log( + "INFO", "Clean Junk", detail=f"Deleted {total_deleted // 1024} KB", operate=True + ) + return total_deleted + + except Exception as e: + log_exception("PYAS_Maintenance.MaintenanceMixin.clean_system_junk:152") + self.write_log("WARN", "clean_system_junk", detail=str(e), operate=True, success=False) + return 0 + + def optimize_memory(self, icon=None, item=None): + def _optimize(): + try: + hwnd = self.user32.GetForegroundWindow() + fg_pid = ctypes.wintypes.DWORD(0) + + if hwnd: + self.user32.GetWindowThreadProcessId(hwnd, ctypes.byref(fg_pid)) + + pids = self.get_process_list_pids() + optimized_count = 0 + + skip_set = { + "dwm.exe", + "explorer.exe", + "csrss.exe", + "smss.exe", + "winlogon.exe", + "lsass.exe", + "services.exe", + "svchost.exe", + "wininit.exe", + "audiodg.exe", + "spoolsv.exe", + "sihost.exe", + "fontdrvhost.exe", + "taskmgr.exe", + } + + for pid in pids: + if pid <= 4 or pid == self.pid_pyas or pid == fg_pid.value: + continue + + name, _ = self.get_exe_info(pid) + + if name and name.lower() in skip_set: + continue + + h = self.kernel32.OpenProcess(0x0100, False, pid) + + if h: + try: + if self.kernel32.SetProcessWorkingSetSize( + h, ctypes.c_size_t(-1), ctypes.c_size_t(-1) + ): + optimized_count += 1 + finally: + self.kernel32.CloseHandle(h) + + title = self._loc( + { + "traditional_switch": "一鍵加速", + "simplified_switch": "一键加速", + "english_switch": "Memory Boost", + "japanese_switch": "メモリ最適化", + "korean_switch": "메모리 최적화", + "french_switch": "Optimiser", + "spanish_switch": "Optimizar", + "hindi_switch": "मेमोरी बूस्ट", + "arabic_switch": "تسريع", + "russian_switch": "Ускорение", + "slovenian_switch": "Optimizacija", + } + ) + + msg = self._loc( + { + "traditional_switch": f"已釋放 {optimized_count} 個背景進程的記憶體", + "simplified_switch": f"已释放 {optimized_count} 个后台进程的内存", + "english_switch": f"Freed memory for {optimized_count} background processes", + "japanese_switch": f"{optimized_count} 個のバックグラウンドプロセスのメモリを解放しました", + "korean_switch": f"{optimized_count}개 백그라운드 프로세스 메모리 확보", + "french_switch": f"Mémoire libérée pour {optimized_count} processus", + "spanish_switch": f"Memoria liberada para {optimized_count} procesos", + "hindi_switch": f"{optimized_count} पृष्ठभूमि प्रक्रियाओं के लिए मेमोरी मुक्त की गई", + "arabic_switch": f"تم تحرير الذاكرة لـ {optimized_count} من العمليات في الخلفية", + "russian_switch": f"Освобождена память {optimized_count} фоновых процессов", + "slovenian_switch": f"Sproščen pomnilnik za {optimized_count} procesov v ozadju", + } + ) + + self.show_notification(title, msg) + self.write_log( + "INFO", + "Memory Boost", + detail=f"Optimized {optimized_count} processes", + operate=True, + ) + + except Exception as e: + log_exception("PYAS_Maintenance.MaintenanceMixin.optimize_memory._optimize:213") + self.write_log("WARN", "optimize_memory", detail=str(e), success=False) + + threading.Thread(target=_optimize, daemon=True).start() diff --git a/PYAS_PE.py b/PYAS_PE.py new file mode 100644 index 0000000..af5aa02 --- /dev/null +++ b/PYAS_PE.py @@ -0,0 +1,283 @@ +from PYAS_Diagnostics import log_exception +import os +import re +import json +import zlib +import numpy +from PYAS_Signature import SignatureScanner as sign_scanner +from PYAS_Features import PEFeatureMixin + + +class PEScanner(PEFeatureMixin): + _STRING_PATTERN = re.compile(b"[\x20-\x7E]{5,}") + + _STRING_FAST_PATH_BYTES = 4194304 + + _STRING_BATCH_BYTES = 4194304 + + _HISTOGRAM_BATCH_WINDOWS = 1024 + + _BYTE_HIST_KEYS = [f"ByteHist_{i:02X}" for i in range(256)] + + _BYTE_ENT_KEYS = [f"ByteEnt_{i:02X}" for i in range(256)] + + _STRING_KEYS = [f"StringHist_{i:02d}" for i in range(95)] + + _SECTION_HASH_KEYS = [f"SectionHash_{i:02d}" for i in range(50)] + + _CHAR_FLAGS = [ + 0x00000020, + 0x00000040, + 0x00000080, + 0x02000000, + 0x20000000, + 0x40000000, + 0x80000000, + ] + + _CHAR_COUNT_KEYS = {flag: f"Char_{flag:08X}_Count" for flag in _CHAR_FLAGS} + + _CHAR_ENT_KEYS = {flag: f"Char_{flag:08X}_MeanEntropy" for flag in _CHAR_FLAGS} + + _C_LOG_C_TABLE = numpy.zeros(2049, dtype=numpy.float64) + + _C_LOG_C_TABLE[1:] = numpy.arange(1, 2049) * numpy.log2(numpy.arange(1, 2049)) + + _OPTIONAL_HEADER_FIELDS = [ + "Magic", + "MajorLinkerVersion", + "MinorLinkerVersion", + "SizeOfCode", + "SizeOfInitializedData", + "SizeOfUninitializedData", + "AddressOfEntryPoint", + "BaseOfCode", + "ImageBase", + "SectionAlignment", + "FileAlignment", + "MajorOperatingSystemVersion", + "MinorOperatingSystemVersion", + "MajorImageVersion", + "MinorImageVersion", + "MajorSubsystemVersion", + "MinorSubsystemVersion", + "SizeOfImage", + "SizeOfHeaders", + "CheckSum", + "Subsystem", + "DllCharacteristics", + "SizeOfStackReserve", + "SizeOfStackCommit", + "SizeOfHeapReserve", + "SizeOfHeapCommit", + "LoaderFlags", + "NumberOfRvaAndSizes", + ] + + _SECTION_DEFAULTS = { + "SectionCount": 0.0, + "SectionMaxEntropy": 0.0, + "SectionMinEntropy": 0.0, + "SectionMeanEntropy": 0.0, + "SectionMaxRawSize": 0.0, + "SectionMinRawSize": 0.0, + "SectionMeanRawSize": 0.0, + "SectionMaxVSize": 0.0, + "SectionMinVSize": 0.0, + "SectionMeanVSize": 0.0, + "ExecutableSections": 0.0, + "WritableSections": 0.0, + "ReadableSections": 0.0, + "SectionException": 0.0, + } + + _CONDITIONAL_DEFAULTS = { + "IsDriver": 0.0, + "HasTlsCallbacks": 0.0, + "IsDebug": 0.0, + "IsPreRelease": 0.0, + "IsPatched": 0.0, + "IsPrivateBuild": 0.0, + "IsSpecialBuild": 0.0, + "ImportCount": 0.0, + "ImportFunctionCount": 0.0, + "ExportCount": 0.0, + "IconCount": 0.0, + "DebugCount": 0.0, + } + + def __init__(self): + self.model = None + self.input_name = None + self.feature_order = [] + self.feature_index = {} + self.dll_hash_dim = 512 + self.api_hash_dim = 4096 + self.dll_hash_pad = 3 + self.api_hash_pad = 4 + self.signer = sign_scanner() + self.signer.init_windll(["wintrust"]) + + def load_path(self, path, callback=None): + for root, _, files in os.walk(path): + for file in files: + full_path = os.path.join(root, file) + + if callback: + callback(full_path) + + if full_path.endswith(".onnx"): + self.load_model(full_path) + + def load_model(self, model_path): + if not os.path.exists(model_path): + return + + try: + import onnxruntime + + model = onnxruntime.InferenceSession(model_path, providers=["CPUExecutionProvider"]) + input_name = model.get_inputs()[0].name + feat_path = os.path.join(os.path.dirname(model_path), "features.json") + feature_order = self.feature_order + + if os.path.exists(feat_path): + with open(feat_path, "r", encoding="utf-8") as stream: + feature_order = json.load(stream) + + if not isinstance(feature_order, list) or any( + not isinstance(feature, str) for feature in feature_order + ): + raise ValueError("Invalid model feature order") + + previous = ( + self.model, + self.input_name, + self.feature_order, + self.feature_index, + self.dll_hash_dim, + self.api_hash_dim, + self.dll_hash_pad, + self.api_hash_pad, + ) + + try: + self.feature_order = feature_order + self._parse_hash_dims() + self.input_name = input_name + self.model = model + except Exception: + ( + self.model, + self.input_name, + self.feature_order, + self.feature_index, + self.dll_hash_dim, + self.api_hash_dim, + self.dll_hash_pad, + self.api_hash_pad, + ) = previous + raise + except Exception: + log_exception("PEScanner.load_model") + + def _parse_hash_dims(self): + self.feature_index = {feature: index for index, feature in enumerate(self.feature_order)} + max_dll = -1 + max_api = -1 + + for feat in self.feature_order: + if feat.startswith("DllHash_"): + try: + val_str = feat.split("_")[1] + max_dll = max(max_dll, int(val_str)) + self.dll_hash_pad = len(val_str) + except Exception: + log_exception("PYAS_PE.PEScanner._parse_hash_dims:127") + pass + elif feat.startswith("ApiHash_"): + try: + val_str = feat.split("_")[1] + max_api = max(max_api, int(val_str)) + self.api_hash_pad = len(val_str) + except Exception: + log_exception("PYAS_PE.PEScanner._parse_hash_dims:134") + pass + + if max_dll >= 0: + self.dll_hash_dim = max_dll + 1 + + if max_api >= 0: + self.api_hash_dim = max_api + 1 + + def pe_scan(self, file_path, enhanced_mode=False): + if not self.model or not self.feature_order: + return False, False + + try: + raw_data = self.extract_features(file_path) + + if not raw_data: + return False, False + + vec = numpy.zeros((1, len(self.feature_order)), dtype=numpy.float32) + base = raw_data.get("Base", {}) + dlls = set(raw_data.get("DLLs", [])) + apis = set(raw_data.get("APIs", [])) + feat_map = self.feature_index + + for k, v in base.items(): + if k in feat_map: + vec[0, feat_map[k]] = v + + for d in dlls: + old_fn = f"Dll_{d}" + + if old_fn in feat_map: + vec[0, feat_map[old_fn]] = 1.0 + + h = zlib.crc32(d.encode("utf-8", "ignore")) % self.dll_hash_dim + new_fn = f"DllHash_{h:0{self.dll_hash_pad}d}" + + if new_fn in feat_map: + vec[0, feat_map[new_fn]] += 1.0 + + for a in apis: + old_fn = f"Api_{a}" + + if old_fn in feat_map: + vec[0, feat_map[old_fn]] = 1.0 + + h = zlib.crc32(a.encode("utf-8", "ignore")) % self.api_hash_dim + new_fn = f"ApiHash_{h:0{self.api_hash_pad}d}" + + if new_fn in feat_map: + vec[0, feat_map[new_fn]] += 1.0 + + outputs = self.model.run(None, {self.input_name: vec}) + + del raw_data, vec + result = outputs[1] + prob = 0.0 + + if isinstance(result, list) and len(result) > 0: + prob_dict = result[0] + + if hasattr(prob_dict, "get"): + prob = float(prob_dict.get(1, prob_dict.get("1", 0.0))) + + elif isinstance(result, numpy.ndarray): + if result.ndim == 2 and result.shape[1] > 1: + prob = float(result[0][1]) + + score = int(prob * 100) + + if score >= 80: + return f"Malware:WinPE/General.{score}!ml", score + elif enhanced_mode and score >= 50: + return f"Suspicious:WinPE/General.{score}!ml", score + + return False, False + except Exception: + log_exception("PYAS_PE.PEScanner.pe_scan:203") + return False, False diff --git a/PYAS_Popup.py b/PYAS_Popup.py new file mode 100644 index 0000000..90fa857 --- /dev/null +++ b/PYAS_Popup.py @@ -0,0 +1,763 @@ +from PYAS_Diagnostics import log_exception +import time +import copy +import ctypes +import ctypes.wintypes +from PYAS_WinAPI import POINT, RECT +import os +import uuid +import math + + +def normalized_path(path): + return os.path.normcase(os.path.normpath(path)) if isinstance(path, str) and path else "" + + +def signature_matches(expected, actual): + if not isinstance(expected, dict) or not isinstance(actual, dict): + return False + + for key in ("class", "style", "exstyle", "children", "owner"): + if key not in expected or expected[key] != actual.get(key): + return False + + for key in ("width", "height"): + size = expected.get(key) + current = actual.get(key) + + if not isinstance(size, (int, float)) or not isinstance(current, (int, float)): + return False + + if ( + not math.isfinite(size) + or not math.isfinite(current) + or size <= 0 + or current <= 0 + or abs(size - current) > max(8, size * 0.08) + ): + return False + + return True + + +def is_main_anchor(window, target): + return ( + window.get("hwnd") != target.get("hwnd") + and window.get("pid") == target.get("pid") + and normalized_path(window.get("path")) == normalized_path(target.get("path")) + and not window.get("fingerprint", {}).get("owner") + and not window.get("fingerprint", {}).get("exstyle", 0) & 0x80 + and window.get("fingerprint", {}).get("style", 0) & 0x00030000 == 0x00030000 + and window.get("fingerprint", {}).get("width", 0) + * window.get("fingerprint", {}).get("height", 0) + > target.get("fingerprint", {}).get("width", 0) + * target.get("fingerprint", {}).get("height", 0) + * 1.25 + ) + + +def secondary_candidate(target): + fingerprint = target.get("fingerprint", {}) + return bool( + fingerprint.get("owner") + or fingerprint.get("exstyle", 0) & 0x80 + or not fingerprint.get("style", 0) & 0x00030000 + ) + + +def build_popup_rule(target, windows): + if not target or not target.get("path") or not target.get("fingerprint", {}).get("class"): + return None + + fingerprint = target["fingerprint"] + peers = [ + window + for window in windows + if window.get("pid") == target.get("pid") + and normalized_path(window.get("path")) == normalized_path(target.get("path")) + and window["hwnd"] != target["hwnd"] + ] + identical = any( + all( + fingerprint.get(key) == window.get("fingerprint", {}).get(key) + for key in ("class", "style", "exstyle", "children", "owner") + ) + for window in peers + ) + anchors = [window for window in peers if is_main_anchor(window, target)] + distinctive = secondary_candidate(target) or bool(anchors) + mode = ( + "adaptive" + if isinstance(fingerprint.get("children"), list) + and signature_matches(fingerprint, fingerprint) + and distinctive + and not identical + and not target.get("modal") + else "title" + ) + return { + "id": uuid.uuid4().hex, + "schema_version": 2, + "exe": target["exe"], + "path": target["path"], + "class": fingerprint["class"], + "title": target.get("title", ""), + "fingerprint": fingerprint, + "match_mode": mode, + "protected": [ + window["fingerprint"] + for window in peers + if not signature_matches(fingerprint, window.get("fingerprint")) + ], + } + + +def popup_rule_matches(rule, target, windows): + if not isinstance(rule, dict) or not target: + return False + + if rule.get("schema_version") != 2: + return bool( + rule.get("exe") == target.get("exe") + and rule.get("class") == target.get("fingerprint", {}).get("class") + and rule.get("title") + and rule["title"] == target.get("title") + and not target.get("modal") + ) + + if ( + not rule.get("id") + or not normalized_path(rule.get("path")) + or normalized_path(rule["path"]) != normalized_path(target.get("path")) + ): + return False + + fingerprint = target.get("fingerprint", {}) + + if not rule.get("class") or rule["class"] != fingerprint.get("class"): + return False + + if any(signature_matches(protected, fingerprint) for protected in rule.get("protected", [])): + return False + + expected = rule.get("fingerprint", {}) + structural_match = signature_matches(expected, fingerprint) + title_match = bool(rule.get("title") and rule["title"] == target.get("title")) + + if rule.get("match_mode", "adaptive") == "title": + if not title_match or not all( + expected.get(key) == fingerprint.get(key) for key in ("style", "exstyle", "owner") + ): + return False + + return not any( + window.get("hwnd") != target.get("hwnd") + and window.get("pid") == target.get("pid") + and normalized_path(window.get("path")) == normalized_path(target.get("path")) + and window.get("title") == target.get("title") + and window.get("fingerprint", {}).get("class") == rule["class"] + and all( + window.get("fingerprint", {}).get(key) == fingerprint.get(key) + for key in ("style", "exstyle", "owner") + ) + for window in windows + ) + + if not structural_match or target.get("modal"): + return False + + peers = [ + window + for window in windows + if window.get("pid") == target.get("pid") + and normalized_path(window.get("path")) == normalized_path(target.get("path")) + and window["hwnd"] != target["hwnd"] + ] + identical = [ + window for window in peers if signature_matches(expected, window.get("fingerprint")) + ] + + if identical: + return title_match and not any( + window.get("title") == target.get("title") for window in identical + ) + + return True + + +class PopupWindowMixin: + def capture_popup_window(self): + self._init_popup_api() + overlay = self.user32.CreateWindowExW( + 0x000800A8, "STATIC", "", 0x98000004, 0, 0, 0, 0, None, None, None, None + ) + + if overlay: + self.user32.SetLayeredWindowAttributes(overlay, 0, 120, 2) + + target_hwnd = None + last_hwnd = None + msg = ctypes.wintypes.MSG() + + start_time = time.monotonic() + + try: + while self.user32.GetAsyncKeyState(0x01) & 0x8000: + if time.monotonic() - start_time >= 30: + return None + + time.sleep(0.01) + + while time.monotonic() - start_time < 30: + if self.user32.GetAsyncKeyState(0x1B) & 0x8000: + break + + pt = POINT() + self.user32.GetCursorPos(ctypes.byref(pt)) + + hwnd = self.user32.WindowFromPoint(pt) + root_hwnd = self.user32.GetAncestor(hwnd, 2) + + if not root_hwnd: + root_hwnd = hwnd + + if root_hwnd and root_hwnd != overlay: + if root_hwnd != last_hwnd: + last_hwnd = root_hwnd + rect = RECT() + self.user32.GetWindowRect(root_hwnd, ctypes.byref(rect)) + self.user32.SetWindowPos( + overlay, + -1, + rect.left, + rect.top, + rect.right - rect.left, + rect.bottom - rect.top, + 0x0050, + ) + + if self.user32.GetAsyncKeyState(0x01) & 0x8000: + target_hwnd = root_hwnd + break + + while self.user32.PeekMessageW(ctypes.byref(msg), 0, 0, 0, 1): + self.user32.TranslateMessage(ctypes.byref(msg)) + self.user32.DispatchMessageW(ctypes.byref(msg)) + + time.sleep(0.02) + finally: + if overlay: + self.user32.DestroyWindow(overlay) + + if not target_hwnd: + return None + + pid = ctypes.c_ulong(0) + self.user32.GetWindowThreadProcessId(target_hwnd, ctypes.byref(pid)) + + if pid.value != self.pid_pyas and pid.value > 4: + length = self.user32.GetWindowTextLengthW(target_hwnd) + title = ctypes.create_unicode_buffer(length + 1) + class_name = ctypes.create_unicode_buffer(256) + + self.user32.GetWindowTextW(target_hwnd, title, length + 1) + self.user32.GetClassNameW(target_hwnd, class_name, 256) + + proc_name, _ = self.get_exe_info(pid.value) + t_str, c_str = str(title.value), str(class_name.value) + + if proc_name and not any( + item.get("exe") == proc_name or item.get("class") == c_str + for item in self.pass_windows + ): + target = self._popup_window_snapshot(target_hwnd) + windows = self._popup_windows({normalized_path(target["path"])}) if target else [] + rule = build_popup_rule(target, windows) + + if not rule: + return {"error": "popup_capture_failed"} + + token = (time.monotonic_ns() & 0x7FFFFFFF) or 1 + selected = {"hwnd": target_hwnd, "pid": target["pid"], "token": token} + + if not self.user32.SetPropW(target_hwnd, "PYAS_Popup_Selected", token): + return {"error": "popup_capture_failed"} + + with self.lock_config: + previous = getattr(self, "_popup_selected", {}) + active = { + item.get("id") + for item in self.pyas_config.get("block_list", []) + if isinstance(item, dict) + } + + for identifier, record in list(previous.items()): + if identifier not in active: + if ( + self.user32.GetPropW(record["hwnd"], "PYAS_Popup_Selected") + == record["token"] + ): + self.user32.RemovePropW(record["hwnd"], "PYAS_Popup_Selected") + + previous.pop(identifier, None) + + previous[rule["id"]] = selected + self._popup_selected = previous + self._captured_popup_rules = {rule["id"]: copy.deepcopy(rule)} + + return rule + + return {"error": "popup_capture_failed"} + + def _init_popup_api(self): + if getattr(self, "_popup_api_ready", False): + return + + self._popup_enum_type = ctypes.WINFUNCTYPE( + ctypes.wintypes.BOOL, ctypes.wintypes.HWND, ctypes.wintypes.LPARAM + ) + self.user32.CreateWindowExW.argtypes = [ + ctypes.wintypes.DWORD, + ctypes.c_wchar_p, + ctypes.c_wchar_p, + ctypes.wintypes.DWORD, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.wintypes.HWND, + ctypes.wintypes.HMENU, + ctypes.wintypes.HINSTANCE, + ctypes.c_void_p, + ] + self.user32.CreateWindowExW.restype = ctypes.wintypes.HWND + self.user32.DestroyWindow.argtypes = [ctypes.wintypes.HWND] + self.user32.DestroyWindow.restype = ctypes.wintypes.BOOL + self.user32.PeekMessageW.argtypes = [ + ctypes.POINTER(ctypes.wintypes.MSG), + ctypes.wintypes.HWND, + ctypes.c_uint, + ctypes.c_uint, + ctypes.c_uint, + ] + self.user32.PeekMessageW.restype = ctypes.wintypes.BOOL + self.user32.TranslateMessage.argtypes = [ctypes.POINTER(ctypes.wintypes.MSG)] + self.user32.TranslateMessage.restype = ctypes.wintypes.BOOL + self.user32.DispatchMessageW.argtypes = [ctypes.POINTER(ctypes.wintypes.MSG)] + self.user32.DispatchMessageW.restype = ctypes.wintypes.LPARAM + + for name in ("EnumWindows", "EnumChildWindows"): + function = getattr(self.user32, name) + function.argtypes = ( + [self._popup_enum_type, ctypes.wintypes.LPARAM] + if name == "EnumWindows" + else [ctypes.wintypes.HWND, self._popup_enum_type, ctypes.wintypes.LPARAM] + ) + function.restype = ctypes.wintypes.BOOL + + self._popup_get_style = getattr( + self.user32, "GetWindowLongPtrW", self.user32.GetWindowLongW + ) + self._popup_get_style.argtypes = [ctypes.wintypes.HWND, ctypes.c_int] + self._popup_get_style.restype = ctypes.c_ssize_t + + for name in ("IsWindow", "IsWindowVisible", "IsWindowEnabled", "IsIconic"): + function = getattr(self.user32, name) + function.argtypes = [ctypes.wintypes.HWND] + function.restype = ctypes.wintypes.BOOL + + self.user32.GetWindow.argtypes = [ctypes.wintypes.HWND, ctypes.c_uint] + self.user32.GetWindow.restype = ctypes.wintypes.HWND + self.user32.GetClientRect.argtypes = [ctypes.wintypes.HWND, ctypes.POINTER(RECT)] + self.user32.GetClientRect.restype = ctypes.wintypes.BOOL + + class WINDOWPLACEMENT(ctypes.Structure): + _fields_ = [ + ("length", ctypes.wintypes.DWORD), + ("flags", ctypes.wintypes.DWORD), + ("showCmd", ctypes.wintypes.DWORD), + ("ptMinPosition", POINT), + ("ptMaxPosition", POINT), + ("rcNormalPosition", RECT), + ] + + self._popup_placement_type = WINDOWPLACEMENT + self.user32.GetWindowPlacement.argtypes = [ + ctypes.wintypes.HWND, + ctypes.POINTER(WINDOWPLACEMENT), + ] + self.user32.GetWindowPlacement.restype = ctypes.wintypes.BOOL + self.user32.GetWindowThreadProcessId.argtypes = [ + ctypes.wintypes.HWND, + ctypes.POINTER(ctypes.wintypes.DWORD), + ] + self.user32.GetWindowThreadProcessId.restype = ctypes.wintypes.DWORD + self.user32.GetWindowTextLengthW.argtypes = [ctypes.wintypes.HWND] + self.user32.GetWindowTextLengthW.restype = ctypes.c_int + self.user32.GetWindowTextW.argtypes = [ctypes.wintypes.HWND, ctypes.c_wchar_p, ctypes.c_int] + self.user32.GetClassNameW.argtypes = [ctypes.wintypes.HWND, ctypes.c_wchar_p, ctypes.c_int] + self.user32.ShowWindowAsync.argtypes = [ctypes.wintypes.HWND, ctypes.c_int] + self.user32.ShowWindowAsync.restype = ctypes.wintypes.BOOL + self.user32.SetPropW.argtypes = [ + ctypes.wintypes.HWND, + ctypes.c_wchar_p, + ctypes.wintypes.HANDLE, + ] + self.user32.SetPropW.restype = ctypes.wintypes.BOOL + self.user32.GetPropW.argtypes = [ctypes.wintypes.HWND, ctypes.c_wchar_p] + self.user32.GetPropW.restype = ctypes.wintypes.HANDLE + self.user32.RemovePropW.argtypes = [ctypes.wintypes.HWND, ctypes.c_wchar_p] + self.user32.RemovePropW.restype = ctypes.wintypes.HANDLE + + if hasattr(self.user32, "GetDpiForWindow"): + self.user32.GetDpiForWindow.argtypes = [ctypes.wintypes.HWND] + self.user32.GetDpiForWindow.restype = ctypes.c_uint + + self._popup_api_ready = True + + def _popup_window_snapshot(self, hwnd, process_info=None): + if not self.user32.IsWindow(hwnd): + return None + + pid = ctypes.wintypes.DWORD() + self.user32.GetWindowThreadProcessId(hwnd, ctypes.byref(pid)) + + if pid.value <= 4 or pid.value == self.pid_pyas: + return None + + process_info = {} if process_info is None else process_info + + if pid.value not in process_info: + process_info[pid.value] = self.get_exe_info(pid.value) + + exe, path = process_info[pid.value] + + if not exe or not path: + return None + + class_name = ctypes.create_unicode_buffer(256) + + if not self.user32.GetClassNameW(hwnd, class_name, 256): + return None + + if any( + item.get("exe") == exe or item.get("class") == class_name.value + for item in self.pass_windows + ): + return None + + rect = RECT() + + if not self.user32.GetClientRect(hwnd, ctypes.byref(rect)): + self.user32.GetWindowRect(hwnd, ctypes.byref(rect)) + + if self.user32.IsIconic(hwnd): + placement = self._popup_placement_type() + placement.length = ctypes.sizeof(placement) + + if self.user32.GetWindowPlacement(hwnd, ctypes.byref(placement)): + rect = placement.rcNormalPosition + + dpi = self.user32.GetDpiForWindow(hwnd) if hasattr(self.user32, "GetDpiForWindow") else 96 + dpi = dpi or 96 + classes = {} + child_count = 0 + + def enum_child(child, parameter): + nonlocal child_count + child_count += 1 + + if child_count > 128: + return False + + name = ctypes.create_unicode_buffer(256) + + if not self.user32.GetClassNameW(child, name, 256): + child_count = 129 + return False + + classes[name.value] = classes.get(name.value, 0) + 1 + return True + + callback = self._popup_enum_type(enum_child) + self.user32.EnumChildWindows(hwnd, callback, 0) + children = ( + None + if child_count > 128 + else [[name, count] for name, count in sorted(classes.items())] + ) + owner_hwnd = self.user32.GetWindow(hwnd, 4) + owner = None + + if owner_hwnd: + owner_pid = ctypes.wintypes.DWORD() + self.user32.GetWindowThreadProcessId(owner_hwnd, ctypes.byref(owner_pid)) + + if owner_pid.value not in process_info: + process_info[owner_pid.value] = self.get_exe_info(owner_pid.value) + + _, owner_path = process_info[owner_pid.value] + owner_class = ctypes.create_unicode_buffer(256) + self.user32.GetClassNameW(owner_hwnd, owner_class, 256) + owner = {"path": normalized_path(owner_path), "class": owner_class.value} + + title = ctypes.create_unicode_buffer(min(self.user32.GetWindowTextLengthW(hwnd), 4096) + 1) + self.user32.GetWindowTextW(hwnd, title, len(title)) + return { + "hwnd": hwnd, + "pid": pid.value, + "exe": exe, + "path": path, + "title": title.value, + "visible": bool(self.user32.IsWindowVisible(hwnd)), + "minimized": bool(self.user32.IsIconic(hwnd)), + "modal": bool(owner_hwnd and not self.user32.IsWindowEnabled(owner_hwnd)), + "fingerprint": { + "class": class_name.value, + "style": self._popup_get_style(hwnd, -16) & 0x80CF0000, + "exstyle": self._popup_get_style(hwnd, -20) & 0x08040088, + "width": round((rect.right - rect.left) * 96 / dpi), + "height": round((rect.bottom - rect.top) * 96 / dpi), + "children": children, + "owner": owner, + }, + } + + def _popup_windows(self, paths=None): + self._init_popup_api() + windows = [] + process_info = {} + + def enum_window(hwnd, parameter): + try: + if paths is not None: + pid = ctypes.wintypes.DWORD() + self.user32.GetWindowThreadProcessId(hwnd, ctypes.byref(pid)) + + if pid.value not in process_info: + process_info[pid.value] = self.get_exe_info(pid.value) + + if normalized_path(process_info[pid.value][1]) not in paths: + return True + + snapshot = self._popup_window_snapshot(hwnd, process_info) + + if snapshot: + windows.append(snapshot) + except Exception: + log_exception("PYAS_Popup.PopupWindowMixin._popup_windows.enum_window:246") + return True + + return True + + callback = self._popup_enum_type(enum_window) + self.user32.EnumWindows(callback, 0) + return windows + + def add_popup_rule(self, rule): + if not isinstance(rule, dict) or rule.get("schema_version") != 2 or not rule.get("id"): + return False + + with self.lock_config: + pending = getattr(self, "_captured_popup_rules", {}) + captured = pending.get(rule["id"]) + + if captured != rule: + return any(item == rule for item in self.pyas_config.get("block_list", [])) + + target_list = self.pyas_config.setdefault("block_list", []) + duplicate = next( + ( + item + for item in target_list + if isinstance(item, dict) + and item.get("schema_version") == 2 + and normalized_path(item.get("path")) == normalized_path(rule.get("path")) + and item.get("fingerprint") == rule.get("fingerprint") + and item.get("title") == rule.get("title") + and item.get("match_mode") == rule.get("match_mode") + ), + None, + ) + + if duplicate: + selected = getattr(self, "_popup_selected", {}).pop(rule["id"], None) + + if selected: + self._popup_selected[duplicate["id"]] = selected + else: + target_list.append(copy.deepcopy(captured)) + self.save_config() + self.write_log( + "INFO", + "Config Update", + detail=f"Popup rule added: {rule['exe']} [{rule['id']}]", + ) + + pending.pop(rule["id"], None) + return True + + def _popup_matches(self, rule, target, windows): + with self.lock_config: + selected = getattr(self, "_popup_selected", {}).get(rule.get("id")) + + if ( + selected + and selected["hwnd"] == target.get("hwnd") + and selected["pid"] == target.get("pid") + ): + if self.user32.GetPropW(selected["hwnd"], "PYAS_Popup_Selected") == selected["token"]: + return normalized_path(rule.get("path")) == normalized_path(target.get("path")) + + return popup_rule_matches(rule, target, windows) + + def remove_popup_rules(self, identifiers): + with self.lock_config: + rules = self.pyas_config.get("block_list", []) + identifiers = set(identifiers or []) + remaining = [ + rule + for index, rule in enumerate(rules) + if (rule.get("id") or f"legacy:{index}") not in identifiers + ] + + if len(remaining) == len(rules): + return False + + selected = getattr(self, "_popup_selected", {}) + + for identifier in identifiers: + record = selected.pop(identifier, None) + + if ( + record + and self.user32.GetPropW(record["hwnd"], "PYAS_Popup_Selected") + == record["token"] + ): + self.user32.RemovePropW(record["hwnd"], "PYAS_Popup_Selected") + + self.pyas_config["block_list"] = remaining + self.write_log( + "INFO", + "Config Update", + detail=f"Popup rules removed: {len(rules) - len(remaining)}", + ) + self.save_config() + return True + + def _restore_popup_windows(self, hidden, active_ids): + for hwnd, record in list(hidden.items()): + if self.user32.GetPropW(hwnd, "PYAS_Popup_Blocker") != record["token"]: + hidden.pop(hwnd, None) + elif record["rule_id"] not in active_ids: + self.user32.RemovePropW(hwnd, "PYAS_Popup_Blocker") + self.user32.ShowWindowAsync(hwnd, 4) + hidden.pop(hwnd, None) + + def _hide_popup_window(self, rule, target, windows, hidden): + hwnd = target["hwnd"] + current = self._popup_window_snapshot(hwnd) + + if ( + not current + or current["pid"] != target["pid"] + or not self._popup_matches(rule, current, windows) + ): + return False + + record = hidden.get(hwnd) + new_record = ( + not record or self.user32.GetPropW(hwnd, "PYAS_Popup_Blocker") != record["token"] + ) + + if new_record: + token = (time.monotonic_ns() & 0x7FFFFFFF) or 1 + + if not self.user32.SetPropW(hwnd, "PYAS_Popup_Blocker", token): + return False + + record = {"token": token, "rule_id": rule["id"]} + hidden[hwnd] = record + + if not self.user32.ShowWindowAsync(hwnd, 0): + if new_record: + self.user32.RemovePropW(hwnd, "PYAS_Popup_Blocker") + hidden.pop(hwnd, None) + + return False + + return True + + def popup_intercept_thread(self): + hidden = {} + last_warning = {} + + while True: + try: + time.sleep(0.5) + self._init_popup_api() + + with self.lock_config: + rules = copy.deepcopy(self.pyas_config.get("block_list", [])) + + valid_rules = [] + + for index, rule in enumerate(rules): + if isinstance(rule, dict) and ( + rule.get("schema_version") == 2 + and rule.get("id") + or rule.get("exe") + and rule.get("class") + and rule.get("title") + ): + rule.setdefault("id", f"legacy:{index}") + valid_rules.append(rule) + + self._restore_popup_windows(hidden, {rule["id"] for rule in valid_rules}) + + if not valid_rules: + continue + + paths = {normalized_path(rule.get("path")) for rule in valid_rules} + windows = self._popup_windows(paths if all(paths) else None) + + for target in windows: + if not target.get("visible") or target.get("minimized"): + continue + + for rule in valid_rules: + if self._popup_matches(rule, target, windows): + first_hide = target["hwnd"] not in hidden + + if self._hide_popup_window(rule, target, windows, hidden): + if first_hide: + self.write_log( + "INFO", + "Popup Blocker", + source=target["path"], + detail=f"Hidden popup window: {target['title']}", + pid=target["pid"], + ) + elif first_hide: + key = (rule["id"], target["hwnd"], target["pid"]) + now = time.monotonic() + + if now - last_warning.get(key, -30) >= 30: + self.write_log( + "WARN", + "Popup Blocker", + source=target["path"], + detail="Unable to hide selected window", + pid=target["pid"], + success=False, + ) + last_warning[key] = now + + break + except Exception as e: + log_exception("PYAS_Popup.PopupWindowMixin.popup_intercept_thread:369") + now = time.monotonic() + + if now - last_warning.get("error", -30) >= 30: + self.write_log("WARN", "Popup Blocker", detail=str(e), success=False) + last_warning["error"] = now diff --git a/PYAS_Process.py b/PYAS_Process.py new file mode 100644 index 0000000..7ea5c86 --- /dev/null +++ b/PYAS_Process.py @@ -0,0 +1,342 @@ +from PYAS_Diagnostics import log_exception +import os +import re +import io +import time +import ctypes +import ctypes.wintypes +from PYAS_WinAPI import ( + IO_COUNTERS, + MIB_TCPROW_OWNER_PID, + PROCESSENTRY32W, + PROCESS_BASIC_INFORMATION, + PYAS_PROCESS_TERMINATE_ACCESS, + UNICODE_STRING, +) + + +class ProcessMixin: + def device_path_to_drive(self, path): + if not path: + return "" + + for d in range(65, 91): + drive = f"{chr(d)}:" + buf = ctypes.create_unicode_buffer(1024) + + if self.kernel32.QueryDosDeviceW(drive, buf, 1024) and path.startswith(buf.value): + return path.replace(buf.value, drive, 1) + + return path + + def get_exe_info(self, pid): + name, file_path = None, None + h = self.kernel32.OpenProcess(0x1000, False, pid) + + if h: + try: + buf = ctypes.create_unicode_buffer(1024) + size = ctypes.wintypes.DWORD(1024) + + if self.kernel32.QueryFullProcessImageNameW(h, 0, buf, ctypes.byref(size)): + file_path = self.norm_path(buf.value) + + if file_path: + name = os.path.basename(file_path) + + elif self.psapi.GetProcessImageFileNameW(h, buf, 1024): + file_path = self.norm_path(self.device_path_to_drive(buf.value)) + + if file_path: + name = os.path.basename(file_path) + + finally: + self.kernel32.CloseHandle(h) + + return name, file_path + + def get_process_file(self, h_process): + buf = ctypes.create_unicode_buffer(1024) + + if self.psapi.GetProcessImageFileNameW(h_process, buf, 1024): + return self.norm_path(self.device_path_to_drive(buf.value)) + + return "" + + def get_process_cmdline(self, h): + try: + pbi = PROCESS_BASIC_INFORMATION() + returned = ctypes.wintypes.ULONG() + + if ( + self.ntdll.NtQueryInformationProcess( + h, 0, ctypes.byref(pbi), ctypes.sizeof(pbi), ctypes.byref(returned) + ) + != 0 + ): + return "" + + if returned.value != ctypes.sizeof(pbi) or not pbi.PebBaseAddress: + return "" + + def read_exact(address, destination): + transferred = ctypes.c_size_t() + capacity = ctypes.sizeof(destination) + return ( + bool( + self.kernel32.ReadProcessMemory( + h, + ctypes.c_void_p(address), + ctypes.byref(destination), + capacity, + ctypes.byref(transferred), + ) + ) + and transferred.value == capacity + ) + + pointer_size = ctypes.sizeof(ctypes.c_void_p) + parameters = ctypes.c_void_p() + + if ( + not read_exact( + int(pbi.PebBaseAddress) + (0x20 if pointer_size == 8 else 0x10), parameters + ) + or not parameters.value + ): + return "" + + descriptor = UNICODE_STRING() + + if not read_exact(parameters.value + (0x70 if pointer_size == 8 else 0x40), descriptor): + return "" + + length = int(descriptor.Length) + + if not descriptor.Buffer or not length: + return "" + + if length % 2 or length > descriptor.MaximumLength: + raise ValueError("Invalid remote UTF-16 command-line length") + + buffer = (ctypes.c_ubyte * length)() + + if not read_exact(int(descriptor.Buffer), buffer): + return "" + + return bytes(buffer).decode("utf-16-le", errors="surrogatepass") + except Exception: + log_exception("ProcessMixin.get_process_cmdline") + return "" + + def extract_paths_from_cmdline(self, cmdline): + if not cmdline: + return [] + + found = [] + + for m in re.finditer(r'"([^"]+)"|\'([^\']+)\'', cmdline): + path = m.group(1) or m.group(2) + + if re.match(r"^[A-Za-z]:\\|^\\\\", path): + found.append(path) + + if not found: + m = re.search( + r'([A-Za-z]:\\[^\*?"<>\|]+\.(?:exe|dll|bat|cmd|vbs|sys|com|pif))', + cmdline, + re.IGNORECASE, + ) + + if m: + found.append(m.group(1)) + + if not found: + argc = ctypes.c_int(0) + argv = self.shell32.CommandLineToArgvW(cmdline, ctypes.byref(argc)) + + if argv: + args = [argv[i] for i in range(argc.value)] + self.kernel32.LocalFree(ctypes.cast(argv, ctypes.c_void_p)) + patterns = [ + r'([A-Za-z]:\\[^"\']+)', + r'(\\\\[^"\']+)', + r'(\.\\[^"\']+)', + r'(\./[^"\']+)', + r'([A-Za-z]:/[^"\']+)', + r"([^\s]*\\[^\s]+)", + ] + + for arg in args: + for p in patterns: + for match in re.finditer(p, arg): + found.append(match.group(1).strip('"').strip("'")) + + return list(dict.fromkeys([p.strip('"').strip("'") for p in found])) + + def _enum_processes(self): + pe = PROCESSENTRY32W() + pe.dwSize = ctypes.sizeof(PROCESSENTRY32W) + snapshot = self.kernel32.CreateToolhelp32Snapshot(0x00000002, 0) + + if snapshot in (-1, 0xFFFFFFFF, 0xFFFFFFFFFFFFFFFF): + return + + try: + if self.kernel32.Process32FirstW(snapshot, ctypes.byref(pe)): + while True: + yield pe.th32ProcessID, pe.szExeFile + + if not self.kernel32.Process32NextW(snapshot, ctypes.byref(pe)): + break + finally: + self.kernel32.CloseHandle(snapshot) + + def get_process_list(self): + result = [] + + for pid, exe_name in self._enum_processes(): + name, file_path = None, None + + if pid > 4: + name, file_path = self.get_exe_info(pid) + + if not name: + try: + name = exe_name + except Exception: + log_exception("PYAS_Process.ProcessMixin.get_process_list:140") + name = "Unknown" + + result.append({"pid": pid, "name": name, "path": file_path or "None"}) + + return result + + def get_process_list_pids(self): + return {pid for pid, _ in self._enum_processes()} + + def kill_process(self, pid, expected_path=None): + try: + h = self.kernel32.OpenProcess(PYAS_PROCESS_TERMINATE_ACCESS, False, pid) + + if h: + try: + return self._terminate_process_handle(h, pid=pid, expected_path=expected_path) + finally: + self.kernel32.CloseHandle(h) + + except Exception as e: + log_exception("PYAS_Process.ProcessMixin.kill_process:159") + self.write_log( + "WARN", "kill_process", pid=pid, detail=str(e), operate=True, success=False + ) + + return False + + def get_connections_list(self): + connections = set() + + try: + size = ctypes.wintypes.DWORD() + + if self.iphlpapi.GetExtendedTcpTable(None, ctypes.byref(size), True, 2, 5, 0) != 122: + return connections + + buf = ctypes.create_string_buffer(size.value) + + if self.iphlpapi.GetExtendedTcpTable(buf, ctypes.byref(size), True, 2, 5, 0) != 0: + return connections + + num_entries = ctypes.cast(buf, ctypes.POINTER(ctypes.wintypes.DWORD)).contents.value + + for i in range(num_entries): + row = MIB_TCPROW_OWNER_PID.from_address( + ctypes.addressof(buf) + + ctypes.sizeof(ctypes.wintypes.DWORD) + + i * ctypes.sizeof(MIB_TCPROW_OWNER_PID) + ) + connections.add((row.dwOwningPid, row.dwRemoteAddr, row.dwRemotePort)) + + except Exception: + log_exception("PYAS_Process.ProcessMixin.get_connections_list:179") + pass + + return connections + + def get_traffic_list(self): + conns = self.get_connections_list() + conn_map = {} + + for pid, _, _ in conns: + conn_map[pid] = conn_map.get(pid, 0) + 1 + + current_io = {} + current_time = time.time() + + with self.lock_io: + time_diff = current_time - self.last_io_time + + if time_diff <= 0: + time_diff = 1 + + self.last_io_time = current_time + old_counters = self.last_io_counters.copy() + + result = [] + exist_process = self.get_process_list_pids() + + for pid in exist_process: + name, file_path = "Unknown", "" + + if pid > 4: + name, file_path = self.get_exe_info(pid) + else: + name = "System" + + down_speed, up_speed = 0, 0 + h = self.kernel32.OpenProcess(0x1000, False, pid) + + if h: + try: + io = IO_COUNTERS() + + if self.kernel32.GetProcessIoCounters(h, ctypes.byref(io)): + current_io[pid] = (io.ReadTransferCount, io.WriteTransferCount) + + if pid in old_counters: + old_read, old_write = old_counters[pid] + down_speed = max(0, (io.ReadTransferCount - old_read) / time_diff) + up_speed = max(0, (io.WriteTransferCount - old_write) / time_diff) + + finally: + self.kernel32.CloseHandle(h) + + count = conn_map.get(pid, 0) + + if count > 0 or down_speed > 0 or up_speed > 0: + result.append( + { + "pid": pid, + "name": name, + "path": file_path or "None", + "down": int(down_speed), + "up": int(up_speed), + "conn": count, + } + ) + + with self.lock_io: + self.last_io_counters = current_io + + return result + + def _terminate_process_handle(self, handle, pid=None, expected_path=None): + file_path = self.norm_path(self.get_process_file(handle)) + + if self.path_equal(file_path, self.file_pyas): + return False + + if expected_path and not self.path_equal(file_path, expected_path): + return False + + return bool(self.kernel32.TerminateProcess(handle, 0)) diff --git a/PYAS_Protect.py b/PYAS_Protect.py index 6d4fa58..e2f452a 100644 --- a/PYAS_Protect.py +++ b/PYAS_Protect.py @@ -1,116 +1,103 @@ -import os, time, shutil, msvcrt, winreg, threading, subprocess -import ctypes, ctypes.wintypes +from PYAS_Diagnostics import log_exception +import os +import time +import shutil +import msvcrt +import winreg +import threading +import subprocess +import copy +import ctypes +import ctypes.wintypes +from PYAS_Popup import build_popup_rule, normalized_path, popup_rule_matches +from PYAS_WinAPI import ( + FILE_NOTIFY_INFORMATION, + LUID, + MEMORY_BASIC_INFORMATION, + POINT, + PYAS_FULL_MESSAGE, + PYAS_USER_MESSAGE, + RECT, + SERVICE_STATUS_PROCESS, + TOKEN_PRIVILEGES, +) +from PYAS_WinAPI import ( + COPYDATASTRUCT, + ERROR_LOCK_VIOLATION, + ERROR_NOT_FOUND, + ERROR_OPERATION_ABORTED, + ERROR_SHARING_VIOLATION, + FILE_ATTRIBUTE_NORMAL, + FILE_LOCK_RETRY_SECONDS, + FILE_SCAN_DEBOUNCE_SECONDS, + FILE_SCAN_RETRY_SECONDS, + FILE_SHARE_READ, + FILE_SHARE_WRITE, + FILTER_MESSAGE_HEADER, + FLT_PORT_FLAG_SYNC_HANDLE, + GENERIC_READ, + GUID, + HRESULT_IO_PENDING, + INTERNAL_STORAGE_BUS_TYPES, + INVALID_HANDLE_VALUE, + IOCTL_STORAGE_GET_HOTPLUG_INFO, + IOCTL_STORAGE_QUERY_PROPERTY, + IO_COUNTERS, + LUID_AND_ATTRIBUTES, + MIB_TCPROW_OWNER_PID, + OPEN_EXISTING, + OVERLAPPED, + PROCESSENTRY32W, + PROCESS_BASIC_INFORMATION, + PROCESS_QUERY_INFORMATION, + PROCESS_QUERY_LIMITED_INFORMATION, + PROCESS_SUSPEND_RESUME, + PROCESS_TERMINATE, + PROCESS_VM_READ, + PROPERTY_STANDARD_QUERY, + PYAS_CONNECTION_CONTEXT, + PYAS_CONNECTION_MAGIC, + PYAS_CONNECTION_VERSION, + PYAS_MESSAGE, + PYAS_PROCESS_NETWORK_ACCESS, + PYAS_PROCESS_SCAN_ACCESS, + PYAS_PROCESS_TERMINATE_ACCESS, + SHQUERYRBINFO, + STORAGE_DESCRIPTOR_HEADER, + STORAGE_DEVICE_DESCRIPTOR, + STORAGE_DEVICE_PROPERTY, + STORAGE_HOTPLUG_INFO, + STORAGE_PROPERTY_QUERY, + UNICODE_STRING, + WAIT_OBJECT_0, + WAIT_TIMEOUT, + WINTRUST_DATA, + WINTRUST_DATA_UNION, + WINTRUST_FILE_INFO, +) +from PYAS_Driver import DriverMixin +from PYAS_System import SystemMixin +from PYAS_Realtime import RealtimeMixin +from PYAS_Popup import PopupWindowMixin -from PYAS_Tools import FILE_NOTIFY_INFORMATION, LUID, MEMORY_BASIC_INFORMATION, POINT, PYAS_FULL_MESSAGE, PYAS_USER_MESSAGE, RECT, SERVICE_STATUS_PROCESS, TOKEN_PRIVILEGES -#################################################################################################### - -FLT_PORT_FLAG_SYNC_HANDLE = 0x00000001 -HRESULT_IO_PENDING = 0x800703E5 -ERROR_OPERATION_ABORTED = 995 -ERROR_NOT_FOUND = 1168 -ERROR_SHARING_VIOLATION = 32 -ERROR_LOCK_VIOLATION = 33 -WAIT_OBJECT_0 = 0 -WAIT_TIMEOUT = 258 -INVALID_HANDLE_VALUE = ctypes.c_void_p(-1).value -GENERIC_READ = 0x80000000 -FILE_SHARE_READ = 0x00000001 -FILE_SHARE_WRITE = 0x00000002 -OPEN_EXISTING = 3 -FILE_ATTRIBUTE_NORMAL = 0x00000080 -IOCTL_STORAGE_QUERY_PROPERTY = 0x002D1400 -IOCTL_STORAGE_GET_HOTPLUG_INFO = 0x002D0C14 -STORAGE_DEVICE_PROPERTY = 0 -PROPERTY_STANDARD_QUERY = 0 -INTERNAL_STORAGE_BUS_TYPES = frozenset({1, 3, 8, 10, 11, 17, 18, 19}) -FILE_SCAN_DEBOUNCE_SECONDS = 0.35 -FILE_SCAN_RETRY_SECONDS = 0.25 -FILE_LOCK_RETRY_SECONDS = 0.5 -PYAS_CONNECTION_MAGIC = 0x53415950 -PYAS_CONNECTION_VERSION = 1 -PROCESS_TERMINATE = 0x0001 -PROCESS_VM_READ = 0x0010 -PROCESS_QUERY_INFORMATION = 0x0400 -PROCESS_SUSPEND_RESUME = 0x0800 -PROCESS_QUERY_LIMITED_INFORMATION = 0x1000 -PYAS_PROCESS_SCAN_ACCESS = PROCESS_TERMINATE | PROCESS_VM_READ | PROCESS_QUERY_INFORMATION | PROCESS_SUSPEND_RESUME | PROCESS_QUERY_LIMITED_INFORMATION -PYAS_PROCESS_NETWORK_ACCESS = PROCESS_TERMINATE | PROCESS_QUERY_INFORMATION | PROCESS_QUERY_LIMITED_INFORMATION - -#################################################################################################### - -class PYAS_CONNECTION_CONTEXT(ctypes.Structure): - _fields_ = [ - ("Size", ctypes.wintypes.ULONG), - ("Version", ctypes.wintypes.ULONG), - ("Magic", ctypes.wintypes.ULONG), - ("ProcessId", ctypes.wintypes.ULONG) - ] - -class OVERLAPPED(ctypes.Structure): - _fields_ = [ - ("Internal", ctypes.c_size_t), - ("InternalHigh", ctypes.c_size_t), - ("Offset", ctypes.wintypes.DWORD), - ("OffsetHigh", ctypes.wintypes.DWORD), - ("hEvent", ctypes.wintypes.HANDLE) - ] - -class STORAGE_PROPERTY_QUERY(ctypes.Structure): - _fields_ = [ - ("PropertyId", ctypes.c_int), - ("QueryType", ctypes.c_int), - ("AdditionalParameters", ctypes.c_ubyte * 1) - ] - -class STORAGE_DESCRIPTOR_HEADER(ctypes.Structure): - _fields_ = [ - ("Version", ctypes.wintypes.DWORD), - ("Size", ctypes.wintypes.DWORD) - ] - -class STORAGE_DEVICE_DESCRIPTOR(ctypes.Structure): - _fields_ = [ - ("Version", ctypes.wintypes.DWORD), - ("Size", ctypes.wintypes.DWORD), - ("DeviceType", ctypes.c_ubyte), - ("DeviceTypeModifier", ctypes.c_ubyte), - ("RemovableMedia", ctypes.c_ubyte), - ("CommandQueueing", ctypes.c_ubyte), - ("VendorIdOffset", ctypes.wintypes.DWORD), - ("ProductIdOffset", ctypes.wintypes.DWORD), - ("ProductRevisionOffset", ctypes.wintypes.DWORD), - ("SerialNumberOffset", ctypes.wintypes.DWORD), - ("BusType", ctypes.c_int), - ("RawPropertiesLength", ctypes.wintypes.DWORD) - ] - -class STORAGE_HOTPLUG_INFO(ctypes.Structure): - _fields_ = [ - ("Size", ctypes.wintypes.DWORD), - ("MediaRemovable", ctypes.c_ubyte), - ("MediaHotplug", ctypes.c_ubyte), - ("DeviceHotplug", ctypes.c_ubyte), - ("WriteCacheEnableOverride", ctypes.c_ubyte) - ] - -#################################################################################################### - -class ProtectMixin: +class ProtectMixin(DriverMixin, SystemMixin, RealtimeMixin, PopupWindowMixin): def is_in_whitelist(self, file_path): p = self.norm_path(file_path, must_exist=False) - if not p: + + if not p: return False p_norm = os.path.normcase(p) - with self.lock_config: + with self.lock_config: whitelist = self.pyas_config.get("white_list", []) for item in whitelist: if isinstance(item, dict): wl_path = item.get("file", "") wl_norm = self.norm_path(wl_path, must_exist=False) + if not wl_norm: continue @@ -132,6 +119,7 @@ class ProtectMixin: def _driver_whitelist_patterns(self, file_path, is_directory=None): normalized_path = self.norm_path(file_path, must_exist=False) + if not normalized_path: return [] @@ -139,11 +127,14 @@ class ProtectMixin: is_directory = os.path.isdir(normalized_path) drive_letter = os.path.splitdrive(normalized_path)[0] - driver_path = "*" + normalized_path[len(drive_letter):] if drive_letter else normalized_path + driver_path = ( + "*" + normalized_path[len(drive_letter) :] if drive_letter else normalized_path + ) patterns = [driver_path] if is_directory: descendant_pattern = driver_path.rstrip("\\/") + "\\*" + if descendant_pattern not in patterns: patterns.append(descendant_pattern) @@ -155,10 +146,12 @@ class ProtectMixin: return False patterns = self._driver_whitelist_patterns(file_path, is_directory) + if not patterns: return False success = True + for driver_path in patterns: msg = PYAS_USER_MESSAGE() msg.Command = 1 if is_add else 2 @@ -166,22 +159,35 @@ class ProtectMixin: try: msg.Path = driver_path - if self.fltlib.FilterSendMessage(self.driver_port, ctypes.byref(msg), ctypes.sizeof(msg), None, 0, ctypes.byref(bytes_returned)) != 0: + + if ( + self.fltlib.FilterSendMessage( + self.driver_port, + ctypes.byref(msg), + ctypes.sizeof(msg), + None, + 0, + ctypes.byref(bytes_returned), + ) + != 0 + ): success = False except Exception: + log_exception("PYAS_Protect.ProtectMixin.sync_driver_whitelist:84") success = False return success -#################################################################################################### - def _try_lock_file(self, file_path): fd = None + try: fd = os.open(file_path, os.O_RDWR | os.O_BINARY) + try: size = os.path.getsize(file_path) except Exception: + log_exception("PYAS_Protect.ProtectMixin._try_lock_file:95") size = 1 lock_size = size if size > 0 else 1 @@ -189,11 +195,15 @@ class ProtectMixin: self.virus_lock[file_path] = (fd, lock_size) return True, None except Exception as e: + log_exception("PYAS_Protect.ProtectMixin._try_lock_file:102") + if fd is not None: try: os.close(fd) except Exception: + log_exception("PYAS_Protect.ProtectMixin._try_lock_file:106") pass + return False, e def _queue_file_lock(self, file_path): @@ -206,6 +216,7 @@ class ProtectMixin: with self.lock_file_ops: if file_path in self.virus_lock: return + success, _ = self._try_lock_file(file_path) if not success: @@ -217,6 +228,7 @@ class ProtectMixin: return False paths_to_lock = [] + if os.path.isdir(target_path): for root, _, files in os.walk(target_path): for file_name in files: @@ -225,23 +237,29 @@ class ProtectMixin: paths_to_lock.append(target_path) failed_locks = [] + with self.lock_file_ops: for file_path in paths_to_lock: if file_path in self.virus_lock: continue success, error = self._try_lock_file(file_path) + if not success: failed_locks.append((file_path, error)) for file_path, error in failed_locks: self._queue_file_lock(file_path) + if not quiet: - self.write_log("INFO", "File Lock Deferred", source=file_path, detail=str(error)) + self.write_log( + "INFO", "File Lock Deferred", source=file_path, detail=str(error) + ) return not failed_locks self._cancel_pending_file_tasks("lock", target_path) + with self.lock_file_ops: target_norm = os.path.normcase(target_path) target_dir = target_norm + os.sep @@ -249,21 +267,29 @@ class ProtectMixin: for locked_path in self.virus_lock: locked_norm = os.path.normcase(locked_path) + if locked_norm == target_norm or locked_norm.startswith(target_dir): keys_to_unlock.append(locked_path) for file_path in keys_to_unlock: fd, lock_size = self.virus_lock[file_path] + try: msvcrt.locking(fd, msvcrt.LK_UNLCK, lock_size) except Exception as e: + log_exception("PYAS_Protect.ProtectMixin.lock_file:170") + if not quiet: - self.write_log("WARN", "unlock_file", source=file_path, detail=str(e), success=False) + self.write_log( + "WARN", "unlock_file", source=file_path, detail=str(e), success=False + ) finally: try: os.close(fd) except Exception: + log_exception("PYAS_Protect.ProtectMixin.lock_file:176") pass + del self.virus_lock[file_path] return True @@ -279,6 +305,7 @@ class ProtectMixin: for item in quarantine_list: file_path = item.get("file") + if not file_path or not os.path.exists(file_path): continue @@ -288,1842 +315,7 @@ class ProtectMixin: self.lock_file(file_path, True, quiet=True) except Exception: + log_exception("PYAS_Protect.ProtectMixin.relock_file:201") pass + time.sleep(5) - -#################################################################################################### - - 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: - 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: - 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: - 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: - 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: - break - - except Exception: - 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: - 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: - break - - except Exception: - 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: - break - - except Exception as e: - 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: - pass - - shutil.rmtree(os.path.join(theme_dir, "CachedFiles"), ignore_errors=True) - self.user32.SystemParametersInfoW(20, 0, wallpaper, 3) - - except Exception as e: - 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: - self.write_log("WARN", "execute_system_repair", detail=str(e), operate=True, success=False) - return False - -#################################################################################################### - - def protect_proc_thread(self): - with self.lock_proc: - self.exist_process = self.get_process_list_pids() - - while True: - with self.lock_config: - if not self.pyas_config.get("process_switch", False): - break - try: - time.sleep(0.1) - cur = self.get_process_list_pids() - - with self.lock_proc: - new_pids = cur - self.exist_process - self.exist_process = cur - - for pid in new_pids: - h = self.kernel32.OpenProcess(PYAS_PROCESS_SCAN_ACCESS, False, pid) - if h: - with self.lock_config: - should_suspend = self.pyas_config.get("suspend_switch", True) - - if should_suspend: - self.ntdll.NtSuspendProcess(h) - with self.lock_proc: - self.suspended_procs.add(h) - try: - self.proc_pool.submit(self.handle_new_process, pid, h, should_suspend) - - except Exception: - if should_suspend: - self.ntdll.NtResumeProcess(h) - with self.lock_proc: - self.suspended_procs.discard(h) - - self.kernel32.CloseHandle(h) - except Exception as e: - self.write_log("WARN", "protect_proc_thread", detail=str(e), success=False) - - def handle_new_process(self, pid, h=None, suspended=None): - if suspended is None: - with self.lock_config: - suspended = self.pyas_config.get("suspend_switch", True) - - if not h: - h = self.kernel32.OpenProcess(PYAS_PROCESS_SCAN_ACCESS, False, pid) - if not h: - return - - if suspended: - self.ntdll.NtSuspendProcess(h) - with self.lock_proc: - if not hasattr(self, 'suspended_procs'): - self.suspended_procs = set() - - self.suspended_procs.add(h) - try: - cmdline = self.get_process_cmdline(h) - process_file = self.get_process_file(h) - if "-scan" in cmdline and self.path_equal(process_file, self.file_pyas): - return - - raw_targets = self.extract_paths_from_cmdline(cmdline) - if process_file: - raw_targets.append(process_file) - - all_targets = [] - for p in raw_targets: - np = self.norm_path(self.device_path_to_drive(p)) - if np and np not in all_targets: - all_targets.append(np) - - with self.lock_config: - ext_filter = self.pyas_config.get("suffix_switch", True) - suffix = self.pyas_config.get("suffix", []) - load_switch = self.pyas_config.get("load_switch", True) - - scan_targets = [] - for file_path in all_targets: - if not file_path or not os.path.isfile(file_path): - continue - if ext_filter and os.path.splitext(file_path)[-1].lower() not in suffix: - continue - if self.is_in_whitelist(file_path): - continue - if process_file and process_file.lower().endswith("explorer.exe") and "/select" in cmdline.lower(): - continue - scan_targets.append(file_path) - - if not scan_targets: - return - - virus_found = False - for file_path in scan_targets: - result = self.safe_scan_engine(file_path) - self.cloud_check(file_path) - - if result: - self.kernel32.TerminateProcess(h, 0) - self.write_log("BLOCK", "Process Block", pid=pid, source=file_path, file_hash=self.calc_file_hash(file_path)) - virus_found = True - break - - if not virus_found and load_switch: - hidden_virus_path = self.scan_process_memory(pid, h) - if hidden_virus_path: - self.kernel32.TerminateProcess(h, 0) - self.write_log("BLOCK", "Process DLL Block", pid=pid, source=hidden_virus_path, file_hash=self.calc_file_hash(hidden_virus_path)) - - finally: - if suspended: - try: - self.ntdll.NtResumeProcess(h) - except Exception: - pass - with self.lock_proc: - self.suspended_procs.discard(h) - try: - self.kernel32.CloseHandle(h) - except Exception: - pass - -#################################################################################################### - - def scan_process_memory(self, pid, h_process): - scanned_paths = set() - address = 0 - mbi = MEMORY_BASIC_INFORMATION() - system_dir = self.path_system.lower() - - buf = ctypes.create_unicode_buffer(1024) - max_address = 0x7FFFFFFFFFFF if ctypes.sizeof(ctypes.c_void_p) == 8 else 0x7FFFFFFF - - with self.lock_config: - ext_filter = self.pyas_config.get("suffix_switch", True) - suffix = self.pyas_config.get("suffix", []) - - try: - while address < max_address and self.kernel32.VirtualQueryEx(h_process, ctypes.c_void_p(address), ctypes.byref(mbi), ctypes.sizeof(mbi)): - if mbi.State == 0x1000 and mbi.Type == 0x1000000: - if self.psapi.GetMappedFileNameW(h_process, ctypes.c_void_p(address), buf, 1024): - - raw_path = buf.value - if raw_path.startswith("\\"): - raw_path = self.device_path_to_drive(raw_path) - - file_path = self.norm_path(raw_path) - if file_path and file_path not in scanned_paths and self.path_system not in file_path: - - scanned_paths.add(file_path) - file_path_lower = file_path.lower() - - if not file_path_lower.startswith(system_dir) and not self.is_in_whitelist(file_path): - ext = os.path.splitext(file_path_lower)[-1] - - if ext != ".exe" and (not ext_filter or ext in suffix): - if self.safe_scan_engine(file_path): - return file_path - - self.cloud_check(file_path) - - if mbi.RegionSize == 0: - break - address += mbi.RegionSize - - finally: - scanned_paths.clear() - - return None - -#################################################################################################### - - def protect_file_thread(self): - with self.lock_file_ops: - if getattr(self, 'h_dir_file', None): - return - - hDir = self.kernel32.CreateFileW(self.path_user, 0x0001, 0x00000007, None, 3, 0x02000000, None) - if not hDir or hDir == -1: - return - - with self.lock_file_ops: - self.h_dir_file = hDir - - try: - buffer = ctypes.create_string_buffer(262144) - temp_prefix = os.path.normcase(self.path_temp) - if not temp_prefix.endswith(os.sep): - temp_prefix += os.sep - - while True: - with self.lock_config: - if not self.pyas_config.get("document_switch", False): - break - try: - bytes_returned = ctypes.wintypes.DWORD() - res = self.kernel32.ReadDirectoryChangesW(self.h_dir_file, buffer, ctypes.sizeof(buffer), True, 0x0000001F, ctypes.byref(bytes_returned), None, None) - if not res or bytes_returned.value == 0: - break - - offset = 0 - while True: - notify = FILE_NOTIFY_INFORMATION.from_buffer(buffer, offset) - raw_filename = notify.FileName[:notify.FileNameLength // 2] - - if raw_filename and notify.Action in [1, 3, 5]: - file_path = self.norm_path(os.path.join(self.path_user, raw_filename), must_exist=True) - - if file_path and not self.is_in_whitelist(file_path): - norm_path = os.path.normcase(file_path) - - if not (norm_path.startswith(temp_prefix) and norm_path[len(temp_prefix):].startswith("_mei")): - with self.lock_config: - ext_filter = self.pyas_config.get("suffix_switch", True) - suffix = self.pyas_config.get("suffix", []) - - if not ext_filter or os.path.splitext(file_path)[-1].lower() in suffix: - self._queue_file_scan(file_path) - - if notify.NextEntryOffset == 0: - break - - offset += notify.NextEntryOffset - - except Exception: - break - finally: - with self.lock_file_ops: - if getattr(self, 'h_dir_file', None): - try: - self.kernel32.CloseHandle(self.h_dir_file) - except Exception: - pass - self.h_dir_file = None - - def _cancel_pending_file_tasks(self, task_name=None, target_path=None): - target_norm = os.path.normcase(target_path) if target_path else None - target_dir = target_norm + os.sep if target_norm else None - - with self.lock_file_ops: - task_keys = [] - for task_key in self.file_task_timers: - name, path_key = task_key - if task_name and name != task_name: - continue - if target_norm and path_key != target_norm and not path_key.startswith(target_dir): - continue - task_keys.append(task_key) - - timers = [self.file_task_timers.pop(task_key) for task_key in task_keys] - for task_key in task_keys: - self.file_task_generations.pop(task_key, None) - - for timer in timers: - timer.cancel() - - def _cancel_pending_file_scans(self): - self._cancel_pending_file_tasks("scan") - - def _schedule_file_task(self, task_name, file_path, callback, delay): - task_key = (task_name, os.path.normcase(file_path)) - - with self.lock_file_ops: - previous_timer = self.file_task_timers.get(task_key) - if previous_timer: - previous_timer.cancel() - - generation = self.file_task_generations.get(task_key, 0) + 1 - timer = threading.Timer(delay, self._run_file_task, args=(task_key, file_path, generation, callback)) - timer.daemon = True - self.file_task_timers[task_key] = timer - self.file_task_generations[task_key] = generation - - timer.start() - - def _run_file_task(self, task_key, file_path, generation, callback): - with self.lock_file_ops: - if self.file_task_generations.get(task_key) != generation: - return - if self.file_task_timers.get(task_key) is threading.current_thread(): - self.file_task_timers.pop(task_key, None) - - try: - callback(file_path) - finally: - with self.lock_file_ops: - if self.file_task_generations.get(task_key) == generation: - self.file_task_timers.pop(task_key, None) - self.file_task_generations.pop(task_key, None) - - def _queue_file_scan(self, file_path, delay=FILE_SCAN_DEBOUNCE_SECONDS): - self._schedule_file_task("scan", file_path, self.handle_new_file, delay) - - def _acquire_file_scan_gate(self, file_path): - ctypes.set_last_error(0) - handle = self.kernel32.CreateFileW( - file_path, - GENERIC_READ, - FILE_SHARE_READ, - None, - OPEN_EXISTING, - FILE_ATTRIBUTE_NORMAL, - None - ) - - if handle == INVALID_HANDLE_VALUE: - return None, ctypes.get_last_error() - - return handle, 0 - - def handle_new_file(self, file_path): - try: - with self.lock_config: - if not self.pyas_config.get("document_switch", False): - return - - gate_handle, error = self._acquire_file_scan_gate(file_path) - if not gate_handle: - if error in (ERROR_SHARING_VIOLATION, ERROR_LOCK_VIOLATION): - self._queue_file_scan(file_path, FILE_SCAN_RETRY_SECONDS) - return - - try: - with self.lock_file_ops: - if file_path in self.virus_lock: - return - - result = self.safe_scan_engine(file_path) - self.cloud_check(file_path) - - if result: - if self.manage_named_list("quarantine", [file_path], action="add", lock_func=self.lock_file) > 0: - self.write_log("BLOCK", "File Block", source=file_path, file_hash=self.calc_file_hash(file_path)) - finally: - self.kernel32.CloseHandle(gate_handle) - except Exception: - pass - -#################################################################################################### - - def protect_net_thread(self): - with self.lock_net: - self.exist_connections = set() - - while True: - with self.lock_config: - if not self.pyas_config.get("network_switch", False): break - - try: - time.sleep(0.5) - conns = self.get_connections_list() - - with self.lock_net: - new_conns = conns - self.exist_connections - self.exist_connections = conns - - for key in new_conns: - self.protect_pool.submit(self.handle_new_connection, key) - - except Exception as e: - self.write_log("WARN", "protect_net_thread", detail=str(e), success=False) - - with self.lock_net: - self.exist_connections = set() - - def handle_new_connection(self, key): - pid, remote_addr, remote_port = key - try: - h = self.kernel32.OpenProcess(PYAS_PROCESS_NETWORK_ACCESS, False, pid) - if not h: - return - - try: - remote_ip = f"{remote_addr & 0xFF}.{(remote_addr >> 8) & 0xFF}.{(remote_addr >> 16) & 0xFF}.{(remote_addr >> 24) & 0xFF}" - file_path = self.norm_path(self.get_process_file(h)) - - if file_path and not self.is_in_whitelist(file_path) and hasattr(self.heuristic, "network") and remote_ip in self.heuristic.network: - self.kernel32.TerminateProcess(h, 0) - self.write_log("BLOCK", "Network Block", pid=pid, source=file_path, target=remote_ip, file_hash=self.calc_file_hash(file_path)) - - finally: - self.kernel32.CloseHandle(h) - - except Exception as e: - self.write_log("WARN", "handle_new_connection", pid=pid, detail=str(e), success=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: - 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: - pass - -#################################################################################################### - - def capture_popup_window(self): - overlay = self.user32.CreateWindowExW(0x000800A8, "STATIC", "", 0x90000004, 0, 0, 0, 0, None, None, None, None) - self.user32.SetLayeredWindowAttributes(overlay, 0, 120, 2) - - target_hwnd = None - last_hwnd = None - msg = ctypes.wintypes.MSG() - - while self.user32.GetAsyncKeyState(0x01) & 0x8000: - time.sleep(0.01) - - try: - start_time = time.time() - while time.time() - start_time < 30: - if self.user32.GetAsyncKeyState(0x1B) & 0x8000: - break - - pt = POINT() - self.user32.GetCursorPos(ctypes.byref(pt)) - - hwnd = self.user32.WindowFromPoint(pt) - root_hwnd = self.user32.GetAncestor(hwnd, 2) - if not root_hwnd: - root_hwnd = hwnd - - if root_hwnd and root_hwnd != overlay: - if root_hwnd != last_hwnd: - last_hwnd = root_hwnd - rect = RECT() - self.user32.GetWindowRect(root_hwnd, ctypes.byref(rect)) - self.user32.SetWindowPos(overlay, -1, rect.left, rect.top, rect.right - rect.left, rect.bottom - rect.top, 0x0050) - - if self.user32.GetAsyncKeyState(0x01) & 0x8000: - target_hwnd = root_hwnd - break - - while self.user32.PeekMessageW(ctypes.byref(msg), 0, 0, 0, 1): - self.user32.TranslateMessage(ctypes.byref(msg)) - self.user32.DispatchMessageW(ctypes.byref(msg)) - - time.sleep(0.02) - finally: - self.user32.DestroyWindow(overlay) - - if not target_hwnd: - return None - - pid = ctypes.c_ulong(0) - self.user32.GetWindowThreadProcessId(target_hwnd, ctypes.byref(pid)) - - if pid.value != self.pid_pyas and pid.value > 4: - length = self.user32.GetWindowTextLengthW(target_hwnd) - title = ctypes.create_unicode_buffer(length + 1) - class_name = ctypes.create_unicode_buffer(256) - - self.user32.GetWindowTextW(target_hwnd, title, length + 1) - self.user32.GetClassNameW(target_hwnd, class_name, 256) - - proc_name, _ = self.get_exe_info(pid.value) - t_str, c_str = str(title.value), str(class_name.value) - - if proc_name and not any(item.get("exe") == proc_name or item.get("class") == c_str for item in self.pass_windows): - return {"exe": proc_name, "class": c_str, "title": t_str} - - return None - - def add_popup_rule(self, rule): - if not rule: - return False - - with self.lock_config: - target_list = self.pyas_config.setdefault("block_list", []) - if any(item.get("exe") == rule.get("exe") and item.get("class") == rule.get("class") and item.get("title") == rule.get("title") for item in target_list): - return False - - target_list.append(rule) - self.write_log("INFO", "Config Update", detail=f"List [block_list] add: {rule}") - self.save_config() - return True - - def popup_intercept_thread(self): - hwnd_list = [] - - def enum_windows_callback(hWnd, lParam): - if self.user32.IsWindowVisible(hWnd): - hwnd_list.append(hWnd) - return True - - WNDENUMPROC = ctypes.WINFUNCTYPE(ctypes.c_bool, ctypes.c_int, ctypes.c_int) - self._enum_cb_keepalive = WNDENUMPROC(enum_windows_callback) - - while True: - try: - time.sleep(0.5) - with self.lock_config: - rules = self.pyas_config.get("block_list", []) - if not rules: - continue - - hwnd_list.clear() - self.user32.EnumWindows(self._enum_cb_keepalive, 0) - - for hWnd in list(hwnd_list): - pid = ctypes.c_ulong(0) - self.user32.GetWindowThreadProcessId(hWnd, ctypes.byref(pid)) - if pid.value <= 4: - continue - - proc_name, _ = self.get_exe_info(pid.value) - length = self.user32.GetWindowTextLengthW(hWnd) - title, class_name = ctypes.create_unicode_buffer(length + 1), ctypes.create_unicode_buffer(256) - - self.user32.GetWindowTextW(hWnd, title, length + 1) - self.user32.GetClassNameW(hWnd, class_name, 256) - t_str, c_str = str(title.value), str(class_name.value) - - if any(item.get("exe") == proc_name or item.get("class") == c_str for item in self.pass_windows): - continue - - is_block = False - for item in rules: - if item.get("exe") and item.get("exe") == proc_name and (item.get("class") == c_str or item.get("title") == t_str): - is_block = True - break - - if is_block: - for msg in [0x0010, 0x0002, 0x0012, 0x0112]: - self.user32.SendMessageW(hWnd, msg, 0xF060, 0) - if proc_name: - self.kill_process(pid.value) - - except Exception: - pass - -#################################################################################################### - - def _query_service_state_handle(self, service): - status = SERVICE_STATUS_PROCESS() - needed = ctypes.wintypes.DWORD(0) - buffer = ctypes.cast(ctypes.byref(status), ctypes.POINTER(ctypes.c_ubyte)) - - if not self.advapi32.QueryServiceStatusEx(service, 0, buffer, ctypes.sizeof(status), ctypes.byref(needed)): - return None, ctypes.get_last_error() - - return int(status.dwCurrentState), 0 - - def _query_driver_service_state(self): - scm = self.advapi32.OpenSCManagerW(None, None, 0x0001) - if not scm: - return None - - service = None - try: - service = self.advapi32.OpenServiceW(scm, "PYAS_Driver", 0x0004) - if not service: - return None - - state, _ = self._query_service_state_handle(service) - return state - finally: - if service: - self.advapi32.CloseServiceHandle(service) - self.advapi32.CloseServiceHandle(scm) - - def _get_driver_client_image_path(self): - buffer = ctypes.create_unicode_buffer(32768) - size = ctypes.wintypes.DWORD(len(buffer)) - if not self.kernel32.QueryFullProcessImageNameW( - self.kernel32.GetCurrentProcess(), - 0, - buffer, - ctypes.byref(size) - ): - return None - - image_path = os.path.abspath(buffer.value) - drive, tail = os.path.splitdrive(image_path) - if not drive or not tail: - return None - - device_buffer = ctypes.create_unicode_buffer(32768) - if not self.kernel32.QueryDosDeviceW(drive, device_buffer, len(device_buffer)): - return None - - return device_buffer.value.rstrip("\\") + tail - - def _configure_driver_identity(self): - image_path = self._get_driver_client_image_path() - if not image_path: - return False, ctypes.get_last_error() or 1 - - try: - access = winreg.KEY_SET_VALUE | getattr(winreg, "KEY_WOW64_64KEY", 0) - with winreg.CreateKeyEx( - winreg.HKEY_LOCAL_MACHINE, - r"SYSTEM\CurrentControlSet\Services\PYAS_Driver", - 0, - access - ) as key: - winreg.SetValueEx(key, "ClientImagePath", 0, winreg.REG_SZ, image_path) - return True, 0 - except OSError as e: - return False, int(getattr(e, "winerror", 0) or 1) - - def _ensure_driver_service(self): - desired_service_access = 0x0002 | 0x0004 | 0x0010 | 0x00010000 - dependencies = ctypes.create_unicode_buffer("FltMgr\0\0") - dependencies_ptr = ctypes.cast(dependencies, ctypes.wintypes.LPCWSTR) - last_error = 0 - - for attempt in range(20): - scm = self.advapi32.OpenSCManagerW(None, None, 0x0001 | 0x0002) - if not scm: - return False, ctypes.get_last_error() - - service = None - try: - ctypes.set_last_error(0) - service = self.advapi32.CreateServiceW( - scm, "PYAS_Driver", "PYAS_Driver", desired_service_access, - 0x00000001, 0x00000003, 0x00000001, self.path_drivers, - "FSFilter Activity Monitor", None, dependencies_ptr, None, None - ) - - if not service: - last_error = ctypes.get_last_error() - if last_error in (1073, 1078): - service = self.advapi32.OpenServiceW(scm, "PYAS_Driver", desired_service_access) - if not service: - last_error = ctypes.get_last_error() - - elif last_error == 1072: - pass - - if service: - ctypes.set_last_error(0) - changed = self.advapi32.ChangeServiceConfigW( - service, 0x00000001, 0x00000003, 0x00000001, self.path_drivers, - "FSFilter Activity Monitor", None, dependencies_ptr, None, None, None - ) - if not changed: - last_error = ctypes.get_last_error() - return False, last_error - - return True, 0 - finally: - if service: - self.advapi32.CloseServiceHandle(service) - self.advapi32.CloseServiceHandle(scm) - - if last_error != 1072: - break - if attempt < 19: - time.sleep(0.05) - - return False, last_error or 1 - - def _start_driver_service(self): - scm = self.advapi32.OpenSCManagerW(None, None, 0x0001) - if not scm: - return False, ctypes.get_last_error() - - service = None - try: - service = self.advapi32.OpenServiceW(scm, "PYAS_Driver", 0x0004 | 0x0010) - if not service: - return False, ctypes.get_last_error() - - state, error = self._query_service_state_handle(service) - if state in (2, 4): - return True, 0 - if error: - return False, error - - ctypes.set_last_error(0) - if self.advapi32.StartServiceW(service, 0, None): - return True, 0 - - error = ctypes.get_last_error() - if error == 1056: - return True, 0 - - return False, error - finally: - if service: - self.advapi32.CloseServiceHandle(service) - self.advapi32.CloseServiceHandle(scm) - - def uninstall_system_driver(self): - if not self.stop_system_driver(): - return False, 1051 - return self._delete_driver_service() - - def _delete_driver_service(self, timeout=15.0): - deadline = time.monotonic() + timeout - last_error = 0 - delete_requested = False - - while time.monotonic() < deadline: - scm = self.advapi32.OpenSCManagerW(None, None, 0x0001) - if not scm: - return False, ctypes.get_last_error() - - service = None - try: - ctypes.set_last_error(0) - service = self.advapi32.OpenServiceW(scm, "PYAS_Driver", 0x0004 | 0x00010000) - if not service: - last_error = ctypes.get_last_error() - if last_error == 1060: - return True, 0 - if last_error != 1072: - return False, last_error - else: - state, status_error = self._query_service_state_handle(service) - if status_error: - return False, status_error - - if state not in (2, 3, 4) and not delete_requested: - ctypes.set_last_error(0) - if self.advapi32.DeleteService(service): - delete_requested = True - last_error = 1072 - else: - last_error = ctypes.get_last_error() - if last_error == 1060: - return True, 0 - if last_error == 1072: - delete_requested = True - else: - return False, last_error - elif state in (2, 3, 4): - last_error = 1051 - finally: - if service: - self.advapi32.CloseServiceHandle(service) - self.advapi32.CloseServiceHandle(scm) - - time.sleep(0.1) - - return False, last_error or 1072 - - def install_system_driver(self): - try: - ensured, service_error = self._ensure_driver_service() - if not ensured: - self.write_log("WARN", "Driver Service", detail=f"CreateServiceW/ChangeServiceConfigW failed: 0x{service_error & 0xFFFFFFFF:08X}", success=False) - return False - - configured, identity_error = self._configure_driver_identity() - if not configured: - self.write_log("WARN", "Driver Identity", detail=f"ClientImagePath configuration failed: 0x{identity_error & 0xFFFFFFFF:08X}", success=False) - return False - - self._reg_write(winreg.HKEY_LOCAL_MACHINE, r"SYSTEM\CurrentControlSet\Services\PYAS_Driver\Instances", "DefaultInstance", winreg.REG_SZ, "PYAS Instance") - self._reg_write(winreg.HKEY_LOCAL_MACHINE, r"SYSTEM\CurrentControlSet\Services\PYAS_Driver\Instances\PYAS Instance", "Altitude", winreg.REG_SZ, "320000") - self._reg_write(winreg.HKEY_LOCAL_MACHINE, r"SYSTEM\CurrentControlSet\Services\PYAS_Driver\Instances\PYAS Instance", "Flags", winreg.REG_DWORD, 0) - - started, start_error = self._start_driver_service() - if not started: - self.write_log("WARN", "Driver Service", detail=f"StartServiceW failed: 0x{start_error & 0xFFFFFFFF:08X}", success=False) - return False - - final_state = None - last_start_attempt = time.monotonic() - deadline = time.monotonic() + 3.0 - - while time.monotonic() < deadline: - service_state = self._query_driver_service_state() - if service_state == 4: - final_state = self._query_driver_state() - if final_state == 2: - return True - if final_state in (4, 5): - break - elif service_state == 1: - now = time.monotonic() - if now - last_start_attempt >= 0.1: - started, start_error = self._start_driver_service() - last_start_attempt = now - if not started: - break - elif service_state is None: - break - - time.sleep(0.02) - - if self.check_system_driver(): - self.stop_system_driver() - - detail = "Driver entered unload-retry state during startup" if final_state == 4 else "Driver did not reach running state" - self.write_log("WARN", "Driver Protection", detail=detail, success=False) - return False - - except Exception as e: - self.write_log("WARN", "install_system_driver", detail=str(e), success=False) - return False - - def _driver_handle_value(self, handle): - if handle is None: - return None - - value = getattr(handle, "value", handle) - if value is None: - return None - - try: - return int(value) - except Exception: - return None - - def _connect_driver_port(self, asynchronous=False): - context = PYAS_CONNECTION_CONTEXT() - context.Size = ctypes.sizeof(context) - context.Version = PYAS_CONNECTION_VERSION - context.Magic = PYAS_CONNECTION_MAGIC - context.ProcessId = os.getpid() - - temp_port = ctypes.wintypes.HANDLE() - status = self.fltlib.FilterConnectCommunicationPort( - "\\PYAS_Output_Pipe", - 0 if asynchronous else FLT_PORT_FLAG_SYNC_HANDLE, - ctypes.byref(context), - ctypes.sizeof(context), - None, - ctypes.byref(temp_port) - ) - if status != 0: - return None - return temp_port - - def _detach_driver_port(self, expected_port=None): - expected_value = self._driver_handle_value(expected_port) - - with self.lock_driver: - current_port = self.driver_port - current_value = self._driver_handle_value(current_port) - if current_value is None: - return None - - if expected_port is not None and current_value != expected_value: - return None - - self.driver_port = None - return current_port - - def _close_driver_port(self, expected_port=None): - port = self._detach_driver_port(expected_port) - if not port: - return False - - try: - self.kernel32.CloseHandle(port) - except Exception: - pass - return True - - def _driver_listener_should_run(self): - if self.driver_stop_event.is_set(): - return False - - with self.lock_config: - return bool(self.pyas_config.get("driver_switch", False)) - - def start_driver_listener(self, wait_ready=True): - with self.lock_driver: - current = self.driver_listener_thread - if current and current.is_alive(): - return self.driver_listener_ready_event.is_set() and not self.driver_stop_event.is_set() - - self.driver_stop_event.clear() - self.driver_listener_ready_event.clear() - self.driver_listener_failed_event.clear() - thread = threading.Thread(target=self.pipe_server_thread, daemon=True) - self.driver_listener_thread = thread - thread.start() - - if not wait_ready: - return True - - deadline = time.monotonic() + 5.0 - while time.monotonic() < deadline: - if self.driver_listener_ready_event.wait(0.02): - return True - if self.driver_listener_failed_event.is_set() or not thread.is_alive(): - break - - self._stop_driver_listener() - return False - - def _stop_driver_listener(self): - self.driver_stop_event.set() - - with self.lock_driver: - thread = self.driver_listener_thread - port = self.driver_port - - if port: - try: - self.kernel32.CancelIoEx(port, None) - except Exception: - pass - - if thread and thread is not threading.current_thread(): - thread.join(timeout=2.0) - - alive = bool(thread and thread.is_alive()) - if not alive: - with self.lock_driver: - if self.driver_listener_thread is thread: - self.driver_listener_thread = None - self.driver_listener_ready_event.clear() - self._close_driver_port() - - return not alive - - def _resume_driver_listener(self): - if not self.check_system_driver(): - return False - - with self.lock_config: - enabled = bool(self.pyas_config.get("driver_switch", False)) - - if not enabled: - return False - - self._close_driver_port() - return self.start_driver_listener(wait_ready=True) - - def _set_driver_unload_authorization(self, enabled): - msg = PYAS_USER_MESSAGE() - msg.Command = 5 if enabled else 6 - msg.Path = "" - - for _ in range(3): - if not self._ensure_driver_port(): - time.sleep(0.1) - continue - - with self.lock_driver: - current_port = self.driver_port - if not current_port: - continue - - bytes_returned = ctypes.wintypes.DWORD(0) - try: - status = self.fltlib.FilterSendMessage( - current_port, - ctypes.byref(msg), - ctypes.sizeof(msg), - None, - 0, - ctypes.byref(bytes_returned) - ) - except Exception as e: - self.write_log("WARN", "Driver Unload Authorization", detail=str(e), success=False) - status = -1 - - if status == 0: - return True - - self._close_driver_port(current_port) - time.sleep(0.1) - - return False - - def _ensure_driver_port(self): - for _ in range(20): - with self.lock_driver: - if self.driver_port: - return True - - temp_port = self._connect_driver_port() - if temp_port: - with self.lock_driver: - if not self.driver_port: - self.driver_port = temp_port - return True - - try: - self.kernel32.CloseHandle(temp_port) - except Exception: - pass - return True - - time.sleep(0.05) - - return False - - def _query_driver_state(self): - for use_existing in (True, False): - local_port = None - port = None - - if use_existing: - with self.lock_driver: - if self.driver_port: - port = self.driver_port - if not port: - continue - else: - local_port = self._connect_driver_port() - if not local_port: - return None - port = local_port - - msg = PYAS_USER_MESSAGE() - msg.Command = 7 - msg.Path = "" - state = ctypes.wintypes.DWORD(0) - bytes_returned = ctypes.wintypes.DWORD(0) - - try: - status = self.fltlib.FilterSendMessage( - port, - ctypes.byref(msg), - ctypes.sizeof(msg), - ctypes.byref(state), - ctypes.sizeof(state), - ctypes.byref(bytes_returned) - ) - if status == 0 and bytes_returned.value == ctypes.sizeof(state): - return int(state.value) - except Exception: - pass - finally: - if local_port: - try: - self.kernel32.CloseHandle(local_port) - except Exception: - pass - - if use_existing: - self._close_driver_port(port) - - return None - - def _enable_token_privilege(self, privilege_name): - token = ctypes.wintypes.HANDLE() - desired_access = 0x0020 | 0x0008 - - if not self.advapi32.OpenProcessToken(self.kernel32.GetCurrentProcess(), desired_access, ctypes.byref(token)): - return None, ctypes.get_last_error() - - luid = LUID() - if not self.advapi32.LookupPrivilegeValueW(None, privilege_name, ctypes.byref(luid)): - error = ctypes.get_last_error() - self.kernel32.CloseHandle(token) - return None, error - - new_state = TOKEN_PRIVILEGES() - new_state.PrivilegeCount = 1 - new_state.Privileges[0].Luid = luid - new_state.Privileges[0].Attributes = 0x00000002 - - previous_state = TOKEN_PRIVILEGES() - return_length = ctypes.wintypes.DWORD(0) - ctypes.set_last_error(0) - - adjusted = self.advapi32.AdjustTokenPrivileges( - token, - False, - ctypes.byref(new_state), - ctypes.sizeof(previous_state), - ctypes.byref(previous_state), - ctypes.byref(return_length) - ) - error = ctypes.get_last_error() - - if not adjusted or error == 1300: - self.kernel32.CloseHandle(token) - return None, error or 1 - - return (token, previous_state, return_length.value), 0 - - def _restore_token_privilege(self, privilege_state): - if not privilege_state: - return - - token, previous_state, previous_length = privilege_state - try: - if previous_length: - self.advapi32.AdjustTokenPrivileges(token, False, ctypes.byref(previous_state), 0, None, None) - finally: - self.kernel32.CloseHandle(token) - - def _filter_unload_with_privilege(self, filter_name): - result = { - "status": None, - "privilege_error": 0, - "authorization_failed": False, - "exception": None, - "attempts": 0 - } - privilege_state = None - unload_authorized = False - - self.driver_unload_worker = None - self.driver_unload_result = result - - try: - privilege_state, privilege_error = self._enable_token_privilege("SeLoadDriverPrivilege") - if not privilege_state: - result["privilege_error"] = privilege_error - return result - - retryable_statuses = {0x80070522, 0x80070005, 0x800700AA} - - for delay in (0.0, 0.05, 0.15): - if delay: - time.sleep(delay) - - if not self.check_system_driver(): - result["status"] = 0 - unload_authorized = False - return result - - if not self._set_driver_unload_authorization(True): - continue - - unload_authorized = True - result["attempts"] += 1 - unload_status = self.fltlib.FilterUnload(filter_name) - result["status"] = unload_status - - if unload_status == 0: - unload_authorized = False - return result - - revoked = False - if self.check_system_driver(): - try: - revoked = self._set_driver_unload_authorization(False) - except Exception: - revoked = False - unload_authorized = not revoked - - if (unload_status & 0xFFFFFFFF) not in retryable_statuses: - return result - - if result["status"] is None: - result["authorization_failed"] = True - - except Exception as e: - result["exception"] = str(e) - finally: - if unload_authorized and self.check_system_driver(): - try: - self._set_driver_unload_authorization(False) - except Exception: - pass - - self._restore_token_privilege(privilege_state) - - return result - - def stop_system_driver(self): - with self.lock_driver_unload: - runtime_unloaded = False - unload_confirmed = False - try: - if not self._stop_driver_listener(): - self.write_log("WARN", "Driver Protection", detail="Driver listener cancellation did not complete", success=False) - return False - - if not self.check_system_driver(): - runtime_unloaded = True - unload_confirmed = True - else: - if not self._ensure_driver_port(): - if not self.check_system_driver(): - runtime_unloaded = True - unload_confirmed = True - else: - self.write_log("WARN", "Driver Protection", detail="Control port unavailable", success=False) - return False - - if not runtime_unloaded: - unload_result = self._filter_unload_with_privilege("PYAS_Driver") - - if unload_result["privilege_error"]: - error = unload_result["privilege_error"] - self.write_log("WARN", "Driver Protection", detail=f"SeLoadDriverPrivilege unavailable: 0x{error & 0xFFFFFFFF:08X}", success=False) - return False - - if unload_result["authorization_failed"]: - self.write_log("WARN", "Driver Protection", detail="Unload authorization rejected", success=False) - return False - - if unload_result["exception"]: - self.write_log("WARN", "stop_system_driver", detail=unload_result["exception"], success=False) - return False - - unload_status = unload_result["status"] - if unload_status == 0: - runtime_unloaded = True - unload_confirmed = True - elif not self.check_system_driver(): - runtime_unloaded = True - unload_confirmed = True - else: - attempts = unload_result["attempts"] - self.write_log("WARN", "Driver Protection", detail=f"FilterUnload failed after {attempts} attempt(s): 0x{unload_status & 0xFFFFFFFF:08X}", success=False) - return False - - self._close_driver_port() - - if unload_confirmed: - deadline = time.monotonic() + 0.5 - while time.monotonic() < deadline and self.check_system_driver(): - time.sleep(0.01) - - if self.check_system_driver(): - self.write_log("INFO", "Driver Protection", detail="Filter unloaded; SCM state is still settling") - - return runtime_unloaded - except Exception as e: - self.write_log("WARN", "stop_system_driver", detail=str(e), success=False) - return False - finally: - if not runtime_unloaded and self.check_system_driver(): - self._resume_driver_listener() - - def check_system_driver(self): - runtime_state = self._query_driver_state() - if runtime_state in (1, 2, 3, 4): - return True - - service_state = self._query_driver_service_state() - return service_state in (2, 3, 4) - - def clear_driver_rules(self): - with self.lock_driver: - if not self.driver_port: - return False - - msg = PYAS_USER_MESSAGE() - msg.Command = 4 - msg.Path = "" - bytes_returned = ctypes.wintypes.DWORD(0) - - try: - return self.fltlib.FilterSendMessage( - self.driver_port, - ctypes.byref(msg), - ctypes.sizeof(msg), - None, - 0, - ctypes.byref(bytes_returned) - ) == 0 - except Exception: - return False - - def load_driver_rule_file(self, json_path): - with self.lock_driver: - if not self.driver_port: - return False - - norm_path = os.path.abspath(json_path) - if not os.path.exists(norm_path): - return False - - nt_path = f"\\??\\{norm_path}" - - msg = PYAS_USER_MESSAGE() - msg.Command = 3 - msg.Path = nt_path - bytes_returned = ctypes.wintypes.DWORD(0) - - try: - return self.fltlib.FilterSendMessage( - self.driver_port, - ctypes.byref(msg), - ctypes.sizeof(msg), - None, - 0, - ctypes.byref(bytes_returned) - ) == 0 - except Exception: - return False - - def _cancel_driver_receive(self, port, overlapped): - try: - ctypes.set_last_error(0) - cancelled = self.kernel32.CancelIoEx(port, ctypes.byref(overlapped)) - error = ctypes.get_last_error() - if not cancelled and error not in (0, ERROR_NOT_FOUND): - return False - - wait_status = self.kernel32.WaitForSingleObject(overlapped.hEvent, 1000) - if wait_status != WAIT_OBJECT_0: - return False - - transferred = ctypes.wintypes.DWORD(0) - self.kernel32.GetOverlappedResult( - port, - ctypes.byref(overlapped), - ctypes.byref(transferred), - False - ) - return True - except Exception: - return False - - def _receive_driver_message(self, port): - message = PYAS_FULL_MESSAGE() - overlapped = OVERLAPPED() - event_handle = self.kernel32.CreateEventW(None, True, False, None) - if not event_handle: - return "error", None - - overlapped.hEvent = event_handle - try: - status = self.fltlib.FilterGetMessage( - port, - ctypes.byref(message), - ctypes.sizeof(PYAS_FULL_MESSAGE), - ctypes.byref(overlapped) - ) - - if status == 0: - return "message", message - - if (status & 0xFFFFFFFF) != HRESULT_IO_PENDING: - return "disconnect", None - - while True: - if self.driver_stop_event.is_set(): - self._cancel_driver_receive(port, overlapped) - return "stop", None - - wait_status = self.kernel32.WaitForSingleObject(event_handle, 20) - if wait_status == WAIT_TIMEOUT: - continue - if wait_status != WAIT_OBJECT_0: - self._cancel_driver_receive(port, overlapped) - return "error", None - - transferred = ctypes.wintypes.DWORD(0) - if self.kernel32.GetOverlappedResult( - port, - ctypes.byref(overlapped), - ctypes.byref(transferred), - False - ): - return "message", message - - error = ctypes.get_last_error() - if error == ERROR_OPERATION_ABORTED and self.driver_stop_event.is_set(): - return "stop", None - return "disconnect", None - finally: - self.kernel32.CloseHandle(event_handle) - - def pipe_server_thread(self): - owned_port = None - ready = False - try: - while self._driver_listener_should_run(): - temp_port = self._connect_driver_port(asynchronous=True) - if not temp_port: - self.driver_stop_event.wait(0.02) - continue - - if not self._driver_listener_should_run(): - self.kernel32.CloseHandle(temp_port) - break - - with self.lock_driver: - if self.driver_port: - accepted = False - else: - self.driver_port = temp_port - owned_port = temp_port - accepted = True - - if not accepted: - self.kernel32.CloseHandle(temp_port) - self.driver_stop_event.wait(0.02) - continue - - rule_files = [] - if os.path.exists(self.path_rules): - rule_files = sorted( - os.path.join(self.path_rules, name) - for name in os.listdir(self.path_rules) - if name.lower().endswith(".json") and name.lower() != "rules_driver_p1.json" - ) - - load_failed = False - for rule_file in rule_files: - if self.driver_stop_event.is_set() or not self.load_driver_rule_file(rule_file): - load_failed = True - break - - if load_failed: - self.driver_listener_failed_event.set() - break - - with self.lock_config: - whitelist = list(self.pyas_config.get("white_list", [])) - - for item in whitelist: - if self.driver_stop_event.is_set(): - break - if isinstance(item, dict) and item.get("file"): - self.sync_driver_whitelist(item["file"], True, item.get("is_dir")) - - if self.driver_stop_event.is_set(): - break - - ready = True - self.driver_listener_ready_event.set() - - while self._driver_listener_should_run(): - with self.lock_driver: - current_port = self.driver_port - owns_current = self._driver_handle_value(current_port) == self._driver_handle_value(owned_port) - - if not owns_current: - break - - receive_state, message = self._receive_driver_message(owned_port) - if receive_state != "message": - break - - code = message.Data.MessageCode - pid = message.Data.ProcessId - target = message.Data.Path - exe_info = self.get_exe_info(pid) - safe_source = exe_info[1] if exe_info and exe_info[1] else "Unknown" - - if not self.is_in_whitelist(safe_source): - self.write_log( - "BLOCK", - "Driver Block", - detail="Driver protection event", - pid=pid, - source=safe_source, - target=target, - code=code, - file_hash=self.calc_file_hash(safe_source), - operate=None, - success=True - ) - - self._close_driver_port(owned_port) - owned_port = None - ready = False - self.driver_listener_ready_event.clear() - - if self.driver_stop_event.is_set(): - break - self.driver_stop_event.wait(0.02) - - except Exception as e: - self.driver_listener_failed_event.set() - if not self.driver_stop_event.is_set(): - self.write_log("WARN", "pipe_server_thread", detail=str(e), success=False) - - finally: - if owned_port: - try: - self.kernel32.CancelIoEx(owned_port, None) - except Exception: - pass - self._close_driver_port(owned_port) - - if not ready and not self.driver_stop_event.is_set(): - self.driver_listener_failed_event.set() - - self.driver_listener_ready_event.clear() - with self.lock_driver: - if self.driver_listener_thread is threading.current_thread(): - self.driver_listener_thread = None diff --git a/PYAS_Realtime.py b/PYAS_Realtime.py new file mode 100644 index 0000000..a7daae7 --- /dev/null +++ b/PYAS_Realtime.py @@ -0,0 +1,555 @@ +from PYAS_WinAPI import iter_file_notifications +from PYAS_Scheduler import TaskScheduler +from PYAS_Diagnostics import log_exception +from PYAS_Diagnostics import observe_future +import os +import time +import threading +import ctypes +import ctypes.wintypes +from PYAS_WinAPI import FILE_NOTIFY_INFORMATION, MEMORY_BASIC_INFORMATION +from PYAS_WinAPI import ( + ERROR_LOCK_VIOLATION, + ERROR_SHARING_VIOLATION, + FILE_ATTRIBUTE_NORMAL, + FILE_SCAN_DEBOUNCE_SECONDS, + FILE_SCAN_RETRY_SECONDS, + FILE_SHARE_READ, + GENERIC_READ, + INVALID_HANDLE_VALUE, + OPEN_EXISTING, + PYAS_PROCESS_NETWORK_ACCESS, + PYAS_PROCESS_SCAN_ACCESS, +) + + +class RealtimeMixin: + def protect_proc_thread(self): + with self.lock_proc: + self.exist_process = self.get_process_list_pids() + + while True: + with self.lock_config: + if not self.pyas_config.get("process_switch", False): + break + + try: + time.sleep(0.1) + cur = self.get_process_list_pids() + + with self.lock_proc: + new_pids = cur - self.exist_process + self.exist_process = cur + + for pid in new_pids: + h = self.kernel32.OpenProcess(PYAS_PROCESS_SCAN_ACCESS, False, pid) + + if h: + with self.lock_config: + should_suspend = self.pyas_config.get("suspend_switch", True) + + if should_suspend: + self.ntdll.NtSuspendProcess(h) + + with self.lock_proc: + self.suspended_procs.add(h) + + try: + observe_future( + self.proc_pool.submit( + self.handle_new_process, pid, h, should_suspend + ), + "RealtimeMixin.handle_new_process", + ) + + except Exception: + log_exception("PYAS_Realtime.RealtimeMixin.protect_proc_thread:39") + + if should_suspend: + self.ntdll.NtResumeProcess(h) + + with self.lock_proc: + self.suspended_procs.discard(h) + + self.kernel32.CloseHandle(h) + except Exception as e: + log_exception("PYAS_Realtime.RealtimeMixin.protect_proc_thread:46") + self.write_log("WARN", "protect_proc_thread", detail=str(e), success=False) + + def handle_new_process(self, pid, h=None, suspended=None): + if suspended is None: + with self.lock_config: + suspended = self.pyas_config.get("suspend_switch", True) + + if not h: + h = self.kernel32.OpenProcess(PYAS_PROCESS_SCAN_ACCESS, False, pid) + + if not h: + return + + if suspended: + self.ntdll.NtSuspendProcess(h) + + with self.lock_proc: + if not hasattr(self, "suspended_procs"): + self.suspended_procs = set() + + self.suspended_procs.add(h) + + try: + cmdline = self.get_process_cmdline(h) + process_file = self.get_process_file(h) + + if "-scan" in cmdline and self.path_equal(process_file, self.file_pyas): + return + + raw_targets = self.extract_paths_from_cmdline(cmdline) + + if process_file: + raw_targets.append(process_file) + + all_targets = [] + + for p in raw_targets: + np = self.norm_path(self.device_path_to_drive(p)) + + if np and np not in all_targets: + all_targets.append(np) + + with self.lock_config: + ext_filter = self.pyas_config.get("suffix_switch", True) + suffix = self.pyas_config.get("suffix", []) + load_switch = self.pyas_config.get("load_switch", True) + + scan_targets = [] + + for file_path in all_targets: + if not file_path or not os.path.isfile(file_path): + continue + + if ext_filter and os.path.splitext(file_path)[-1].lower() not in suffix: + continue + + if self.is_in_whitelist(file_path): + continue + + if ( + process_file + and process_file.lower().endswith("explorer.exe") + and "/select" in cmdline.lower() + ): + continue + + scan_targets.append(file_path) + + if not scan_targets: + return + + virus_found = False + + for file_path in scan_targets: + result = self.safe_scan_engine(file_path) + self.cloud_check(file_path) + + if result: + self._terminate_process_handle(h, pid=pid) + self.write_log( + "BLOCK", + "Process Block", + pid=pid, + source=file_path, + file_hash=self.calc_file_hash(file_path), + ) + virus_found = True + break + + if not virus_found and load_switch: + hidden_virus_path = self.scan_process_memory(pid, h) + + if hidden_virus_path: + self._terminate_process_handle(h, pid=pid) + self.write_log( + "BLOCK", + "Process DLL Block", + pid=pid, + source=hidden_virus_path, + file_hash=self.calc_file_hash(hidden_virus_path), + ) + + finally: + if suspended: + try: + self.ntdll.NtResumeProcess(h) + except Exception: + log_exception("PYAS_Realtime.RealtimeMixin.handle_new_process:123") + pass + + with self.lock_proc: + self.suspended_procs.discard(h) + + try: + self.kernel32.CloseHandle(h) + except Exception: + log_exception("PYAS_Realtime.RealtimeMixin.handle_new_process:129") + pass + + def scan_process_memory(self, pid, h_process): + scanned_paths = set() + address = 0 + mbi = MEMORY_BASIC_INFORMATION() + system_dir = self.path_system.lower() + + buf = ctypes.create_unicode_buffer(1024) + max_address = 0x7FFFFFFFFFFF if ctypes.sizeof(ctypes.c_void_p) == 8 else 0x7FFFFFFF + + with self.lock_config: + ext_filter = self.pyas_config.get("suffix_switch", True) + suffix = self.pyas_config.get("suffix", []) + + try: + while address < max_address and self.kernel32.VirtualQueryEx( + h_process, ctypes.c_void_p(address), ctypes.byref(mbi), ctypes.sizeof(mbi) + ): + if mbi.State == 0x1000 and mbi.Type == 0x1000000: + if self.psapi.GetMappedFileNameW( + h_process, ctypes.c_void_p(address), buf, 1024 + ): + + raw_path = buf.value + + if raw_path.startswith("\\"): + raw_path = self.device_path_to_drive(raw_path) + + file_path = self.norm_path(raw_path) + + if ( + file_path + and file_path not in scanned_paths + and self.path_system not in file_path + ): + + scanned_paths.add(file_path) + file_path_lower = file_path.lower() + + if not file_path_lower.startswith( + system_dir + ) and not self.is_in_whitelist(file_path): + ext = os.path.splitext(file_path_lower)[-1] + + if ext != ".exe" and (not ext_filter or ext in suffix): + if self.safe_scan_engine(file_path): + return file_path + + self.cloud_check(file_path) + + if mbi.RegionSize == 0: + break + + address += mbi.RegionSize + + finally: + scanned_paths.clear() + + return None + + def protect_file_thread(self): + with self.lock_file_ops: + if getattr(self, "h_dir_file", None): + return + + hDir = self.kernel32.CreateFileW( + self.path_user, 0x0001, 0x00000007, None, 3, 0x02000000, None + ) + + if not hDir or hDir == -1: + return + + with self.lock_file_ops: + self.h_dir_file = hDir + + try: + buffer = ctypes.create_string_buffer(262144) + temp_prefix = os.path.normcase(self.path_temp) + + if not temp_prefix.endswith(os.sep): + temp_prefix += os.sep + + while True: + with self.lock_config: + if not self.pyas_config.get("document_switch", False): + break + + try: + bytes_returned = ctypes.wintypes.DWORD() + res = self.kernel32.ReadDirectoryChangesW( + self.h_dir_file, + buffer, + ctypes.sizeof(buffer), + True, + 0x0000001F, + ctypes.byref(bytes_returned), + None, + None, + ) + + if not res: + raise OSError(ctypes.get_last_error(), "Directory watcher read failed") + + if bytes_returned.value == 0: + self._schedule_file_task( + "recovery", self.path_user, lambda _: self._recover_file_tasks(), 0 + ) + continue + + for action, raw_filename in iter_file_notifications( + buffer, bytes_returned.value + ): + if raw_filename and action in [1, 3, 5]: + file_path = self.norm_path( + os.path.join(self.path_user, raw_filename), must_exist=True + ) + + if file_path and not self.is_in_whitelist(file_path): + norm_path = os.path.normcase(file_path) + + if not ( + norm_path.startswith(temp_prefix) + and norm_path[len(temp_prefix) :].startswith("_mei") + ): + with self.lock_config: + ext_filter = self.pyas_config.get("suffix_switch", True) + suffix = self.pyas_config.get("suffix", []) + + if ( + not ext_filter + or os.path.splitext(file_path)[-1].lower() in suffix + ): + self._queue_file_scan(file_path) + + except Exception: + log_exception("PYAS_Realtime.RealtimeMixin.protect_file_thread:230") + + with self.lock_config: + enabled = self.pyas_config.get("document_switch", False) + + if not enabled or self.closing: + break + + self._schedule_file_task( + "recovery", self.path_user, lambda _: self._recover_file_tasks(), 0 + ) + time.sleep(0.25) + finally: + with self.lock_file_ops: + if getattr(self, "h_dir_file", None): + try: + self.kernel32.CloseHandle(self.h_dir_file) + except Exception: + log_exception("PYAS_Realtime.RealtimeMixin.protect_file_thread:237") + pass + + self.h_dir_file = None + + def _cancel_pending_file_tasks(self, task_name=None, target_path=None): + scheduler = getattr(self, "file_scheduler", None) + + if scheduler is None: + return + + target = os.path.normcase(target_path) if target_path else None + prefix = target + os.sep if target else None + scheduler.cancel( + lambda key: (not task_name or key[0] == task_name) + and (not target or key[1] == target or key[1].startswith(prefix)) + ) + + def _cancel_pending_file_scans(self): + self._cancel_pending_file_tasks("scan") + + def _schedule_file_task(self, task_name, file_path, callback, delay): + with self.lock_file_ops: + if self.closing: + return False + + if getattr(self, "file_scheduler", None) is None: + self.file_scheduler = TaskScheduler(self._recover_file_tasks) + + scheduler = self.file_scheduler + + return scheduler.schedule( + (task_name, os.path.normcase(file_path)), + delay, + callback, + (file_path,), + critical=task_name in ("delete", "lock", "verify"), + ) + + def _recover_file_tasks(self): + self.write_log( + "WARN", + "File Event Recovery", + detail="Pending task capacity reached; recovering by rescan", + success=False, + ) + + with self.lock_config: + quarantined = list(self.pyas_config.get("quarantine", [])) + + for entry in quarantined: + if self.closing: + return + + path = entry.get("file") if isinstance(entry, dict) else entry + + if path: + self.lock_file(path, True, quiet=True) + + for file_path, _ in self.yield_files(self.path_user): + with self.lock_config: + enabled = self.pyas_config.get("document_switch", False) + ext_filter = self.pyas_config.get("suffix_switch", True) + suffix = self.pyas_config.get("suffix", []) + + if self.closing or not enabled: + break + + if self.is_in_whitelist(file_path): + continue + + norm_path = os.path.normcase(file_path) + temp_prefix = os.path.normcase(self.path_temp).rstrip(os.sep) + os.sep + + if norm_path.startswith(temp_prefix) and norm_path[len(temp_prefix) :].startswith( + "_mei" + ): + continue + + if not ext_filter or os.path.splitext(norm_path)[-1] in suffix: + self.handle_new_file(file_path) + + def _queue_file_scan(self, file_path, delay=FILE_SCAN_DEBOUNCE_SECONDS): + self._schedule_file_task("scan", file_path, self.handle_new_file, delay) + + def _acquire_file_scan_gate(self, file_path): + ctypes.set_last_error(0) + handle = self.kernel32.CreateFileW( + file_path, + GENERIC_READ, + FILE_SHARE_READ, + None, + OPEN_EXISTING, + FILE_ATTRIBUTE_NORMAL, + None, + ) + + if not handle or handle == INVALID_HANDLE_VALUE: + return None, ctypes.get_last_error() + + return handle, 0 + + def handle_new_file(self, file_path): + try: + with self.lock_config: + if not self.pyas_config.get("document_switch", False): + return + + gate_handle, error = self._acquire_file_scan_gate(file_path) + + if not gate_handle: + if error in (ERROR_SHARING_VIOLATION, ERROR_LOCK_VIOLATION): + self._queue_file_scan(file_path, FILE_SCAN_RETRY_SECONDS) + + return + + try: + with self.lock_file_ops: + if file_path in self.virus_lock: + return + + result = self.safe_scan_engine(file_path) + self.cloud_check(file_path) + + if result: + if ( + self.manage_named_list( + "quarantine", [file_path], action="add", lock_func=self.lock_file + ) + > 0 + ): + self.write_log( + "BLOCK", + "File Block", + source=file_path, + file_hash=self.calc_file_hash(file_path), + ) + finally: + self.kernel32.CloseHandle(gate_handle) + except Exception: + log_exception("PYAS_Realtime.RealtimeMixin.handle_new_file:341") + pass + + def protect_net_thread(self): + with self.lock_net: + self.exist_connections = set() + + while True: + with self.lock_config: + if not self.pyas_config.get("network_switch", False): + break + + try: + time.sleep(0.5) + conns = self.get_connections_list() + + with self.lock_net: + new_conns = conns - self.exist_connections + self.exist_connections = conns + + for key in new_conns: + observe_future( + self.protect_pool.submit(self.handle_new_connection, key), + "RealtimeMixin.handle_new_connection", + ) + + except Exception as e: + log_exception("PYAS_Realtime.RealtimeMixin.protect_net_thread:363") + self.write_log("WARN", "protect_net_thread", detail=str(e), success=False) + + with self.lock_net: + self.exist_connections = set() + + def handle_new_connection(self, key): + pid, remote_addr, remote_port = key + + try: + h = self.kernel32.OpenProcess(PYAS_PROCESS_NETWORK_ACCESS, False, pid) + + if not h: + return + + try: + remote_ip = f"{remote_addr & 0xFF}.{(remote_addr >> 8) & 0xFF}.{(remote_addr >> 16) & 0xFF}.{(remote_addr >> 24) & 0xFF}" + file_path = self.norm_path(self.get_process_file(h)) + + if ( + file_path + and not self.is_in_whitelist(file_path) + and hasattr(self.heuristic, "network") + and remote_ip in self.heuristic.network + ): + self._terminate_process_handle(h, pid=pid) + self.write_log( + "BLOCK", + "Network Block", + pid=pid, + source=file_path, + target=remote_ip, + file_hash=self.calc_file_hash(file_path), + ) + + finally: + self.kernel32.CloseHandle(h) + + except Exception as e: + log_exception("PYAS_Realtime.RealtimeMixin.handle_new_connection:387") + self.write_log("WARN", "handle_new_connection", pid=pid, detail=str(e), success=False) diff --git a/PYAS_Rules.py b/PYAS_Rules.py new file mode 100644 index 0000000..ae48400 --- /dev/null +++ b/PYAS_Rules.py @@ -0,0 +1,107 @@ +from PYAS_Diagnostics import log_exception +import os +import yara + + +class RuleScanner: + def __init__(self): + self.rules = None + self.network = [] + + def load_path(self, path, callback=None): + yara_files = {} + + for root, _, files in os.walk(path): + for file in files: + + full_path = os.path.join(root, file) + + if callback: + callback(full_path) + + ext = os.path.splitext(file)[1].lower() + + if ext in (".yara", ".yar"): + namespace = os.path.relpath(full_path, path).replace(os.sep, "_") + yara_files[namespace] = full_path + + elif ext in (".yc", ".yrc"): + self.load_compiled_rule(full_path) + + elif ext in (".ip", ".txt"): + self.load_network_list(full_path) + + if yara_files: + self.compile_all_rules(yara_files) + + def load_compiled_rule(self, file): + try: + self.rules = yara.load(file) + except Exception: + log_exception("PYAS_Rules.RuleScanner.load_compiled_rule:38") + pass + + def load_network_list(self, file): + try: + with open(file, "r", encoding="utf-8", errors="ignore") as f: + self.network.extend(line.strip() for line in f if line.strip()) + except Exception: + log_exception("PYAS_Rules.RuleScanner.load_network_list:45") + pass + + def compile_all_rules(self, file_map): + try: + self.rules = yara.compile(filepaths=file_map) + except Exception: + log_exception("PYAS_Rules.RuleScanner.compile_all_rules:51") + pass + + def yara_scan(self, file_path): + try: + if not self.rules: + return False, False + + matches = self.rules.match(filepath=file_path) + + if matches: + rule_name = str(matches[0]) + + try: + label = rule_name.split("_")[0] + platform = rule_name.split("_")[1] + family = rule_name.split("_")[2] + + return f"{label}:{platform}/{family}.{len(matches)}!yr", 100 + except Exception: + log_exception("PYAS_Rules.RuleScanner.yara_scan:70") + return rule_name, 100 + + return False, False + except Exception: + log_exception("PYAS_Rules.RuleScanner.yara_scan:74") + return False, False + + def yara_mem_scan(self, pid): + try: + if not self.rules: + return False, False + + matches = self.rules.match(pid=pid) + + if matches: + rule_name = str(matches[0]) + + try: + label = rule_name.split("_")[0] + platform = rule_name.split("_")[1] + family = rule_name.split("_")[2] + + return f"{label}:{platform}/{family}.{len(matches)}!ym", 100 + except Exception: + log_exception("PYAS_Rules.RuleScanner.yara_mem_scan:91") + return rule_name, 100 + + return False, False + except Exception: + log_exception("PYAS_Rules.RuleScanner.yara_mem_scan:95") + return False, False diff --git a/PYAS_Runtime.py b/PYAS_Runtime.py new file mode 100644 index 0000000..bc4770b --- /dev/null +++ b/PYAS_Runtime.py @@ -0,0 +1,255 @@ +from PYAS_Diagnostics import log_exception, report_log_path +import os +import sys +import time +import queue +import threading +from concurrent.futures import ThreadPoolExecutor +from PYAS_Engine import rule_scanner, pe_scanner, cloud_scanner +from PYAS_Version import VERSION + + +class RuntimeMixin: + def init_environ(self): + self.python = sys.executable + + if getattr(sys, "frozen", False): + self.file_pyas = self.norm_path(sys.executable) + else: + self.file_pyas = self.norm_path(os.path.abspath(sys.argv[0])) + + self.args_pyas = sys.argv[1:] + self.path_pyas = os.path.dirname(self.file_pyas) + self.pid_pyas = int(os.getpid()) + + self.path_appdata = os.environ.get("APPDATA") + self.path_localappdata = os.environ.get("LOCALAPPDATA") or self.path_appdata + self.path_name = os.environ.get("USERNAME") + self.path_temp = os.environ.get( + "TEMP", f"C:\\Users\\{self.path_name}\\AppData\\Local\\Temp" + ) + self.path_config = os.environ.get("ALLUSERSPROFILE", "C:\\ProgramData") + self.path_system = os.environ.get("SYSTEMROOT", "C:\\Windows") + self.path_user = os.environ.get("USERPROFILE", f"C:\\Users\\{self.path_name}") + + self.path_systemp = os.path.join(self.path_system, "Temp") + self.file_config = os.path.join(self.path_config, "PYAS", "Config.json") + self.file_log = report_log_path() + self.path_webview = os.path.join( + self.path_localappdata or self.path_temp, "PYAS", "WebView2" + ) + self.file_webview_log = self.file_log + self.path_properties = os.path.join(self.path_pyas, "Engine", "Properties") + self.path_heuristic = os.path.join(self.path_pyas, "Engine", "Heuristic") + self.path_protect = os.path.join(self.path_pyas, "Plugins", "Filter") + self.path_rules = os.path.join(self.path_pyas, "Plugins", "Rules") + self.path_drivers = os.path.join(self.path_protect, "PYAS_Driver.sys") + + def init_variables(self): + self.lock_workers = threading.RLock() + self.feature_workers = {} + self.feature_restarts = set() + + self.heuristic = rule_scanner() + self.properties = pe_scanner() + self.cloud = cloud_scanner() + self.cloud_queue = queue.Queue() + + self.ui_queue = queue.Queue() + self.start_daemon_thread(self.ui_dispatcher_thread) + + self.driver_port = None + self.driver_stop_event = threading.Event() + self.driver_listener_ready_event = threading.Event() + self.driver_listener_failed_event = threading.Event() + self.driver_listener_thread = None + + self.ui_ready_event = threading.Event() + self.engine_initialized = False + + self.scan_running = False + self.scan_preparing = False + self.scan_stop_requested = False + self.scan_finished = False + + self.virus_lock = {} + self.virus_results = [] + + self.scan_count = 0 + self.scan_events = {} + self.hash_cache = {} + self.file_scheduler = None + + self.mbr_backup = {} + + self.cloud_pending = set() + self.cloud_cancel_event = threading.Event() + + self.autostart_mode = "unknown" + + self.last_io_counters = {} + self.last_io_time = time.time() + self.suspended_procs = set() + + self.lock_driver = threading.RLock() + self.lock_driver_unload = threading.Lock() + self.driver_unload_worker = None + self.driver_unload_result = None + self.closing = False + + self.lock_update = threading.RLock() + self.lock_proc = threading.RLock() + self.lock_net = threading.RLock() + self.lock_io = threading.RLock() + + self.pyas_default = { + "version": VERSION, + "api_host": "https://pyas-security.com/", + "api_key": "fBRZxYS1UxykM-qzNOlKOEl63WILzlvgNMn6QfsG6FXCAAIktCrOPTAfY5_hEyuZ", + "suffix": [ + ".exe", + ".dll", + ".sys", + ".ocx", + ".scr", + ".efi", + ".acm", + ".ax", + ".cpl", + ".drv", + ".com", + ".mui", + ".pyd", + ".wfx", + ".api", + ".awx", + ".rll", + ".winmd", + ".bat", + ".cmd", + ".ps1", + ".vbs", + ".vbe", + ".wsf", + ".reg", + ".html", + ".js", + ".jse", + ".jsp", + ".php", + ".hta", + ".lnk", + ".py", + ".sh", + ".url", + ".rtf", + ".ini", + ], + "size": 256 * 1024 * 1024, + "language": "english_switch", + "theme": "system_switch", + "first_launch": True, + "process_switch": False, + "suspend_switch": True, + "load_switch": True, + "document_switch": False, + "system_switch": False, + "driver_switch": False, + "network_switch": False, + "extension_switch": False, + "sensitive_switch": False, + "cloud_switch": False, + "suffix_switch": True, + "autostart_switch": True, + "context_switch": True, + "custom_rule": [], + "white_list": [], + "quarantine": [], + "block_list": [], + } + + self.pass_windows = [ + {"exe": "System Idle Process", "class": "", "title": ""}, + {"exe": "", "class": "Windows.UI.Core.CoreWindow", "title": ""}, + {"exe": "explorer.exe", "class": "", "title": ""}, + ] + + self.scan_pool = ThreadPoolExecutor(max_workers=2) + self.protect_pool = ThreadPoolExecutor(max_workers=8) + self.proc_pool = ThreadPoolExecutor(max_workers=16) + self.start_daemon_thread(self.log_flush_thread) + + for _ in range(2): + self.start_daemon_thread(self.cloud_worker) + + def ui_dispatcher_thread(self): + while not getattr(self, "closing", False): + batch = [] + + try: + batch.append(self.ui_queue.get(timeout=0.1)) + + while len(batch) < 50: + try: + batch.append(self.ui_queue.get_nowait()) + except queue.Empty: + break + + if self._window: + self._window.evaluate_js("".join(batch)) + except queue.Empty: + continue + except Exception: + log_exception("PYAS_Runtime.RuntimeMixin.ui_dispatcher_thread:146") + pass + finally: + for _ in batch: + self.ui_queue.task_done() + + def start_feature_thread(self, target, switch): + with self.lock_workers: + existing = self.feature_workers.get(switch) + + if existing and existing.is_alive(): + self.feature_restarts.add(switch) + return existing + + def worker(): + try: + while True: + try: + target() + except Exception: + log_exception(f"RuntimeMixin.{target.__name__}") + + with self.lock_config: + enabled = self.pyas_config.get(switch, False) + + if not enabled or self.closing: + break + + self.write_log( + "WARN", + "Protection Recovery", + detail=f"Restarting {switch}", + success=False, + ) + time.sleep(0.5) + finally: + with self.lock_workers: + if self.feature_workers.get(switch) is threading.current_thread(): + self.feature_workers.pop(switch, None) + + restart = switch in self.feature_restarts + self.feature_restarts.discard(switch) + + with self.lock_config: + enabled = self.pyas_config.get(switch, False) + + if restart and enabled and not self.closing: + self.start_feature_thread(target, switch) + + thread = threading.Thread(target=worker, daemon=True, name=f"PYAS.{target.__name__}") + self.feature_workers[switch] = thread + thread.start() + return thread diff --git a/PYAS_Scanner.py b/PYAS_Scanner.py index 85e64eb..7ca0e6a 100644 --- a/PYAS_Scanner.py +++ b/PYAS_Scanner.py @@ -1,11 +1,18 @@ -import os, gc, json, time, stat, threading -import ctypes, ctypes.wintypes - +from PYAS_Diagnostics import log_exception +from PYAS_Diagnostics import observe_future, exception_sequence +import os +import json +import time +import stat +import threading +from itertools import islice +import ctypes +import ctypes.wintypes from PYAS_Tools import MEMORY_BASIC_INFORMATION +from PYAS_Cloud import CloudQueueMixin -#################################################################################################### -class ScannerMixin: +class ScannerMixin(CloudQueueMixin): def init_engine_thread(self): try: self.start_daemon_thread(self.backup_mbr) @@ -23,8 +30,9 @@ class ScannerMixin: if "-scan" in self.args_pyas: try: idx = self.args_pyas.index("-scan") - self.trigger_context_scan(self.args_pyas[idx+1]) + self.trigger_context_scan(self.args_pyas[idx + 1]) except Exception: + log_exception("PYAS_Scanner.ScannerMixin.init_engine_thread:31") pass with self.lock_config: @@ -34,40 +42,61 @@ class ScannerMixin: autostart_enabled = self.pyas_config.get("autostart_switch", True) self.register_context_menu(context_enabled) - self.manage_autostart(autostart_enabled) + + with self.lock_update: + with self.lock_config: + autostart_enabled = self.pyas_config.get("autostart_switch", True) + + if not self.manage_autostart(autostart_enabled) and autostart_enabled: + with self.lock_config: + self.pyas_config["autostart_switch"] = False + self.save_config() + + self.ui_queue.put( + "(()=>{const ats=document.querySelector('[data-config-key=autostart_switch]');if(ats)ats.checked=false;})();" + ) if not first_launch: with self.lock_config: if self.pyas_config.get("process_switch"): - self.start_daemon_thread(self.protect_proc_thread) + self.start_feature_thread(self.protect_proc_thread, "process_switch") + if self.pyas_config.get("document_switch"): - self.start_daemon_thread(self.protect_file_thread) + self.start_feature_thread(self.protect_file_thread, "document_switch") + if self.pyas_config.get("system_switch"): - self.start_daemon_thread(self.protect_system_thread) + self.start_feature_thread(self.protect_system_thread, "system_switch") + if self.pyas_config.get("network_switch"): - self.start_daemon_thread(self.protect_net_thread) + self.start_feature_thread(self.protect_net_thread, "network_switch") if driver_enabled: if self.install_system_driver() and self.start_driver_listener(wait_ready=True): pass else: self.stop_system_driver() + with self.lock_config: self.pyas_config["driver_switch"] = False self.save_config() - self.write_log("WARN", "System", detail="Driver Protection Failed to Start", success=False) + self.write_log( + "WARN", + "System", + detail="Driver Protection Failed to Start", + success=False, + ) except Exception as e: + log_exception("PYAS_Scanner.ScannerMixin.init_engine_thread:72") self.write_log("WARN", "init_engine_thread", detail=str(e), success=False) -#################################################################################################### - def yield_files(self, targets): if isinstance(targets, str): if os.path.isdir(targets): for root, dirs, files in os.walk(targets): dirs[:] = [d for d in dirs if not self._is_reparse_point(os.path.join(root, d))] + for f in files: yield self.norm_path(os.path.join(root, f)), False @@ -82,16 +111,16 @@ class ScannerMixin: try: if os.path.islink(path): return True - if hasattr(os.path, 'isjunction') and os.path.isjunction(path): + + if hasattr(os.path, "isjunction") and os.path.isjunction(path): return True st = os.lstat(path) - return bool(getattr(st, 'st_file_attributes', 0) & 0x400) + return bool(getattr(st, "st_file_attributes", 0) & 0x400) except OSError: + log_exception("PYAS_Scanner.ScannerMixin._is_reparse_point:101") return False -#################################################################################################### - def scan_engine(self, file_path): with self.lock_config: ext_switch = self.pyas_config.get("extension_switch", False) @@ -99,72 +128,122 @@ class ScannerMixin: try: pe_label, _ = self.properties.pe_scan(file_path, enhanced_mode=sen_switch) + if pe_label: return pe_label except Exception: + log_exception("PYAS_Scanner.ScannerMixin.scan_engine:115") pass try: if ext_switch: yara_label, _ = self.heuristic.yara_scan(file_path) + if yara_label: return yara_label except Exception: + log_exception("PYAS_Scanner.ScannerMixin.scan_engine:123") pass return False def safe_scan_engine(self, file_path): norm_path = self.norm_path(file_path) + if not norm_path: return False - file_hash = self.calc_file_hash(norm_path) - cache_key = file_hash if file_hash else os.path.normcase(norm_path) + gate, error = self._acquire_file_scan_gate(norm_path) - with self.lock_file_ops: - if cache_key in self.hash_cache: - return self.hash_cache[cache_key] - - if cache_key in self.scan_events: - event = self.scan_events[cache_key] - needs_scan = False - else: - event = threading.Event() - self.scan_events[cache_key] = event - needs_scan = True - - if not needs_scan: - event.wait(timeout=60) - with self.lock_file_ops: - return self.hash_cache.get(cache_key, False) + if not gate: + self.write_log("WARN", "Scan Deferred", source=norm_path, code=error, success=False) + self._schedule_file_task("verify", norm_path, self._retry_verified_scan, 0.5) + return self.scan_engine(norm_path) try: - result = self.scan_engine(norm_path) + file_hash = self.calc_file_hash(norm_path) + + if not file_hash: + return self.scan_engine(norm_path) + + with self.lock_config: + policy = ( + bool(self.pyas_config.get("extension_switch", False)), + bool(self.pyas_config.get("sensitive_switch", False)), + ) + + cache_key = (file_hash, policy) + with self.lock_file_ops: - if len(self.hash_cache) > 10000: - for k in list(self.hash_cache.keys())[:1000]: - del self.hash_cache[k] + if cache_key in self.hash_cache: + return self.hash_cache[cache_key] - self.hash_cache[cache_key] = result - return result + event = self.scan_events.get(cache_key) + needs_scan = event is None + if needs_scan: + event = threading.Event() + self.scan_events[cache_key] = event + + if not needs_scan: + event.wait(timeout=60) + + with self.lock_file_ops: + if cache_key in self.hash_cache: + return self.hash_cache[cache_key] + + return self.scan_engine(norm_path) + + try: + before_exceptions = exception_sequence() + result = self.scan_engine(norm_path) + scan_succeeded = exception_sequence() == before_exceptions + + with self.lock_config: + current_policy = ( + bool(self.pyas_config.get("extension_switch", False)), + bool(self.pyas_config.get("sensitive_switch", False)), + ) + + if scan_succeeded and current_policy == policy: + with self.lock_file_ops: + if len(self.hash_cache) >= 100000: + for key in list(islice(self.hash_cache, 10000)): + del self.hash_cache[key] + + self.hash_cache[cache_key] = result + + return result + finally: + with self.lock_file_ops: + self.scan_events.pop(cache_key, None) + event.set() finally: - with self.lock_file_ops: - if cache_key in self.scan_events: - self.scan_events[cache_key].set() - del self.scan_events[cache_key] + self.kernel32.CloseHandle(gate) - gc.collect() + def _retry_verified_scan(self, file_path): + if self.closing or not os.path.isfile(file_path): + return -#################################################################################################### + result = self.safe_scan_engine(file_path) + + if result: + self.manage_named_list( + "quarantine", [file_path], action="add", lock_func=self.lock_file + ) + self.write_log( + "BLOCK", + "Deferred Scan Block", + source=file_path, + file_hash=self.calc_file_hash(file_path), + ) def _scan_result_messages(self, count=None, scanned=None, elapsed=None): if count is None or scanned is None or elapsed is None: with self.lock_virus: - count = len(getattr(self, 'virus_results', [])) if count is None else count - scanned = getattr(self, 'scan_count', 0) if scanned is None else scanned - scan_start = getattr(self, 'scan_start', time.time()) + count = len(getattr(self, "virus_results", [])) if count is None else count + scanned = getattr(self, "scan_count", 0) if scanned is None else scanned + scan_start = getattr(self, "scan_start", time.time()) elapsed = int(time.time() - scan_start) if elapsed is None else elapsed return { @@ -178,18 +257,18 @@ class ScannerMixin: "hindi_switch": f"{count} वायरस मिले, {scanned} फ़ाइलें स्कैन की गईं, समय {elapsed}s", "arabic_switch": f"تم العثور على {count} فيروسات، تم فحص {scanned} ملفات، الوقت {elapsed} ثانية", "russian_switch": f"Найдено {count} вирусов, проверено {scanned} файлов, время {elapsed}с", - "slovenian_switch": f"Najdenih {count} virusov, skeniranih {scanned} datotek, čas {elapsed}s" + "slovenian_switch": f"Najdenih {count} virusov, skeniranih {scanned} datotek, čas {elapsed}s", } def _is_scan_cancel_requested(self): with self.lock_virus: - return getattr(self, 'scan_stop_requested', False) + return getattr(self, "scan_stop_requested", False) def _finish_scan_cancelled(self, messages=None): with self.lock_virus: - count = len(getattr(self, 'virus_results', [])) - scanned = getattr(self, 'scan_count', 0) - scan_start = getattr(self, 'scan_start', time.time()) + count = len(getattr(self, "virus_results", [])) + scanned = getattr(self, "scan_count", 0) + scan_start = getattr(self, "scan_start", time.time()) elapsed = int(time.time() - scan_start) self.scan_running = False self.scan_preparing = False @@ -206,7 +285,7 @@ class ScannerMixin: def _begin_scan_preparation(self): with self.lock_virus: - if getattr(self, 'scan_running', False) or getattr(self, 'scan_preparing', False): + if getattr(self, "scan_running", False) or getattr(self, "scan_preparing", False): return False self.scan_preparing = True @@ -219,10 +298,12 @@ class ScannerMixin: def start_scan(self, targets, from_preparing=False): with self.lock_virus: - if getattr(self, 'scan_running', False) or (getattr(self, 'scan_preparing', False) and not from_preparing): + if getattr(self, "scan_running", False) or ( + getattr(self, "scan_preparing", False) and not from_preparing + ): return False - if getattr(self, 'scan_stop_requested', False): + if getattr(self, "scan_stop_requested", False): self.scan_running = False self.scan_preparing = False self.scan_finished = True @@ -236,22 +317,37 @@ class ScannerMixin: self.scan_count = 0 self.scan_start = time.time() - self.scan_pool.submit(self.scan_worker, targets) - return True + try: + observe_future( + self.scan_pool.submit(self.scan_worker, targets), "ScannerMixin.scan_worker" + ) + return True + except Exception: + log_exception("PYAS_Scanner.ScannerMixin.start_scan:253") + + with self.lock_virus: + self.scan_running = False + self.scan_preparing = False + self.scan_finished = True + + return False def scan_worker(self, targets): last_update = 0.0 + try: for file_path, is_explicit in self.yield_files(targets): with self.lock_virus: - if not self.scan_running or getattr(self, 'scan_stop_requested', False): + if not self.scan_running or getattr(self, "scan_stop_requested", False): break norm_path = self.norm_path(file_path) + if not norm_path: continue current_time = time.time() + if current_time - last_update >= 0.05: if self._window: js_cmd = f"if(window.updateScanProgress) window.updateScanProgress({json.dumps(norm_path.replace(os.sep, '/'))});" @@ -260,6 +356,7 @@ class ScannerMixin: last_update = current_time was_locked = False + try: with self.lock_file_ops: if norm_path in self.virus_lock: @@ -270,35 +367,49 @@ class ScannerMixin: continue with self.lock_virus: - if not self.scan_running or getattr(self, 'scan_stop_requested', False): + if not self.scan_running or getattr(self, "scan_stop_requested", False): break + self.scan_count += 1 with self.lock_config: ext_filter = self.pyas_config.get("suffix_switch", True) suffix = self.pyas_config.get("suffix", []) - if not is_explicit and ext_filter and os.path.splitext(norm_path)[-1].lower() not in suffix: + if ( + not is_explicit + and ext_filter + and os.path.splitext(norm_path)[-1].lower() not in suffix + ): continue result = self.safe_scan_engine(norm_path) + with self.lock_virus: - if not self.scan_running or getattr(self, 'scan_stop_requested', False): + if not self.scan_running or getattr(self, "scan_stop_requested", False): break if result: with self.lock_virus: self.virus_results.append(norm_path) - if self._window: + if self._window: js_cmd = f"if(window.addVirusResult) window.addVirusResult({json.dumps(result)}, {json.dumps(norm_path.replace(os.sep, '/'))});" self.ui_queue.put(js_cmd) - self.write_log("SCAN", "Virus Detected", source=norm_path, file_hash=self.calc_file_hash(norm_path)) + self.write_log( + "SCAN", + "Virus Detected", + source=norm_path, + file_hash=self.calc_file_hash(norm_path), + ) self.cloud_check(norm_path) except Exception as e: - self.write_log("WARN", "Scan Engine", source=norm_path, detail=str(e), success=False) + log_exception("PYAS_Scanner.ScannerMixin.scan_worker:318") + self.write_log( + "WARN", "Scan Engine", source=norm_path, detail=str(e), success=False + ) finally: if was_locked: @@ -306,30 +417,36 @@ class ScannerMixin: finally: with self.lock_virus: - cancelled = getattr(self, 'scan_stop_requested', False) + cancelled = getattr(self, "scan_stop_requested", False) self.scan_running = False self.scan_preparing = False self.scan_finished = True self.scan_stop_requested = False - count, scanned, elapsed = len(self.virus_results), self.scan_count, int(time.time() - self.scan_start) + count, scanned, elapsed = ( + len(self.virus_results), + self.scan_count, + int(time.time() - self.scan_start), + ) messages = self._scan_result_messages(count, scanned, elapsed) log_detail = f"Found {count} viruses, scanned {scanned} files, time {elapsed}s" if self._window: - js_cmd = f"if(window.finishScan) window.finishScan({json.dumps(messages)}, {count});" + js_cmd = ( + f"if(window.finishScan) window.finishScan({json.dumps(messages)}, {count});" + ) self.ui_queue.put(js_cmd) self.write_log("INFO", "Scan Completed", detail=log_detail) -#################################################################################################### - def stop_scan(self): with self.lock_virus: - active = getattr(self, 'scan_running', False) or getattr(self, 'scan_preparing', False) + active = getattr(self, "scan_running", False) or getattr(self, "scan_preparing", False) + if active: self.scan_stop_requested = True self.scan_running = False + return active def trigger_scan(self, method): @@ -345,6 +462,7 @@ class ScannerMixin: return self._finish_scan_cancelled() fp = os.path.join(self.path_user, folder) + if os.path.exists(fp): targets.append(fp) @@ -362,32 +480,51 @@ class ScannerMixin: if self._is_scan_cancel_requested(): return self._finish_scan_cancelled() - if proc["path"] and proc["path"] != "None": + if proc["path"] and proc["path"] != "None": targets.append(proc["path"]) pid = proc["pid"] + if pid <= 4: continue h_process = self.kernel32.OpenProcess(0x1000, False, pid) + if h_process: try: address = 0 mbi = MEMORY_BASIC_INFORMATION() - while address < max_address and not self._is_scan_cancel_requested() and self.kernel32.VirtualQueryEx(h_process, ctypes.c_void_p(address), ctypes.byref(mbi), ctypes.sizeof(mbi)): + + while ( + address < max_address + and not self._is_scan_cancel_requested() + and self.kernel32.VirtualQueryEx( + h_process, + ctypes.c_void_p(address), + ctypes.byref(mbi), + ctypes.sizeof(mbi), + ) + ): if mbi.State == 0x1000 and mbi.Type == 0x1000000: - if self.psapi.GetMappedFileNameW(h_process, ctypes.c_void_p(address), buf, 1024): + if self.psapi.GetMappedFileNameW( + h_process, ctypes.c_void_p(address), buf, 1024 + ): raw_path = buf.value + if raw_path.startswith("\\"): raw_path = self.device_path_to_drive(raw_path) file_path = self.norm_path(raw_path) - if file_path and not file_path.lower().startswith(system_dir): + + if file_path and not file_path.lower().startswith( + system_dir + ): targets.append(file_path) if mbi.RegionSize == 0: break + address += mbi.RegionSize finally: self.kernel32.CloseHandle(h_process) @@ -402,107 +539,143 @@ class ScannerMixin: elif method == "file": targets = self.select_files() + if self._is_scan_cancel_requested(): return self._finish_scan_cancelled() if targets: if not self.start_scan(targets, from_preparing=True): return self._finish_scan_cancelled() + return True return self._finish_scan_cancelled() elif method == "path": targets = self.select_folder() + if self._is_scan_cancel_requested(): return self._finish_scan_cancelled() if targets: if not self.start_scan(targets, from_preparing=True): return self._finish_scan_cancelled() + return True return self._finish_scan_cancelled() elif method == "full": targets = [f"{chr(d)}:/" for d in range(65, 91) if os.path.exists(f"{chr(d)}:/")] - if self._is_scan_cancel_requested() or not self.start_scan(targets, from_preparing=True): + + if self._is_scan_cancel_requested() or not self.start_scan( + targets, from_preparing=True + ): return self._finish_scan_cancelled() + return True return self._finish_scan_cancelled() except Exception as e: + log_exception("PYAS_Scanner.ScannerMixin.trigger_scan:453") self.write_log("WARN", "trigger_scan", detail=str(e), success=False) return self._finish_scan_cancelled() -#################################################################################################### - def _queue_virus_delete(self, file_path): self._schedule_file_task("delete", file_path, self._retry_virus_delete, 0.5) def _retry_virus_delete(self, file_path): file_hash = self.calc_file_hash(file_path) + try: with self.lock_file_ops: if file_path in self.virus_lock: self.lock_file(file_path, False) target_key = os.path.normcase(self.norm_path(file_path, must_exist=False) or file_path) + for proc in self.get_process_list(): proc_path = proc.get("path") + if not proc_path or proc_path == "None": continue - proc_key = os.path.normcase(self.norm_path(proc_path, must_exist=False) or proc_path) + + proc_key = os.path.normcase( + self.norm_path(proc_path, must_exist=False) or proc_path + ) + if proc_key == target_key: - self.kill_process(proc["pid"]) + self.kill_process(proc["pid"], expected_path=file_path) try: os.chmod(file_path, stat.S_IWRITE) except Exception: + log_exception("PYAS_Scanner.ScannerMixin._retry_virus_delete:480") pass + os.remove(file_path) except FileNotFoundError: pass except Exception as e: + log_exception("PYAS_Scanner.ScannerMixin._retry_virus_delete:486") self.lock_file(file_path, True, quiet=True) self._queue_virus_delete(file_path) self.write_log("INFO", "Virus Delete Deferred", source=file_path, detail=str(e)) return self.remove_list_items("quarantine", [file_path]) + with self.lock_virus: target_key = os.path.normcase(self.norm_path(file_path, must_exist=False) or file_path) self.virus_results = [ - result for result in self.virus_results - if os.path.normcase(self.norm_path(result, must_exist=False) or result) != target_key + result + for result in self.virus_results + if os.path.normcase(self.norm_path(result, must_exist=False) or result) + != target_key ] - self.write_log("INFO", "Virus Delete", source=file_path, file_hash=file_hash, operate=True, success=True) + + self.write_log( + "INFO", + "Virus Delete", + source=file_path, + file_hash=file_hash, + operate=True, + success=True, + ) def solve_scan(self, file_paths): deleted_paths = [] proc_map = {} + for proc in self.get_process_list(): if proc["path"] and proc["path"] != "None": norm_p = self.norm_path(proc["path"], must_exist=False) + if norm_p: proc_map.setdefault(os.path.normcase(norm_p), []).append(proc["pid"]) last_update = 0.0 + with self.lock_virus: deleted_set = set() + for raw_path in file_paths: path = self.norm_path(raw_path, must_exist=False) + if not path: continue current_time = time.time() + if current_time - last_update >= 0.05: if self._window: - self._window.evaluate_js(f"if(window.updateDeleteProgress) window.updateDeleteProgress({json.dumps(path.replace(os.sep, '/'))});") + self._window.evaluate_js( + f"if(window.updateDeleteProgress) window.updateDeleteProgress({json.dumps(path.replace(os.sep, '/'))});" + ) last_update = current_time @@ -511,19 +684,23 @@ class ScannerMixin: self.lock_file(path, False) path_key = os.path.normcase(path) + if path_key in proc_map: for pid in proc_map[path_key]: - self.kill_process(pid) + self.kill_process(pid, expected_path=path) try: os.chmod(path, stat.S_IWRITE) except Exception: + log_exception("PYAS_Scanner.ScannerMixin.solve_scan:538") pass file_hash = self.calc_file_hash(path) + try: os.remove(path) except Exception as e: + log_exception("PYAS_Scanner.ScannerMixin.solve_scan:544") self.lock_file(path, True, quiet=True) self._queue_virus_delete(path) self.write_log("INFO", "Virus Delete Deferred", source=path, detail=str(e)) @@ -532,11 +709,27 @@ class ScannerMixin: self.remove_list_items("quarantine", [path]) deleted_paths.append(raw_path) deleted_set.add(path) - self.write_log("INFO", "Virus Delete", source=path, file_hash=file_hash, operate=True, success=True) + self.write_log( + "INFO", + "Virus Delete", + source=path, + file_hash=file_hash, + operate=True, + success=True, + ) except Exception as e: + log_exception("PYAS_Scanner.ScannerMixin.solve_scan:555") self.lock_file(path, True) - self.write_log("SCAN", "Virus Delete", source=path, file_hash=self.calc_file_hash(path), detail=str(e), operate=True, success=False) + self.write_log( + "SCAN", + "Virus Delete", + source=path, + file_hash=self.calc_file_hash(path), + detail=str(e), + operate=True, + success=False, + ) if deleted_set: self.virus_results = [p for p in self.virus_results if p not in deleted_set] @@ -554,10 +747,13 @@ class ScannerMixin: "hindi_switch": f"शेष वायरस {remaining}, हटाए गए फ़ाइलें {len(deleted_paths)}.", "arabic_switch": f"الفيروسات المتبقية {remaining}، تم حذف {len(deleted_paths)} ملفات.", "russian_switch": f"Осталось вирусов: {remaining}, удалено файлов: {len(deleted_paths)}.", - "slovenian_switch": f"Preostalih virusov: {remaining}, izbrisanih datotek: {len(deleted_paths)}." + "slovenian_switch": f"Preostalih virusov: {remaining}, izbrisanih datotek: {len(deleted_paths)}.", } + if self._window: - self._window.evaluate_js(f"if(window.finishScan) window.finishScan({json.dumps(messages)}, {remaining});") + self._window.evaluate_js( + f"if(window.finishScan) window.finishScan({json.dumps(messages)}, {remaining});" + ) return deleted_paths @@ -566,77 +762,3 @@ class ScannerMixin: norm_paths = set(self.norm_path(p, must_exist=False) for p in paths) self.virus_results = [p for p in self.virus_results if p not in norm_paths] return len(self.virus_results) - -#################################################################################################### - - def cloud_check(self, file_path): - norm_path = self.norm_path(file_path) - if not norm_path: - return - - cache_key = os.path.normcase(norm_path) - with self.lock_file_ops: - if cache_key in self.cloud_pending: - return - - self.cloud_pending.add(cache_key) - - self.cloud_queue.put(norm_path) - - def cloud_worker(self): - while True: - try: - file_path = self.cloud_queue.get() - cache_key = os.path.normcase(file_path) - - try: - self.perform_cloud_scan(file_path) - finally: - with self.lock_file_ops: - self.cloud_pending.discard(cache_key) - - self.cloud_queue.task_done() - except Exception: - pass - - def perform_cloud_scan(self, file_path): - was_locked = False - try: - with self.lock_config: - if not self.pyas_config.get("cloud_switch", True): - return False - - api_host, api_key, max_size = self.pyas_config.get("api_host"), self.pyas_config.get("api_key"), self.pyas_config.get("size", 256 * 1024 * 1024) - - if not os.path.exists(file_path) or not os.path.isfile(file_path): - return False - - with self.lock_file_ops: - if file_path in self.virus_lock: - self.lock_file(file_path, False) - was_locked = True - - if os.path.getsize(file_path) > max_size: - if was_locked: - self.lock_file(file_path, True) - return False - - file_hash = self.calc_file_hash(file_path) - success, sha256 = self.cloud.upload_file(file_path, api_host, api_key, file_hash=file_hash) - if was_locked: - self.lock_file(file_path, True) - was_locked = False - - if not success: - self.write_log("WARN", "Cloud API", source=file_path, detail="Failed", success=False) - - except Exception as e: - self.write_log("WARN", "perform_cloud_scan", detail=str(e), success=False) - - finally: - if was_locked: - try: - self.lock_file(file_path, True) - except Exception: - pass - return False diff --git a/PYAS_Scheduler.py b/PYAS_Scheduler.py new file mode 100644 index 0000000..a762408 --- /dev/null +++ b/PYAS_Scheduler.py @@ -0,0 +1,160 @@ +import heapq +import itertools +import threading +import time +from concurrent.futures import ThreadPoolExecutor +from PYAS_Diagnostics import log_exception + + +class TaskScheduler: + def __init__(self, on_overflow, max_pending=2048, max_workers=4): + self.max_pending = max_pending + self.max_workers = max_workers + self.on_overflow = on_overflow + + self.condition = threading.Condition() + self.pending = {} + self.heap = [] + self.running = set() + self.sequence = itertools.count() + + self.overflow = False + self.closed = False + + self.pool = ThreadPoolExecutor(max_workers=max_workers, thread_name_prefix="PYAS.File") + + self.thread = threading.Thread( + target=self._dispatch, name="PYAS.FileScheduler", daemon=True + ) + + self.thread.start() + + def schedule(self, key, delay, callback, args=(), critical=False): + with self.condition: + if self.closed: + return False + + worker = threading.current_thread().name.startswith("PYAS.File_") + + limit = self.max_pending + (self.max_workers if worker and critical else 0) + + while key not in self.pending and len(self.pending) >= limit: + if not critical or worker: + self.overflow = True + self.condition.notify_all() + return False + + self.condition.wait() + + if self.closed: + return False + + token = next(self.sequence) + item = (time.monotonic() + max(0, delay), token, key, callback, args) + self.pending[key] = item + heapq.heappush(self.heap, item[:3]) + + if len(self.heap) > max(64, 2 * len(self.pending)): + self.heap = [item[:3] for item in self.pending.values()] + heapq.heapify(self.heap) + + self.condition.notify_all() + return True + + def cancel(self, predicate=None): + with self.condition: + for key in list(self.pending): + if predicate is None or predicate(key): + self.pending.pop(key, None) + + self.heap = [item[:3] for item in self.pending.values()] + heapq.heapify(self.heap) + + if predicate is None: + self.overflow = False + + self.condition.notify_all() + + def _dispatch(self): + while True: + with self.condition: + if self.closed: + return + + if len(self.running) >= self.max_workers: + self.condition.wait() + continue + + if ( + self.overflow + and len(self.pending) < max(1, self.max_pending // 2) + and "__recovery__" not in self.running + ): + self.overflow = False + item = (0, next(self.sequence), "__recovery__", self.on_overflow, ()) + else: + blocked = [] + item = None + + while self.heap: + due, token, key = self.heap[0] + current = self.pending.get(key) + + if current is None or current[1] != token: + heapq.heappop(self.heap) + continue + + if key in self.running: + blocked.append(heapq.heappop(self.heap)) + continue + + if due > time.monotonic(): + break + + heapq.heappop(self.heap) + item = self.pending.pop(key) + break + + for entry in blocked: + heapq.heappush(self.heap, entry) + + if item is None: + self.condition.wait(timeout=0.1 if self.heap else None) + continue + + self.running.add(item[2]) + self.condition.notify_all() + + try: + self.pool.submit(self._run, item) + except Exception: + log_exception("TaskScheduler.submit") + + with self.condition: + self.running.discard(item[2]) + self.overflow = True + self.condition.wait(timeout=0.5) + + def _run(self, item): + try: + with self.condition: + if self.closed: + return + + item[3](*item[4]) + except Exception: + log_exception("TaskScheduler.callback") + finally: + with self.condition: + self.running.discard(item[2]) + self.condition.notify_all() + + def close(self, wait=False): + with self.condition: + self.closed = True + self.pending.clear() + self.heap.clear() + self.condition.notify_all() + + self.thread.join(timeout=1) + self.pool.shutdown(wait=wait, cancel_futures=True) diff --git a/PYAS_Signature.py b/PYAS_Signature.py new file mode 100644 index 0000000..82e64fa --- /dev/null +++ b/PYAS_Signature.py @@ -0,0 +1,69 @@ +from PYAS_Diagnostics import log_exception +import os +import ctypes +import ctypes.wintypes +from PYAS_WinAPI import GUID, WINTRUST_DATA, WINTRUST_DATA_UNION, WINTRUST_FILE_INFO + + +class SignatureScanner: + def __init__(self): + self.verify = GUID( + 0x00AAC56B, + 0xCD44, + 0x11D0, + (ctypes.c_ubyte * 8)(0x8C, 0xC2, 0x00, 0xC0, 0x4F, 0xC2, 0x95, 0xEE), + ) + + def init_windll(self, path): + for name in path: + try: + setattr(self, name.lower(), ctypes.WinDLL(name, use_last_error=True)) + except Exception: + log_exception("PYAS_Signature.SignatureScanner.init_windll:14") + pass + + try: + self.WinVerifyTrust = self.wintrust.WinVerifyTrust + self.WinVerifyTrust.restype = ctypes.wintypes.LONG + self.WinVerifyTrust.argtypes = [ + ctypes.wintypes.HWND, + ctypes.POINTER(GUID), + ctypes.c_void_p, + ] + except Exception: + log_exception("PYAS_Signature.SignatureScanner.init_windll:21") + pass + + def sign_verify(self, file_path): + if os.name != "nt": + return False + + try: + fi = WINTRUST_FILE_INFO() + fi.cbStruct = ctypes.sizeof(WINTRUST_FILE_INFO) + fi.pcwszFilePath = os.path.abspath(file_path) + fi.hFile = None + fi.pgKnownSubject = None + + wt_union = WINTRUST_DATA_UNION() + wt_union.pFile = ctypes.pointer(fi) + + td = WINTRUST_DATA() + td.cbStruct = ctypes.sizeof(WINTRUST_DATA) + td.pPolicyCallbackData = None + td.pSIPClientData = None + td.dwUIChoice = 2 + td.fdwRevocationChecks = 0 + td.dwUnionChoice = 1 + td.u = wt_union + td.dwStateAction = 0 + td.hWVTStateData = None + td.pwszURLReference = None + td.dwProvFlags = 0 + td.dwUIContext = 0 + td.pSignatureSettings = None + + return self.WinVerifyTrust(None, ctypes.byref(self.verify), ctypes.byref(td)) == 0 + except Exception: + log_exception("PYAS_Signature.SignatureScanner.sign_verify:55") + return False diff --git a/PYAS_Startup.py b/PYAS_Startup.py new file mode 100644 index 0000000..42673fd --- /dev/null +++ b/PYAS_Startup.py @@ -0,0 +1,201 @@ +from PYAS_Diagnostics import log_exception +import ctypes +import importlib.metadata +import logging +import os +import platform +import subprocess +import sys +import tempfile +import threading +import uuid +from PYAS_Diagnostics import configure_report_logging + + +def configure_startup_logging(): + return configure_report_logging() + + +def bootstrap_startup(): + log_path = configure_startup_logging() + logger = logging.getLogger("PYAS.Startup") + + for stream in ("stdout", "stderr"): + if getattr(sys, stream) is None: + setattr(sys, stream, open(os.devnull, "w")) + + def exception_hook(exc_type, exc_value, traceback): + logger.critical("Unhandled startup exception", exc_info=(exc_type, exc_value, traceback)) + + if not any( + arg in sys.argv + for arg in ("-hide", "-h", "-quit", "-driver-unload", "-driver-uninstall") + ): + try: + ctypes.windll.user32.MessageBoxW( + None, + f"PYAS failed to start.\n\n{exc_value}\n\nDiagnostic log: {log_path or 'unavailable'}", + "PYAS Security", + 0x10, + ) + except Exception: + logger.exception("Could not display startup error") + + sys.excepthook = exception_hook + threading.excepthook = lambda args: logger.error( + "Unhandled thread exception: %s", + args.thread.name, + exc_info=(args.exc_type, args.exc_value, args.exc_traceback), + ) + logger.info( + "Starting PYAS; Python=%s; platform=%s; frozen=%s; recovery=%s", + sys.version.split()[0], + platform.platform(), + getattr(sys, "frozen", False), + os.environ.get("PYAS_WEBVIEW_RECOVERY", "0"), + ) + + for package in ("pywebview", "pythonnet", "pyinstaller"): + try: + logger.info("Dependency %s=%s", package, importlib.metadata.version(package)) + except importlib.metadata.PackageNotFoundError: + logger.info("Dependency metadata unavailable: %s", package) + + return log_path + + +def prepare_webview_profile(preferred_path): + logger = logging.getLogger("PYAS.Startup") + candidates = [preferred_path, os.path.join(tempfile.gettempdir(), "PYAS", "WebView2")] + + for path in dict.fromkeys(candidates): + try: + os.makedirs(path, exist_ok=True) + + with tempfile.TemporaryFile(dir=path) as probe: + probe.write(b"PYAS") + probe.flush() + probe.seek(0) + + if probe.read() != b"PYAS": + raise OSError("Profile readback failed") + + logger.info("WebView2 profile ready: %s", path) + return path + except OSError: + logger.exception("WebView2 profile is not writable: %s", path) + + raise OSError("No writable WebView2 user data folder is available") + + +def check_webview_dependencies(webview): + from webview.util import interop_dll_path + import sysconfig + + logger = logging.getLogger("PYAS.Startup") + architecture = ( + "win-x86" + if ctypes.sizeof(ctypes.c_void_p) == 4 + else ("win-arm64" if "arm64" in sysconfig.get_platform().lower() else "win-x64") + ) + loader_directory = interop_dll_path(architecture) + loader_path = os.path.join(loader_directory, "WebView2Loader.dll") + loader = ctypes.WinDLL(loader_path) + get_version = loader.GetAvailableCoreWebView2BrowserVersionString + get_version.argtypes = [ctypes.c_wchar_p, ctypes.POINTER(ctypes.c_void_p)] + get_version.restype = ctypes.c_long + version_pointer = ctypes.c_void_p() + runtime_path = webview.settings.get("WEBVIEW2_RUNTIME_PATH") or None + result = get_version(runtime_path, ctypes.byref(version_pointer)) + + try: + if result < 0 or not version_pointer.value: + raise RuntimeError( + f"WebView2 Runtime detection failed: HRESULT 0x{result & 0xFFFFFFFF:08X}" + ) + + version = ctypes.wstring_at(version_pointer) + logger.info("WebView2 Runtime=%s; loader=%s", version, loader_path) + finally: + if version_pointer.value: + ole32 = ctypes.WinDLL("ole32") + ole32.CoTaskMemFree.argtypes = [ctypes.c_void_p] + ole32.CoTaskMemFree.restype = None + ole32.CoTaskMemFree(version_pointer) + + os.environ["PATH"] = loader_directory + os.pathsep + os.environ.get("PATH", "") + import clr + + clr.AddReference("System.Windows.Forms") + + for name in ("Microsoft.Web.WebView2.Core.dll", "Microsoft.Web.WebView2.WinForms.dll"): + dll_path = interop_dll_path(name) + clr.AddReference(dll_path) + logger.info("WebView2 managed dependency loaded: %s", dll_path) + + return version + + +def wait_for_recovery_parent(): + parent = os.environ.pop("PYAS_WEBVIEW_RECOVERY_PARENT", None) + + if not parent: + return + + kernel32 = ctypes.WinDLL("kernel32", use_last_error=True) + kernel32.OpenProcess.argtypes = [ctypes.c_ulong, ctypes.c_int, ctypes.c_ulong] + kernel32.OpenProcess.restype = ctypes.c_void_p + kernel32.WaitForSingleObject.argtypes = [ctypes.c_void_p, ctypes.c_ulong] + kernel32.WaitForSingleObject.restype = ctypes.c_ulong + kernel32.CloseHandle.argtypes = [ctypes.c_void_p] + kernel32.CloseHandle.restype = ctypes.c_int + handle = kernel32.OpenProcess(0x00100000, False, int(parent)) + + if handle: + try: + if kernel32.WaitForSingleObject(handle, 10000) == 0x102: + raise RuntimeError("Previous PYAS instance did not exit for WebView2 recovery") + finally: + kernel32.CloseHandle(handle) + + +def restart_with_new_webview_profile(profile_path, script_path): + logger = logging.getLogger("PYAS.Startup") + + if os.environ.get("PYAS_WEBVIEW_RECOVERY") == "1": + return False + + try: + recovery_path = prepare_webview_profile(profile_path + "-Recovery-" + uuid.uuid4().hex) + environment = os.environ.copy() + environment.update( + { + "PYAS_WEBVIEW_RECOVERY": "1", + "PYAS_WEBVIEW_PROFILE": recovery_path, + "PYAS_WEBVIEW_RECOVERY_PARENT": str(os.getpid()), + "PYINSTALLER_RESET_ENVIRONMENT": "1", + } + ) + command = [sys.executable] + + if not getattr(sys, "frozen", False): + command.append(os.path.abspath(script_path)) + + command.extend(sys.argv[1:]) + working_directory = ( + os.path.dirname(sys.executable) + if getattr(sys, "frozen", False) + else os.path.dirname(os.path.abspath(script_path)) + ) + subprocess.Popen( + command, + env=environment, + cwd=working_directory, + creationflags=0x08000000, + close_fds=True, + ) + logger.warning("Restarting once with a new WebView2 profile: %s", recovery_path) + return True + except Exception: + logger.exception("Could not restart for WebView2 profile recovery") + return False diff --git a/PYAS_Storage.py b/PYAS_Storage.py new file mode 100644 index 0000000..0d523b0 --- /dev/null +++ b/PYAS_Storage.py @@ -0,0 +1,26 @@ +import json +import os +import tempfile + + +def atomic_write_json(path, data): + directory = os.path.dirname(os.path.abspath(path)) + os.makedirs(directory, exist_ok=True) + descriptor, temporary_path = tempfile.mkstemp(prefix=".pyas-", suffix=".tmp", dir=directory) + + try: + with os.fdopen(descriptor, "w", encoding="utf-8") as stream: + descriptor = None + + json.dump(data, stream, indent=4, ensure_ascii=False) + + stream.flush() + os.fsync(stream.fileno()) + + os.replace(temporary_path, path) + finally: + if descriptor is not None: + os.close(descriptor) + + if os.path.exists(temporary_path): + os.unlink(temporary_path) diff --git a/PYAS_System.py b/PYAS_System.py new file mode 100644 index 0000000..d91102f --- /dev/null +++ b/PYAS_System.py @@ -0,0 +1,596 @@ +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 diff --git a/PYAS_Threats.py b/PYAS_Threats.py new file mode 100644 index 0000000..2494f77 --- /dev/null +++ b/PYAS_Threats.py @@ -0,0 +1,254 @@ +from PYAS_Diagnostics import log_exception +import os +import time +import shutil + + +class ThreatMixin: + def manage_named_list(self, list_key, files, action="add", lock_func=None): + if list_key == "quarantine" and lock_func is None: + lock_func = self.lock_file + + norm_paths = self.norm_path(files or [], must_exist=False) + + if isinstance(norm_paths, str): + norm_paths = [norm_paths] if norm_paths else [] + + if not norm_paths: + return 0 + + if list_key == "custom_rule": + acted_items = [] + os.makedirs(self.path_rules, exist_ok=True) + + for file_path in norm_paths: + if action == "add": + if os.path.exists(file_path) and file_path.lower().endswith(".json"): + dest = os.path.join(self.path_rules, os.path.basename(file_path)) + + if file_path != dest: + try: + shutil.copy2(file_path, dest) + acted_items.append(dest) + + except Exception: + log_exception("PYAS_Threats.ThreatMixin.manage_named_list:34") + pass + else: + acted_items.append(dest) + + elif action == "remove": + try: + if os.path.exists(file_path): + os.remove(file_path) + + acted_items.append(file_path) + except Exception: + log_exception("PYAS_Threats.ThreatMixin.manage_named_list:45") + pass + + if acted_items and getattr(self, "driver_port", None): + self.clear_driver_rules() + + for f in os.listdir(self.path_rules): + if f.lower().endswith(".json"): + self.load_driver_rule_file(os.path.join(self.path_rules, f)) + + return len(acted_items) + + if action == "add": + if list_key == "white_list": + self.remove_list_items("quarantine", norm_paths) + + elif list_key == "quarantine": + self.remove_list_items("white_list", norm_paths) + + acted_items = [] + + with self.lock_config: + target_list = self.pyas_config.setdefault(list_key, []) + + if action == "add": + for path in norm_paths: + path_case = os.path.normcase(path) + exists = False + + for item in target_list: + val = item.get("file", "") if isinstance(item, dict) else item + np = self.norm_path(val, must_exist=False) + + if np and os.path.normcase(np) == path_case: + exists = True + break + + if not exists: + if lock_func: + lock_func(path, True) + + new_item = {"file": path, "time": time.time()} + is_directory = None + + if list_key == "white_list": + is_directory = os.path.isdir(path) + new_item["is_dir"] = is_directory + + target_list.append(new_item) + acted_items.append(path) + + if list_key == "white_list": + self.sync_driver_whitelist(path, True, is_directory) + + elif action == "remove": + norm_paths_case = {os.path.normcase(p) for p in norm_paths} + new_list = [] + + for item in target_list: + val = item.get("file", "") if isinstance(item, dict) else item + + if val: + np = self.norm_path(val, must_exist=False) + + if np and os.path.normcase(np) in norm_paths_case: + if lock_func: + lock_func(val, False) + + acted_items.append(val) + + if list_key == "white_list": + is_directory = ( + item.get("is_dir") if isinstance(item, dict) else None + ) + self.sync_driver_whitelist(val, False, is_directory) + + continue + + new_list.append(item) + + target_list[:] = new_list + + if acted_items: + self.write_log( + "INFO", "Config Update", detail=f"List [{list_key}] {action}: {acted_items}" + ) + self.save_config() + + return len(acted_items) + + def remove_list_items(self, list_key, paths_to_remove): + if list_key == "custom_rule": + return self.manage_named_list(list_key, paths_to_remove, action="remove") > 0 + + with self.lock_config: + target_list = self.pyas_config.get(list_key, []) + original_len = len(target_list) + + norm_paths_to_remove = set() + + for p in paths_to_remove: + np = self.norm_path(p, must_exist=False) + + if np: + norm_paths_to_remove.add(os.path.normcase(np)) + + if list_key == "quarantine": + for item in target_list: + val = item.get("file") if isinstance(item, dict) else item + + if val: + np = self.norm_path(val, must_exist=False) + + if np and os.path.normcase(np) in norm_paths_to_remove: + self.lock_file(val, False) + + new_list = [] + removed_items = [] + + for item in target_list: + val = ( + item.get("file") or item.get("exe") or item.get("title") + if isinstance(item, dict) + else item + ) + + if val: + np = self.norm_path(val, must_exist=False) + + if np and os.path.normcase(np) in norm_paths_to_remove: + removed_items.append(val) + + if list_key == "white_list": + is_directory = item.get("is_dir") if isinstance(item, dict) else None + self.sync_driver_whitelist(val, False, is_directory) + + continue + + new_list.append(item) + + self.pyas_config[list_key] = new_list + + if len(self.pyas_config[list_key]) < original_len: + self.write_log( + "INFO", "Config Update", detail=f"List [{list_key}] remove: {removed_items}" + ) + self.save_config() + return True + + return False + + def extract_list_items(self, paths, dest_dir): + if not dest_dir or not os.path.exists(dest_dir): + return False + + extracted_count = 0 + + for raw_path in paths: + src = self.norm_path(raw_path, must_exist=True) + + if not src: + continue + + base_name = os.path.basename(src) + target_path = os.path.join(dest_dir, base_name) + + if os.path.exists(target_path): + name, ext = os.path.splitext(base_name) + counter = 1 + + while os.path.exists(target_path): + target_path = os.path.join(dest_dir, f"{name} ({counter}){ext}") + counter += 1 + + was_locked = False + src_norm = os.path.normcase(src) + src_dir = src_norm + os.sep + + with self.lock_file_ops: + for locked_path in self.virus_lock: + locked_norm = os.path.normcase(locked_path) + + if locked_norm == src_norm or locked_norm.startswith(src_dir): + was_locked = True + break + + if was_locked: + self.lock_file(src, False) + + try: + if os.path.isdir(src): + shutil.copytree(src, target_path) + else: + shutil.copy2(src, target_path) + + extracted_count += 1 + self.write_log("INFO", "File Extract", source=src, target=target_path, operate=True) + + except Exception as e: + log_exception("PYAS_Threats.ThreatMixin.extract_list_items:210") + self.write_log( + "WARN", "extract_list_items", source=src, detail=str(e), success=False + ) + finally: + if was_locked: + self.lock_file(src, True) + + return extracted_count > 0 diff --git a/PYAS_Tools.py b/PYAS_Tools.py index 846bc09..4386a94 100644 --- a/PYAS_Tools.py +++ b/PYAS_Tools.py @@ -1,190 +1,110 @@ -import os, re, io, csv, json, time, shutil, threading -import pefile, hashlib, winreg, requests, subprocess -import ctypes, ctypes.wintypes - +from PYAS_Diagnostics import log_exception +import os +import re +import io +import csv +import json +import time +import shutil +import threading +import sys +import pefile +import hashlib +import winreg +import requests +import subprocess +import ctypes +import ctypes.wintypes from concurrent.futures import ThreadPoolExecutor +from PYAS_WinAPI import ( + COPYDATASTRUCT, + ERROR_LOCK_VIOLATION, + ERROR_NOT_FOUND, + ERROR_OPERATION_ABORTED, + ERROR_SHARING_VIOLATION, + FILE_ATTRIBUTE_NORMAL, + FILE_LOCK_RETRY_SECONDS, + FILE_NOTIFY_INFORMATION, + FILE_SCAN_DEBOUNCE_SECONDS, + FILE_SCAN_RETRY_SECONDS, + FILE_SHARE_READ, + FILE_SHARE_WRITE, + FILTER_MESSAGE_HEADER, + FLT_PORT_FLAG_SYNC_HANDLE, + GENERIC_READ, + GUID, + HRESULT_IO_PENDING, + INTERNAL_STORAGE_BUS_TYPES, + INVALID_HANDLE_VALUE, + IOCTL_STORAGE_GET_HOTPLUG_INFO, + IOCTL_STORAGE_QUERY_PROPERTY, + IO_COUNTERS, + LUID, + LUID_AND_ATTRIBUTES, + MEMORY_BASIC_INFORMATION, + MIB_TCPROW_OWNER_PID, + OPEN_EXISTING, + OVERLAPPED, + POINT, + PROCESSENTRY32W, + PROCESS_BASIC_INFORMATION, + PROCESS_QUERY_INFORMATION, + PROCESS_QUERY_LIMITED_INFORMATION, + PROCESS_SUSPEND_RESUME, + PROCESS_TERMINATE, + PROCESS_VM_READ, + PROPERTY_STANDARD_QUERY, + PYAS_CONNECTION_CONTEXT, + PYAS_CONNECTION_MAGIC, + PYAS_CONNECTION_VERSION, + PYAS_FULL_MESSAGE, + PYAS_MESSAGE, + PYAS_PROCESS_NETWORK_ACCESS, + PYAS_PROCESS_SCAN_ACCESS, + PYAS_PROCESS_TERMINATE_ACCESS, + PYAS_USER_MESSAGE, + RECT, + SERVICE_STATUS_PROCESS, + SHQUERYRBINFO, + STORAGE_DESCRIPTOR_HEADER, + STORAGE_DEVICE_DESCRIPTOR, + STORAGE_DEVICE_PROPERTY, + STORAGE_HOTPLUG_INFO, + STORAGE_PROPERTY_QUERY, + TOKEN_PRIVILEGES, + UNICODE_STRING, + WAIT_OBJECT_0, + WAIT_TIMEOUT, + WINTRUST_DATA, + WINTRUST_DATA_UNION, + WINTRUST_FILE_INFO, +) +from PYAS_Autostart import AutostartMixin +from PYAS_Process import ProcessMixin +from PYAS_Maintenance import MaintenanceMixin +from PYAS_Threats import ThreatMixin -#################################################################################################### -PROCESS_TERMINATE = 0x0001 -PROCESS_QUERY_INFORMATION = 0x0400 -PROCESS_QUERY_LIMITED_INFORMATION = 0x1000 -PYAS_PROCESS_TERMINATE_ACCESS = PROCESS_TERMINATE | PROCESS_QUERY_INFORMATION | PROCESS_QUERY_LIMITED_INFORMATION - -#################################################################################################### - -class PROCESSENTRY32W(ctypes.Structure): - _fields_ = [ - ("dwSize", ctypes.wintypes.DWORD), - ("cntUsage", ctypes.wintypes.DWORD), - ("th32ProcessID", ctypes.wintypes.DWORD), - ("th32DefaultHeapID", ctypes.wintypes.LPVOID), - ("th32ModuleID", ctypes.wintypes.DWORD), - ("cntThreads", ctypes.wintypes.DWORD), - ("th32ParentProcessID", ctypes.wintypes.DWORD), - ("pcPriClassBase", ctypes.wintypes.LONG), - ("dwFlags", ctypes.wintypes.DWORD), - ("szExeFile", ctypes.wintypes.WCHAR * 260) - ] - -class MIB_TCPROW_OWNER_PID(ctypes.Structure): - _fields_ = [ - ("dwState", ctypes.wintypes.DWORD), - ("dwLocalAddr", ctypes.wintypes.DWORD), - ("dwLocalPort", ctypes.wintypes.DWORD), - ("dwRemoteAddr", ctypes.wintypes.DWORD), - ("dwRemotePort", ctypes.wintypes.DWORD), - ("dwOwningPid", ctypes.wintypes.DWORD) - ] - -class FILE_NOTIFY_INFORMATION(ctypes.Structure): - _fields_ = [ - ("NextEntryOffset", ctypes.wintypes.DWORD), - ("Action", ctypes.wintypes.DWORD), - ("FileNameLength", ctypes.wintypes.DWORD), - ("FileName", ctypes.wintypes.WCHAR * 1024) - ] - -class PROCESS_BASIC_INFORMATION(ctypes.Structure): - _fields_ = [ - ("Reserved1", ctypes.wintypes.LPVOID), - ("PebBaseAddress", ctypes.wintypes.LPVOID), - ("Reserved2", ctypes.wintypes.LPVOID * 2), - ("UniqueProcessId", ctypes.wintypes.LPVOID), - ("Reserved3", ctypes.wintypes.LPVOID) - ] - -class SHQUERYRBINFO(ctypes.Structure): - _fields_ = [ - ("cbSize", ctypes.wintypes.DWORD), - ("i64Size", ctypes.c_int64), - ("i64NumItems", ctypes.c_int64) - ] - -class UNICODE_STRING(ctypes.Structure): - _fields_ = [ - ("Length", ctypes.wintypes.USHORT), - ("MaximumLength", ctypes.wintypes.USHORT), - ("Buffer", ctypes.c_void_p) - ] - -class FILTER_MESSAGE_HEADER(ctypes.Structure): - _fields_ = [ - ("ReplyLength", ctypes.wintypes.ULONG), - ("MessageId", ctypes.c_uint64) - ] - -class PYAS_MESSAGE(ctypes.Structure): - _fields_ = [ - ("MessageCode", ctypes.wintypes.ULONG), - ("ProcessId", ctypes.wintypes.ULONG), - ("Path", ctypes.wintypes.WCHAR * 1024) - ] - -class PYAS_FULL_MESSAGE(ctypes.Structure): - _fields_ = [ - ("Header", FILTER_MESSAGE_HEADER), - ("Data", PYAS_MESSAGE) - ] - -class PYAS_USER_MESSAGE(ctypes.Structure): - _fields_ = [ - ("Command", ctypes.wintypes.ULONG), - ("Path", ctypes.wintypes.WCHAR * 1024) - ] - -class COPYDATASTRUCT(ctypes.Structure): - _fields_ = [ - ("dwData", ctypes.c_size_t), - ("cbData", ctypes.wintypes.DWORD), - ("lpData", ctypes.c_void_p) - ] - -class IO_COUNTERS(ctypes.Structure): - _fields_ = [ - ("ReadOperationCount", ctypes.c_ulonglong), - ("WriteOperationCount", ctypes.c_ulonglong), - ("OtherOperationCount", ctypes.c_ulonglong), - ("ReadTransferCount", ctypes.c_ulonglong), - ("WriteTransferCount", ctypes.c_ulonglong), - ("OtherTransferCount", ctypes.c_ulonglong) - ] - -class MEMORY_BASIC_INFORMATION(ctypes.Structure): - _fields_ = [ - ("BaseAddress", ctypes.c_void_p), - ("AllocationBase", ctypes.c_void_p), - ("AllocationProtect", ctypes.wintypes.DWORD), - ("RegionSize", ctypes.c_size_t), - ("State", ctypes.wintypes.DWORD), - ("Protect", ctypes.wintypes.DWORD), - ("Type", ctypes.wintypes.DWORD) - ] - -class LUID(ctypes.Structure): - _fields_ = [ - ("LowPart", ctypes.wintypes.DWORD), - ("HighPart", ctypes.wintypes.LONG) - ] - -class LUID_AND_ATTRIBUTES(ctypes.Structure): - _fields_ = [ - ("Luid", LUID), - ("Attributes", ctypes.wintypes.DWORD) - ] - -class TOKEN_PRIVILEGES(ctypes.Structure): - _fields_ = [ - ("PrivilegeCount", ctypes.wintypes.DWORD), - ("Privileges", LUID_AND_ATTRIBUTES * 1) - ] - -class SERVICE_STATUS_PROCESS(ctypes.Structure): - _fields_ = [ - ("dwServiceType", ctypes.wintypes.DWORD), - ("dwCurrentState", ctypes.wintypes.DWORD), - ("dwControlsAccepted", ctypes.wintypes.DWORD), - ("dwWin32ExitCode", ctypes.wintypes.DWORD), - ("dwServiceSpecificExitCode", ctypes.wintypes.DWORD), - ("dwCheckPoint", ctypes.wintypes.DWORD), - ("dwWaitHint", ctypes.wintypes.DWORD), - ("dwProcessId", ctypes.wintypes.DWORD), - ("dwServiceFlags", ctypes.wintypes.DWORD) - ] - -class POINT(ctypes.Structure): - _fields_ = [ - ("x", ctypes.c_long), - ("y", ctypes.c_long) - ] - -class RECT(ctypes.Structure): - _fields_ = [ - ("left", ctypes.c_long), - ("top", ctypes.c_long), - ("right", ctypes.c_long), - ("bottom", ctypes.c_long) - ] - -#################################################################################################### - -class ToolsMixin: +class ToolsMixin(AutostartMixin, ProcessMixin, MaintenanceMixin, ThreatMixin): def _find_windows_tool(self, *names): windows_dir = os.environ.get("SystemRoot") or os.environ.get("WINDIR") or r"C:\Windows" search_dirs = [ os.path.join(windows_dir, "Sysnative"), os.path.join(windows_dir, "System32"), - windows_dir + windows_dir, ] for name in names: executable_name = name if name.lower().endswith(".exe") else name + ".exe" + for directory in search_dirs: candidate = os.path.join(directory, executable_name) + if os.path.isfile(candidate): return candidate candidate = shutil.which(name) + if candidate: return candidate @@ -192,59 +112,59 @@ class ToolsMixin: def _run_windows_tool(self, tool_names, arguments, **kwargs): executable = self._find_windows_tool(*tool_names) + if not executable: return None options = {"creationflags": 0x08000000} options.update(kwargs) + try: return subprocess.run([executable, *arguments], **options) except (FileNotFoundError, OSError): + log_exception("PYAS_Tools.ToolsMixin._run_windows_tool:43") return None def _run_powershell(self, command, **kwargs): return self._run_windows_tool( - ("powershell.exe", "pwsh.exe"), - ["-NoProfile", "-NonInteractive", "-ExecutionPolicy", "Bypass", "-WindowStyle", "Hidden", "-Command", command], - **kwargs + (r"WindowsPowerShell\v1.0\powershell.exe", "powershell.exe", "pwsh.exe"), + [ + "-NoProfile", + "-NonInteractive", + "-ExecutionPolicy", + "Bypass", + "-WindowStyle", + "Hidden", + "-Command", + command, + ], + **kwargs, ) - def _set_registry_autostart(self, enable): - run_key = r"Software\Microsoft\Windows\CurrentVersion\Run" - value_name = "PYAS_Security" - if enable: - command = f'"{self.file_pyas}" -hide' - return self._reg_write(winreg.HKEY_CURRENT_USER, run_key, value_name, winreg.REG_SZ, command) - - try: - with winreg.OpenKey(winreg.HKEY_CURRENT_USER, run_key, 0, winreg.KEY_SET_VALUE) as reg: - try: - winreg.DeleteValue(reg, value_name) - except FileNotFoundError: - pass - return True - except FileNotFoundError: - return True - except Exception: - return False - def _clear_event_log(self, log_name): result = self._run_windows_tool( ("wevtutil.exe",), ["cl", log_name], stdout=subprocess.DEVNULL, - stderr=subprocess.DEVNULL + stderr=subprocess.DEVNULL, ) + if result is not None and result.returncode == 0: return True try: wevtapi = ctypes.WinDLL("wevtapi", use_last_error=True) - wevtapi.EvtClearLog.argtypes = [ctypes.wintypes.HANDLE, ctypes.wintypes.LPCWSTR, ctypes.wintypes.LPCWSTR, ctypes.wintypes.DWORD] + wevtapi.EvtClearLog.argtypes = [ + ctypes.wintypes.HANDLE, + ctypes.wintypes.LPCWSTR, + ctypes.wintypes.LPCWSTR, + ctypes.wintypes.DWORD, + ] wevtapi.EvtClearLog.restype = ctypes.wintypes.BOOL return bool(wevtapi.EvtClearLog(None, log_name, None, 0)) except Exception: + log_exception("PYAS_Tools.ToolsMixin._clear_event_log:69") return False def _manage_service_native(self, service_name, action): @@ -257,21 +177,32 @@ class ToolsMixin: try: advapi32 = ctypes.WinDLL("advapi32", use_last_error=True) - advapi32.OpenSCManagerW.argtypes = [ctypes.wintypes.LPCWSTR, ctypes.wintypes.LPCWSTR, ctypes.wintypes.DWORD] + advapi32.OpenSCManagerW.argtypes = [ + ctypes.wintypes.LPCWSTR, + ctypes.wintypes.LPCWSTR, + ctypes.wintypes.DWORD, + ] advapi32.OpenSCManagerW.restype = ctypes.wintypes.HANDLE - advapi32.OpenServiceW.argtypes = [ctypes.wintypes.HANDLE, ctypes.wintypes.LPCWSTR, ctypes.wintypes.DWORD] + advapi32.OpenServiceW.argtypes = [ + ctypes.wintypes.HANDLE, + ctypes.wintypes.LPCWSTR, + ctypes.wintypes.DWORD, + ] advapi32.OpenServiceW.restype = ctypes.wintypes.HANDLE advapi32.CloseServiceHandle.argtypes = [ctypes.wintypes.HANDLE] advapi32.CloseServiceHandle.restype = ctypes.wintypes.BOOL manager = advapi32.OpenSCManagerW(None, None, SC_MANAGER_CONNECT) + if not manager: return False service = None + try: desired_access = DELETE if action == "delete" else SERVICE_CHANGE_CONFIG service = advapi32.OpenServiceW(manager, service_name, desired_access) + if not service: return False @@ -292,43 +223,46 @@ class ToolsMixin: ctypes.wintypes.LPCWSTR, ctypes.wintypes.LPCWSTR, ctypes.wintypes.LPCWSTR, - ctypes.wintypes.LPCWSTR + ctypes.wintypes.LPCWSTR, ] advapi32.ChangeServiceConfigW.restype = ctypes.wintypes.BOOL - return bool(advapi32.ChangeServiceConfigW( - service, - SERVICE_NO_CHANGE, - start_type, - SERVICE_NO_CHANGE, - None, - None, - None, - None, - None, - None, - None - )) + return bool( + advapi32.ChangeServiceConfigW( + service, + SERVICE_NO_CHANGE, + start_type, + SERVICE_NO_CHANGE, + None, + None, + None, + None, + None, + None, + None, + ) + ) finally: if service: advapi32.CloseServiceHandle(service) + advapi32.CloseServiceHandle(manager) except Exception: + log_exception("PYAS_Tools.ToolsMixin._manage_service_native:137") return False def _manage_scheduled_task(self, task_path, action): schtasks_action = "/delete" if action == "delete" else "/change" arguments = [schtasks_action, "/tn", task_path] + if action == "delete": arguments.append("/f") else: arguments.append("/enable" if action == "enable" else "/disable") result = self._run_windows_tool( - ("schtasks.exe",), - arguments, - capture_output=True, - text=True + ("schtasks.exe",), arguments, capture_output=True, text=True ) + if result is not None and result.returncode == 0: return True @@ -336,8 +270,9 @@ class ToolsMixin: operation = { "delete": "Unregister-ScheduledTask -InputObject $Task -Confirm:$false", "enable": "Enable-ScheduledTask -InputObject $Task | Out-Null", - "disable": "Disable-ScheduledTask -InputObject $Task | Out-Null" + "disable": "Disable-ScheduledTask -InputObject $Task | Out-Null", }.get(action) + if not operation: return False @@ -357,6 +292,7 @@ class ToolsMixin: return val except Exception: + log_exception("PYAS_Tools.ToolsMixin._reg_read:181") return None def _reg_write(self, root, path, value_name, value_type, value): @@ -366,9 +302,11 @@ class ToolsMixin: winreg.SetValue(reg, "", value_type, value) else: winreg.SetValueEx(reg, value_name, 0, value_type, value) - return True + actual, actual_type = winreg.QueryValueEx(reg, value_name or "") + return actual == value and actual_type == value_type except Exception: + log_exception("ToolsMixin._reg_write") return False def _reg_delete(self, root, path, value_name=None): @@ -381,306 +319,17 @@ class ToolsMixin: return True except Exception: + log_exception("PYAS_Tools.ToolsMixin._reg_delete:205") return False -#################################################################################################### - - def manage_autostart(self, enable): - task_name = "PYAS_Security_ATS" - executable = os.path.normpath(self.file_pyas) - escaped_executable = executable.replace("'", "''") - errors = [] - - if enable: - task_command = f'"{executable}" -hide' - result = self._run_windows_tool( - ("schtasks.exe",), - ["/Create", "/TN", task_name, "/TR", task_command, "/SC", "ONLOGON", "/RL", "HIGHEST", "/F"], - capture_output=True, - text=True - ) - if result and result.returncode == 0: - self._set_registry_autostart(False) - return True - if result: - errors.append((result.stderr or result.stdout or f"schtasks exit {result.returncode}").strip()) - - command = ( - f"$Action = New-ScheduledTaskAction -Execute '{escaped_executable}' -Argument '-hide'; " - "$Trigger = New-ScheduledTaskTrigger -AtLogOn; " - "$Settings = New-ScheduledTaskSettingsSet -AllowStartIfOnBatteries -DontStopIfGoingOnBatteries; " - f"Register-ScheduledTask -TaskName '{task_name}' -Action $Action -Trigger $Trigger -Settings $Settings -RunLevel Highest -Force | Out-Null" - ) - result = self._run_powershell(command, capture_output=True, text=True) - if result and result.returncode == 0: - self._set_registry_autostart(False) - return True - if result: - errors.append((result.stderr or result.stdout or f"PowerShell exit {result.returncode}").strip()) - - if self._set_registry_autostart(True): - self.write_log("INFO", "manage_autostart", detail="Scheduled Task unavailable; using HKCU Run fallback", success=True) - return True - else: - task_removed = False - query_result = self._run_windows_tool( - ("schtasks.exe",), - ["/Query", "/TN", task_name], - capture_output=True, - text=True - ) - result = None - if query_result is not None and query_result.returncode != 0: - task_removed = True - elif query_result is not None: - result = self._run_windows_tool( - ("schtasks.exe",), - ["/Delete", "/TN", task_name, "/F"], - capture_output=True, - text=True - ) - task_removed = bool(result and result.returncode == 0) - - if query_result is None or not task_removed: - command = ( - f"$Task = Get-ScheduledTask -TaskName '{task_name}' -ErrorAction SilentlyContinue; " - "if ($Task) { $Task | Unregister-ScheduledTask -Confirm:$false }" - ) - powershell_result = self._run_powershell(command, capture_output=True, text=True) - if powershell_result is not None: - task_removed = powershell_result.returncode == 0 - result = powershell_result - - if result and result.returncode != 0: - errors.append((result.stderr or result.stdout or f"task removal exit {result.returncode}").strip()) - - registry_removed = self._set_registry_autostart(False) - if registry_removed and (task_removed or query_result is None and result is None): - return True - - detail = "; ".join(error for error in errors if error) or "No supported autostart method was available" - self.write_log("WARN", "manage_autostart", detail=detail, success=False) - return False - -#################################################################################################### - - def manage_named_list(self, list_key, files, action="add", lock_func=None): - if list_key == "quarantine" and lock_func is None: - lock_func = self.lock_file - - norm_paths = self.norm_path(files or [], must_exist=False) - if isinstance(norm_paths, str): - norm_paths = [norm_paths] if norm_paths else [] - - if not norm_paths: - return 0 - - if list_key == "custom_rule": - acted_items = [] - os.makedirs(self.path_rules, exist_ok=True) - - for file_path in norm_paths: - if action == "add": - if os.path.exists(file_path) and file_path.lower().endswith(".json"): - dest = os.path.join(self.path_rules, os.path.basename(file_path)) - - if file_path != dest: - try: - shutil.copy2(file_path, dest) - acted_items.append(dest) - - except Exception: - pass - else: - acted_items.append(dest) - - elif action == "remove": - try: - if os.path.exists(file_path): - os.remove(file_path) - - acted_items.append(file_path) - except Exception: - pass - - if acted_items and getattr(self, 'driver_port', None): - self.clear_driver_rules() - for f in os.listdir(self.path_rules): - if f.lower().endswith(".json"): - self.load_driver_rule_file(os.path.join(self.path_rules, f)) - - return len(acted_items) - - if action == "add": - if list_key == "white_list": - self.remove_list_items("quarantine", norm_paths) - - elif list_key == "quarantine": - self.remove_list_items("white_list", norm_paths) - - acted_items = [] - with self.lock_config: - target_list = self.pyas_config.setdefault(list_key, []) - if action == "add": - for path in norm_paths: - path_case = os.path.normcase(path) - exists = False - - for item in target_list: - val = item.get("file", "") if isinstance(item, dict) else item - np = self.norm_path(val, must_exist=False) - - if np and os.path.normcase(np) == path_case: - exists = True - break - - if not exists: - if lock_func: - lock_func(path, True) - - new_item = {"file": path, "time": time.time()} - is_directory = None - if list_key == "white_list": - is_directory = os.path.isdir(path) - new_item["is_dir"] = is_directory - - target_list.append(new_item) - acted_items.append(path) - if list_key == "white_list": - self.sync_driver_whitelist(path, True, is_directory) - - elif action == "remove": - norm_paths_case = {os.path.normcase(p) for p in norm_paths} - new_list = [] - - for item in target_list: - val = item.get("file", "") if isinstance(item, dict) else item - if val: - np = self.norm_path(val, must_exist=False) - if np and os.path.normcase(np) in norm_paths_case: - if lock_func: - lock_func(val, False) - - acted_items.append(val) - if list_key == "white_list": - is_directory = item.get("is_dir") if isinstance(item, dict) else None - self.sync_driver_whitelist(val, False, is_directory) - - continue - new_list.append(item) - target_list[:] = new_list - - if acted_items: - self.write_log("INFO", "Config Update", detail=f"List [{list_key}] {action}: {acted_items}") - self.save_config() - - return len(acted_items) - -#################################################################################################### - - def remove_list_items(self, list_key, paths_to_remove): - if list_key == "custom_rule": - return self.manage_named_list(list_key, paths_to_remove, action="remove") > 0 - - with self.lock_config: - target_list = self.pyas_config.get(list_key, []) - original_len = len(target_list) - - norm_paths_to_remove = set() - for p in paths_to_remove: - np = self.norm_path(p, must_exist=False) - if np: - norm_paths_to_remove.add(os.path.normcase(np)) - - if list_key == "quarantine": - for item in target_list: - val = item.get("file") if isinstance(item, dict) else item - if val: - np = self.norm_path(val, must_exist=False) - if np and os.path.normcase(np) in norm_paths_to_remove: - self.lock_file(val, False) - - new_list = [] - removed_items = [] - for item in target_list: - val = item.get("file") or item.get("exe") or item.get("title") if isinstance(item, dict) else item - if val: - np = self.norm_path(val, must_exist=False) - if np and os.path.normcase(np) in norm_paths_to_remove: - removed_items.append(val) - if list_key == "white_list": - is_directory = item.get("is_dir") if isinstance(item, dict) else None - self.sync_driver_whitelist(val, False, is_directory) - - continue - new_list.append(item) - - self.pyas_config[list_key] = new_list - - if len(self.pyas_config[list_key]) < original_len: - self.write_log("INFO", "Config Update", detail=f"List [{list_key}] remove: {removed_items}") - self.save_config() - return True - - return False - -#################################################################################################### - - def extract_list_items(self, paths, dest_dir): - if not dest_dir or not os.path.exists(dest_dir): - return False - - extracted_count = 0 - for raw_path in paths: - src = self.norm_path(raw_path, must_exist=True) - if not src: - continue - - base_name = os.path.basename(src) - target_path = os.path.join(dest_dir, base_name) - - if os.path.exists(target_path): - name, ext = os.path.splitext(base_name) - counter = 1 - - while os.path.exists(target_path): - target_path = os.path.join(dest_dir, f"{name} ({counter}){ext}") - counter += 1 - - was_locked = False - src_norm = os.path.normcase(src) - src_dir = src_norm + os.sep - - with self.lock_file_ops: - for locked_path in self.virus_lock: - locked_norm = os.path.normcase(locked_path) - if locked_norm == src_norm or locked_norm.startswith(src_dir): - was_locked = True - break - - if was_locked: - self.lock_file(src, False) - - try: - if os.path.isdir(src): - shutil.copytree(src, target_path) - else: - shutil.copy2(src, target_path) - extracted_count += 1 - self.write_log("INFO", "File Extract", source=src, target=target_path, operate=True) - - except Exception as e: - self.write_log("WARN", "extract_list_items", source=src, detail=str(e), success=False) - finally: - if was_locked: - self.lock_file(src, True) - - return extracted_count > 0 - -#################################################################################################### - def start_daemon_thread(self, target, *args, **kwargs): - t = threading.Thread(target=target, args=args, kwargs=kwargs, daemon=True) + t = threading.Thread( + target=target, + args=args, + kwargs=kwargs, + daemon=True, + name=f"PYAS.{getattr(target, '__name__', 'worker')}", + ) t.start() return t @@ -694,6 +343,7 @@ class ToolsMixin: return ap if (not must_exist or os.path.exists(ap)) else None except Exception: + log_exception("PYAS_Tools.ToolsMixin.norm_path:223") return None return path @@ -702,29 +352,39 @@ class ToolsMixin: pa, pb = self.norm_path(a, must_exist=False), self.norm_path(b, must_exist=False) return os.path.normcase(pa) == os.path.normcase(pb) if pa and pb else False -#################################################################################################### - def get_file_version(self, file_path): + pe = None + try: pe = pefile.PE(file_path, fast_load=True) - pe.parse_data_directories(directories=[pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_RESOURCE']]) + pe.parse_data_directories( + directories=[pefile.DIRECTORY_ENTRY["IMAGE_DIRECTORY_ENTRY_RESOURCE"]] + ) - if hasattr(pe, 'FileInfo'): - for fl in pe.FileInfo: - for info in fl: - if getattr(info, 'name', '') in ('StringFileInfo', b'StringFileInfo'): + for file_info in getattr(pe, "FileInfo", []): + for info in file_info: + if getattr(info, "name", "") in ("StringFileInfo", b"StringFileInfo"): + for table in getattr(info, "StringTable", []): + for key, value in table.entries.items(): + key = ( + key.decode("utf-8", "ignore") + if isinstance(key, bytes) + else str(key) + ) - for st in getattr(info, 'StringTable', []): - for k, val in st.entries.items(): - k_str = k.decode('utf-8', 'ignore') if isinstance(k, bytes) else str(k) - v_str = val.decode('utf-8', 'ignore') if isinstance(val, bytes) else str(val) - if k_str == 'FileVersion': - pe.close() - return v_str.strip() - - pe.close() + if key == "FileVersion": + value = ( + value.decode("utf-8", "ignore") + if isinstance(value, bytes) + else str(value) + ) + return value.strip() except Exception: + log_exception("PYAS_Tools.ToolsMixin.get_file_version:246") pass + finally: + if pe is not None: + pe.close() return "0.0.0.0" @@ -735,673 +395,41 @@ class ToolsMixin: return val1 >= val2 except Exception: + log_exception("PYAS_Tools.ToolsMixin.compare_versions:259") return False def check_update(self): try: current = self.pyas_config.get("version", "0.0.0") - j = requests.get("https://api.github.com/repos/87owo/PYAS/releases/latest", headers={"Accept": "application/vnd.github+json", "User-Agent": "PYAS"}, timeout=10).json() + j = requests.get( + "https://api.github.com/repos/87owo/PYAS/releases/latest", + headers={"Accept": "application/vnd.github+json", "User-Agent": "PYAS"}, + timeout=10, + ).json() latest = str(j.get("tag_name") or j.get("name") or "").strip() page = j.get("html_url") or "https://github.com/87owo/PYAS/releases" if latest: if not self.compare_versions(current, latest) and current != latest: return {"has_update": True, "latest": latest, "current": current, "url": page} + return {"has_update": False, "latest": latest, "current": current, "url": page} except Exception: + log_exception("PYAS_Tools.ToolsMixin.check_update:274") pass - return {"error": True} -#################################################################################################### + return {"error": True} def calc_file_hash(self, file_path, block_size=65536): try: h = hashlib.sha256() + with open(file_path, "rb") as f: - for chunk in iter(lambda: f.read(block_size), b""): h.update(chunk) + for chunk in iter(lambda: f.read(block_size), b""): + h.update(chunk) return h.hexdigest() except Exception: + log_exception("PYAS_Tools.ToolsMixin.calc_file_hash:285") return None - -#################################################################################################### - - def device_path_to_drive(self, path): - if not path: - return "" - - for d in range(65, 91): - drive = f"{chr(d)}:" - buf = ctypes.create_unicode_buffer(1024) - - if self.kernel32.QueryDosDeviceW(drive, buf, 1024) and path.startswith(buf.value): - return path.replace(buf.value, drive, 1) - - return path - - def get_exe_info(self, pid): - name, file_path = None, None - h = self.kernel32.OpenProcess(0x1000, False, pid) - if h: - try: - buf = ctypes.create_unicode_buffer(1024) - size = ctypes.wintypes.DWORD(1024) - - if self.kernel32.QueryFullProcessImageNameW(h, 0, buf, ctypes.byref(size)): - file_path = self.norm_path(buf.value) - if file_path: - name = os.path.basename(file_path) - - elif self.psapi.GetProcessImageFileNameW(h, buf, 1024): - file_path = self.norm_path(self.device_path_to_drive(buf.value)) - if file_path: - name = os.path.basename(file_path) - - finally: - self.kernel32.CloseHandle(h) - - return name, file_path - - def get_process_file(self, h_process): - buf = ctypes.create_unicode_buffer(1024) - if self.psapi.GetProcessImageFileNameW(h_process, buf, 1024): - return self.norm_path(self.device_path_to_drive(buf.value)) - - return "" - -#################################################################################################### - - def get_process_cmdline(self, h): - pbi = PROCESS_BASIC_INFORMATION() - retlen = ctypes.wintypes.ULONG(0) - if self.ntdll.NtQueryInformationProcess(h, 0, ctypes.byref(pbi), ctypes.sizeof(pbi), ctypes.byref(retlen)) != 0: - return "" - - pointer_size = ctypes.sizeof(ctypes.c_void_p) - addr_pp = (int(pbi.PebBaseAddress) if pbi.PebBaseAddress else 0) + (0x20 if pointer_size == 8 else 0x10) - if not addr_pp: - return "" - - read_buf = (ctypes.c_ubyte * pointer_size)() - lpBytesRead = ctypes.c_size_t(0) - if not self.kernel32.ReadProcessMemory(h, ctypes.c_void_p(addr_pp), read_buf, pointer_size, ctypes.byref(lpBytesRead)): - return "" - - proc_params_address = ctypes.c_void_p.from_buffer_copy(read_buf).value - if not proc_params_address: - return "" - - us = UNICODE_STRING() - if not self.kernel32.ReadProcessMemory(h, ctypes.c_void_p(int(proc_params_address) + (0x70 if pointer_size == 8 else 0x40)), ctypes.byref(us), ctypes.sizeof(us), ctypes.byref(lpBytesRead)): - return "" - if not us.Buffer or us.Length == 0: - return "" - - buf = (ctypes.c_wchar * int(us.Length // 2))() - if not self.kernel32.ReadProcessMemory(h, ctypes.c_void_p(int(us.Buffer)), buf, us.Length, ctypes.byref(lpBytesRead)): - return "" - - return "".join(buf) - - def extract_paths_from_cmdline(self, cmdline): - if not cmdline: - return [] - - found = [] - for m in re.finditer(r'"([^"]+)"|\'([^\']+)\'', cmdline): - path = m.group(1) or m.group(2) - if re.match(r'^[A-Za-z]:\\|^\\\\', path): - found.append(path) - - if not found: - m = re.search(r'([A-Za-z]:\\[^\*?"<>\|]+\.(?:exe|dll|bat|cmd|vbs|sys|com|pif))', cmdline, re.IGNORECASE) - if m: - found.append(m.group(1)) - - if not found: - argc = ctypes.c_int(0) - argv = self.shell32.CommandLineToArgvW(cmdline, ctypes.byref(argc)) - if argv: - args = [argv[i] for i in range(argc.value)] - self.kernel32.LocalFree(ctypes.cast(argv, ctypes.c_void_p)) - patterns = [r'([A-Za-z]:\\[^"\']+)', r'(\\\\[^"\']+)', r'(\.\\[^"\']+)', r'(\./[^"\']+)', r'([A-Za-z]:/[^"\']+)', r'([^\s]*\\[^\s]+)'] - - for arg in args: - for p in patterns: - for match in re.finditer(p, arg): found.append(match.group(1).strip('"').strip("'")) - - return list(dict.fromkeys([p.strip('"').strip("'") for p in found])) - -#################################################################################################### - - def _enum_processes(self): - pe = PROCESSENTRY32W() - pe.dwSize = ctypes.sizeof(PROCESSENTRY32W) - snapshot = self.kernel32.CreateToolhelp32Snapshot(0x00000002, 0) - - if snapshot in (-1, 0xFFFFFFFF, 0xFFFFFFFFFFFFFFFF): - return - - try: - if self.kernel32.Process32FirstW(snapshot, ctypes.byref(pe)): - while True: - yield pe.th32ProcessID, pe.szExeFile - if not self.kernel32.Process32NextW(snapshot, ctypes.byref(pe)): - break - finally: - self.kernel32.CloseHandle(snapshot) - - def get_process_list(self): - result = [] - for pid, exe_name in self._enum_processes(): - name, file_path = None, None - if pid > 4: - name, file_path = self.get_exe_info(pid) - - if not name: - try: - name = exe_name - except Exception: - name = "Unknown" - - result.append({"pid": pid, "name": name, "path": file_path or "None"}) - - return result - - def get_process_list_pids(self): - return {pid for pid, _ in self._enum_processes()} - -#################################################################################################### - - def kill_process(self, pid, expected_path=None): - try: - h = self.kernel32.OpenProcess(PYAS_PROCESS_TERMINATE_ACCESS, False, pid) - if h: - try: - file_path = self.norm_path(self.get_process_file(h)) - if self.path_equal(file_path, self.file_pyas): - return False - if expected_path and not self.path_equal(file_path, expected_path): - return False - - return bool(self.kernel32.TerminateProcess(h, 0)) - finally: - self.kernel32.CloseHandle(h) - - except Exception as e: - self.write_log("WARN", "kill_process", pid=pid, detail=str(e), operate=True, success=False) - return False - -#################################################################################################### - - def get_connections_list(self): - connections = set() - try: - size = ctypes.wintypes.DWORD() - if self.iphlpapi.GetExtendedTcpTable(None, ctypes.byref(size), True, 2, 5, 0) != 122: - return connections - - buf = ctypes.create_string_buffer(size.value) - if self.iphlpapi.GetExtendedTcpTable(buf, ctypes.byref(size), True, 2, 5, 0) != 0: - return connections - - num_entries = ctypes.cast(buf, ctypes.POINTER(ctypes.wintypes.DWORD)).contents.value - for i in range(num_entries): - row = MIB_TCPROW_OWNER_PID.from_address(ctypes.addressof(buf) + ctypes.sizeof(ctypes.wintypes.DWORD) + i * ctypes.sizeof(MIB_TCPROW_OWNER_PID)) - connections.add((row.dwOwningPid, row.dwRemoteAddr, row.dwRemotePort)) - - except Exception: - pass - return connections - - def get_traffic_list(self): - conns = self.get_connections_list() - conn_map = {} - for pid, _, _ in conns: - conn_map[pid] = conn_map.get(pid, 0) + 1 - - current_io = {} - current_time = time.time() - - with self.lock_io: - time_diff = current_time - self.last_io_time - if time_diff <= 0: - time_diff = 1 - - self.last_io_time = current_time - old_counters = self.last_io_counters.copy() - - result = [] - exist_process = self.get_process_list_pids() - - for pid in exist_process: - name, file_path = "Unknown", "" - if pid > 4: - name, file_path = self.get_exe_info(pid) - else: - name = "System" - - down_speed, up_speed = 0, 0 - h = self.kernel32.OpenProcess(0x1000, False, pid) - if h: - try: - io = IO_COUNTERS() - if self.kernel32.GetProcessIoCounters(h, ctypes.byref(io)): - current_io[pid] = (io.ReadTransferCount, io.WriteTransferCount) - - if pid in old_counters: - old_read, old_write = old_counters[pid] - down_speed = max(0, (io.ReadTransferCount - old_read) / time_diff) - up_speed = max(0, (io.WriteTransferCount - old_write) / time_diff) - - finally: - self.kernel32.CloseHandle(h) - - count = conn_map.get(pid, 0) - if count > 0 or down_speed > 0 or up_speed > 0: - result.append({"pid": pid, "name": name, "path": file_path or "None", "down": int(down_speed), "up": int(up_speed), "conn": count}) - - with self.lock_io: - self.last_io_counters = current_io - return result - -#################################################################################################### - - def _traverse_delete(self, path): - deleted = 0 - if not path or not os.path.exists(path): - return deleted - - try: - items = os.listdir(path) - except Exception: - return deleted - - for fd in items: - file = os.path.join(path, fd) - try: - if os.path.isdir(file): - if self._is_reparse_point(file): - continue - - deleted += self._traverse_delete(file) - try: - os.rmdir(file) - except Exception: - pass - - else: - if file.lower().endswith(('.sys', '.dll', '.exe', '.ini', '.dat')): - continue - - size = os.path.getsize(file) - os.remove(file) - deleted += size - - except Exception: - continue - - return deleted - - def _get_junk_dirs(self): - return [ - self.path_temp, - self.path_systemp, - os.path.join(self.path_system, "SoftwareDistribution", "Download") - ] - - def _yield_log_files(self): - log_dir = os.path.join(self.path_system, "Logs") - if os.path.exists(log_dir): - for root, _, files in os.walk(log_dir): - for file in files: - if file.lower().endswith(('.log', '.etl', '.evtx')): - yield os.path.join(root, file) - - def _get_recycle_bin_size(self): - try: - info = SHQUERYRBINFO() - info.cbSize = ctypes.sizeof(SHQUERYRBINFO) - - if self.shell32.SHQueryRecycleBinW(None, ctypes.byref(info)) == 0 and info.i64Size > 0: - return info.i64Size - except Exception: - pass - - return 0 - - def scan_system_junk(self): - junk_list = [] - try: - for path in self._get_junk_dirs(): - if os.path.exists(path): - for root, _, files in os.walk(path): - for file in files: - if not file.lower().endswith(('.sys', '.dll', '.exe', '.ini', '.dat')): - try: - fp = os.path.join(root, file) - junk_list.append({"path": fp, "size": os.path.getsize(fp)}) - except Exception: - pass - - for fp in self._yield_log_files(): - try: - junk_list.append({"path": fp, "size": os.path.getsize(fp)}) - except Exception: - pass - - rb_size = self._get_recycle_bin_size() - if rb_size > 0: - junk_list.append({"path": "Recycle Bin", "size": rb_size}) - - return junk_list - - except Exception as e: - self.write_log("WARN", "scan_system_junk", detail=str(e), success=False) - return [] - - def clean_system_junk(self, paths_to_delete=None): - total_deleted = 0 - try: - if paths_to_delete is not None: - for path in paths_to_delete: - if path == "Recycle Bin": - rb_size = self._get_recycle_bin_size() - if rb_size > 0 and self.shell32.SHEmptyRecycleBinW(None, None, 7) == 0: - total_deleted += rb_size - continue - - if path.lower().endswith(('.sys', '.dll', '.exe', '.ini', '.dat')): - continue - - try: - size = os.path.getsize(path) - os.remove(path) - total_deleted += size - except Exception: - pass - else: - for path in self._get_junk_dirs(): - total_deleted += self._traverse_delete(path) - - for fp in self._yield_log_files(): - try: - size = os.path.getsize(fp) - os.remove(fp) - total_deleted += size - except Exception: - pass - - rb_size = self._get_recycle_bin_size() - if rb_size > 0 and self.shell32.SHEmptyRecycleBinW(None, None, 7) == 0: - total_deleted += rb_size - - try: - for log_type in ["Application", "Security", "Setup", "System"]: - log_path = os.path.join(self.path_system, "System32", "winevt", "Logs", f"{log_type}.evtx") - - size = os.path.getsize(log_path) if os.path.exists(log_path) else 0 - if self._clear_event_log(log_type): - total_deleted += size - - except Exception: - pass - - self.write_log("INFO", "Clean Junk", detail=f"Deleted {total_deleted // 1024} KB", operate=True) - return total_deleted - - except Exception as e: - self.write_log("WARN", "clean_system_junk", detail=str(e), operate=True, success=False) - return 0 - -#################################################################################################### - - def optimize_memory(self, icon=None, item=None): - def _optimize(): - try: - hwnd = self.user32.GetForegroundWindow() - fg_pid = ctypes.wintypes.DWORD(0) - if hwnd: - self.user32.GetWindowThreadProcessId(hwnd, ctypes.byref(fg_pid)) - - pids = self.get_process_list_pids() - optimized_count = 0 - - skip_set = { - "dwm.exe", "explorer.exe", "csrss.exe", "smss.exe", "winlogon.exe", - "lsass.exe", "services.exe", "svchost.exe", "wininit.exe", "audiodg.exe", - "spoolsv.exe", "sihost.exe", "fontdrvhost.exe", "taskmgr.exe" - } - - for pid in pids: - if pid <= 4 or pid == self.pid_pyas or pid == fg_pid.value: - continue - - name, _ = self.get_exe_info(pid) - if name and name.lower() in skip_set: - continue - - h = self.kernel32.OpenProcess(0x0100, False, pid) - if h: - try: - if self.kernel32.SetProcessWorkingSetSize(h, ctypes.c_size_t(-1), ctypes.c_size_t(-1)): - optimized_count += 1 - finally: - self.kernel32.CloseHandle(h) - - title = self._loc({ - "traditional_switch": "一鍵加速", "simplified_switch": "一键加速", "english_switch": "Memory Boost", - "japanese_switch": "メモリ最適化", "korean_switch": "메모리 최적화", "french_switch": "Optimiser", - "spanish_switch": "Optimizar", "hindi_switch": "मेमोरी बूस्ट", "arabic_switch": "تسريع", - "russian_switch": "Ускорение", "slovenian_switch": "Optimizacija" - }) - - msg = self._loc({ - "traditional_switch": f"已釋放 {optimized_count} 個背景進程的記憶體", - "simplified_switch": f"已释放 {optimized_count} 个后台进程的内存", - "english_switch": f"Freed memory for {optimized_count} background processes", - "japanese_switch": f"{optimized_count} 個のバックグラウンドプロセスのメモリを解放しました", - "korean_switch": f"{optimized_count}개 백그라운드 프로세스 메모리 확보", - "french_switch": f"Mémoire libérée pour {optimized_count} processus", - "spanish_switch": f"Memoria liberada para {optimized_count} procesos", - "hindi_switch": f"{optimized_count} पृष्ठभूमि प्रक्रियाओं के लिए मेमोरी मुक्त की गई", - "arabic_switch": f"تم تحرير الذاكرة لـ {optimized_count} من العمليات في الخلفية", - "russian_switch": f"Освобождена память {optimized_count} фоновых процессов", - "slovenian_switch": f"Sproščen pomnilnik za {optimized_count} procesov v ozadju" - }) - - self.show_notification(title, msg) - self.write_log("INFO", "Memory Boost", detail=f"Optimized {optimized_count} processes", operate=True) - - except Exception as e: - self.write_log("WARN", "optimize_memory", detail=str(e), success=False) - - threading.Thread(target=_optimize, daemon=True).start() - -#################################################################################################### - - def get_startup_list(self): - items = [] - - def get_reg_startup(): - res = [] - paths = [ - (winreg.HKEY_LOCAL_MACHINE, r"SOFTWARE\Microsoft\Windows\CurrentVersion\Run", "enabled"), - (winreg.HKEY_CURRENT_USER, r"SOFTWARE\Microsoft\Windows\CurrentVersion\Run", "enabled"), - (winreg.HKEY_LOCAL_MACHINE, r"SOFTWARE\Microsoft\Windows\CurrentVersion\Run_Disabled", "disabled"), - (winreg.HKEY_CURRENT_USER, r"SOFTWARE\Microsoft\Windows\CurrentVersion\Run_Disabled", "disabled") - ] - - for root, path, status in paths: - try: - with winreg.OpenKey(root, path, 0, winreg.KEY_READ) as reg: - i = 0 - while True: - try: - name, val, _ = winreg.EnumValue(reg, i) - file_path = self.extract_paths_from_cmdline(val) - fp = file_path[0] if file_path else val - res.append({"id": f"reg|{root}|{path}|{name}", "type": "type_reg", "name": name, "status": status, "path": fp}) - i += 1 - - except OSError: - break - except Exception: - pass - return res - - def get_services(): - res = [] - try: - with winreg.OpenKey(winreg.HKEY_LOCAL_MACHINE, r"SYSTEM\CurrentControlSet\Services", 0, winreg.KEY_READ) as reg: - i = 0 - while True: - try: - name = winreg.EnumKey(reg, i) - i += 1 - try: - with winreg.OpenKey(reg, name, 0, winreg.KEY_READ) as sub: - stype, _ = winreg.QueryValueEx(sub, "Type") - if stype not in (16, 32, 256): - continue - - start, _ = winreg.QueryValueEx(sub, "Start") - if start not in (2, 3, 4): - continue - - path, _ = winreg.QueryValueEx(sub, "ImagePath") - if path and "svchost" not in path.lower() and "windows" not in path.lower(): - fp = self.extract_paths_from_cmdline(path) - fpp = fp[0] if fp else path - status = "enabled" if start == 2 else "disabled" - res.append({"id": f"srv|{name}", "type": "type_srv", "name": name, "status": status, "path": fpp}) - - except Exception: - pass - except OSError: - break - - except Exception: - pass - return res - - def get_tasks(): - res = [] - try: - proc = self._run_windows_tool( - ("schtasks.exe",), - ["/query", "/fo", "csv", "/v"], - capture_output=True, - text=True - ) - if proc is not None and proc.returncode == 0: - reader = csv.reader(io.StringIO(proc.stdout)) - next(reader, None) - - for row in reader: - if len(row) > 8 and row[1].strip() and row[8].strip() and row[8] != "N/A": - name = row[1].split('\\')[-1] - path = row[8] - status = row[3] - - if path and "windows" not in path.lower(): - fp = self.extract_paths_from_cmdline(path) - fpp = fp[0] if fp else path - res.append({"id": f"tsk|{row[1]}", "type": "type_tsk", "name": name, "status": "disabled" if status.lower() == "disabled" else "enabled", "path": fpp}) - return res - - command = ( - "[Console]::OutputEncoding = [Text.UTF8Encoding]::new(); " - "Get-ScheduledTask -ErrorAction SilentlyContinue | ForEach-Object { " - "$Action = @($_.Actions)[0]; " - "[pscustomobject]@{Name=$_.TaskName; FullName=($_.TaskPath + $_.TaskName); " - "State=[string]$_.State; Execute=$Action.Execute; Arguments=$Action.Arguments} " - "} | ConvertTo-Json -Compress" - ) - proc = self._run_powershell(command, capture_output=True, text=True, encoding="utf-8", errors="replace") - if not proc or proc.returncode != 0 or not proc.stdout.strip(): - return res - - task_data = json.loads(proc.stdout) - if isinstance(task_data, dict): - task_data = [task_data] - - for task in task_data: - name = str(task.get("Name") or "").strip() - full_name = str(task.get("FullName") or name).strip() - execute = str(task.get("Execute") or "").strip() - arguments = str(task.get("Arguments") or "").strip() - command_line = " ".join(part for part in (execute, arguments) if part) - if not name or not command_line or "windows" in command_line.lower(): - continue - - paths = self.extract_paths_from_cmdline(command_line) - path = paths[0] if paths else execute or command_line - status = "disabled" if str(task.get("State") or "").lower() == "disabled" else "enabled" - res.append({"id": f"tsk|{full_name}", "type": "type_tsk", "name": name, "status": status, "path": path}) - - except Exception: - pass - return res - - with ThreadPoolExecutor(max_workers=3) as executor: - f1 = executor.submit(get_reg_startup) - f2 = executor.submit(get_services) - f3 = executor.submit(get_tasks) - - items.extend(f1.result()) - items.extend(f2.result()) - items.extend(f3.result()) - - return items - - def manage_startup(self, items, action): - def process_item(item_id): - try: - parts = item_id.split('|', 3) - stype = parts[0] - - if stype == "reg": - root, old_path, name = int(parts[1]), parts[2], parts[3] - if action == "delete": - return self._reg_delete(root, old_path, name) - - else: - new_path = old_path.replace("_Disabled", "") if action == "enable" else (old_path if "_Disabled" in old_path else old_path + "_Disabled") - if old_path != new_path: - val = self._reg_read(root, old_path, name) - if val: - return self._reg_write(root, new_path, name, winreg.REG_SZ, val) and self._reg_delete(root, old_path, name) - return True - - elif stype == "srv": - name = parts[1] - if action == "delete": - result = self._run_windows_tool(("sc.exe",), ["delete", name], capture_output=True) - - else: - mode = "auto" if action == "enable" else "disabled" - result = self._run_windows_tool(("sc.exe",), ["config", name, "start=", mode], capture_output=True) - - if result is not None and result.returncode == 0: - return True - return self._manage_service_native(name, action) - - elif stype == "tsk": - name = item_id.partition('|')[2] - return self._manage_scheduled_task(name, action) - - return False - - except Exception as e: - self.write_log("WARN", "manage_startup", detail=str(e), success=False) - return False - - with ThreadPoolExecutor(max_workers=5) as executor: - results = list(executor.map(process_item, items)) - - success = all(results) if results else True - if not success: - self.write_log("WARN", "manage_startup", detail="One or more startup items could not be updated", success=False) - return success diff --git a/PYAS_Version.py b/PYAS_Version.py index 43e2fe2..7b793f4 100644 --- a/PYAS_Version.py +++ b/PYAS_Version.py @@ -1,12 +1,21 @@ -import pyinstaller_versionfile - -pyinstaller_versionfile.create_versionfile( - output_file="versionfile.txt", - version='3.7.0.0', - company_name="PYAS Security", - file_description="PYAS Security Antivirus", - internal_name="PYAS", - legal_copyright="PYAS Security", - original_filename="PYAS.exe", - product_name="PYAS" -) +VERSION = "3.7.1" +FILE_VERSION = VERSION + ".0" + + +def create_version_file(output_file="versionfile.txt"): + import pyinstaller_versionfile + + pyinstaller_versionfile.create_versionfile( + output_file=output_file, + version=FILE_VERSION, + company_name="PYAS Security", + file_description="PYAS Security Antivirus", + internal_name="PYAS", + legal_copyright="PYAS Security", + original_filename="PYAS.exe", + product_name="PYAS", + ) + + +if __name__ == "__main__": + create_version_file() diff --git a/PYAS_WinAPI.py b/PYAS_WinAPI.py new file mode 100644 index 0000000..4cadd14 --- /dev/null +++ b/PYAS_WinAPI.py @@ -0,0 +1,649 @@ +import struct +from PYAS_Diagnostics import log_exception +import ctypes +import ctypes.wintypes + +PROCESS_TERMINATE = 0x0001 +PROCESS_QUERY_INFORMATION = 0x0400 +PROCESS_QUERY_LIMITED_INFORMATION = 0x1000 +PYAS_PROCESS_TERMINATE_ACCESS = ( + PROCESS_TERMINATE | PROCESS_QUERY_INFORMATION | PROCESS_QUERY_LIMITED_INFORMATION +) + + +class PROCESSENTRY32W(ctypes.Structure): + _fields_ = [ + ("dwSize", ctypes.wintypes.DWORD), + ("cntUsage", ctypes.wintypes.DWORD), + ("th32ProcessID", ctypes.wintypes.DWORD), + ("th32DefaultHeapID", ctypes.wintypes.LPVOID), + ("th32ModuleID", ctypes.wintypes.DWORD), + ("cntThreads", ctypes.wintypes.DWORD), + ("th32ParentProcessID", ctypes.wintypes.DWORD), + ("pcPriClassBase", ctypes.wintypes.LONG), + ("dwFlags", ctypes.wintypes.DWORD), + ("szExeFile", ctypes.wintypes.WCHAR * 260), + ] + + +class MIB_TCPROW_OWNER_PID(ctypes.Structure): + _fields_ = [ + ("dwState", ctypes.wintypes.DWORD), + ("dwLocalAddr", ctypes.wintypes.DWORD), + ("dwLocalPort", ctypes.wintypes.DWORD), + ("dwRemoteAddr", ctypes.wintypes.DWORD), + ("dwRemotePort", ctypes.wintypes.DWORD), + ("dwOwningPid", ctypes.wintypes.DWORD), + ] + + +class FILE_NOTIFY_INFORMATION(ctypes.Structure): + _fields_ = [ + ("NextEntryOffset", ctypes.wintypes.DWORD), + ("Action", ctypes.wintypes.DWORD), + ("FileNameLength", ctypes.wintypes.DWORD), + ("FileName", ctypes.wintypes.WCHAR * 1024), + ] + + +class PROCESS_BASIC_INFORMATION(ctypes.Structure): + _fields_ = [ + ("Reserved1", ctypes.wintypes.LPVOID), + ("PebBaseAddress", ctypes.wintypes.LPVOID), + ("Reserved2", ctypes.wintypes.LPVOID * 2), + ("UniqueProcessId", ctypes.wintypes.LPVOID), + ("Reserved3", ctypes.wintypes.LPVOID), + ] + + +class SHQUERYRBINFO(ctypes.Structure): + _fields_ = [ + ("cbSize", ctypes.wintypes.DWORD), + ("i64Size", ctypes.c_int64), + ("i64NumItems", ctypes.c_int64), + ] + + +class UNICODE_STRING(ctypes.Structure): + _fields_ = [ + ("Length", ctypes.wintypes.USHORT), + ("MaximumLength", ctypes.wintypes.USHORT), + ("Buffer", ctypes.c_void_p), + ] + + +class FILTER_MESSAGE_HEADER(ctypes.Structure): + _fields_ = [("ReplyLength", ctypes.wintypes.ULONG), ("MessageId", ctypes.c_uint64)] + + +class PYAS_MESSAGE(ctypes.Structure): + _fields_ = [ + ("MessageCode", ctypes.wintypes.ULONG), + ("ProcessId", ctypes.wintypes.ULONG), + ("Path", ctypes.wintypes.WCHAR * 1024), + ] + + +class PYAS_FULL_MESSAGE(ctypes.Structure): + _fields_ = [("Header", FILTER_MESSAGE_HEADER), ("Data", PYAS_MESSAGE)] + + +class PYAS_USER_MESSAGE(ctypes.Structure): + _fields_ = [("Command", ctypes.wintypes.ULONG), ("Path", ctypes.wintypes.WCHAR * 1024)] + + +class COPYDATASTRUCT(ctypes.Structure): + _fields_ = [ + ("dwData", ctypes.c_size_t), + ("cbData", ctypes.wintypes.DWORD), + ("lpData", ctypes.c_void_p), + ] + + +class IO_COUNTERS(ctypes.Structure): + _fields_ = [ + ("ReadOperationCount", ctypes.c_ulonglong), + ("WriteOperationCount", ctypes.c_ulonglong), + ("OtherOperationCount", ctypes.c_ulonglong), + ("ReadTransferCount", ctypes.c_ulonglong), + ("WriteTransferCount", ctypes.c_ulonglong), + ("OtherTransferCount", ctypes.c_ulonglong), + ] + + +class MEMORY_BASIC_INFORMATION(ctypes.Structure): + _fields_ = [ + ("BaseAddress", ctypes.c_void_p), + ("AllocationBase", ctypes.c_void_p), + ("AllocationProtect", ctypes.wintypes.DWORD), + ("RegionSize", ctypes.c_size_t), + ("State", ctypes.wintypes.DWORD), + ("Protect", ctypes.wintypes.DWORD), + ("Type", ctypes.wintypes.DWORD), + ] + + +class LUID(ctypes.Structure): + _fields_ = [("LowPart", ctypes.wintypes.DWORD), ("HighPart", ctypes.wintypes.LONG)] + + +class LUID_AND_ATTRIBUTES(ctypes.Structure): + _fields_ = [("Luid", LUID), ("Attributes", ctypes.wintypes.DWORD)] + + +class TOKEN_PRIVILEGES(ctypes.Structure): + _fields_ = [("PrivilegeCount", ctypes.wintypes.DWORD), ("Privileges", LUID_AND_ATTRIBUTES * 1)] + + +class SERVICE_STATUS_PROCESS(ctypes.Structure): + _fields_ = [ + ("dwServiceType", ctypes.wintypes.DWORD), + ("dwCurrentState", ctypes.wintypes.DWORD), + ("dwControlsAccepted", ctypes.wintypes.DWORD), + ("dwWin32ExitCode", ctypes.wintypes.DWORD), + ("dwServiceSpecificExitCode", ctypes.wintypes.DWORD), + ("dwCheckPoint", ctypes.wintypes.DWORD), + ("dwWaitHint", ctypes.wintypes.DWORD), + ("dwProcessId", ctypes.wintypes.DWORD), + ("dwServiceFlags", ctypes.wintypes.DWORD), + ] + + +class POINT(ctypes.Structure): + _fields_ = [("x", ctypes.c_long), ("y", ctypes.c_long)] + + +class RECT(ctypes.Structure): + _fields_ = [ + ("left", ctypes.c_long), + ("top", ctypes.c_long), + ("right", ctypes.c_long), + ("bottom", ctypes.c_long), + ] + + +FLT_PORT_FLAG_SYNC_HANDLE = 0x00000001 +HRESULT_IO_PENDING = 0x800703E5 +ERROR_OPERATION_ABORTED = 995 +ERROR_NOT_FOUND = 1168 +ERROR_SHARING_VIOLATION = 32 +ERROR_LOCK_VIOLATION = 33 +WAIT_OBJECT_0 = 0 +WAIT_TIMEOUT = 258 +INVALID_HANDLE_VALUE = ctypes.c_void_p(-1).value +GENERIC_READ = 0x80000000 +FILE_SHARE_READ = 0x00000001 +FILE_SHARE_WRITE = 0x00000002 +OPEN_EXISTING = 3 +FILE_ATTRIBUTE_NORMAL = 0x00000080 +IOCTL_STORAGE_QUERY_PROPERTY = 0x002D1400 +IOCTL_STORAGE_GET_HOTPLUG_INFO = 0x002D0C14 +STORAGE_DEVICE_PROPERTY = 0 +PROPERTY_STANDARD_QUERY = 0 +INTERNAL_STORAGE_BUS_TYPES = frozenset({1, 3, 8, 10, 11, 17, 18, 19}) +FILE_SCAN_DEBOUNCE_SECONDS = 0.35 +FILE_SCAN_RETRY_SECONDS = 0.25 +FILE_LOCK_RETRY_SECONDS = 0.5 +PYAS_CONNECTION_MAGIC = 0x53415950 +PYAS_CONNECTION_VERSION = 1 +PROCESS_VM_READ = 0x0010 +PROCESS_SUSPEND_RESUME = 0x0800 +PYAS_PROCESS_SCAN_ACCESS = ( + PROCESS_TERMINATE + | PROCESS_VM_READ + | PROCESS_QUERY_INFORMATION + | PROCESS_SUSPEND_RESUME + | PROCESS_QUERY_LIMITED_INFORMATION +) +PYAS_PROCESS_NETWORK_ACCESS = ( + PROCESS_TERMINATE | PROCESS_QUERY_INFORMATION | PROCESS_QUERY_LIMITED_INFORMATION +) + + +class PYAS_CONNECTION_CONTEXT(ctypes.Structure): + _fields_ = [ + ("Size", ctypes.wintypes.ULONG), + ("Version", ctypes.wintypes.ULONG), + ("Magic", ctypes.wintypes.ULONG), + ("ProcessId", ctypes.wintypes.ULONG), + ] + + +class OVERLAPPED(ctypes.Structure): + _fields_ = [ + ("Internal", ctypes.c_size_t), + ("InternalHigh", ctypes.c_size_t), + ("Offset", ctypes.wintypes.DWORD), + ("OffsetHigh", ctypes.wintypes.DWORD), + ("hEvent", ctypes.wintypes.HANDLE), + ] + + +class STORAGE_PROPERTY_QUERY(ctypes.Structure): + _fields_ = [ + ("PropertyId", ctypes.c_int), + ("QueryType", ctypes.c_int), + ("AdditionalParameters", ctypes.c_ubyte * 1), + ] + + +class STORAGE_DESCRIPTOR_HEADER(ctypes.Structure): + _fields_ = [("Version", ctypes.wintypes.DWORD), ("Size", ctypes.wintypes.DWORD)] + + +class STORAGE_DEVICE_DESCRIPTOR(ctypes.Structure): + _fields_ = [ + ("Version", ctypes.wintypes.DWORD), + ("Size", ctypes.wintypes.DWORD), + ("DeviceType", ctypes.c_ubyte), + ("DeviceTypeModifier", ctypes.c_ubyte), + ("RemovableMedia", ctypes.c_ubyte), + ("CommandQueueing", ctypes.c_ubyte), + ("VendorIdOffset", ctypes.wintypes.DWORD), + ("ProductIdOffset", ctypes.wintypes.DWORD), + ("ProductRevisionOffset", ctypes.wintypes.DWORD), + ("SerialNumberOffset", ctypes.wintypes.DWORD), + ("BusType", ctypes.c_int), + ("RawPropertiesLength", ctypes.wintypes.DWORD), + ] + + +class STORAGE_HOTPLUG_INFO(ctypes.Structure): + _fields_ = [ + ("Size", ctypes.wintypes.DWORD), + ("MediaRemovable", ctypes.c_ubyte), + ("MediaHotplug", ctypes.c_ubyte), + ("DeviceHotplug", ctypes.c_ubyte), + ("WriteCacheEnableOverride", ctypes.c_ubyte), + ] + + +class GUID(ctypes.Structure): + _fields_ = [ + ("Data1", ctypes.wintypes.DWORD), + ("Data2", ctypes.wintypes.WORD), + ("Data3", ctypes.wintypes.WORD), + ("Data4", ctypes.c_ubyte * 8), + ] + + +class WINTRUST_FILE_INFO(ctypes.Structure): + _fields_ = [ + ("cbStruct", ctypes.wintypes.DWORD), + ("pcwszFilePath", ctypes.wintypes.LPCWSTR), + ("hFile", ctypes.wintypes.HANDLE), + ("pgKnownSubject", ctypes.wintypes.LPVOID), + ] + + +class WINTRUST_DATA_UNION(ctypes.Union): + _fields_ = [ + ("pFile", ctypes.POINTER(WINTRUST_FILE_INFO)), + ("pCatalog", ctypes.wintypes.LPVOID), + ("pBlob", ctypes.wintypes.LPVOID), + ("pSgnr", ctypes.wintypes.LPVOID), + ("pCert", ctypes.wintypes.LPVOID), + ] + + +class WINTRUST_DATA(ctypes.Structure): + _fields_ = [ + ("cbStruct", ctypes.wintypes.DWORD), + ("pPolicyCallbackData", ctypes.wintypes.LPVOID), + ("pSIPClientData", ctypes.wintypes.LPVOID), + ("dwUIChoice", ctypes.wintypes.DWORD), + ("fdwRevocationChecks", ctypes.wintypes.DWORD), + ("dwUnionChoice", ctypes.wintypes.DWORD), + ("u", WINTRUST_DATA_UNION), + ("dwStateAction", ctypes.wintypes.DWORD), + ("hWVTStateData", ctypes.wintypes.HANDLE), + ("pwszURLReference", ctypes.wintypes.LPCWSTR), + ("dwProvFlags", ctypes.wintypes.DWORD), + ("dwUIContext", ctypes.wintypes.DWORD), + ("pSignatureSettings", ctypes.wintypes.LPVOID), + ] + + +class WindowsMixin: + def init_windll(self): + for name in [ + "ntdll", + "Psapi", + "user32", + "kernel32", + "iphlpapi", + "shell32", + "fltlib", + "advapi32", + ]: + try: + setattr(self, name.lower(), ctypes.WinDLL(name, use_last_error=True)) + except Exception as e: + log_exception("PYAS_WinAPI.WindowsMixin.init_windll:323") + self.write_log("WARN", "init_windll", detail=str(e), success=False) + + self.user32.FindWindowW.argtypes = [ctypes.c_wchar_p, ctypes.c_wchar_p] + self.user32.FindWindowW.restype = ctypes.wintypes.HWND + self.user32.ShowWindow.argtypes = [ctypes.wintypes.HWND, ctypes.c_int] + self.user32.ShowWindow.restype = ctypes.wintypes.BOOL + self.user32.SendMessageTimeoutW.argtypes = [ + ctypes.wintypes.HWND, + ctypes.wintypes.UINT, + ctypes.wintypes.WPARAM, + ctypes.c_void_p, + ctypes.wintypes.UINT, + ctypes.wintypes.UINT, + ctypes.c_void_p, + ] + self.user32.SendMessageTimeoutW.restype = ctypes.wintypes.LPARAM + + self.user32.WindowFromPoint.argtypes = [POINT] + self.user32.WindowFromPoint.restype = ctypes.wintypes.HWND + self.user32.GetAncestor.argtypes = [ctypes.wintypes.HWND, ctypes.c_uint] + self.user32.GetAncestor.restype = ctypes.wintypes.HWND + self.user32.GetCursorPos.argtypes = [ctypes.POINTER(POINT)] + self.user32.GetCursorPos.restype = ctypes.wintypes.BOOL + self.user32.GetWindowRect.argtypes = [ctypes.wintypes.HWND, ctypes.POINTER(RECT)] + self.user32.GetWindowRect.restype = ctypes.wintypes.BOOL + self.user32.SetWindowPos.argtypes = [ + ctypes.wintypes.HWND, + ctypes.wintypes.HWND, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.c_int, + ctypes.wintypes.UINT, + ] + self.user32.SetWindowPos.restype = ctypes.wintypes.BOOL + self.user32.SetLayeredWindowAttributes.argtypes = [ + ctypes.wintypes.HWND, + ctypes.wintypes.DWORD, + ctypes.c_byte, + ctypes.wintypes.DWORD, + ] + self.user32.SetLayeredWindowAttributes.restype = ctypes.wintypes.BOOL + self.user32.CreateWindowExW.restype = ctypes.wintypes.HWND + + self.ntdll.NtQueryInformationProcess.argtypes = [ + ctypes.wintypes.HANDLE, + ctypes.wintypes.ULONG, + ctypes.c_void_p, + ctypes.wintypes.ULONG, + ctypes.POINTER(ctypes.wintypes.ULONG), + ] + self.ntdll.NtQueryInformationProcess.restype = ctypes.wintypes.ULONG + self.ntdll.NtSuspendProcess.argtypes = [ctypes.wintypes.HANDLE] + self.ntdll.NtSuspendProcess.restype = ctypes.c_ulong + self.ntdll.NtResumeProcess.argtypes = [ctypes.wintypes.HANDLE] + self.ntdll.NtResumeProcess.restype = ctypes.c_ulong + + self.shell32.CommandLineToArgvW.argtypes = [ + ctypes.wintypes.LPCWSTR, + ctypes.POINTER(ctypes.c_int), + ] + self.shell32.CommandLineToArgvW.restype = ctypes.POINTER(ctypes.wintypes.LPWSTR) + self.shell32.SHEmptyRecycleBinW.argtypes = [ + ctypes.wintypes.HWND, + ctypes.c_wchar_p, + ctypes.wintypes.DWORD, + ] + self.shell32.SHEmptyRecycleBinW.restype = ctypes.c_long + self.shell32.SHQueryRecycleBinW.argtypes = [ctypes.c_wchar_p, ctypes.POINTER(SHQUERYRBINFO)] + self.shell32.SHQueryRecycleBinW.restype = ctypes.c_long + + self.fltlib.FilterConnectCommunicationPort.argtypes = [ + ctypes.wintypes.LPCWSTR, + ctypes.wintypes.DWORD, + ctypes.c_void_p, + ctypes.wintypes.DWORD, + ctypes.c_void_p, + ctypes.POINTER(ctypes.wintypes.HANDLE), + ] + self.fltlib.FilterConnectCommunicationPort.restype = ctypes.c_long + self.fltlib.FilterGetMessage.argtypes = [ + ctypes.wintypes.HANDLE, + ctypes.c_void_p, + ctypes.wintypes.DWORD, + ctypes.c_void_p, + ] + self.fltlib.FilterGetMessage.restype = ctypes.c_long + self.fltlib.FilterSendMessage.argtypes = [ + ctypes.wintypes.HANDLE, + ctypes.c_void_p, + ctypes.wintypes.DWORD, + ctypes.c_void_p, + ctypes.wintypes.DWORD, + ctypes.POINTER(ctypes.wintypes.DWORD), + ] + self.fltlib.FilterSendMessage.restype = ctypes.c_long + self.fltlib.FilterUnload.argtypes = [ctypes.wintypes.LPCWSTR] + self.fltlib.FilterUnload.restype = ctypes.c_long + + self.advapi32.OpenProcessToken.argtypes = [ + ctypes.wintypes.HANDLE, + ctypes.wintypes.DWORD, + ctypes.POINTER(ctypes.wintypes.HANDLE), + ] + self.advapi32.OpenProcessToken.restype = ctypes.wintypes.BOOL + self.advapi32.LookupPrivilegeValueW.argtypes = [ + ctypes.wintypes.LPCWSTR, + ctypes.wintypes.LPCWSTR, + ctypes.POINTER(LUID), + ] + self.advapi32.LookupPrivilegeValueW.restype = ctypes.wintypes.BOOL + self.advapi32.AdjustTokenPrivileges.argtypes = [ + ctypes.wintypes.HANDLE, + ctypes.wintypes.BOOL, + ctypes.POINTER(TOKEN_PRIVILEGES), + ctypes.wintypes.DWORD, + ctypes.POINTER(TOKEN_PRIVILEGES), + ctypes.POINTER(ctypes.wintypes.DWORD), + ] + self.advapi32.AdjustTokenPrivileges.restype = ctypes.wintypes.BOOL + self.advapi32.OpenSCManagerW.argtypes = [ + ctypes.wintypes.LPCWSTR, + ctypes.wintypes.LPCWSTR, + ctypes.wintypes.DWORD, + ] + self.advapi32.OpenSCManagerW.restype = ctypes.wintypes.HANDLE + self.advapi32.CreateServiceW.argtypes = [ + ctypes.wintypes.HANDLE, + ctypes.wintypes.LPCWSTR, + ctypes.wintypes.LPCWSTR, + ctypes.wintypes.DWORD, + ctypes.wintypes.DWORD, + ctypes.wintypes.DWORD, + ctypes.wintypes.DWORD, + ctypes.wintypes.LPCWSTR, + ctypes.wintypes.LPCWSTR, + ctypes.POINTER(ctypes.wintypes.DWORD), + ctypes.wintypes.LPCWSTR, + ctypes.wintypes.LPCWSTR, + ctypes.wintypes.LPCWSTR, + ] + self.advapi32.CreateServiceW.restype = ctypes.wintypes.HANDLE + self.advapi32.OpenServiceW.argtypes = [ + ctypes.wintypes.HANDLE, + ctypes.wintypes.LPCWSTR, + ctypes.wintypes.DWORD, + ] + self.advapi32.OpenServiceW.restype = ctypes.wintypes.HANDLE + self.advapi32.ChangeServiceConfigW.argtypes = [ + ctypes.wintypes.HANDLE, + ctypes.wintypes.DWORD, + ctypes.wintypes.DWORD, + ctypes.wintypes.DWORD, + ctypes.wintypes.LPCWSTR, + ctypes.wintypes.LPCWSTR, + ctypes.POINTER(ctypes.wintypes.DWORD), + ctypes.wintypes.LPCWSTR, + ctypes.wintypes.LPCWSTR, + ctypes.wintypes.LPCWSTR, + ctypes.wintypes.LPCWSTR, + ] + self.advapi32.ChangeServiceConfigW.restype = ctypes.wintypes.BOOL + self.advapi32.StartServiceW.argtypes = [ + ctypes.wintypes.HANDLE, + ctypes.wintypes.DWORD, + ctypes.POINTER(ctypes.wintypes.LPCWSTR), + ] + self.advapi32.StartServiceW.restype = ctypes.wintypes.BOOL + self.advapi32.QueryServiceStatusEx.argtypes = [ + ctypes.wintypes.HANDLE, + ctypes.c_int, + ctypes.POINTER(ctypes.c_ubyte), + ctypes.wintypes.DWORD, + ctypes.POINTER(ctypes.wintypes.DWORD), + ] + self.advapi32.QueryServiceStatusEx.restype = ctypes.wintypes.BOOL + self.advapi32.DeleteService.argtypes = [ctypes.wintypes.HANDLE] + self.advapi32.DeleteService.restype = ctypes.wintypes.BOOL + self.advapi32.CloseServiceHandle.argtypes = [ctypes.wintypes.HANDLE] + self.advapi32.CloseServiceHandle.restype = ctypes.wintypes.BOOL + + self.kernel32.GetCurrentProcess.argtypes = [] + self.kernel32.GetCurrentProcess.restype = ctypes.wintypes.HANDLE + self.kernel32.CreateToolhelp32Snapshot.restype = ctypes.wintypes.HANDLE + self.kernel32.CreateToolhelp32Snapshot.argtypes = [ + ctypes.wintypes.DWORD, + ctypes.wintypes.DWORD, + ] + self.kernel32.Process32FirstW.restype = ctypes.wintypes.BOOL + self.kernel32.Process32FirstW.argtypes = [ctypes.wintypes.HANDLE, ctypes.c_void_p] + self.kernel32.Process32NextW.restype = ctypes.wintypes.BOOL + self.kernel32.Process32NextW.argtypes = [ctypes.wintypes.HANDLE, ctypes.c_void_p] + self.kernel32.OpenProcess.restype = ctypes.wintypes.HANDLE + self.kernel32.OpenProcess.argtypes = [ + ctypes.wintypes.DWORD, + ctypes.wintypes.BOOL, + ctypes.wintypes.DWORD, + ] + self.kernel32.CloseHandle.restype = ctypes.wintypes.BOOL + self.kernel32.CloseHandle.argtypes = [ctypes.wintypes.HANDLE] + self.kernel32.CreateEventW.argtypes = [ + ctypes.c_void_p, + ctypes.wintypes.BOOL, + ctypes.wintypes.BOOL, + ctypes.wintypes.LPCWSTR, + ] + self.kernel32.CreateEventW.restype = ctypes.wintypes.HANDLE + self.kernel32.WaitForSingleObject.argtypes = [ctypes.wintypes.HANDLE, ctypes.wintypes.DWORD] + self.kernel32.WaitForSingleObject.restype = ctypes.wintypes.DWORD + self.kernel32.CancelIoEx.argtypes = [ctypes.wintypes.HANDLE, ctypes.c_void_p] + self.kernel32.CancelIoEx.restype = ctypes.wintypes.BOOL + self.kernel32.GetOverlappedResult.argtypes = [ + ctypes.wintypes.HANDLE, + ctypes.c_void_p, + ctypes.POINTER(ctypes.wintypes.DWORD), + ctypes.wintypes.BOOL, + ] + self.kernel32.GetOverlappedResult.restype = ctypes.wintypes.BOOL + self.kernel32.TerminateProcess.restype = ctypes.wintypes.BOOL + self.kernel32.TerminateProcess.argtypes = [ctypes.wintypes.HANDLE, ctypes.c_uint] + self.kernel32.CreateMutexW.restype = ctypes.wintypes.HANDLE + self.kernel32.CreateMutexW.argtypes = [ + ctypes.c_void_p, + ctypes.wintypes.BOOL, + ctypes.c_wchar_p, + ] + self.kernel32.CreateFileW.restype = ctypes.wintypes.HANDLE + self.kernel32.CreateFileW.argtypes = [ + ctypes.c_wchar_p, + ctypes.wintypes.DWORD, + ctypes.wintypes.DWORD, + ctypes.c_void_p, + ctypes.wintypes.DWORD, + ctypes.wintypes.DWORD, + ctypes.wintypes.HANDLE, + ] + self.kernel32.DeviceIoControl.restype = ctypes.wintypes.BOOL + self.kernel32.DeviceIoControl.argtypes = [ + ctypes.wintypes.HANDLE, + ctypes.wintypes.DWORD, + ctypes.c_void_p, + ctypes.wintypes.DWORD, + ctypes.c_void_p, + ctypes.wintypes.DWORD, + ctypes.POINTER(ctypes.wintypes.DWORD), + ctypes.c_void_p, + ] + self.kernel32.ReadProcessMemory.restype = ctypes.wintypes.BOOL + self.kernel32.ReadProcessMemory.argtypes = [ + ctypes.wintypes.HANDLE, + ctypes.c_void_p, + ctypes.c_void_p, + ctypes.c_size_t, + ctypes.POINTER(ctypes.c_size_t), + ] + self.kernel32.QueryFullProcessImageNameW.restype = ctypes.wintypes.BOOL + self.kernel32.QueryFullProcessImageNameW.argtypes = [ + ctypes.wintypes.HANDLE, + ctypes.wintypes.DWORD, + ctypes.wintypes.LPWSTR, + ctypes.POINTER(ctypes.wintypes.DWORD), + ] + self.kernel32.QueryDosDeviceW.restype = ctypes.wintypes.DWORD + self.kernel32.QueryDosDeviceW.argtypes = [ + ctypes.wintypes.LPCWSTR, + ctypes.wintypes.LPWSTR, + ctypes.wintypes.DWORD, + ] + self.kernel32.GetProcessIoCounters.argtypes = [ + ctypes.wintypes.HANDLE, + ctypes.POINTER(IO_COUNTERS), + ] + self.kernel32.GetProcessIoCounters.restype = ctypes.wintypes.BOOL + self.kernel32.VirtualQueryEx.argtypes = [ + ctypes.wintypes.HANDLE, + ctypes.c_void_p, + ctypes.POINTER(MEMORY_BASIC_INFORMATION), + ctypes.c_size_t, + ] + self.kernel32.VirtualQueryEx.restype = ctypes.c_size_t + self.kernel32.SetProcessWorkingSetSize.argtypes = [ + ctypes.wintypes.HANDLE, + ctypes.c_size_t, + ctypes.c_size_t, + ] + self.kernel32.SetProcessWorkingSetSize.restype = ctypes.wintypes.BOOL + + self.psapi.GetMappedFileNameW.argtypes = [ + ctypes.wintypes.HANDLE, + ctypes.c_void_p, + ctypes.wintypes.LPWSTR, + ctypes.wintypes.DWORD, + ] + self.psapi.GetMappedFileNameW.restype = ctypes.wintypes.DWORD + + +def iter_file_notifications(buffer, size): + view = memoryview(buffer).cast("B") + + if size < 0 or size > len(view): + raise ValueError("Invalid file notification buffer size") + + offset = 0 + + while offset < size: + if offset + 12 > size: + raise ValueError("Truncated file notification header") + + next_offset, action, length = struct.unpack_from(" size: + raise ValueError("Invalid file notification length") + + if next_offset and ( + next_offset < 12 + length or next_offset % 4 or offset + next_offset >= size + ): + raise ValueError("Invalid file notification offset") + + filename = bytes(view[offset + 12 : offset + 12 + length]).decode( + "utf-16-le", errors="surrogatepass" + ) + yield action, filename + + if not next_offset: + break + + offset += next_offset diff --git a/README.md b/README.md index fa5018e..43fa8e0 100644 --- a/README.md +++ b/README.md @@ -19,6 +19,8 @@ ## Overview +Current desktop source version: **3.7.1** (`PYAS_Version.py`; Windows file version `3.7.1.0`). The source version does not imply that a matching installer has already been published. + PYAS Security is a source-available Windows endpoint security project that combines multiple detection and protection layers in one desktop application. It brings together local PE-file machine-learning inference, YARA-based matching, digital-signature inspection, optional cloud analysis, real-time user-mode monitoring, and a native Windows minifilter driver. The project is designed both as a usable security application and as an engineering platform for studying malware detection, Windows internals, rule-driven prevention, model training, and online file analysis. @@ -42,7 +44,7 @@ The project is designed both as a usable security application and as an engineer | Static rules | File and process-memory matching | YARA rules compiled and loaded by `rule_scanner` | | PE machine learning | Feature extraction and local classification | `pefile`, NumPy, ONNX Runtime | | Signature trust | Authenticode verification | Windows WinVerifyTrust APIs | -| Cloud analysis | Upload, polling, rescan, and result retrieval | HTTPS API with chunked upload support | +| Cloud analysis | Optional queued submission; client methods for polling and rescanning | `PYAS_Cloud.py`, HTTPS and chunked upload | | Process protection | New-process inspection, optional suspension, termination controls | Win32/NT APIs and worker pools | | File protection | Directory change monitoring, debounced real-time scans, file locking | Windows file notification APIs | | Memory protection | Process-memory YARA scanning and kernel callbacks | User-mode scanner plus driver rules | @@ -50,119 +52,124 @@ The project is designed both as a usable security application and as an engineer | Kernel enforcement | File, process/thread, registry, memory, image-load, and boot/disk controls | Native C++ minifilter driver | | Recovery and maintenance | Quarantine, system repair, startup management, cleanup, MBR backup/check | Python application services | -Protection modules are individually configurable. Several active-protection switches are disabled by default so users can enable only the controls appropriate for their environment. +Protection modules are individually configurable. Cloud submission is disabled by default. Several active-protection switches are also disabled by default so users can enable only the controls appropriate for their environment. ## Detection pipeline ```mermaid flowchart LR A[File or process event] --> B{Scope and policy checks} - B -->|Allowed| C[Hash and metadata] B -->|Excluded| Z[Skip] - C --> D[Digital signature verification] - C --> E[YARA rule scan] - C --> F[PE feature extraction] - F --> G[ONNX model inference] - D --> H[Local verdict aggregation] - E --> H - G --> H - H --> I{Cloud analysis enabled?} - I -->|Yes| J[Hash lookup / chunked upload] - J --> K[Cloud report] - I -->|No| L[Local result] - K --> M[Alert and response] - L --> M + B -->|Allowed| C[Shared scan entry: stable file gate and hash] + C --> D{Content and policy cache hit?} + D -->|Yes| H[Cached local verdict] + D -->|No| E[Local engines: signature / YARA / PE and ONNX] + E --> F[Local verdict and eligible cache publication] + F --> H + H --> M[Alert and configured response] M --> N[Quarantine / delete / allow] + H -. optional submission .-> I{Cloud switch enabled?} + I -->|Yes| J[Cloud queue: hash lookup / chunked upload] + I -->|No| K[No upload] ``` +`PYAS_Scanner.py` coordinates local scanning and remediation. The common scan entry holds a read handle that denies concurrent writes/deletes while hashing and classifying. If the handle cannot be acquired, classification still runs and a retry is scheduled, but the result is not published to the shared cache. + +`PYAS_Cloud.py` contains both `CloudScanner` and `CloudQueueMixin`. The desktop cloud worker currently submits files; it does not automatically call `get_result()` or merge a remote verdict into local remediation. Polling and rescan methods remain available to explicit client callers. + ## System architecture +The desktop application currently consists of **27 `PYAS*.py` modules**. `WindowAPI` in `PYAS.py` combines `_MainMixin`, `ScannerMixin`, `ToolsMixin`, and `ProtectMixin` on one application instance. The mixins share the existing configuration, state, queues, and locks rather than creating separate service instances. + ```mermaid flowchart TB - subgraph UI[Presentation layer] - WEB[HTML / CSS / JavaScript] - WV[WebView2 desktop shell] - TRAY[System tray and Windows shell integration] - end + UI[Interface: HTML / CSS / JavaScript] <--> WV[WebView2] + WV <--> CORE[PYAS.py: WindowAPI / tray / window messages] + START[PYAS_Startup.py: diagnostics and profile recovery] --> WV + CORE --> MAIN[Main mixins: Runtime / Config / Logs / WinAPI] + CORE --> SCAN[PYAS_Scanner.py: ScannerMixin] + CORE --> TOOLS[PYAS_Tools.py: ToolsMixin] + CORE --> PROTECT[PYAS_Protect.py: ProtectMixin] - subgraph APP[Python application layer] - CORE[PYAS.py orchestration] - SCAN[ScannerMixin] - PROTECT[ProtectMixin] - TOOLS[ToolsMixin] - ENGINE[PYAS_Engine.py] - end + MAIN --> ENGINE[PYAS_Engine.py: compatible engine exports] + SCAN --> LOCAL[Rules / Signature / PE] + ENGINE --> LOCAL + LOCAL --> FEATURES[PYAS_Features.py] + LOCAL --> ASSETS[Engine: YARA assets and ONNX models] + SCAN --> CLOUD[PYAS_Cloud.py: client and cloud queue] + CLOUD -. optional HTTPS submission .-> ONLINE[Analyze: separate online analysis service] - subgraph DETECT[Detection engines] - YARA[YARA rules] - PE[PE feature extractor] - ONNX[ONNX Runtime models] - SIGN[Authenticode verification] - CLOUD[PYAS Cloud API] - end + TOOLS --> UTIL[Autostart / Process / Maintenance / Threats] + PROTECT --> RT[PYAS_Realtime.py] + PROTECT --> SYS[PYAS_System.py] + PROTECT --> POP[PYAS_Popup.py: matching and window operations] + PROTECT --> DRIVER[PYAS_Driver.py] + RT --> SCHED[PYAS_Scheduler.py: bounded delayed work] + DRIVER <--> PORT[Filter Manager communication port] + PORT <--> KERNEL[Plugins/Filter: native minifilter] + KERNEL --> POLICY[Plugins/Rules: protection policies] - subgraph KERNEL[Windows kernel layer] - PORT[Filter Manager communication port] - RULES[Dynamic rule engine and trust cache] - FILE[File-system minifilter] - PROC[Process and thread callbacks] - REG[Registry callbacks] - MEM[Memory and image-load controls] - BOOT[Boot and disk I/O controls] - end - - subgraph ONLINE[Online analysis platform] - API[Flask application and REST API] - QUEUE[Priority analysis workers] - DB[(PostgreSQL)] - STORAGE[(Sample and report storage)] - end - - WEB <--> WV - WV <--> CORE - TRAY <--> CORE - CORE --> SCAN - CORE --> PROTECT - CORE --> TOOLS - SCAN --> ENGINE - PROTECT --> ENGINE - ENGINE --> YARA - ENGINE --> PE - PE --> ONNX - ENGINE --> SIGN - ENGINE -. optional HTTPS .-> CLOUD - PROTECT <--> PORT - PORT <--> RULES - RULES --> FILE - RULES --> PROC - RULES --> REG - RULES --> MEM - RULES --> BOOT - CLOUD --> API - API --> QUEUE - QUEUE --> DB - QUEUE --> STORAGE + MAIN --> STORE[PYAS_Storage.py: atomic JSON writes] + CORE -. exception reporting .-> DIAG[PYAS_Diagnostics.py] ``` +- **Application state:** `PYAS_Runtime.py` owns initialization, UI dispatch, and feature-worker lifecycle. `PYAS_Config.py` applies switches and persists settings after successful operations; failures restore the prior state and attempt the corresponding rollback. Driver disable state is committed only after unloading succeeds. +- **Protection:** `PYAS_Realtime.py` handles process, file, and network monitoring; `PYAS_System.py` handles system protection and repair; `PYAS_Driver.py` owns driver lifecycle and communication; `PYAS_Popup.py` combines popup matching with native window operations. +- **Shared infrastructure:** `PYAS_Process.py` centralizes process enumeration and operations; `PYAS_WinAPI.py` centralizes Win32/NT declarations, native structures, and file-notification decoding. `PYAS_Scheduler.py` uses one scheduler and at most four workers, with a default 2048 pending-work limit, bounded critical retry capacity, and recovery scans on overflow. +- **Storage and diagnostics:** configuration and reports use atomic JSON replacement via `PYAS_Storage.py`. `PYAS_Diagnostics.py` supplies consistent exception context and tracebacks, including background futures and workers, with bounded repeat suppression. The UI text log trims from 100,000 UTF-16 code units to approximately 90,000; tables render at most 1,000 rows, while log retention/export allows 10,000 records. The scan cache allows 100,000 entries and evicts 10,000 at capacity. +- **Compatibility:** `PYAS_Engine.py` continues to export `sign_scanner`, `rule_scanner`, `pe_scanner`, and `cloud_scanner` as aliases for the implementation classes. Cloud queue logic resides in `PYAS_Cloud.py`; native popup operations reside in `PYAS_Popup.py`; Windows API declarations reside in `PYAS_WinAPI.py`. + +See [ARCHITECTURE.md](ARCHITECTURE.md) for state ownership, rollback behavior, performance limits, diagnostics, and validation details. + ## Repository layout ```text PYAS/ -├── PYAS.py # Desktop entry point, UI bridge, configuration and lifecycle -├── PYAS_Engine.py # YARA, PE/ML, signature and cloud scanning engines -├── PYAS_Scanner.py # Scan scheduling, workers, results and remediation -├── PYAS_Protect.py # Real-time protection and driver lifecycle/communication -├── PYAS_Tools.py # Windows utilities, startup, process, network and cleanup tools -├── PYAS_Version.py # Windows executable version metadata +├── PYAS.py # Desktop entry point, WindowAPI, tray, messages and lifecycle +├── PYAS_Startup.py # WebView2 startup checks, diagnostics and profile recovery +├── PYAS_Runtime.py # Environment/state initialization, UI queue and feature workers +├── PYAS_Config.py # Settings, switch application, persistence and rollback +├── PYAS_Logs.py # Activity reports, bounded history, export and flushing +├── PYAS_WinAPI.py # Shared Win32/NT constants, structures and API declarations +├── PYAS_Storage.py # Atomic JSON persistence +├── PYAS_Diagnostics.py # Exception logging, repeat suppression and future observation +├── PYAS_Scheduler.py # Bounded delayed work, debounce and overflow recovery +├── PYAS_Tools.py # Windows utility facade and general system operations +├── PYAS_Autostart.py # SID-bound elevated startup tasks and startup management +├── PYAS_Process.py # Shared process enumeration, command lines and termination +├── PYAS_Maintenance.py # Cleanup and memory-maintenance utilities +├── PYAS_Threats.py # Named threat lists and threat extraction/removal +├── PYAS_Protect.py # Protection facade, allowlists and file locking +├── PYAS_Driver.py # Driver install/unload, listener and communication protocol +├── PYAS_System.py # System monitoring/repair and MBR protection +├── PYAS_Realtime.py # Process, file and network protection workers +├── PYAS_Popup.py # Popup fingerprints, matching and native window operations +├── PYAS_Scanner.py # Scan coordination, shared verdict cache and remediation +├── PYAS_Engine.py # Compatible exports for the local/cloud engine classes +├── PYAS_Signature.py # Authenticode signature verification +├── PYAS_Rules.py # YARA and heuristic rule scanning +├── PYAS_PE.py # PE model loading, feature ordering and ONNX classification +├── PYAS_Features.py # PE structural/statistical feature extraction +├── PYAS_Cloud.py # Cloud HTTP client, session reuse, queue and cancellation +├── PYAS_Version.py # Source version and Windows executable metadata generation +├── ARCHITECTURE.md # Detailed architecture and validation notes +├── README.md # Project overview and source layout ├── Engine/ -│ ├── Heuristic/ # YARA signatures and rule assets -│ └── Properties/ # PE feature models and training/inference utilities -├── Interface/ # WebView2 HTML, CSS, JavaScript and icons -└── Plugins/ - ├── Filter/ # Native Windows minifilter driver source - └── Rules/ # Driver protection policies +│ ├── Heuristic/ # Desktop YARA signatures and rule assets +│ └── Properties/ # PE feature models and training/inference utilities +├── Interface/ # WebView2 HTML, CSS, JavaScript and icons +├── Plugins/ +│ ├── Filter/ # Native Windows minifilter components +│ └── Rules/ # Driver protection policies +├── Install/ +│ └── Inno.Installer.Source/ # Inno Setup installer source +├── Analyze/ # Separate online analysis platform +├── Experimental/ # Experimental/research components +└── tests/ # Python regression and JavaScript UI tests ``` +This lists all desktop Python modules and the main project directories. Additional engine assets, generated files, caches, local backups, and duplicate development directories are not part of this architecture inventory. Packaging must include all 27 desktop modules and the required runtime assets. + ## Getting started ### Option 1 — Install a packaged release @@ -238,7 +245,7 @@ PYAS stores machine-wide configuration and reports under: %ProgramData%\PYAS\Report.json ``` -WebView2 user data and its diagnostic log are stored under the current user's local application-data directory. Configuration includes scan limits, enabled protection layers, language and theme preferences, custom rule references, allow/block lists, and quarantine metadata. +WebView2 user data is stored under the current user's local application-data directory. Startup, WebView, background exceptions, and activity records share `Report.json` (up to 10,000 records). New records include the desktop source `version`; exception tracebacks are stored in `detail`. If the machine-wide log directory is not writable, logging uses `PYAS/Report.json` in the temporary directory for that session. Configuration includes scan limits, enabled protection layers, language and theme preferences, custom rule references, allow/block lists, and quarantine metadata. Before reporting a configuration issue, reproduce it with default settings when safe to do so and remove sensitive paths or file information from logs. @@ -258,23 +265,23 @@ The service can be containerized with Docker Compose, but the checked-in deploym ## Python API example -`PYAS_Cloud.py` provides a client for the online analysis API: +`PYAS_Cloud.py` provides `CloudScanner`, the desktop application's cloud client. The example below explicitly requests the existing client's verdict for a previously submitted SHA-256; it does not upload a new file: ```python -from PYAS_Cloud import PYAS_Client +from PYAS_Cloud import CloudScanner -client = PYAS_Client( +client = CloudScanner() +verdict = client.get_result( + sha256="SHA256_OF_A_PREVIOUSLY_SUBMITTED_FILE", + api_host="https://pyas-security.com", api_key="YOUR_API_KEY", - hosts="https://pyas-security.com", + max_retries=6, + interval=10, ) - -sha256 = client.upload_file(r"C:\samples\candidate.exe") -if sha256 and client.wait_for_analysis(sha256): - report = client.get_report(sha256) - print(report) +print(verdict) ``` -Only submit files that you are authorized to upload. Samples may contain confidential or personal information. +`get_result()` returns a supported malicious classification string or `False`; `False` is not a complete remote report or a guarantee that a sample is safe. The desktop's queued submission path does not call this method automatically. Only submit files that you are authorized to upload; samples may contain confidential or personal information. ## Building the driver @@ -346,4 +353,4 @@ The repository contains component-specific license files and bundled third-party Built for Windows security research, layered endpoint defense, and transparent experimentation. - + \ No newline at end of file