PYAS_Scanner.py

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(CloudQueueMixin):
    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:
                    log_exception("PYAS_Scanner.ScannerMixin.init_engine_thread:31")
                    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)

            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_feature_thread(self.protect_proc_thread, "process_switch")

                    if self.pyas_config.get("document_switch"):
                        self.start_feature_thread(self.protect_file_thread, "document_switch")

                    if self.pyas_config.get("system_switch"):
                        self.start_feature_thread(self.protect_system_thread, "system_switch")

                    if self.pyas_config.get("network_switch"):
                        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,
                        )

        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

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

        gate, error = self._acquire_file_scan_gate(norm_path)

        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:
            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 cache_key in self.hash_cache:
                    return self.hash_cache[cache_key]

                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:
            self.kernel32.CloseHandle(gate)

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

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

                if proc_key == target_key:
                    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
            ]

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

            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)