diff --git a/Engine/Properties/convert.py b/Engine/Properties/convert.py index bb91b18..fbd9ba7 100644 --- a/Engine/Properties/convert.py +++ b/Engine/Properties/convert.py @@ -1,876 +1,1101 @@ -import re, os, sys, math, time, zlib, mmap, datetime, hashlib -import ctypes, ctypes.wintypes, pefile, multiprocessing -import numpy as np - -#################################################################################################### - -try: - import orjson - JSON_DUMPS = lambda x: orjson.dumps(x).decode('utf-8') - JSON_LOADS = orjson.loads - HAS_ORJSON = True - -except ImportError: - import json - JSON_DUMPS = json.dumps - JSON_LOADS = json.loads - HAS_ORJSON = False - -#################################################################################################### - -JSONL_PATH = "pe_features.jsonl" -BATCH_SIZE = 1000 -MAX_FILE_SIZE = 4 * 1024 * 1024 * 1024 -MAX_WORKERS = os.cpu_count() - 2 - -#################################################################################################### - -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)) - -#################################################################################################### - -def verify_signature(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 - - abs_path = os.path.abspath(file_path) - - fi = WINTRUST_FILE_INFO() - fi.cbStruct = ctypes.sizeof(WINTRUST_FILE_INFO) - fi.pcwszFilePath = abs_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 - -#################################################################################################### - -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): - if not data: - return 0.0 - - arr = np.frombuffer(data, dtype=np.uint8) - sz = len(arr) - counts = np.bincount(arr, minlength=256) - counts = counts[counts > 0] - - return float(np.log2(sz) - np.sum(counts * np.log2(counts)) / sz) - - @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): - 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 = [] - 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: - data = pe.get_data(entry_lang.data.struct.OffsetToData, entry_lang.data.struct.Size) - entropies.append(cls._calc_entropy(data)) - except Exception: - continue - - 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_bytes, fsize): - res = {"HasSignature": 0.0, "SignatureCount": 0.0} - if not hasattr(pe, 'OPTIONAL_HEADER') or not hasattr(pe.OPTIONAL_HEADER, 'DATA_DIRECTORY'): - return res - - sec_idx = 4 - directories = pe.OPTIONAL_HEADER.DATA_DIRECTORY - if len(directories) <= sec_idx: - return res - - sec_dir = directories[sec_idx] - if sec_dir.Size == 0 or sec_dir.VirtualAddress == 0: - return res - - res["HasSignature"] = 1.0 - offset = sec_dir.VirtualAddress - size = sec_dir.Size - - if offset + size > fsize or offset < 0 or size < 0: - return res - - count = 0 - curr = offset - end = offset + size - - while curr + 8 <= end: - dw_length = int.from_bytes(file_bytes[curr:curr+4], 'little') - if dw_length < 8: - break - - count += 1 - curr += (dw_length + 7) & ~7 - - res["SignatureCount"] = float(count) - return res - - @classmethod - def extract(cls, file_bytes, file_path, fsize): - if fsize == 0: - return None - - base, dlls, apis, pe = cls._base_defaults(), set(), set(), None - try: - pe = pefile.PE(data=file_bytes, 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 - - base['TrustSigned'] = float(verify_signature(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 f in fields: base[f] = getattr(op, f, 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 in 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 - - try: - s_data = section.get_data() - s_entropy = cls._calc_entropy(s_data) - except Exception: - s_entropy = 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(cls._extract_strings(file_bytes)) - base.update(cls._extract_histograms(file_bytes)) - base.update(cls._extract_overlay_features(pe, fsize)) - base.update(cls._extract_rich_header(pe)) - base.update(cls._extract_ep_anomalies(pe)) - base.update(cls._extract_advanced_resources(pe)) - base.update(cls._extract_load_config(pe)) - base.update(cls._extract_security_directory(pe, file_bytes, fsize)) - - for k in base: - base[k] = cls._safe_float(base[k]) - - return {"Base": base, "DLLs": list(dlls), "APIs": list(apis)} - - except Exception: - return None - finally: - if pe: pe.close() - -#################################################################################################### - -def load_existing_hashes(): - hashes = set() - if not os.path.exists(JSONL_PATH): - return hashes - - pattern = re.compile(rb'"FileHash"\s*:\s*"([a-fA-F0-9]{64})"') - with open(JSONL_PATH, 'rb') as f: - for line in f: - match = pattern.search(line) - if match: - hashes.add(match.group(1).decode('ascii')) - - return hashes - -GLOBAL_HASHES = set() - -def worker_init(shared_hashes): - global GLOBAL_HASHES - GLOBAL_HASHES = shared_hashes - -def _extract_worker_wrapper(args): - return _extract_worker(*args) - -def _extract_worker(file_path, label): - mm = None - f = None - - try: - fsize = os.path.getsize(file_path) - if fsize == 0 or fsize > MAX_FILE_SIZE: - return 'error', None - - f = open(file_path, "rb") - mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) - sha256 = hashlib.sha256(mm).hexdigest() - - if sha256 in GLOBAL_HASHES: - return 'skipped', None - - res = FeatureExtractor.extract(mm, file_path, fsize) - if res: - final_res = {"Label": label, "FileHash": sha256, **res} - return 'success', (sha256, JSON_DUMPS(final_res) + "\n") - - return 'error', None - except Exception: - return 'error', None - finally: - if mm: - mm.close() - if f: - f.close() - -#################################################################################################### - -def scan_and_save(path, label, existing_hashes): - print(f"\n[*] Scanning path: {path} (Label={label})") - files = [path] if os.path.isfile(path) else [os.path.join(r, f) for r, _, fs in os.walk(path) for f in fs] - - total = len(files) - if total == 0: - print("[-] No files found.") - return - - print(f"[*] Found {total} files. Starting {MAX_WORKERS}-process extraction...\n") - if HAS_ORJSON: - print("[+] orjson acceleration enabled.") - - count, errors, skipped, processed_files = 0, 0, 0, 0 - start_time = time.time() - batch_buffer = [] - - pool = multiprocessing.Pool(processes=MAX_WORKERS, initializer=worker_init, initargs=(existing_hashes,)) - tasks = [(fpath, label) for fpath in files] - - try: - for status, data in pool.imap_unordered(_extract_worker_wrapper, tasks): - processed_files += 1 - - if status == 'error': - errors += 1 - elif status == 'skipped': - skipped += 1 - elif status == 'success': - sha256, json_str = data - - if sha256 not in existing_hashes: - existing_hashes.add(sha256) - batch_buffer.append(json_str) - count += 1 - else: - skipped += 1 - - if len(batch_buffer) >= BATCH_SIZE: - with open(JSONL_PATH, 'a', encoding='utf-8') as f_out: - f_out.writelines(batch_buffer) - batch_buffer.clear() - - elapsed = time.time() - start_time - rate = processed_files / elapsed if elapsed > 0 else 0 - sys.stdout.write(f"\r[{processed_files}/{total}] Added: {count} | Skipped: {skipped} | Errors: {errors} | Speed: {rate:.1f} files/s") - sys.stdout.flush() - - pool.close() - pool.join() - - except KeyboardInterrupt: - print("\n\n[-] User aborted. Force terminating all child processes...") - pool.terminate() - pool.join() - - except Exception as e: - print(f"\n[-] Critical error: {e}") - pool.terminate() - pool.join() - - finally: - if batch_buffer: - try: - with open(JSONL_PATH, 'a', encoding='utf-8') as f_out: - f_out.writelines(batch_buffer) - except Exception as save_err: - print(f"\n[-] Failed to save remaining buffer: {save_err}") - - print(f"\n\n[+] Scan complete. Total added: {count}, Total skipped: {skipped}") - -#################################################################################################### - -if __name__ == "__main__": - multiprocessing.freeze_support() - - print("\n---------------- PE Dataset Builder v4.0 ----------------\n") - - existing_hashes = load_existing_hashes() - print(f"[+] Found {len(existing_hashes)} existing hashes.\n") - - try: - while True: - print("-" * 57) - target_path = input("\n[*] Enter path (file or directory): ").strip().strip('"').strip("'") - if target_path.lower() in ['exit', 'q']: - break - - if not target_path or not os.path.exists(target_path): - print("[-] Invalid path.") - continue - - while True: - label_input = input("[*] Enter label (0=Safe, 1=Malware): ").strip() - if label_input.lower() in ['exit', 'q']: - sys.exit() - if label_input in ['0', '1']: - break - - print("[-] Please enter 0 or 1.") - - scan_and_save(target_path, int(label_input), existing_hashes) - - except KeyboardInterrupt: - print("\n\n[-] Aborted.") +import re, os, sys, math, time, zlib, datetime, hashlib +import ctypes, ctypes.wintypes, pefile, multiprocessing +import numpy as np + +#################################################################################################### + +try: + import orjson + JSON_DUMPS = lambda x: orjson.dumps(x).decode('utf-8') + JSON_LOADS = orjson.loads + HAS_ORJSON = True + +except ImportError: + import json + JSON_DUMPS = json.dumps + JSON_LOADS = json.loads + HAS_ORJSON = False + +#################################################################################################### + +JSONL_PATH = "pe_features.jsonl" +BATCH_SIZE = 1000 +MAX_FILE_SIZE = 4 * 1024 * 1024 * 1024 +MAX_WORKERS = os.cpu_count() - 2 + +#################################################################################################### + +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)) + +#################################################################################################### + +def verify_signature(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 + + abs_path = os.path.abspath(file_path) + + fi = WINTRUST_FILE_INFO() + fi.cbStruct = ctypes.sizeof(WINTRUST_FILE_INFO) + fi.pcwszFilePath = abs_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_bytes, file_path, fsize): + if fsize == 0: + return None + + base, dlls, apis, pe = cls._base_defaults(), set(), set(), None + try: + 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(verify_signature(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 f in fields: base[f] = getattr(op, f, 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 Exception: + return None + finally: + if pe: pe.close() + +#################################################################################################### + +def load_existing_hashes(): + hashes = set() + if not os.path.exists(JSONL_PATH): + return hashes + + pattern = re.compile(rb'"FileHash"\s*:\s*"([a-fA-F0-9]{64})"') + with open(JSONL_PATH, 'rb') as f: + for line in f: + match = pattern.search(line) + if match: + hashes.add(match.group(1).decode('ascii')) + + return hashes + +GLOBAL_HASHES = set() + +def worker_init(shared_hashes): + global GLOBAL_HASHES + GLOBAL_HASHES = shared_hashes + +def _extract_worker_wrapper(args): + return _extract_worker(*args) + +def _extract_worker(file_path, label): + try: + fsize = os.path.getsize(file_path) + if fsize == 0 or fsize > MAX_FILE_SIZE: + return 'error', None + + digest = hashlib.sha256() + with open(file_path, "rb") as source: + while True: + chunk = source.read(4194304) + if not chunk: + break + digest.update(chunk) + sha256 = digest.hexdigest() + + if sha256 in GLOBAL_HASHES: + return 'skipped', None + + res = FeatureExtractor.extract(None, file_path, fsize) + if res: + final_res = {"Label": label, "FileHash": sha256, **res} + return 'success', (sha256, JSON_DUMPS(final_res) + "\n") + return 'error', None + except Exception: + return 'error', None + +#################################################################################################### + +def scan_and_save(path, label, existing_hashes): + print(f"\n[*] Scanning path: {path} (Label={label})") + files = [path] if os.path.isfile(path) else [os.path.join(r, f) for r, _, fs in os.walk(path) for f in fs] + + total = len(files) + if total == 0: + print("[-] No files found.") + return + + print(f"[*] Found {total} files. Starting {MAX_WORKERS}-process extraction...\n") + if HAS_ORJSON: + print("[+] orjson acceleration enabled.") + + count, errors, skipped, processed_files = 0, 0, 0, 0 + start_time = time.time() + batch_buffer = [] + + pool = multiprocessing.Pool(processes=MAX_WORKERS, initializer=worker_init, initargs=(existing_hashes,)) + tasks = [(fpath, label) for fpath in files] + + try: + for status, data in pool.imap_unordered(_extract_worker_wrapper, tasks): + processed_files += 1 + + if status == 'error': + errors += 1 + elif status == 'skipped': + skipped += 1 + elif status == 'success': + sha256, json_str = data + + if sha256 not in existing_hashes: + existing_hashes.add(sha256) + batch_buffer.append(json_str) + count += 1 + else: + skipped += 1 + + if len(batch_buffer) >= BATCH_SIZE: + with open(JSONL_PATH, 'a', encoding='utf-8') as f_out: + f_out.writelines(batch_buffer) + batch_buffer.clear() + + elapsed = time.time() - start_time + rate = processed_files / elapsed if elapsed > 0 else 0 + sys.stdout.write(f"\r[{processed_files}/{total}] Added: {count} | Skipped: {skipped} | Errors: {errors} | Speed: {rate:.1f} files/s") + sys.stdout.flush() + + pool.close() + pool.join() + + except KeyboardInterrupt: + print("\n\n[-] User aborted. Force terminating all child processes...") + pool.terminate() + pool.join() + + except Exception as e: + print(f"\n[-] Critical error: {e}") + pool.terminate() + pool.join() + + finally: + if batch_buffer: + try: + with open(JSONL_PATH, 'a', encoding='utf-8') as f_out: + f_out.writelines(batch_buffer) + except Exception as save_err: + print(f"\n[-] Failed to save remaining buffer: {save_err}") + + print(f"\n\n[+] Scan complete. Total added: {count}, Total skipped: {skipped}") + +#################################################################################################### + +if __name__ == "__main__": + multiprocessing.freeze_support() + + print("\n---------------- PE Dataset Builder v4.0 ----------------\n") + + existing_hashes = load_existing_hashes() + print(f"[+] Found {len(existing_hashes)} existing hashes.\n") + + try: + while True: + print("-" * 57) + target_path = input("\n[*] Enter path (file or directory): ").strip().strip('"').strip("'") + if target_path.lower() in ['exit', 'q']: + break + + if not target_path or not os.path.exists(target_path): + print("[-] Invalid path.") + continue + + while True: + label_input = input("[*] Enter label (0=Safe, 1=Malware): ").strip() + if label_input.lower() in ['exit', 'q']: + sys.exit() + if label_input in ['0', '1']: + break + + print("[-] Please enter 0 or 1.") + + scan_and_save(target_path, int(label_input), existing_hashes) + + except KeyboardInterrupt: + print("\n\n[-] Aborted.") diff --git a/Engine/Properties/scanner.py b/Engine/Properties/scanner.py index 1d749c3..aaafec4 100644 --- a/Engine/Properties/scanner.py +++ b/Engine/Properties/scanner.py @@ -1,942 +1,1156 @@ -import os, re, sys, math, json, zlib, mmap, pefile, datetime -import ctypes, ctypes.wintypes -import numpy as np -import onnxruntime as ort - -#################################################################################################### - -MODEL_FILE = "Pefile_General_T1.onnx" -FEATURE_FILE = "features.json" -MAX_FILE_SIZE = 4 * 1024 * 1024 * 1024 - -#################################################################################################### - -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 - -#################################################################################################### - -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): - if not data: - return 0.0 - - arr = np.frombuffer(data, dtype=np.uint8) - sz = len(arr) - counts = np.bincount(arr, minlength=256) - counts = counts[counts > 0] - - return float(np.log2(sz) - np.sum(counts * np.log2(counts)) / sz) - - @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): - 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 = [] - 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: - data = pe.get_data(entry_lang.data.struct.OffsetToData, entry_lang.data.struct.Size) - entropies.append(cls._calc_entropy(data)) - except Exception: - continue - - 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_bytes, fsize): - res = {"HasSignature": 0.0, "SignatureCount": 0.0} - if not hasattr(pe, 'OPTIONAL_HEADER') or not hasattr(pe.OPTIONAL_HEADER, 'DATA_DIRECTORY'): - return res - - sec_idx = 4 - directories = pe.OPTIONAL_HEADER.DATA_DIRECTORY - if len(directories) <= sec_idx: - return res - - sec_dir = directories[sec_idx] - if sec_dir.Size == 0 or sec_dir.VirtualAddress == 0: - return res - - res["HasSignature"] = 1.0 - offset = sec_dir.VirtualAddress - size = sec_dir.Size - - if offset + size > fsize or offset < 0 or size < 0: - return res - - count = 0 - curr = offset - end = offset + size - - while curr + 8 <= end: - dw_length = int.from_bytes(file_bytes[curr:curr+4], 'little') - if dw_length < 8: - break - - count += 1 - curr += (dw_length + 7) & ~7 - - res["SignatureCount"] = float(count) - return res - - @classmethod - def extract(cls, file_path): - pe = None - mm = None - source_file = None - - try: - fsize = os.path.getsize(file_path) - if fsize == 0 or fsize > MAX_FILE_SIZE: - return None - - source_file = open(file_path, "rb") - mm = mmap.mmap(source_file.fileno(), 0, access=mmap.ACCESS_READ) - file_bytes = mm - - base, dlls, apis = cls._base_defaults(), set(), set() - pe = pefile.PE(data=file_bytes, 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 - - 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 in 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 - - try: - s_data = section.get_data() - s_entropy = cls._calc_entropy(s_data) - except Exception: - s_entropy = 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(cls._extract_strings(file_bytes)) - base.update(cls._extract_histograms(file_bytes)) - base.update(cls._extract_overlay_features(pe, fsize)) - base.update(cls._extract_rich_header(pe)) - base.update(cls._extract_ep_anomalies(pe)) - base.update(cls._extract_advanced_resources(pe)) - base.update(cls._extract_load_config(pe)) - base.update(cls._extract_security_directory(pe, file_bytes, 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 - if mm is not None: - try: - mm.close() - except Exception: - pass - if source_file is not None: - try: - source_file.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: - files.append(os.path.join(r, f)) - - if not files: - print("[-] No valid PE files found.") - return - - log_buffer = [] - - def _log(message): - print(message) - log_buffer.append(message) - - _log(f"[*] Scanning {len(files)} files...\n") - _log(f"{'RESULT':<10} | {'PROB':<8} | {'FILE'}") - _log("-" * 80) - - safe_count = 0 - mal_count = 0 - unsupport_count = 0 - error_count = 0 - - try: - for fpath in files: - try: - data = FeatureExtractor.extract(fpath) - if not data: - _log(f"{'UNSUPPORT':<10} | {'-':<8} | {fpath}") - unsupport_count += 1 - continue - - prob = predictor.predict(data) - is_malware = prob > 0.5 - - label = "MALWARE" if is_malware else "SAFE" - if is_malware: - mal_count += 1 - else: - safe_count += 1 - - _log(f"{label:<10} | {prob:.4f} | {fpath}") - - except Exception as e: - _log(f"{'FAIL':<10} | {'-':<8} | {fpath} ({str(e)})") - error_count += 1 - - except KeyboardInterrupt: - _log("\n[-] Scan aborted by user (Ctrl+C).\n") - - _log("-" * 80) - _log(f"\n[*] Summary: Safe={safe_count}, Malware={mal_count}, Unsupport={unsupport_count}, Error={error_count}, Total={len(files)}") - - log_filename = datetime.datetime.now().strftime("%Y%m%d_%H%M%S.log") - try: - with open(log_filename, "w", encoding="utf-8") as f: - f.write("\n".join(log_buffer) + "\n") - print(f"[*] Log saved successfully: {log_filename}") - except Exception as e: - print(f"[-] Failed to save log: {e}") - -#################################################################################################### - -if __name__ == "__main__": - base_dir = os.path.dirname(sys.executable) if getattr(sys, 'frozen', False) else os.path.dirname(os.path.abspath(__file__)) - - custom_model = MODEL_FILE - targets = [] - - args = sys.argv[1:] - i = 0 - while i < len(args): - if args[i] in ("-m", "-model", "--model") and i + 1 < len(args): - custom_model = args[i+1] - i += 2 - else: - targets.append(args[i]) - i += 1 - - print("\n-------------------------- PE Malware Predictor v4.0 --------------------------\n") - print(f"[*] © 2020-2026 87owo (PYAS Security)") - print(f"[*] Official Website: https://github.com/87owo/PYAS") - print(f"[*] Loading model: {custom_model}") - - model_path = os.path.join(base_dir, custom_model) if not os.path.isabs(custom_model) else custom_model - feature_path = os.path.join(base_dir, FEATURE_FILE) - - try: - predictor = ModelPredictor(model_path, feature_path) - except Exception as e: - print(f"[-] Critical Error: {e}") - sys.exit(1) - - if targets: - for path in targets: - target = path.strip('"').strip("'") - if os.path.exists(target): - print() - print("-" * 80) - scan_target(target, predictor) - else: - print(f"[-] Path does not exist: {target}") - else: - while True: - try: - print() - print("-" * 80) - path = input("\n[?] 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("[-] Path does not exist.") - continue - - scan_target(path, predictor) - - except KeyboardInterrupt: - break - except Exception as e: - print(f"[-] Error: {e}") +import os, re, sys, math, json, zlib, pefile, datetime +import ctypes, ctypes.wintypes +import numpy as np +import onnxruntime as ort + +#################################################################################################### + +MODEL_FILE = "Pefile_General_T1.onnx" +FEATURE_FILE = "features.json" +MAX_FILE_SIZE = 4 * 1024 * 1024 * 1024 + +#################################################################################################### + +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: + files.append(os.path.join(r, f)) + + if not files: + print("[-] No valid PE files found.") + return + + log_buffer = [] + + def _log(message): + print(message) + log_buffer.append(message) + + _log(f"[*] Scanning {len(files)} files...\n") + _log(f"{'RESULT':<10} | {'PROB':<8} | {'FILE'}") + _log("-" * 80) + + safe_count = 0 + mal_count = 0 + unsupport_count = 0 + error_count = 0 + + try: + for fpath in files: + try: + data = FeatureExtractor.extract(fpath) + if not data: + _log(f"{'UNSUPPORT':<10} | {'-':<8} | {fpath}") + unsupport_count += 1 + continue + + prob = predictor.predict(data) + is_malware = prob > 0.5 + + label = "MALWARE" if is_malware else "SAFE" + if is_malware: + mal_count += 1 + else: + safe_count += 1 + + _log(f"{label:<10} | {prob:.4f} | {fpath}") + + except Exception as e: + _log(f"{'FAIL':<10} | {'-':<8} | {fpath} ({str(e)})") + error_count += 1 + + except KeyboardInterrupt: + _log("\n[-] Scan aborted by user (Ctrl+C).\n") + + _log("-" * 80) + _log(f"\n[*] Summary: Safe={safe_count}, Malware={mal_count}, Unsupport={unsupport_count}, Error={error_count}, Total={len(files)}") + + log_filename = datetime.datetime.now().strftime("%Y%m%d_%H%M%S.log") + try: + with open(log_filename, "w", encoding="utf-8") as f: + f.write("\n".join(log_buffer) + "\n") + print(f"[*] Log saved successfully: {log_filename}") + except Exception as e: + print(f"[-] Failed to save log: {e}") + +#################################################################################################### + +if __name__ == "__main__": + base_dir = os.path.dirname(sys.executable) if getattr(sys, 'frozen', False) else os.path.dirname(os.path.abspath(__file__)) + + custom_model = MODEL_FILE + targets = [] + + args = sys.argv[1:] + i = 0 + while i < len(args): + if args[i] in ("-m", "-model", "--model") and i + 1 < len(args): + custom_model = args[i+1] + i += 2 + else: + targets.append(args[i]) + i += 1 + + print("\n-------------------------- PE Malware Predictor v4.0 --------------------------\n") + print(f"[*] © 2020-2026 87owo (PYAS Security)") + print(f"[*] Official Website: https://github.com/87owo/PYAS") + print(f"[*] Loading model: {custom_model}") + + model_path = os.path.join(base_dir, custom_model) if not os.path.isabs(custom_model) else custom_model + feature_path = os.path.join(base_dir, FEATURE_FILE) + + try: + predictor = ModelPredictor(model_path, feature_path) + except Exception as e: + print(f"[-] Critical Error: {e}") + sys.exit(1) + + if targets: + for path in targets: + target = path.strip('"').strip("'") + if os.path.exists(target): + print() + print("-" * 80) + scan_target(target, predictor) + else: + print(f"[-] Path does not exist: {target}") + else: + while True: + try: + print() + print("-" * 80) + path = input("\n[?] 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("[-] Path does not exist.") + continue + + scan_target(path, predictor) + + except KeyboardInterrupt: + break + except Exception as e: + print(f"[-] Error: {e}")