From 814baadbdf6d8e6a41dd3eaafffb1a497dc71089 Mon Sep 17 00:00:00 2001 From: 87owo <85057800+87owo@users.noreply.github.com> Date: Tue, 15 Sep 2026 00:09:47 +0800 Subject: [PATCH] Add files via upload --- Engine/Properties/convert.py | 2 +- Engine/Properties/predict_onnx.py | 1093 +++++++++++++++++++++++++++++ 2 files changed, 1094 insertions(+), 1 deletion(-) create mode 100644 Engine/Properties/predict_onnx.py diff --git a/Engine/Properties/convert.py b/Engine/Properties/convert.py index fbd9ba7..9659262 100644 --- a/Engine/Properties/convert.py +++ b/Engine/Properties/convert.py @@ -21,7 +21,7 @@ except ImportError: JSONL_PATH = "pe_features.jsonl" BATCH_SIZE = 1000 MAX_FILE_SIZE = 4 * 1024 * 1024 * 1024 -MAX_WORKERS = os.cpu_count() - 2 +MAX_WORKERS = max(1, os.cpu_count() - 2) #################################################################################################### diff --git a/Engine/Properties/predict_onnx.py b/Engine/Properties/predict_onnx.py new file mode 100644 index 0000000..b79bf37 --- /dev/null +++ b/Engine/Properties/predict_onnx.py @@ -0,0 +1,1093 @@ +import os, re, sys, math, json, zlib, pefile, datetime, argparse +import ctypes, ctypes.wintypes +import numpy as np +import onnxruntime as ort + +#################################################################################################### + +FEATURE_FILE = "features.json" +MAX_FILE_SIZE = 4 * 1024 * 1024 * 1024 +TARGET_EXTENSIONS = { + '.exe', '.dll', '.sys', '.ocx', '.scr', '.efi', '.acm', '.ax', '.cpl', '.drv', '.com', '.mui', '.pyd' +} + +#################################################################################################### + +class WINTRUST_FILE_INFO(ctypes.Structure): + _fields_ = [ + ("cbStruct", ctypes.wintypes.DWORD), + ("pcwszFilePath", ctypes.wintypes.LPCWSTR), + ("hFile", ctypes.wintypes.HANDLE), + ("pgKnownSubject", ctypes.wintypes.LPVOID) + ] + +class WINTRUST_DATA_UNION(ctypes.Union): + _fields_ = [ + ("pFile", ctypes.POINTER(WINTRUST_FILE_INFO)), + ("pCatalog", ctypes.wintypes.LPVOID), + ("pBlob", ctypes.wintypes.LPVOID), + ("pSgnr", ctypes.wintypes.LPVOID), + ("pCert", ctypes.wintypes.LPVOID) + ] + +class WINTRUST_DATA(ctypes.Structure): + _fields_ = [ + ("cbStruct", ctypes.wintypes.DWORD), + ("pPolicyCallbackData", ctypes.wintypes.LPVOID), + ("pSIPClientData", ctypes.wintypes.LPVOID), + ("dwUIChoice", ctypes.wintypes.DWORD), + ("fdwRevocationChecks", ctypes.wintypes.DWORD), + ("dwUnionChoice", ctypes.wintypes.DWORD), + ("u", WINTRUST_DATA_UNION), + ("dwStateAction", ctypes.wintypes.DWORD), + ("hWVTStateData", ctypes.wintypes.HANDLE), + ("pwszURLReference", ctypes.wintypes.LPCWSTR), + ("dwProvFlags", ctypes.wintypes.DWORD), + ("dwUIContext", ctypes.wintypes.DWORD), + ("pSignatureSettings", ctypes.wintypes.LPVOID) + ] + +class GUID(ctypes.Structure): + _fields_ = [ + ("Data1", ctypes.wintypes.DWORD), + ("Data2", ctypes.wintypes.WORD), + ("Data3", ctypes.wintypes.WORD), + ("Data4", ctypes.c_ubyte * 8) + ] + +V2_GUID = GUID(0x00AAC56B, 0xCD44, 0x11D0, (ctypes.c_ubyte * 8)(0x8C, 0xC2, 0x00, 0xC0, 0x4F, 0xC2, 0x95, 0xEE)) + +#################################################################################################### + +class WinTrust: + @staticmethod + def verify(file_path): + if os.name != 'nt': + return 0 + try: + wintrust = ctypes.windll.wintrust + wintrust.WinVerifyTrust.argtypes = [ctypes.wintypes.HWND, ctypes.POINTER(GUID), ctypes.wintypes.LPVOID] + wintrust.WinVerifyTrust.restype = ctypes.wintypes.LONG + + fi = WINTRUST_FILE_INFO() + fi.cbStruct = ctypes.sizeof(WINTRUST_FILE_INFO) + fi.pcwszFilePath = os.path.abspath(file_path) + fi.hFile = None + fi.pgKnownSubject = None + + wt_union = WINTRUST_DATA_UNION() + wt_union.pFile = ctypes.pointer(fi) + + td = WINTRUST_DATA() + td.cbStruct = ctypes.sizeof(WINTRUST_DATA) + td.pPolicyCallbackData = None + td.pSIPClientData = None + td.dwUIChoice = 2 + td.fdwRevocationChecks = 0 + td.dwUnionChoice = 1 + td.u = wt_union + td.dwStateAction = 0 + td.hWVTStateData = None + td.pwszURLReference = None + td.dwProvFlags = 0 + td.dwUIContext = 0 + td.pSignatureSettings = None + + return 1 if wintrust.WinVerifyTrust(None, ctypes.byref(V2_GUID), ctypes.byref(td)) == 0 else 0 + except Exception: + return 0 + +#################################################################################################### + +def _entropy_from_counts(counts, size): + if size <= 0: + return 0.0 + nonzero = counts[counts > 0] + return float(np.log2(size) - np.sum(nonzero * np.log2(nonzero)) / size) + + +def _bounded_entropy(data, batch_size=1048576): + size = len(data) + if size == 0: + return 0.0 + counts = np.zeros(256, dtype=np.int64) + view = memoryview(data) + try: + for start in range(0, size, batch_size): + chunk = np.frombuffer( + view[start:min(start + batch_size, size)], dtype=np.uint8 + ) + counts += np.bincount(chunk, minlength=256) + finally: + view.release() + return _entropy_from_counts(counts, size) + + +def _accumulate_string_chunk(chunk, array, state, pattern): + separators = (array < 0x20) | (array > 0x7E) + if not separators.any(): + state["pending_counts"] += np.bincount(array, minlength=256) + state["pending_length"] += len(array) + return + + first_separator = int(np.argmax(separators)) + last_separator = len(separators) - 1 - int(np.argmax(separators[::-1])) + if first_separator: + state["pending_counts"] += np.bincount( + array[:first_separator], minlength=256 + ) + state["pending_length"] += first_separator + if state["pending_length"] >= 5: + state["count"] += 1 + state["total_length"] += state["pending_length"] + state["counts"] += state["pending_counts"] + state["pending_length"] = 0 + state["pending_counts"].fill(0) + + middle_start = first_separator + 1 + middle_end = last_separator + if middle_end > middle_start: + matches = pattern.findall(chunk[middle_start:middle_end]) + if matches: + combined = b"".join(matches) + state["count"] += len(matches) + state["total_length"] += len(combined) + state["counts"] += np.bincount( + np.frombuffer(combined, dtype=np.uint8), minlength=256 + ) + + trailing_start = last_separator + 1 + if trailing_start < len(array): + trailing = array[trailing_start:] + state["pending_counts"] += np.bincount(trailing, minlength=256) + state["pending_length"] = len(trailing) + + +def _accumulate_entropy_windows(data, entropy_histogram, c_log_c_table): + block_count = len(data) // 1024 + if block_count < 2: + return data + + complete_size = block_count * 1024 + blocks = np.frombuffer( + data, dtype=np.uint8, count=complete_size + ).reshape(block_count, 1024) + encoded = blocks.astype(np.int32) + encoded += np.arange(block_count, dtype=np.int32)[:, None] * 256 + block_counts = np.bincount( + encoded.ravel(), minlength=block_count * 256 + ).reshape(block_count, 256) + window_counts = block_counts[:-1] + block_counts[1:] + sum_c_log_c = np.sum(c_log_c_table[window_counts], axis=1) + entropies = 11.0 - sum_c_log_c / 2048.0 + entropy_bins = (entropies * 2.0).astype(np.int32) + np.clip(entropy_bins, 0, 15, out=entropy_bins) + byte_bins = window_counts.reshape(block_count - 1, 16, 16).sum(axis=2) + flat_indices = (entropy_bins[:, None] * 16 + np.arange(16)).ravel() + np.add.at(entropy_histogram, flat_indices, byte_bins.ravel()) + return data[(block_count - 1) * 1024:] + + +def _collect_resource_ranges(pe, file_size): + ranges = [] + if not hasattr(pe, "DIRECTORY_ENTRY_RESOURCE"): + return ranges + for entry_type in pe.DIRECTORY_ENTRY_RESOURCE.entries: + for entry_id in getattr( + getattr(entry_type, "directory", None), "entries", [] + ): + for entry_lang in getattr( + getattr(entry_id, "directory", None), "entries", [] + ): + if not hasattr(entry_lang, "data"): + continue + try: + struct = entry_lang.data.struct + offset = int(pe.get_offset_from_rva(int(struct.OffsetToData))) + size = int(struct.Size) + start = max(0, offset) + end = min(file_size, start + max(0, size)) + ranges.append((start, max(start, end))) + except Exception: + ranges.append((0, 0)) + return ranges + + +def _extract_stream_features(file_path, file_size, ranges, extractor): + normalized = [] + for start, end in ranges: + start = max(0, min(file_size, int(start))) + end = max(start, min(file_size, int(end))) + normalized.append((start, end)) + + byte_counts = np.zeros(256, dtype=np.int64) + range_counts = [np.zeros(256, dtype=np.int64) for _ in normalized] + entropy_histogram = np.zeros(256, dtype=np.float64) + string_state = { + "count": 0, + "total_length": 0, + "counts": np.zeros(256, dtype=np.int64), + "pending_length": 0, + "pending_counts": np.zeros(256, dtype=np.int64), + } + histogram_tail = b"" + position = 0 + + with open(file_path, "rb") as source: + while True: + chunk = source.read(4194304) + if not chunk: + break + array = np.frombuffer(chunk, dtype=np.uint8) + _accumulate_string_chunk( + chunk, array, string_state, extractor._STRING_PATTERN + ) + histogram_tail = _accumulate_entropy_windows( + histogram_tail + chunk, + entropy_histogram, + extractor._C_LOG_C_TABLE, + ) + + chunk_end = position + len(chunk) + boundaries = {0, len(chunk)} + for start, end in normalized: + if position < start < chunk_end: + boundaries.add(start - position) + if position < end < chunk_end: + boundaries.add(end - position) + + ordered = sorted(boundaries) + for index in range(len(ordered) - 1): + local_start = ordered[index] + local_end = ordered[index + 1] + if local_end <= local_start: + continue + counts = np.bincount( + array[local_start:local_end], minlength=256 + ) + byte_counts += counts + absolute_start = position + local_start + absolute_end = position + local_end + for range_index, (start, end) in enumerate(normalized): + if start <= absolute_start and absolute_end <= end: + range_counts[range_index] += counts + position = chunk_end + + if string_state["pending_length"] >= 5: + string_state["count"] += 1 + string_state["total_length"] += string_state["pending_length"] + string_state["counts"] += string_state["pending_counts"] + + string_count = string_state["count"] + total_string_length = string_state["total_length"] + result = { + "StringCount": float(string_count), + "StringMeanLength": ( + float(total_string_length) / string_count if string_count else 0.0 + ), + "StringEntropy": _entropy_from_counts( + string_state["counts"], total_string_length + ), + } + result.update(dict.fromkeys(extractor._STRING_KEYS, 0.0)) + if total_string_length: + valid_counts = string_state["counts"][0x20:0x7F] + result.update(dict(zip( + extractor._STRING_KEYS, + valid_counts.astype(np.float64) / total_string_length, + ))) + + result["FileEntropy"] = _entropy_from_counts(byte_counts, file_size) + result.update(dict(zip( + extractor._BYTE_HIST_KEYS, + byte_counts.astype(np.float64) / file_size, + ))) + entropy_sum = np.sum(entropy_histogram) + if entropy_sum: + entropy_histogram /= entropy_sum + result.update(dict(zip(extractor._BYTE_ENT_KEYS, entropy_histogram))) + + range_entropies = [ + _entropy_from_counts(counts, end - start) + for counts, (start, end) in zip(range_counts, normalized) + ] + return result, range_entropies + + +class FeatureExtractor: + _STRING_PATTERN = re.compile(b'[\x20-\x7E]{5,}') + _STRING_FAST_PATH_BYTES = 4194304 + _STRING_BATCH_BYTES = 4194304 + _HISTOGRAM_BATCH_WINDOWS = 1024 + + _BYTE_HIST_KEYS = [f'ByteHist_{i:02X}' for i in range(256)] + _BYTE_ENT_KEYS = [f'ByteEnt_{i:02X}' for i in range(256)] + _STRING_KEYS = [f'StringHist_{i:02d}' for i in range(95)] + _SECTION_HASH_KEYS = [f'SectionHash_{i:02d}' for i in range(50)] + _CHAR_FLAGS = [0x00000020, 0x00000040, 0x00000080, 0x02000000, 0x20000000, 0x40000000, 0x80000000] + _CHAR_COUNT_KEYS = {flag: f'Char_{flag:08X}_Count' for flag in _CHAR_FLAGS} + _CHAR_ENT_KEYS = {flag: f'Char_{flag:08X}_MeanEntropy' for flag in _CHAR_FLAGS} + + _C_LOG_C_TABLE = np.zeros(2049, dtype=np.float64) + _C_LOG_C_TABLE[1:] = np.arange(1, 2049) * np.log2(np.arange(1, 2049)) + + _OPTIONAL_HEADER_FIELDS = [ + 'Magic', 'MajorLinkerVersion', 'MinorLinkerVersion', 'SizeOfCode', + 'SizeOfInitializedData', 'SizeOfUninitializedData', 'AddressOfEntryPoint', + 'BaseOfCode', 'ImageBase', 'SectionAlignment', 'FileAlignment', + 'MajorOperatingSystemVersion', 'MinorOperatingSystemVersion', + 'MajorImageVersion', 'MinorImageVersion', 'MajorSubsystemVersion', + 'MinorSubsystemVersion', 'SizeOfImage', 'SizeOfHeaders', 'CheckSum', + 'Subsystem', 'DllCharacteristics', 'SizeOfStackReserve', 'SizeOfStackCommit', + 'SizeOfHeapReserve', 'SizeOfHeapCommit', 'LoaderFlags', 'NumberOfRvaAndSizes' + ] + + _SECTION_DEFAULTS = { + 'SectionCount': 0.0, + 'SectionMaxEntropy': 0.0, + 'SectionMinEntropy': 0.0, + 'SectionMeanEntropy': 0.0, + 'SectionMaxRawSize': 0.0, + 'SectionMinRawSize': 0.0, + 'SectionMeanRawSize': 0.0, + 'SectionMaxVSize': 0.0, + 'SectionMinVSize': 0.0, + 'SectionMeanVSize': 0.0, + 'ExecutableSections': 0.0, + 'WritableSections': 0.0, + 'ReadableSections': 0.0, + 'SectionException': 0.0 + } + + _CONDITIONAL_DEFAULTS = { + 'IsDriver': 0.0, + 'HasTlsCallbacks': 0.0, + 'IsDebug': 0.0, + 'IsPreRelease': 0.0, + 'IsPatched': 0.0, + 'IsPrivateBuild': 0.0, + 'IsSpecialBuild': 0.0, + 'ImportCount': 0.0, + 'ImportFunctionCount': 0.0, + 'ExportCount': 0.0, + 'IconCount': 0.0, + 'DebugCount': 0.0 + } + + @staticmethod + def _safe_float(val): + try: + f = float(val) + return f if not (math.isinf(f) or math.isnan(f)) else 0.0 + + except Exception: + return 0.0 + + @classmethod + def _base_defaults(cls): + base = dict.fromkeys(cls._OPTIONAL_HEADER_FIELDS, 0.0) + base.update(cls._SECTION_DEFAULTS) + base.update(cls._CONDITIONAL_DEFAULTS) + + for i in range(16): + base[f'DataDirectory_{i}_Size'] = 0.0 + base[f'DataDirectory_{i}_VA'] = 0.0 + + return base + + @staticmethod + def _calc_entropy(data): + return _bounded_entropy(data) + + @classmethod + def _extract_strings(cls, file_bytes): + res = {"StringCount": 0.0, "StringMeanLength": 0.0, "StringEntropy": 0.0} + res.update(dict.fromkeys(cls._STRING_KEYS, 0.0)) + + if len(file_bytes) <= cls._STRING_FAST_PATH_BYTES: + matches = cls._STRING_PATTERN.findall(file_bytes) + if not matches: + return res + + combined_data = b''.join(matches) + count = len(matches) + total_len = len(combined_data) + global_counts = np.bincount(np.frombuffer(combined_data, dtype=np.uint8), minlength=256) + else: + count = 0 + total_len = 0 + global_counts = np.zeros(256, dtype=np.int64) + mv = memoryview(file_bytes) + pos = 0 + size = len(file_bytes) + + while pos < size: + chunk_end = min(pos + cls._STRING_BATCH_BYTES, size) + chunk_arr = np.frombuffer(mv[pos:chunk_end], dtype=np.uint8) + separators = (chunk_arr < 0x20) | (chunk_arr > 0x7E) + + if separators.any(): + last_separator = len(separators) - 1 - int(np.argmax(separators[::-1])) + segment_end = pos + last_separator + 1 + matches = cls._STRING_PATTERN.findall(mv[pos:segment_end]) + + if matches: + combined_data = b''.join(matches) + count += len(matches) + total_len += len(combined_data) + global_counts += np.bincount(np.frombuffer(combined_data, dtype=np.uint8), minlength=256) + + pos = segment_end + continue + + pending_len = 0 + pending_counts = np.zeros(256, dtype=np.int64) + + while pos < size: + chunk_end = min(pos + cls._STRING_BATCH_BYTES, size) + chunk_arr = np.frombuffer(mv[pos:chunk_end], dtype=np.uint8) + separators = (chunk_arr < 0x20) | (chunk_arr > 0x7E) + + if separators.any(): + first_separator = int(np.argmax(separators)) + if first_separator: + pending_counts += np.bincount(chunk_arr[:first_separator], minlength=256) + pending_len += first_separator + pos += first_separator + 1 + break + + pending_counts += np.bincount(chunk_arr, minlength=256) + pending_len += len(chunk_arr) + pos = chunk_end + + if pending_len >= 5: + count += 1 + total_len += pending_len + global_counts += pending_counts + + if count == 0 or total_len == 0: + return res + + res["StringCount"] = float(count) + res["StringMeanLength"] = float(total_len) / count + + counts_nonzero = global_counts[global_counts > 0] + if len(counts_nonzero) > 0: + res["StringEntropy"] = float(np.log2(total_len) - np.sum(counts_nonzero * np.log2(counts_nonzero)) / total_len) + + valid_chars = global_counts[0x20:0x7F] + res.update(dict(zip(cls._STRING_KEYS, valid_chars.astype(np.float64) / total_len))) + + return res + + + + @classmethod + def _extract_histograms(cls, file_bytes): + arr = np.frombuffer(file_bytes, dtype=np.uint8) + sz = len(arr) + res = {} + + if sz == 0: + res['FileEntropy'] = 0.0 + res.update(dict.fromkeys(cls._BYTE_HIST_KEYS, 0.0)) + res.update(dict.fromkeys(cls._BYTE_ENT_KEYS, 0.0)) + return res + + byte_counts = np.zeros(256, dtype=np.int64) + ent_hist = np.zeros(256, dtype=np.float64) + window = 2048 + step = 1024 + + if sz < window: + byte_counts += np.bincount(arr, minlength=256) + p = byte_counts[byte_counts > 0] / float(sz) + entropy = float(-np.sum(p * np.log2(p))) + ent_bin = min(int(entropy * 2.0), 15) + byte_bins = byte_counts.reshape(16, 16).sum(axis=1) + idx_start = ent_bin * 16 + ent_hist[idx_start : idx_start+16] += byte_bins + else: + num_windows = (sz - window) // step + 1 + + for w_start in range(0, num_windows, cls._HISTOGRAM_BATCH_WINDOWS): + w_end = min(w_start + cls._HISTOGRAM_BATCH_WINDOWS, num_windows) + b_num = w_end - w_start + + start_byte = w_start * step + end_byte = start_byte + b_num * step + step + encoded = arr[start_byte:end_byte].astype(np.int32).reshape(b_num + 1, step) + encoded += np.arange(b_num + 1, dtype=np.int32)[:, None] * 256 + + chunk_counts = np.bincount(encoded.ravel(), minlength=(b_num + 1) * 256).reshape(b_num + 1, 256) + del encoded + + byte_counts += np.sum(chunk_counts[:-1], axis=0) + if w_end == num_windows: + byte_counts += chunk_counts[-1] + + window_counts = chunk_counts[:-1] + chunk_counts[1:] + sum_c_log_c = np.sum(cls._C_LOG_C_TABLE[window_counts], axis=1) + entropies = 11.0 - sum_c_log_c / 2048.0 + + ent_bins = (entropies * 2.0).astype(np.int32) + np.clip(ent_bins, 0, 15, out=ent_bins) + + byte_bins = window_counts.reshape(b_num, 16, 16).sum(axis=2) + flat_ent_bins = (ent_bins[:, None] * 16 + np.arange(16)).ravel() + np.add.at(ent_hist, flat_ent_bins, byte_bins.ravel()) + + covered_size = (num_windows + 1) * step + if covered_size < sz: + byte_counts += np.bincount(arr[covered_size:], minlength=256) + + counts_nonzero = byte_counts[byte_counts > 0] + res['FileEntropy'] = float(np.log2(sz) - np.sum(counts_nonzero * np.log2(counts_nonzero)) / sz) + + byte_hist = byte_counts.astype(np.float64) / sz + res.update(dict(zip(cls._BYTE_HIST_KEYS, byte_hist))) + + ent_sum = np.sum(ent_hist) + if ent_sum > 0: + ent_hist /= ent_sum + + res.update(dict(zip(cls._BYTE_ENT_KEYS, ent_hist))) + return res + + + @classmethod + def _extract_overlay_features(cls, pe, fsize): + overlay_offset = pe.get_overlay_data_start_offset() + if not overlay_offset or overlay_offset >= fsize: + return {"HasOverlay": 0.0, "OverlaySize": 0.0, "OverlayRatio": 0.0, "OverlayEntropy": 0.0} + + overlay_data = pe.get_overlay() + if not overlay_data: + return {"HasOverlay": 0.0, "OverlaySize": 0.0, "OverlayRatio": 0.0, "OverlayEntropy": 0.0} + + sz = len(overlay_data) + return { + "HasOverlay": 1.0, + "OverlaySize": float(sz), + "OverlayRatio": float(sz) / float(fsize), + "OverlayEntropy": cls._calc_entropy(overlay_data) + } + + @classmethod + def _extract_rich_header(cls, pe): + if not hasattr(pe, 'RICH_HEADER') or not pe.RICH_HEADER: + return {"HasRichHeader": 0.0, "RichHeaderCount": 0.0} + return { + "HasRichHeader": 1.0, + "RichHeaderCount": float(len(pe.RICH_HEADER.values) // 2) if hasattr(pe.RICH_HEADER, 'values') else 0.0 + } + + @classmethod + def _extract_ep_anomalies(cls, pe): + if not hasattr(pe, 'OPTIONAL_HEADER') or not hasattr(pe, 'sections') or not pe.sections: + return {"EntryPointSectionIndex": -1.0, "EntryPointInExecutable": 0.0, "EntryPointInLastSection": 0.0} + + ep = pe.OPTIONAL_HEADER.AddressOfEntryPoint + ep_section_idx = -1 + ep_in_exec = 0.0 + + for idx, sec in enumerate(pe.sections): + if sec.VirtualAddress <= ep < (sec.VirtualAddress + sec.Misc_VirtualSize): + ep_section_idx = idx + if sec.Characteristics & 0x20000000: + ep_in_exec = 1.0 + break + + return { + "EntryPointSectionIndex": float(ep_section_idx), + "EntryPointInExecutable": ep_in_exec, + "EntryPointInLastSection": 1.0 if ep_section_idx == len(pe.sections) - 1 else 0.0 + } + + @classmethod + def _extract_advanced_resources(cls, pe, range_entropies=None): + res_data = { + "ResourceMaxEntropy": 0.0, + "ResourceMinEntropy": 0.0, + "ResourceMeanEntropy": 0.0, + "ResourceLangCount": 0.0, + "ResourceRCDataCount": 0.0 + } + + if not hasattr(pe, 'DIRECTORY_ENTRY_RESOURCE'): + return res_data + + entropies = [] + entropy_iter = iter(range_entropies or ()) + langs = set() + rcdata_count = 0 + + for entry_type in pe.DIRECTORY_ENTRY_RESOURCE.entries: + if hasattr(entry_type, 'id') and entry_type.id == 10: + rcdata_count += len(getattr(entry_type.directory, 'entries', [])) + + if not hasattr(entry_type, 'directory'): + continue + + for entry_id in entry_type.directory.entries: + if not hasattr(entry_id, 'directory'): + continue + + for entry_lang in entry_id.directory.entries: + if hasattr(entry_lang, 'id'): + langs.add(entry_lang.id) + + if hasattr(entry_lang, 'data'): + try: + entropies.append(float(next(entropy_iter))) + except (StopIteration, TypeError, ValueError): + entropies.append(0.0) + + res_data["ResourceLangCount"] = float(len(langs)) + res_data["ResourceRCDataCount"] = float(rcdata_count) + + if entropies: + res_data["ResourceMaxEntropy"] = max(entropies) + res_data["ResourceMinEntropy"] = min(entropies) + res_data["ResourceMeanEntropy"] = sum(entropies) / len(entropies) + + return res_data + + @classmethod + def _extract_load_config(cls, pe): + cfg = {"HasLoadConfig": 0.0, "HasCFG": 0.0, "HasSEHTable": 0.0} + if hasattr(pe, 'DIRECTORY_ENTRY_LOAD_CONFIG'): + cfg["HasLoadConfig"] = 1.0 + struct = pe.DIRECTORY_ENTRY_LOAD_CONFIG.struct + + if getattr(struct, 'GuardCFFunctionTable', 0) != 0: + cfg["HasCFG"] = 1.0 + + if getattr(struct, 'SEHandlerTable', 0) != 0: + cfg["HasSEHTable"] = 1.0 + + return cfg + + @classmethod + def _extract_security_directory(cls, pe, file_path, fsize): + res = {"HasSignature": 0.0, "SignatureCount": 0.0} + if not hasattr(pe, 'OPTIONAL_HEADER') or not hasattr(pe.OPTIONAL_HEADER, 'DATA_DIRECTORY'): + return res + directories = pe.OPTIONAL_HEADER.DATA_DIRECTORY + if len(directories) <= 4: + return res + sec_dir = directories[4] + offset = int(sec_dir.VirtualAddress) + size = int(sec_dir.Size) + if offset <= 0 or size <= 0 or offset + size > fsize: + return res + + res["HasSignature"] = 1.0 + count = 0 + curr = offset + end = offset + size + with open(file_path, "rb") as source: + while curr + 8 <= end: + source.seek(curr) + header = source.read(4) + if len(header) != 4: + break + length = int.from_bytes(header, 'little') + if length < 8 or curr + length > end: + break + count += 1 + curr += (length + 7) & ~7 + res["SignatureCount"] = float(count) + return res + + @classmethod + def extract(cls, file_path): + pe = None + + try: + fsize = os.path.getsize(file_path) + if fsize == 0 or fsize > MAX_FILE_SIZE: + return None + + base, dlls, apis = cls._base_defaults(), set(), set() + pe = pefile.PE(name=file_path, fast_load=True) + + try: + pe.parse_rich_header() + pe.parse_data_directories(directories=[ + pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_IMPORT'], + pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_EXPORT'], + pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_RESOURCE'], + pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_DEBUG'], + pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_TLS'], + pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_LOAD_CONFIG'] + ]) + except Exception: + pass + + sections = list(getattr(pe, 'sections', [])) + section_ranges = [] + for section in sections: + start = max(0, int(section.get_PointerToRawData_adj())) + end = min(fsize, start + max(0, int(section.SizeOfRawData))) + section_ranges.append((start, max(start, end))) + resource_ranges = _collect_resource_ranges(pe, fsize) + overlay_offset = pe.get_overlay_data_start_offset() + if not overlay_offset or overlay_offset >= fsize: + overlay_offset = fsize + all_ranges = section_ranges + resource_ranges + [ + (int(overlay_offset), fsize) + ] + stream_features, range_entropies = _extract_stream_features( + file_path, fsize, all_ranges, cls + ) + section_entropies = range_entropies[:len(section_ranges)] + resource_start = len(section_ranges) + resource_entropies = range_entropies[ + resource_start:resource_start + len(resource_ranges) + ] + overlay_entropy = range_entropies[-1] + + base['TrustSigned'] = float(WinTrust.verify(file_path)) + base['FileSize'] = float(fsize) + + fh = pe.FILE_HEADER + base['Machine'] = getattr(fh, 'Machine', 0) + base['NumberOfSections'] = getattr(fh, 'NumberOfSections', 0) + base['TimeDateStamp'] = getattr(fh, 'TimeDateStamp', 0) + base['PointerToSymbolTable'] = getattr(fh, 'PointerToSymbolTable', 0) + base['NumberOfSymbols'] = getattr(fh, 'NumberOfSymbols', 0) + base['SizeOfOptionalHeader'] = getattr(fh, 'SizeOfOptionalHeader', 0) + base['Characteristics'] = getattr(fh, 'Characteristics', 0) + + curr_ts = datetime.datetime.utcnow().timestamp() + base['HasInvalidTimestamp'] = 1.0 if (base['TimeDateStamp'] < 631152000 or base['TimeDateStamp'] > curr_ts + 2592000) else 0.0 + base['FileTimeException'] = 1.0 if base['TimeDateStamp'] == 0 else 0.0 + base['Is64Bit'] = 1.0 if base['Machine'] in (0x8664, 0xAA64, 0x0200) else 0.0 + base['IsExe'] = 1.0 if pe.is_exe() else 0.0 + base['IsDll'] = 1.0 if pe.is_dll() else 0.0 + base['ExceptionCount'] = 0.0 + + if hasattr(pe, 'OPTIONAL_HEADER'): + op = pe.OPTIONAL_HEADER + fields = ['Magic', 'MajorLinkerVersion', 'MinorLinkerVersion', 'SizeOfCode', 'SizeOfInitializedData', 'SizeOfUninitializedData', 'AddressOfEntryPoint', 'BaseOfCode', 'ImageBase', 'SectionAlignment', 'FileAlignment', 'MajorOperatingSystemVersion', 'MinorOperatingSystemVersion', 'MajorImageVersion', 'MinorImageVersion', 'MajorSubsystemVersion', 'MinorSubsystemVersion', 'SizeOfImage', 'SizeOfHeaders', 'CheckSum', 'Subsystem', 'DllCharacteristics', 'SizeOfStackReserve', 'SizeOfStackCommit', 'SizeOfHeapReserve', 'SizeOfHeapCommit', 'LoaderFlags', 'NumberOfRvaAndSizes'] + for field_name in fields: + base[field_name] = getattr(op, field_name, 0) + base['IsDriver'] = 1.0 if base.get('Subsystem') == 1 else 0.0 + + if hasattr(op, 'DATA_DIRECTORY'): + for i, directory in enumerate(op.DATA_DIRECTORY): + base[f'DataDirectory_{i}_Size'] = directory.Size + base[f'DataDirectory_{i}_VA'] = directory.VirtualAddress + + exc_idx = pefile.DIRECTORY_ENTRY.get('IMAGE_DIRECTORY_ENTRY_EXCEPTION', 3) + if len(op.DATA_DIRECTORY) > exc_idx and op.DATA_DIRECTORY[exc_idx].Size > 0: + base['ExceptionCount'] = op.DATA_DIRECTORY[exc_idx].Size // 12 if base['Machine'] in (0x8664, 0xAA64) else 0.0 + + res_flags_count = dict.fromkeys(cls._CHAR_COUNT_KEYS.values(), 0.0) + res_flags_ent = dict.fromkeys(cls._CHAR_ENT_KEYS.values(), 0.0) + res_sec_hash = dict.fromkeys(cls._SECTION_HASH_KEYS, 0.0) + + if hasattr(pe, 'sections'): + base['SectionCount'] = len(pe.sections) + entropies, raw_sizes, v_sizes = [], [], [] + exec_sec, write_sec, read_sec, sec_exc = 0, 0, 0, 0 + + for section_index, section in enumerate(pe.sections): + sec_name = section.Name.rstrip(b'\x00') + hash_val = zlib.crc32(sec_name) % 50 + res_sec_hash[cls._SECTION_HASH_KEYS[hash_val]] += 1.0 + + s_entropy = ( + section_entropies[section_index] + if section_index < len(section_entropies) else 0.0 + ) + + entropies.append(s_entropy) + raw_sizes.append(section.SizeOfRawData) + v_sizes.append(section.Misc_VirtualSize) + + if section.Characteristics & 0x20000000: + exec_sec += 1 + if section.Characteristics & 0x80000000: + write_sec += 1 + if section.Characteristics & 0x40000000: + read_sec += 1 + if section.SizeOfRawData + section.PointerToRawData > fsize: + sec_exc = 1 + + for flag in cls._CHAR_FLAGS: + if section.Characteristics & flag: + res_flags_count[cls._CHAR_COUNT_KEYS[flag]] += 1.0 + res_flags_ent[cls._CHAR_ENT_KEYS[flag]] += s_entropy + + base['SectionMaxEntropy'] = max(entropies) if entropies else 0.0 + base['SectionMinEntropy'] = min(entropies) if entropies else 0.0 + base['SectionMeanEntropy'] = sum(entropies) / len(entropies) if entropies else 0.0 + base['SectionMaxRawSize'] = max(raw_sizes) if raw_sizes else 0.0 + base['SectionMinRawSize'] = min(raw_sizes) if raw_sizes else 0.0 + base['SectionMeanRawSize'] = sum(raw_sizes) / len(raw_sizes) if raw_sizes else 0.0 + base['SectionMaxVSize'] = max(v_sizes) if v_sizes else 0.0 + base['SectionMinVSize'] = min(v_sizes) if v_sizes else 0.0 + base['SectionMeanVSize'] = sum(v_sizes) / len(v_sizes) if v_sizes else 0.0 + base['ExecutableSections'] = float(exec_sec) + base['WritableSections'] = float(write_sec) + base['ReadableSections'] = float(read_sec) + base['SectionException'] = float(sec_exc) + + for flag in cls._CHAR_FLAGS: + cnt_key = cls._CHAR_COUNT_KEYS[flag] + if res_flags_count[cnt_key] > 0: + res_flags_ent[cls._CHAR_ENT_KEYS[flag]] /= res_flags_count[cnt_key] + + base.update(res_flags_count) + base.update(res_flags_ent) + base.update(res_sec_hash) + + string_keys = ['FileDescription', 'FileVersion', 'ProductName', 'ProductVersion', 'CompanyName', 'LegalCopyright', 'Comments', 'InternalName', 'LegalTrademarks', 'SpecialBuild', 'PrivateBuild'] + for key in string_keys: + base[f'{key}Length'] = 0.0 + + if hasattr(pe, 'FileInfo'): + for fileinfo_list in pe.FileInfo: + for fileinfo in fileinfo_list: + if getattr(fileinfo, 'name', '') in ('StringFileInfo', b'StringFileInfo'): + for st in getattr(fileinfo, 'StringTable', []): + for key, val in st.entries.items(): + + try: + k = key.decode('utf-8', 'ignore') if isinstance(key, bytes) else str(key) + if k in string_keys: + v = val.decode('utf-8', 'ignore') if isinstance(val, bytes) else str(val) + base[f'{k}Length'] = float(len(v)) + + except Exception: + continue + + if hasattr(pe, 'DIRECTORY_ENTRY_TLS') and hasattr(pe.DIRECTORY_ENTRY_TLS, 'struct'): + base['HasTlsCallbacks'] = 1.0 if getattr(pe.DIRECTORY_ENTRY_TLS.struct, 'AddressOfCallBacks', 0) != 0 else 0.0 + + if hasattr(pe, 'VS_FIXEDFILEINFO') and len(pe.VS_FIXEDFILEINFO) > 0: + flags = getattr(pe.VS_FIXEDFILEINFO[0], 'FileFlags', 0) + base['IsDebug'] = 1.0 if flags & 0x1 else 0.0 + base['IsPreRelease'] = 1.0 if flags & 0x2 else 0.0 + base['IsPatched'] = 1.0 if flags & 0x4 else 0.0 + base['IsPrivateBuild'] = 1.0 if flags & 0x8 else 0.0 + base['IsSpecialBuild'] = 1.0 if flags & 0x20 else 0.0 + + if hasattr(pe, 'DIRECTORY_ENTRY_IMPORT'): + base['ImportCount'] = float(len(pe.DIRECTORY_ENTRY_IMPORT)) + func_count = 0 + for entry in pe.DIRECTORY_ENTRY_IMPORT: + if getattr(entry, 'dll', None): + try: + dlls.add(entry.dll.decode('ascii', 'ignore').lower()) + except Exception: + pass + + for imp in getattr(entry, 'imports', []): + func_count += 1 + if getattr(imp, 'name', None): + try: + apis.add(imp.name.decode('ascii', 'ignore')) + except Exception: + pass + + base['ImportFunctionCount'] = float(func_count) + + base['ExportCount'] = float(len(pe.DIRECTORY_ENTRY_EXPORT.symbols)) if (hasattr(pe, 'DIRECTORY_ENTRY_EXPORT') and hasattr(pe.DIRECTORY_ENTRY_EXPORT, 'symbols')) else 0.0 + + if hasattr(pe, 'DIRECTORY_ENTRY_RESOURCE'): + icon_count = 0 + for entry in pe.DIRECTORY_ENTRY_RESOURCE.entries: + if getattr(entry, 'id', None) == 3 and hasattr(entry, 'directory'): + icon_count += len(getattr(entry.directory, 'entries', [])) + base['IconCount'] = float(icon_count) + else: + base['IconCount'] = 0.0 + + base['DebugCount'] = float(len(pe.DIRECTORY_ENTRY_DEBUG)) if hasattr(pe, 'DIRECTORY_ENTRY_DEBUG') else 0.0 + + base.update(stream_features) + overlay_size = fsize - int(overlay_offset) + base.update({ + "HasOverlay": 1.0 if overlay_size else 0.0, + "OverlaySize": float(overlay_size), + "OverlayRatio": float(overlay_size) / float(fsize) if overlay_size else 0.0, + "OverlayEntropy": overlay_entropy if overlay_size else 0.0 + }) + base.update(cls._extract_rich_header(pe)) + base.update(cls._extract_ep_anomalies(pe)) + base.update(cls._extract_advanced_resources(pe, resource_entropies)) + base.update(cls._extract_load_config(pe)) + base.update(cls._extract_security_directory(pe, file_path, fsize)) + + for k in base: + base[k] = cls._safe_float(base[k]) + + return {"Base": base, "DLLs": list(dlls), "APIs": list(apis)} + + except pefile.PEFormatError: + return None + finally: + if pe is not None: + try: + pe.close() + except Exception: + pass + +#################################################################################################### + +class ModelPredictor: + def __init__(self, model_path, feature_path): + self.model_features = self._load_features(feature_path) + self.sess = ort.InferenceSession(model_path, providers=['CPUExecutionProvider']) + self.input_name = self.sess.get_inputs()[0].name + + expected_dim = self.sess.get_inputs()[0].shape[1] + current_dim = len(self.model_features) + if current_dim != expected_dim: + raise ValueError(f"Dimension error: Model expects {expected_dim}, but feature file has {current_dim}") + + self.feat_map = {feat: i for i, feat in enumerate(self.model_features)} + + self.dll_hash_dim = 512 + self.api_hash_dim = 4096 + self.dll_hash_pad = 3 + self.api_hash_pad = 4 + self._parse_hash_dims() + + def _load_features(self, path): + if not os.path.exists(path): + raise FileNotFoundError(f"Feature file not found: {path}") + + with open(path, 'r') as f: + return json.load(f) + + def _parse_hash_dims(self): + max_dll = -1 + max_api = -1 + + for feat in self.model_features: + if feat.startswith("DllHash_"): + try: + val_str = feat.split("_")[1] + max_dll = max(max_dll, int(val_str)) + self.dll_hash_pad = len(val_str) + except Exception: + pass + elif feat.startswith("ApiHash_"): + try: + val_str = feat.split("_")[1] + max_api = max(max_api, int(val_str)) + self.api_hash_pad = len(val_str) + except Exception: + pass + + if max_dll >= 0: + self.dll_hash_dim = max_dll + 1 + if max_api >= 0: + self.api_hash_dim = max_api + 1 + + def predict(self, raw_data): + vec = np.zeros((1, len(self.model_features)), dtype=np.float32) + base = raw_data.get('Base', {}) + dlls = raw_data.get('DLLs', []) + apis = raw_data.get('APIs', []) + + for k, v in base.items(): + if k in self.feat_map: + vec[0, self.feat_map[k]] = v + + for d in dlls: + h = zlib.crc32(d.encode('utf-8', 'ignore')) % self.dll_hash_dim + feat_name = f"DllHash_{h:0{self.dll_hash_pad}d}" + if feat_name in self.feat_map: + vec[0, self.feat_map[feat_name]] += 1.0 + + for a in apis: + h = zlib.crc32(a.encode('utf-8', 'ignore')) % self.api_hash_dim + feat_name = f"ApiHash_{h:0{self.api_hash_pad}d}" + if feat_name in self.feat_map: + vec[0, self.feat_map[feat_name]] += 1.0 + + outputs = self.sess.run(None, {self.input_name: vec}) + + if len(outputs) > 1: + result = outputs[1] + if isinstance(result, list) and len(result) > 0: + prob_dict = result[0] + if hasattr(prob_dict, 'get'): + return float(prob_dict.get(1, prob_dict.get('1', 0.0))) + elif isinstance(result, np.ndarray): + if result.ndim == 2 and result.shape[1] > 1: + return float(result[0][1]) + + return 0.0 + +#################################################################################################### + +def scan_target(target, predictor): + files = [] + if os.path.isfile(target): + files.append(target) + elif os.path.isdir(target): + for r, _, fs in os.walk(target): + for f in fs: + if os.path.splitext(f)[1].lower() in TARGET_EXTENSIONS: + files.append(os.path.join(r, f)) + + for fpath in files: + try: + data = FeatureExtractor.extract(fpath) + if not data: + print(json.dumps({"result": "error", "prob": 0.0, "file": fpath, "details": "Extraction failed"})) + continue + + prob = predictor.predict(data) + label = "malware" if prob > 0.5 else "safe" + + print(json.dumps({"result": label, "prob": round(prob, 4), "file": fpath})) + + except Exception as e: + print(json.dumps({"result": "error", "prob": 0.0, "file": fpath, "details": str(e)})) + +#################################################################################################### + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("-model", required=True, type=str) + parser.add_argument("targets", nargs="*", type=str) + args = parser.parse_args() + + base_dir = os.path.dirname(sys.executable) if getattr(sys, 'frozen', False) else os.path.dirname(os.path.abspath(__file__)) + model_path = args.model + feature_path = os.path.join(base_dir, FEATURE_FILE) + + try: + predictor = ModelPredictor(model_path, feature_path) + except Exception as e: + print(json.dumps({"result": "fatal", "details": str(e)})) + sys.exit(1) + + if args.targets: + for path in args.targets: + target = path.strip('"').strip("'") + if os.path.exists(target): + scan_target(target, predictor) + else: + print(json.dumps({"result": "error", "file": target, "details": "Path not found"})) + else: + while True: + try: + path = input("Enter File or Folder Path (or 'q' to exit): ").strip().strip('"').strip("'") + if path.lower() in ['q', 'exit']: + break + if not os.path.exists(path): + print(json.dumps({"result": "error", "file": path, "details": "Path not found"})) + continue + scan_target(path, predictor) + except (KeyboardInterrupt, EOFError): + break + except Exception as e: + print(json.dumps({"result": "error", "details": str(e)}))