diff --git a/PYAS.py b/PYAS.py index 08c8105..6fb80a1 100644 --- a/PYAS.py +++ b/PYAS.py @@ -1,6 +1,6 @@ import os, gc, sys, time, json, psutil -import ctypes, struct, msvcrt, pyperclip -import win32gui, win32api, win32con, win32file +import msvcrt, pyperclip, win32file +import win32gui, win32api, win32con from PYAS_Engine import YRScan, DLScan from PYAS_Suffixes import slist, alist from PYAS_Language import translate_dict @@ -924,14 +924,13 @@ class MainWindow_Controller(QMainWindow): if os.path.exists(file): self.send_notify(self.trans(f"正在掃描: ")+file, False) label, level = self.model.dl_scan(file) - if label and label in self.model.detect: - if self.json["high_sensitive"]: - return label - elif level >= self.model.values: - return label - if self.json["extension_kits"]: - if self.rules.yr_scan(file): - return "Virus/Rules" + if label and self.json["high_sensitive"]: + return label + elif label and level >= self.model.values: + return label + label, level = self.rules.yr_scan(file) + if label and self.json["extension_kits"]: + return label return False except: return False @@ -1330,8 +1329,6 @@ 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: ftype = str(f".{file.split('.')[-1]}").lower() if os.path.exists(file) and ftype in slist: @@ -1351,57 +1348,6 @@ 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: - 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 - def lock_process(self, p, lock): if lock: p.suspend() diff --git a/PYAS_Engine.py b/PYAS_Engine.py index 174fedf..c5d478d 100644 --- a/PYAS_Engine.py +++ b/PYAS_Engine.py @@ -26,11 +26,12 @@ class YRScan: with open(file_path, "rb") as f: file_path = f.read() for name, rules in self.rules.items(): - if rules.match(data=file_path): - return True - return False + matchs_rules = rules.match(data=file_path) + if matchs_rules: + return "Virus/Rules", matchs_rules + return False, False except: - return False + return False, False class DLScan: def __init__(self): @@ -51,12 +52,10 @@ class DLScan: except: pass - def dl_scan(self, file_path): + def predict(self, file_data): try: - if isinstance(file_path, str): - file_path = self.check_file_type(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(): @@ -73,23 +72,37 @@ class DLScan: except: return False, False - def check_file_type(self, file_path): + def check_shell(self, name): try: - ftype = str(f".{file_path.split('.')[-1]}").lower() - if ftype in [".exe", ".dll", ".sys"]: - with pefile.PE(file_path, fast_load=True) as pe: - for section in pe.sections: - if (section.Characteristics & pefile.SECTION_CHARACTERISTICS['IMAGE_SCN_MEM_READ'] and - section.Characteristics & pefile.SECTION_CHARACTERISTICS['IMAGE_SCN_MEM_EXECUTE'] and - section.SizeOfRawData > 0 and section.Name.decode().strip('\x00').lower() in [".text"]): - return section.get_data() - elif ftype in [".bat", ".vbs", ".ps1"]: - with open(file_path, 'rb') as f: - return f.read() - return matching_data if matching_data else False + section_name = name.lower() + for shell in ["upx", "vmp", "be", "orpc", "pack", "crypt"]: + if shell in section_name: + return True + return False except: return False + def dl_scan(self, file_path): + try: + match_data = [] + ftype = str(f".{file_path.split('.')[-1]}").lower() + if ftype in [".exe", ".dll", ".sys", ".com", ".scr"]: + with pefile.PE(file_path, fast_load=True) as pe: + for section in pe.sections: + section_name = section.Name.decode().strip('\x00') + if section_name.lower() in [".text"]: + match_data.append(section.get_data()) + elif ftype in [".bat", ".vbs", ".ps1", ".js", ".cmd"]: + with open(file_path, 'rb') as file: + match_data.append(file.read()) + for file_data in match_data: + label, level = self.predict(file_data) + if label and label in self.detect: + return label, level + return False, False + except: + return False, False + def preprocess_image(self, file_data, target_size): file_data = numpy.frombuffer(file_data, dtype=numpy.uint8) data_count = len(file_data)