diff --git a/PYAS.py b/PYAS.py index 40d8bfd..56d4474 100644 --- a/PYAS.py +++ b/PYAS.py @@ -79,6 +79,7 @@ class MainWindow_Controller(QMainWindow): scan_result_signal = Signal(str) scan_reset_signal = Signal() progress_title_signal = Signal(str) + init_log_signal = Signal(str) def __init__(self): super().__init__() @@ -107,9 +108,7 @@ class MainWindow_Controller(QMainWindow): self.sign = sign_scanner() self.sign.init_windll(["wintrust"]) self.model = model_scanner() - self.model.load_path(self.path_models) self.rule = rule_scanner() - self.rule.load_path(self.path_rules) self.mouse_flag = 0 self.mouse_pos = 0 @@ -270,6 +269,16 @@ class MainWindow_Controller(QMainWindow): self.pyas_config[key] = [] return self.pyas_config[key] +#################################################################################################### + + def init_engine_thread(self): + try: + self.rule.load_path(self.path_rules, callback=lambda x: self.init_log_signal.emit(os.path.basename(x))) + self.model.load_path(self.path_models, callback=lambda x: self.init_log_signal.emit(os.path.basename(x))) + self.init_log_signal.emit("引擎加載完成") + except Exception as e: + self.send_message(e, "warn", False) + #################################################################################################### def init_interface(self): @@ -314,6 +323,7 @@ class MainWindow_Controller(QMainWindow): self.scan_result_signal.connect(self.slot_scan_result) self.scan_reset_signal.connect(self.slot_scan_reset) self.progress_title_signal.connect(self.slot_progress_title) + self.init_log_signal.connect(self.slot_init_log) for name, widget in self.widgets.items(): if hasattr(widget, "setCheckable") and widget.isCheckable(): @@ -578,6 +588,7 @@ class MainWindow_Controller(QMainWindow): if "-h" not in param: self.show_button() self.apply_settings() + self.start_daemon_thread(self.init_engine_thread) else: sys.exit(0) except Exception as e: @@ -740,6 +751,10 @@ class MainWindow_Controller(QMainWindow): self.widgets["method_button"].show() self.send_message(msg, "notify", True) + @Slot(str) + def slot_init_log(self, text): + self.send_message(text, "load", True) + #################################################################################################### def init_scan(self): @@ -822,21 +837,18 @@ class MainWindow_Controller(QMainWindow): def scan_engine(self, file_path): try: - if self.sign.sign_verify(file_path): - return False - label, level = self.model.model_scan(file_path) if label and level >= self.pyas_config.get("sensitive", 95): - return f"{label}.{level}" + if not self.sign.sign_verify(file_path): + return f"{label}.{level}" if self.pyas_config.get("extension_switch", False): - label, level = self.rule.yara_scan(file_path) - if label and level: - return f"{label}.{level}" - return False + types, detail = self.rule.yara_scan(file_path) + if types and detail: + return f"{types}.{detail}" except Exception as e: self.send_message(e, "warn", False) - return None + return False #################################################################################################### diff --git a/PYAS_Config.py b/PYAS_Config.py index 0254e25..263a56e 100644 --- a/PYAS_Config.py +++ b/PYAS_Config.py @@ -24,7 +24,7 @@ translate_dict = { "english_switch": {"開啟":"On","關閉":"Off","开启":"On","关闭":"Off",",":", ", "立即刪除":"Delete Now","安全模式":"Safe Mode","官方網站":"Official Website ", - "腳本":"Script ","鏈式":"Chain ","記憶體":"Memory ","加載":"Load ", + "腳本":"Script ","鏈式":"Chain ","記憶體":"Memory ","正在加載":"Loading ","加載":"Load ", "行為":"Behavior ","可疑":"Suspicious ","勒索":"Ransom ","主要":"Main ", "許可條款":"License Terms ","惡意":"Malicious", "開發":"Developers ", "啟用":"Enable ","禁用":"Disable ","正在初始化中":"Initializing ", diff --git a/PYAS_Engine.py b/PYAS_Engine.py index 9e1edb1..719d095 100644 --- a/PYAS_Engine.py +++ b/PYAS_Engine.py @@ -70,36 +70,65 @@ class sign_scanner: class rule_scanner: def __init__(self): - self.rules = {} + self.suffix = {".com", ".dll", ".drv", ".exe", ".ocx", ".scr", ".sys"} + self.rules = None self.network = [] - def load_path(self, path): + def load_path(self, path, callback=None): + yara_files = {} for root, _, files in os.walk(path): for file in files: - self.load_file(os.path.join(root, file)) + 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) - def load_file(self, file): - ext = os.path.splitext(file)[1].lower() + if yara_files: + self.compile_all_rules(yara_files) + + def load_compiled_rule(self, file): try: - if ext in ('.yara', '.yar'): - self.rules[file] = yara.compile(file) - elif ext in ('.yc', '.yrc'): - self.rules[file] = yara.load(file) - elif ext in ('.ip', '.txt'): - 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: - return False + self.rules = yara.load(file) + except Exception as e: + print(e) + + 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 as e: + print(e) + + def compile_all_rules(self, file_map): + try: + self.rules = yara.compile(filepaths=file_map) + except Exception as e: + print(e) def yara_scan(self, file_path): try: - for rules in self.rules.values(): - matches = rules.match(filepath=file_path) - if matches: - rule_name = str(matches[0]) - label = rule_name.split("_")[0] - level = rule_name.split("_")[-1] - return f"Rules/{label}", level + if not self.rules: + return False, False + + fpath_str = str(file_path) + ext = os.path.splitext(fpath_str)[-1].lower() + if ext not in self.suffix: + return False, False + + matches = self.rules.match(filepath=file_path) + if matches: + rule_name = str(matches[0]) + types = rule_name.split("_")[0] + label = rule_name.split("_")[1] + level = rule_name.split("_")[2] + return f"{types}/{label}", level return False, False except Exception: return False, False @@ -108,26 +137,30 @@ class rule_scanner: class model_scanner: def __init__(self): - self.model = None + self.models = [] self.suffix = {".com", ".dll", ".drv", ".exe", ".ocx", ".scr", ".sys"} self.labels = ["Pefile/White", "Pefile/General"] self.detect_set = {"Pefile/General"} self.resize = (224, 224) - def load_path(self, path): + def load_path(self, path, callback=None): for root, _, files in os.walk(path): for file in files: - self.load_file(os.path.join(root, file)) + full_path = os.path.join(root, file) + if callback: + callback(full_path) + self.load_file(full_path) def load_file(self, file): if file.lower().endswith('.onnx'): try: - self.model = onnxruntime.InferenceSession(file) + session = onnxruntime.InferenceSession(file) + self.models.append(session) except Exception: return False def model_scan(self, file_path, full_output=False): - if not self.model: + if not self.models: return (False, False) if not full_output else ([], str(file_path), None) sections = self.extract_sections(file_path) @@ -142,32 +175,35 @@ class model_scanner: if arr.ndim == 3: arr = numpy.expand_dims(arr, axis=-1) - input_meta = self.model.get_inputs()[0] - input_name = input_meta.name - input_shape = input_meta.shape - - if len(input_shape) == 4: - if input_shape[1] in (1, 3) and input_shape[3] not in (1, 3): - arr = arr.transpose(0, 3, 1, 2) - try: - probs = self.model.run(None, {input_name: arr})[0] - except Exception: - return (False, False) if not full_output else ([], str(file_path), None) - results = [] best_malicious = None - for name, pred, img in zip(sections.keys(), probs, images): - idx = int(numpy.argmax(pred)) - label = self.labels[idx] if idx < len(self.labels) else f"Class_{idx}" - conf = round(float(pred[idx]) * 100, 2) - print(f"Section: {name} | Label: {label} | Confidence: {conf}%") + for model in self.models: + input_meta = model.get_inputs()[0] + input_name = input_meta.name + input_shape = input_meta.shape + + curr_arr = arr.copy() + if len(input_shape) == 4: + if input_shape[1] in (1, 3) and input_shape[3] not in (1, 3): + curr_arr = curr_arr.transpose(0, 3, 1, 2) + + try: + probs = model.run(None, {input_name: curr_arr})[0] + except Exception: + continue - if full_output: - results.append((name, label, conf, self.pil_to_base64(img))) - if label in self.detect_set: - if best_malicious is None or conf > best_malicious[1]: - best_malicious = (label, int(conf)) + for name, pred, img in zip(sections.keys(), probs, images): + idx = int(numpy.argmax(pred)) + label = self.labels[idx] if idx < len(self.labels) else f"Class_{idx}" + conf = round(float(pred[idx]) * 100, 2) + + if full_output: + results.append((name, label, conf, self.pil_to_base64(img))) + + if label in self.detect_set: + if best_malicious is None or conf > best_malicious[1]: + best_malicious = (label, int(conf)) if full_output: results.sort(key=lambda x: (x[1] not in self.detect_set, -x[2]))