diff --git a/PYAS.py b/PYAS.py index c277b03..6c699e4 100644 --- a/PYAS.py +++ b/PYAS.py @@ -1,6 +1,6 @@ import os, gc, sys, time, json, psutil -import ctypes, msvcrt, pyperclip, win32file -import win32gui, win32api, win32con +import ctypes, struct, msvcrt, pyperclip +import win32gui, win32api, win32con, win32file from PYAS_Engine import YRScan, DLScan from PYAS_Suffixes import slist, alist from PYAS_Language import translate_dict @@ -19,7 +19,7 @@ class MainWindow_Controller(QMainWindow): self.pyas = sys.argv[0].replace("\\", "/") self.dir = os.path.dirname(self.pyas) self.pyae_version = "AI Engine" - self.pyas_version = "3.1.7" + self.pyas_version = "3.1.8" self.first_startup = True self.init_data_base() self.init_tray_icon() @@ -1299,7 +1299,7 @@ class MainWindow_Controller(QMainWindow): def protect_proc_thread(self): while self.proc_protect: - time.sleep(0.01) + time.sleep(0.1) new_process = set() for p in psutil.process_iter(): new_process.add(p.pid) @@ -1330,6 +1330,8 @@ class MainWindow_Controller(QMainWindow): self.kill_process("病毒攔截", p, file, False) elif self.load_scan(p): self.kill_process("加載攔截", p, file, False) + elif self.memory_scan(p): + self.kill_process("記憶體攔截", p, file, False) elif ":/Windows" not in file and ":/Program" not in file: if os.path.exists(file): self.track_proc = p, file, True @@ -1338,7 +1340,7 @@ class MainWindow_Controller(QMainWindow): def load_scan(self, p): try: - for entry in p.memory_maps(): + for entry in p.memory_maps(grouped=True): file = str(entry.path).replace("\\", "/") if ":/Windows" not in file and file not in self.whitelist: ftype = str(f".{file.split('.')[-1]}").lower() @@ -1348,19 +1350,53 @@ class MainWindow_Controller(QMainWindow): except: return False + def get_executable_section(self, process_handle, base_address): + dos_header = ctypes.create_string_buffer(64) + ctypes.windll.kernel32.ReadProcessMemory(process_handle, + ctypes.c_void_p(base_address), dos_header, 64, None) + nt_headers_offset = struct.unpack(" 0): + return base_address + virtual_address, virtual_size + return None, 0 + + def get_module_base(self, p): + try: + for entry in p.memory_maps(grouped=False): + file = str(entry.path).replace("\\", "/") + ftype = str(f".{file.split('.')[-1]}").lower() + if ftype in [".exe"] and file not in self.whitelist: + return int(entry.addr, 16) + return False + except: + return False + def memory_scan(self, p): try: - memory_bytes = {} - m = ctypes.windll.kernel32.OpenProcess(0x1F0FFF, False, p.pid) - for memory_range in p.memory_maps(grouped=False): - base_addr = int(memory_range.addr, 16) - buffer = ctypes.create_string_buffer(memory_range.rss) - if ctypes.windll.kernel32.ReadProcessMemory(m, ctypes.c_void_p(base_addr), - buffer, memory_range.rss, ctypes.byref(ctypes.c_size_t(0))): - file = str(memory_range.path).replace("\\", "/") - if ":/Windows" not in file and self.start_scan(buffer.raw): - return True - ctypes.windll.kernel32.CloseHandle(m) + base_address = self.get_module_base(p) + process_handle = ctypes.windll.kernel32.OpenProcess(0x1F0FFF, False, p.pid) + text_addr, text_size = self.get_executable_section(process_handle, base_address) + buffer = ctypes.create_string_buffer(text_size) + if ctypes.windll.kernel32.ReadProcessMemory(process_handle, + ctypes.c_void_p(text_addr), buffer, text_size, ctypes.byref(ctypes.c_size_t(0))): + if self.start_scan(buffer.raw): + return True + ctypes.windll.kernel32.CloseHandle(process_handle) return False except: return False diff --git a/PYAS_Engine.py b/PYAS_Engine.py index 8ffcac3..a3cdeb1 100644 --- a/PYAS_Engine.py +++ b/PYAS_Engine.py @@ -1,5 +1,5 @@ import os, io, sys, time, json, yara -import numpy, onnxruntime +import numpy, pefile, onnxruntime from PIL import Image, ImageShow class YRScan: @@ -36,39 +36,38 @@ class DLScan: self.values = 100 def load_model(self, file_path): - ftype = str(f".{file_path.split('.')[-1]}").lower() - if ftype in [".json", ".txt"]: - self.class_names = json.load(open(file_path, 'r')) - self.values = self.class_names['Values'] - self.detect = self.class_names['Detect'] - available_providers = onnxruntime.get_available_providers() - preferred_providers = [ - 'CUDAExecutionProvider', 'ROCmExecutionProvider', - 'OpenVINOExecutionProvider', 'DirectMLExecutionProvider', - 'AzureExecutionProvider', 'CPUExecutionProvider'] - providers = [p for p in preferred_providers if p in available_providers] - for model in self.class_names['Models']: - try: + try: + ftype = str(f".{file_path.split('.')[-1]}").lower() + if ftype in [".json", ".txt"]: + with open(file_path, 'r') as f: + self.class_names = json.load(f) + self.values = self.class_names['Values'] + self.detect = self.class_names['Detect'] + available_providers = onnxruntime.get_available_providers() + preferred_providers = [ + 'CUDAExecutionProvider', 'ROCmExecutionProvider', + 'OpenVINOExecutionProvider', 'DirectMLExecutionProvider', + 'AzureExecutionProvider', 'CPUExecutionProvider'] + providers = [p for p in preferred_providers if p in available_providers] + for model in self.class_names['Models']: model_path = os.path.join(os.path.dirname(file_path), model) try: - self.models[model] = onnxruntime.InferenceSession(model_path, providers=providers) + self.models[model] = onnxruntime.InferenceSession( + model_path, providers=providers) except Exception as e: - providers = ['CPUExecutionProvider'] - self.models[model] = onnxruntime.InferenceSession(model_path, providers=providers) - except: - pass + self.models[model] = onnxruntime.InferenceSession( + model_path, providers=['CPUExecutionProvider']) + except Exception as e: + pass def dl_scan(self, file_path): try: if isinstance(file_path, str): - ftype = str(f".{file_path.split('.')[-1]}").lower() - if ftype in self.class_names['Suffix']: - with open(file_path, 'rb') as f: - file_path = f.read() - else: - return False, False + file_data = self.check_file_type(file_path) + else: + file_data = file_path target_size, sim = tuple(self.class_names['Pixels']), {} - image = self.preprocess_image(file_path, target_size) + image = self.preprocess_image(file_data, target_size) image_array = numpy.array(image).astype('float32') / 255.0 image_expand = numpy.expand_dims(image_array, axis=0) for model_name, model in self.models.items(): @@ -78,19 +77,36 @@ class DLScan: label = self.class_names['Labels'][number].replace("\n", "") sim[label] = sim.setdefault(label, 0) + pre_answer[number] for local_label, sim_sum in sim.items(): - if sim[local_label] > len(self.models) / 2: + if sim_sum > len(self.models) / 2: local_level = sim_sum / len(self.models) return local_label, local_level * 100 return False, False - except: + except Exception as e: return False, False - def preprocess_image(self, file_data, target_size): # New + def check_file_type(self, file_path): + try: + ftype = str(f".{file_path.split('.')[-1]}").lower() + if ftype in [".exe", ".dll", ".sys", ".com"]: + with pefile.PE(file_path, fast_load=True) as pe: + for section in pe.sections: + if (section.Characteristics & pefile.SECTION_CHARACTERISTICS['IMAGE_SCN_MEM_EXECUTE'] and + section.Characteristics & pefile.SECTION_CHARACTERISTICS['IMAGE_SCN_MEM_READ'] and + section.SizeOfRawData > 0 and section.Name.decode().strip('\x00').lower() in [".text"]): + return section.get_data() + elif ftype in [".bat", ".cmd", ".vbs", ".ps1"]: + with open(file_path, 'rb') as f: + return f.read() + return False + except Exception as e: + return False + + def preprocess_image(self, file_data, target_size): file_data = numpy.frombuffer(file_data, dtype=numpy.uint8) data_count = len(file_data) wah = int(numpy.ceil(numpy.sqrt((data_count + 2) // 3))) image_array = numpy.zeros((wah * wah * 3,), dtype=numpy.uint8) image_array[:data_count] = file_data image_array = image_array.reshape((wah, wah, 3)) - image = Image.fromarray(image_array) #ImageShow.show(image) + image = Image.fromarray(image_array) return image.resize(target_size, Image.Resampling.LANCZOS) diff --git a/PYAS_Version.py b/PYAS_Version.py index b5841d1..88171e7 100644 --- a/PYAS_Version.py +++ b/PYAS_Version.py @@ -2,7 +2,7 @@ import pyinstaller_versionfile pyinstaller_versionfile.create_versionfile( output_file="versionfile.txt", -version='3.1.7', +version='3.1.8', company_name="PYAS Security", file_description="Python Antivirus Software", internal_name="PYAS",