PYAS_Scanner.py

import os, gc, json, time, stat, threading
import ctypes, ctypes.wintypes

from PYAS_Tools import MEMORY_BASIC_INFORMATION

####################################################################################################

class ScannerMixin:
    def init_engine_thread(self):
        try:
            self.start_daemon_thread(self.backup_mbr)
            self.start_daemon_thread(self.relock_file)
            self.start_daemon_thread(self.init_whitelist)
            self.start_daemon_thread(self.popup_intercept_thread)

            def log_callback(x):
                self.write_log("INFO", "Load Engine", source=os.path.basename(x))

            self.heuristic.load_path(self.path_heuristic, callback=log_callback)
            self.properties.load_path(self.path_properties, callback=log_callback)
            self.write_log("INFO", "System", detail="Engine Initialization Complete")

            if "-scan" in self.args_pyas:
                try:
                    idx = self.args_pyas.index("-scan")
                    self.trigger_context_scan(self.args_pyas[idx+1])
                except Exception:
                    pass

            with self.lock_config:
                first_launch = self.pyas_config.get("first_launch", True)
                driver_enabled = self.pyas_config.get("driver_switch", False)
                context_enabled = self.pyas_config.get("context_switch", True)
                autostart_enabled = self.pyas_config.get("autostart_switch", True)

            self.register_context_menu(context_enabled)
            self.manage_autostart(autostart_enabled)

            if not first_launch:
                with self.lock_config:
                    if self.pyas_config.get("process_switch"):
                        self.start_daemon_thread(self.protect_proc_thread)
                    if self.pyas_config.get("document_switch"):
                        self.start_daemon_thread(self.protect_file_thread)
                    if self.pyas_config.get("system_switch"):
                        self.start_daemon_thread(self.protect_system_thread)
                    if self.pyas_config.get("network_switch"):
                        self.start_daemon_thread(self.protect_net_thread)

                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)

        except Exception as e:
            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

            elif os.path.isfile(targets):
                yield self.norm_path(targets), True

        elif isinstance(targets, (list, tuple, set)):
            for t in targets:
                yield from self.yield_files(t)

    def _is_reparse_point(self, path):
        try:
            if os.path.islink(path):
                return True
            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)
        except OSError:
            return False

####################################################################################################

    def scan_engine(self, file_path):
        with self.lock_config:
            ext_switch = self.pyas_config.get("extension_switch", False)
            sen_switch = self.pyas_config.get("sensitive_switch", False)

        try:
            pe_label, _ = self.properties.pe_scan(file_path, enhanced_mode=sen_switch)
            if pe_label:
                return pe_label
        except Exception:
            pass

        try:
            if ext_switch:
                yara_label, _ = self.heuristic.yara_scan(file_path)
                if yara_label:
                    return yara_label
        except Exception:
            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)

        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)

        try:
            result = self.scan_engine(norm_path)
            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]

                self.hash_cache[cache_key] = result
            return result

        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]

            gc.collect()

####################################################################################################

    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())
                elapsed = int(time.time() - scan_start) if elapsed is None else elapsed

        return {
            "traditional_switch": f"發現 {count} 個病毒，掃描 {scanned} 個檔案，耗時 {elapsed} 秒",
            "simplified_switch": f"发现 {count} 个病毒，扫描 {scanned} 个文件，耗时 {elapsed} 秒",
            "english_switch": f"Found {count} viruses, scanned {scanned} files, time {elapsed}s",
            "japanese_switch": f"{count} 個のウイルスを発見し、{scanned} 個のファイルをスキャンしました（所要時間 {elapsed} 秒）",
            "korean_switch": f"바이러스 {count}개 발견, 파일 {scanned}개 검사, 소요 시간 {elapsed}초",
            "french_switch": f"{count} virus trouvés, {scanned} fichiers analysés, temps {elapsed}s",
            "spanish_switch": f"Se encontraron {count} virus, {scanned} archivos escaneados, tiempo {elapsed}s",
            "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"
        }

    def _is_scan_cancel_requested(self):
        with self.lock_virus:
            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())
            elapsed = int(time.time() - scan_start)
            self.scan_running = False
            self.scan_preparing = False
            self.scan_finished = True
            self.scan_stop_requested = False

        messages = messages or self._scan_result_messages(count, scanned, elapsed)

        if self._window:
            js_cmd = f"if(window.finishScan) window.finishScan({json.dumps(messages)}, {count});"
            self.ui_queue.put(js_cmd)

        return False

    def _begin_scan_preparation(self):
        with self.lock_virus:
            if getattr(self, 'scan_running', False) or getattr(self, 'scan_preparing', False):
                return False

            self.scan_preparing = True
            self.scan_stop_requested = False
            self.scan_finished = False
            self.virus_results = []
            self.scan_count = 0
            self.scan_start = time.time()
            return True

    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):
                return False

            if getattr(self, 'scan_stop_requested', False):
                self.scan_running = False
                self.scan_preparing = False
                self.scan_finished = True
                return False

            self.scan_preparing = False
            self.scan_running = True
            self.scan_finished = False
            self.scan_stop_requested = False
            self.virus_results = []
            self.scan_count = 0
            self.scan_start = time.time()

        self.scan_pool.submit(self.scan_worker, targets)
        return True

    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):
                        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, '/'))});"
                        self.ui_queue.put(js_cmd)

                    last_update = current_time

                was_locked = False
                try:
                    with self.lock_file_ops:
                        if norm_path in self.virus_lock:
                            self.lock_file(norm_path, False)
                            was_locked = True

                    if self.is_in_whitelist(norm_path):
                        continue

                    with self.lock_virus:
                        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:
                        continue

                    result = self.safe_scan_engine(norm_path)
                    with self.lock_virus:
                        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: 
                            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.cloud_check(norm_path)
                except Exception as e:
                    self.write_log("WARN", "Scan Engine", source=norm_path, detail=str(e), success=False)

                finally:
                    if was_locked:
                        self.lock_file(norm_path, True)

        finally:
            with self.lock_virus:
                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)

            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});"
                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)
            if active:
                self.scan_stop_requested = True
                self.scan_running = False
            return active

    def trigger_scan(self, method):
        targets = []

        if not self._begin_scan_preparation():
            return False

        try:
            if method == "smart":
                for folder in ["Desktop", "Downloads", "AppData"]:
                    if self._is_scan_cancel_requested():
                        return self._finish_scan_cancelled()

                    fp = os.path.join(self.path_user, folder)
                    if os.path.exists(fp):
                        targets.append(fp)

                if self._is_scan_cancel_requested():
                    return self._finish_scan_cancelled()

                if os.path.exists(self.path_config):
                    targets.append(self.path_config)

                buf = ctypes.create_unicode_buffer(1024)
                max_address = 0x7FFFFFFFFFFF if ctypes.sizeof(ctypes.c_void_p) == 8 else 0x7FFFFFFF
                system_dir = self.path_system.lower()

                for proc in self.get_process_list():
                    if self._is_scan_cancel_requested():
                        return self._finish_scan_cancelled()

                    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)):
                                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 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)

                    if self._is_scan_cancel_requested():
                        return self._finish_scan_cancelled()

                if not self.start_scan(list(set(targets)), from_preparing=True):
                    return self._finish_scan_cancelled()

                return True

            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):
                    return self._finish_scan_cancelled()
                return True

            return self._finish_scan_cancelled()

        except Exception as e:
            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)
                if proc_key == target_key:
                    self.kill_process(proc["pid"])

            try:
                os.chmod(file_path, stat.S_IWRITE)
            except Exception:
                pass
            os.remove(file_path)
        except FileNotFoundError:
            pass
        except Exception as e:
            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
            ]
        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, '/'))});")

                    last_update = current_time

                try:
                    if path in self.virus_lock:
                        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)

                    try:
                        os.chmod(path, stat.S_IWRITE)
                    except Exception:
                        pass

                    file_hash = self.calc_file_hash(path)
                    try:
                        os.remove(path)
                    except Exception as e:
                        self.lock_file(path, True, quiet=True)
                        self._queue_virus_delete(path)
                        self.write_log("INFO", "Virus Delete Deferred", source=path, detail=str(e))
                        continue

                    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)

                except Exception as e:
                    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)

            if deleted_set:
                self.virus_results = [p for p in self.virus_results if p not in deleted_set]

            remaining = len(self.virus_results)

        messages = {
            "traditional_switch": f"剩餘 {remaining} 個病毒，已刪除 {len(deleted_paths)} 個檔案。",
            "simplified_switch": f"剩余 {remaining} 个病毒，已删除 {len(deleted_paths)} 个文件。",
            "english_switch": f"Remaining {remaining} viruses, deleted {len(deleted_paths)} files.",
            "japanese_switch": f"残りのウイルス {remaining} 個、削除されたファイル {len(deleted_paths)} 個。",
            "korean_switch": f"남은 바이러스 {remaining}개, 삭제된 파일 {len(deleted_paths)}개.",
            "french_switch": f"Virus restants : {remaining}, fichiers supprimés : {len(deleted_paths)}.",
            "spanish_switch": f"Virus restantes: {remaining}, archivos eliminados: {len(deleted_paths)}.",
            "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)}."
        }
        if self._window:
            self._window.evaluate_js(f"if(window.finishScan) window.finishScan({json.dumps(messages)}, {remaining});")

        return deleted_paths

    def remove_virus_result(self, paths):
        with self.lock_virus:
            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