From ef5793fb15c368f3718c1874f6fa2c1a2fd9fd15 Mon Sep 17 00:00:00 2001 From: 87owo <85057800+87owo@users.noreply.github.com> Date: Tue, 26 May 2026 09:30:27 +0800 Subject: [PATCH] Add files via upload --- Engine/Properties/convert.py | 138 ++++++++++++++++-------------- Engine/Properties/scanner.py | 118 +++++++++++++------------ Engine/Properties/train_server.py | 91 +++++++++++++++----- Engine/Properties/train_user.py | 93 +++++++++++++++----- 4 files changed, 274 insertions(+), 166 deletions(-) diff --git a/Engine/Properties/convert.py b/Engine/Properties/convert.py index b4fb3f7..b6dc22b 100644 --- a/Engine/Properties/convert.py +++ b/Engine/Properties/convert.py @@ -1,6 +1,5 @@ -import os, sys, math, time, datetime, pefile, hashlib -import ctypes, ctypes.wintypes, multiprocessing -import re, zlib +import re, os, sys, math, time, zlib, datetime, hashlib +import ctypes, ctypes.wintypes, pefile, multiprocessing import numpy as np #################################################################################################### @@ -110,6 +109,17 @@ def verify_signature(file_path): class FeatureExtractor: _STRING_PATTERN = re.compile(b'[\x20-\x7E]{5,}') + + _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)) @staticmethod def _safe_float(val): @@ -123,15 +133,15 @@ class FeatureExtractor: 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) - probs = counts[counts > 0] / len(arr) - return float(-np.sum(probs * np.log2(probs))) + 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} - for i in range(95): - res[f'StringHist_{i:02d}'] = 0.0 + res.update(dict.fromkeys(cls._STRING_KEYS, 0.0)) matches = cls._STRING_PATTERN.findall(file_bytes) if not matches: @@ -149,9 +159,7 @@ class FeatureExtractor: counts = np.bincount(arr, minlength=127) char_counts = counts[0x20:0x7F] - for i in range(95): - res[f'StringHist_{i:02d}'] = float(char_counts[i]) / total_len - + res.update(dict(zip(cls._STRING_KEYS, char_counts.astype(np.float64) / total_len))) return res @classmethod @@ -161,15 +169,13 @@ class FeatureExtractor: res = {} if sz == 0: - for i in range(256): - res[f'ByteHist_{i:02X}'] = 0.0 - res[f'ByteEnt_{i:02X}'] = 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.bincount(arr, minlength=256) - byte_hist = byte_counts / sz - for i, v in enumerate(byte_hist): - res[f'ByteHist_{i:02X}'] = float(v) + byte_hist = byte_counts.astype(np.float64) / sz + res.update(dict(zip(cls._BYTE_HIST_KEYS, byte_hist))) ent_hist = np.zeros(256, dtype=np.float64) window = 2048 @@ -177,35 +183,36 @@ class FeatureExtractor: if sz < window: w_counts = byte_counts - p = w_counts[w_counts > 0] / sz - entropy = -np.sum(p * np.log2(p)) + p = w_counts[w_counts > 0] / float(sz) + entropy = float(-np.sum(p * np.log2(p))) ent_bin = min(int(entropy * 2.0), 15) byte_bins = w_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 - batch_size = 50000 + batch_windows = 10000 - for b_start in range(0, num_windows, batch_size): - b_end = min(b_start + batch_size, num_windows) - b_num = b_end - b_start + for w_start in range(0, num_windows, batch_windows): + w_end = min(w_start + batch_windows, num_windows) + b_num = w_end - w_start - b_chunks = arr[b_start * step : (b_end + 1) * step].reshape(b_num + 1, step) + start_byte = w_start * step + end_byte = start_byte + b_num * step + step + b_arr = arr[start_byte:end_byte] - offsets = np.arange(b_num + 1, dtype=np.int32)[:, None] * 256 - flat_idx = (b_chunks.astype(np.int32) + offsets).ravel() - counts_flat = np.bincount(flat_idx, minlength=(b_num + 1) * 256) - chunk_counts = counts_flat.reshape(b_num + 1, 256) + arr_2d = b_arr.reshape(b_num + 1, step) + offsets = (np.arange(b_num + 1, dtype=np.int32)[:, None] * 256) + flat_idx = (arr_2d.astype(np.int32) + offsets).ravel() + chunk_counts = np.bincount(flat_idx, minlength=(b_num + 1) * 256).reshape(b_num + 1, 256) window_counts = chunk_counts[:-1] + chunk_counts[1:] - p = window_counts / float(window) - p_safe = np.where(p > 0, p, 1.0) - entropies = -np.sum(p * np.log2(p_safe), axis=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, None, 15, out=ent_bins) + 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() @@ -215,9 +222,7 @@ class FeatureExtractor: if ent_sum > 0: ent_hist /= ent_sum - for i, v in enumerate(ent_hist): - res[f'ByteEnt_{i:02X}'] = float(v) - + res.update(dict(zip(cls._BYTE_ENT_KEYS, ent_hist))) return res @classmethod @@ -386,9 +391,9 @@ class FeatureExtractor: except Exception: pass - base['TrustSigned'] = verify_signature(file_path) + base['TrustSigned'] = float(verify_signature(file_path)) base['FileEntropy'] = cls._calc_entropy(file_bytes) - base['FileSize'] = fsize + base['FileSize'] = float(fsize) fh = pe.FILE_HEADER base['Machine'] = getattr(fh, 'Machine', 0) @@ -423,13 +428,9 @@ class FeatureExtractor: 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 - char_flags = [0x00000020, 0x00000040, 0x00000080, 0x02000000, 0x20000000, 0x40000000, 0x80000000] - for flag in char_flags: - base[f'Char_{flag:08X}_Count'] = 0.0 - base[f'Char_{flag:08X}_MeanEntropy'] = 0.0 - - for i in range(50): - base[f'SectionHash_{i:02d}'] = 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) @@ -439,7 +440,7 @@ class FeatureExtractor: for section in pe.sections: sec_name = section.Name.rstrip(b'\x00') hash_val = zlib.crc32(sec_name) % 50 - base[f'SectionHash_{hash_val:02d}'] += 1.0 + res_sec_hash[cls._SECTION_HASH_KEYS[hash_val]] += 1.0 try: s_data = section.get_data() @@ -456,10 +457,10 @@ class FeatureExtractor: if section.Characteristics & 0x40000000: read_sec += 1 if section.SizeOfRawData + section.PointerToRawData > fsize: sec_exc = 1 - for flag in char_flags: + for flag in cls._CHAR_FLAGS: if section.Characteristics & flag: - base[f'Char_{flag:08X}_Count'] += 1.0 - base[f'Char_{flag:08X}_MeanEntropy'] += s_entropy + 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 @@ -470,14 +471,19 @@ class FeatureExtractor: 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'] = exec_sec - base['WritableSections'] = write_sec - base['ReadableSections'] = read_sec + base['ExecutableSections'] = float(exec_sec) + base['WritableSections'] = float(write_sec) + base['ReadableSections'] = float(read_sec) base['SectionException'] = float(sec_exc) - for flag in char_flags: - if base[f'Char_{flag:08X}_Count'] > 0: - base[f'Char_{flag:08X}_MeanEntropy'] = base[f'Char_{flag:08X}_MeanEntropy'] / base[f'Char_{flag:08X}_Count'] + 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: @@ -509,7 +515,7 @@ class FeatureExtractor: base['IsSpecialBuild'] = 1.0 if flags & 0x20 else 0.0 if hasattr(pe, 'DIRECTORY_ENTRY_IMPORT'): - base['ImportCount'] = len(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): @@ -520,20 +526,20 @@ class FeatureExtractor: if getattr(imp, 'name', None): try: apis.add(imp.name.decode('ascii', 'ignore')) except Exception: pass - base['ImportFunctionCount'] = func_count + base['ImportFunctionCount'] = float(func_count) - base['ExportCount'] = len(pe.DIRECTORY_ENTRY_EXPORT.symbols) if (hasattr(pe, 'DIRECTORY_ENTRY_EXPORT') and hasattr(pe.DIRECTORY_ENTRY_EXPORT, 'symbols')) else 0.0 + 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'] = icon_count + base['IconCount'] = float(icon_count) else: base['IconCount'] = 0.0 - base['DebugCount'] = len(pe.DIRECTORY_ENTRY_DEBUG) if hasattr(pe, 'DIRECTORY_ENTRY_DEBUG') else 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)) @@ -555,13 +561,15 @@ class FeatureExtractor: def load_existing_hashes(): hashes = set() - if os.path.exists(JSONL_PATH): - with open(JSONL_PATH, 'r', encoding='utf-8') as f: - for line in f: - try: - data = JSON_LOADS(line) - if 'FileHash' in data: hashes.add(data['FileHash']) - except Exception: pass + 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() diff --git a/Engine/Properties/scanner.py b/Engine/Properties/scanner.py index 34991ce..3ec96c3 100644 --- a/Engine/Properties/scanner.py +++ b/Engine/Properties/scanner.py @@ -1,8 +1,7 @@ -import os, sys, math, json, pefile, datetime +import os, re, sys, math, json, zlib, pefile, datetime +import ctypes, ctypes.wintypes import numpy as np import onnxruntime as ort -import ctypes, ctypes.wintypes -import re, zlib #################################################################################################### @@ -100,6 +99,17 @@ class WinTrust: class FeatureExtractor: _STRING_PATTERN = re.compile(b'[\x20-\x7E]{5,}') + + _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)) @staticmethod def _safe_float(val): @@ -113,15 +123,15 @@ class FeatureExtractor: 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) - probs = counts[counts > 0] / len(arr) - return float(-np.sum(probs * np.log2(probs))) + 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} - for i in range(95): - res[f'StringHist_{i:02d}'] = 0.0 + res.update(dict.fromkeys(cls._STRING_KEYS, 0.0)) matches = cls._STRING_PATTERN.findall(file_bytes) if not matches: @@ -139,9 +149,7 @@ class FeatureExtractor: counts = np.bincount(arr, minlength=127) char_counts = counts[0x20:0x7F] - for i in range(95): - res[f'StringHist_{i:02d}'] = float(char_counts[i]) / total_len - + res.update(dict(zip(cls._STRING_KEYS, char_counts.astype(np.float64) / total_len))) return res @classmethod @@ -151,15 +159,13 @@ class FeatureExtractor: res = {} if sz == 0: - for i in range(256): - res[f'ByteHist_{i:02X}'] = 0.0 - res[f'ByteEnt_{i:02X}'] = 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.bincount(arr, minlength=256) - byte_hist = byte_counts / sz - for i, v in enumerate(byte_hist): - res[f'ByteHist_{i:02X}'] = float(v) + byte_hist = byte_counts.astype(np.float64) / sz + res.update(dict(zip(cls._BYTE_HIST_KEYS, byte_hist))) ent_hist = np.zeros(256, dtype=np.float64) window = 2048 @@ -167,35 +173,36 @@ class FeatureExtractor: if sz < window: w_counts = byte_counts - p = w_counts[w_counts > 0] / sz - entropy = -np.sum(p * np.log2(p)) + p = w_counts[w_counts > 0] / float(sz) + entropy = float(-np.sum(p * np.log2(p))) ent_bin = min(int(entropy * 2.0), 15) byte_bins = w_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 - batch_size = 50000 + batch_windows = 10000 - for b_start in range(0, num_windows, batch_size): - b_end = min(b_start + batch_size, num_windows) - b_num = b_end - b_start + for w_start in range(0, num_windows, batch_windows): + w_end = min(w_start + batch_windows, num_windows) + b_num = w_end - w_start - b_chunks = arr[b_start * step : (b_end + 1) * step].reshape(b_num + 1, step) + start_byte = w_start * step + end_byte = start_byte + b_num * step + step + b_arr = arr[start_byte:end_byte] - offsets = np.arange(b_num + 1, dtype=np.int32)[:, None] * 256 - flat_idx = (b_chunks.astype(np.int32) + offsets).ravel() - counts_flat = np.bincount(flat_idx, minlength=(b_num + 1) * 256) - chunk_counts = counts_flat.reshape(b_num + 1, 256) + arr_2d = b_arr.reshape(b_num + 1, step) + offsets = (np.arange(b_num + 1, dtype=np.int32)[:, None] * 256) + flat_idx = (arr_2d.astype(np.int32) + offsets).ravel() + chunk_counts = np.bincount(flat_idx, minlength=(b_num + 1) * 256).reshape(b_num + 1, 256) window_counts = chunk_counts[:-1] + chunk_counts[1:] - p = window_counts / float(window) - p_safe = np.where(p > 0, p, 1.0) - entropies = -np.sum(p * np.log2(p_safe), axis=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, None, 15, out=ent_bins) + 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() @@ -205,9 +212,7 @@ class FeatureExtractor: if ent_sum > 0: ent_hist /= ent_sum - for i, v in enumerate(ent_hist): - res[f'ByteEnt_{i:02X}'] = float(v) - + res.update(dict(zip(cls._BYTE_ENT_KEYS, ent_hist))) return res @classmethod @@ -419,13 +424,9 @@ class FeatureExtractor: 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 - char_flags = [0x00000020, 0x00000040, 0x00000080, 0x02000000, 0x20000000, 0x40000000, 0x80000000] - for flag in char_flags: - base[f'Char_{flag:08X}_Count'] = 0.0 - base[f'Char_{flag:08X}_MeanEntropy'] = 0.0 - - for i in range(50): - base[f'SectionHash_{i:02d}'] = 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) @@ -435,7 +436,7 @@ class FeatureExtractor: for section in pe.sections: sec_name = section.Name.rstrip(b'\x00') hash_val = zlib.crc32(sec_name) % 50 - base[f'SectionHash_{hash_val:02d}'] += 1.0 + res_sec_hash[cls._SECTION_HASH_KEYS[hash_val]] += 1.0 try: s_data = section.get_data() @@ -452,10 +453,10 @@ class FeatureExtractor: if section.Characteristics & 0x40000000: read_sec += 1 if section.SizeOfRawData + section.PointerToRawData > fsize: sec_exc = 1 - for flag in char_flags: + for flag in cls._CHAR_FLAGS: if section.Characteristics & flag: - base[f'Char_{flag:08X}_Count'] += 1.0 - base[f'Char_{flag:08X}_MeanEntropy'] += s_entropy + 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 @@ -466,14 +467,19 @@ class FeatureExtractor: 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'] = exec_sec - base['WritableSections'] = write_sec - base['ReadableSections'] = read_sec + base['ExecutableSections'] = float(exec_sec) + base['WritableSections'] = float(write_sec) + base['ReadableSections'] = float(read_sec) base['SectionException'] = float(sec_exc) - for flag in char_flags: - if base[f'Char_{flag:08X}_Count'] > 0: - base[f'Char_{flag:08X}_MeanEntropy'] = base[f'Char_{flag:08X}_MeanEntropy'] / base[f'Char_{flag:08X}_Count'] + 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: @@ -505,7 +511,7 @@ class FeatureExtractor: base['IsSpecialBuild'] = 1.0 if flags & 0x20 else 0.0 if hasattr(pe, 'DIRECTORY_ENTRY_IMPORT'): - base['ImportCount'] = len(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): @@ -516,20 +522,20 @@ class FeatureExtractor: if getattr(imp, 'name', None): try: apis.add(imp.name.decode('ascii', 'ignore')) except Exception: pass - base['ImportFunctionCount'] = func_count + base['ImportFunctionCount'] = float(func_count) - base['ExportCount'] = len(pe.DIRECTORY_ENTRY_EXPORT.symbols) if (hasattr(pe, 'DIRECTORY_ENTRY_EXPORT') and hasattr(pe.DIRECTORY_ENTRY_EXPORT, 'symbols')) else 0.0 + 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'] = icon_count + base['IconCount'] = float(icon_count) else: base['IconCount'] = 0.0 - base['DebugCount'] = len(pe.DIRECTORY_ENTRY_DEBUG) if hasattr(pe, 'DIRECTORY_ENTRY_DEBUG') else 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)) diff --git a/Engine/Properties/train_server.py b/Engine/Properties/train_server.py index 1e45966..c8c705f 100644 --- a/Engine/Properties/train_server.py +++ b/Engine/Properties/train_server.py @@ -1,4 +1,4 @@ -import os, sys, time, json, onnxmltools, zlib +import os, sys, time, json, zlib, array, onnxmltools, multiprocessing import numpy as np import pandas as pd import lightgbm as lgb @@ -55,14 +55,15 @@ def analyze_schema(jsonl_path): with open(jsonl_path, 'rb') as f: for line in f: - try: - data = JSON_LOADS(line) - base_data = data.get('Base') - if base_data: - base_keys.update(base_data.keys()) - num_lines += 1 - except Exception: - continue + num_lines += 1 + if not base_keys: + try: + data = JSON_LOADS(line) + base_data = data.get('Base') + if base_data: + base_keys.update(base_data.keys()) + except Exception: + pass base_schema = sorted([k for k in base_keys if k.lower().strip() not in EXCLUDE_FEATURES]) dll_schema = [f"DllHash_{i:03d}" for i in range(DLL_HASH_DIM)] @@ -73,28 +74,69 @@ def analyze_schema(jsonl_path): #################################################################################################### def load_data_efficiently(jsonl_path, base_schema, dll_schema, api_schema, num_lines): - print(f"[*] Stage 2: Building sparse feature matrix dynamically (Expected samples: {num_lines})...") + print(f"[*] Stage 2: Building sparse feature matrix dynamically via Multiprocessing...") final_features = base_schema + dll_schema + api_schema num_features = len(final_features) - - base_idx = {feat: i for i, feat in enumerate(base_schema)} offset_dll = len(base_schema) offset_api = offset_dll + DLL_HASH_DIM - rows, cols, vals = [], [], [] - y = np.zeros(num_lines, dtype=np.int8) + file_size = os.path.getsize(jsonl_path) + workers = min(os.cpu_count() or 4, 16) + chunk_size = file_size // workers + + tasks = [] + for i in range(workers): + start = i * chunk_size + end = file_size if i == workers - 1 else (i + 1) * chunk_size + tasks.append((jsonl_path, start, end, base_schema, num_features, offset_dll, offset_api)) + + with multiprocessing.Pool(workers) as pool: + results = pool.map(_chunk_worker, tasks) + + Xs = [r[0] for r in results if r[0].shape[0] > 0] + ys = [r[1] for r in results if r[1].shape[0] > 0] + + X_final = sp.vstack(Xs, format='csr') if Xs else sp.csr_matrix((0, num_features), dtype=np.float32) + y_final = np.concatenate(ys) if ys else np.array([], dtype=np.int8) + + return X_final, pd.Series(y_final), final_features - with open(jsonl_path, 'rb') as f: +def _chunk_worker(args): + filepath, start_byte, end_byte, base_schema, num_features, offset_dll, offset_api = args + + base_idx = {feat: i for i, feat in enumerate(base_schema)} + base_idx_get = base_idx.get + + rows = array.array('i') + cols = array.array('i') + vals = array.array('f') + y_list = array.array('b') + + with open(filepath, 'rb') as f: + if start_byte > 0: + f.seek(start_byte - 1) + if f.read(1) != b'\n': + f.seek(start_byte) + f.readline() + else: + f.seek(start_byte) + else: + f.seek(0) + row_idx = 0 - for line in f: + while f.tell() < end_byte: + line = f.readline() + if not line: + break + try: data = JSON_LOADS(line) - y[row_idx] = data.get('Label', 0) + y_list.append(data.get('Label', 0)) base_data = data.get('Base') if base_data: for k, v in base_data.items(): - idx = base_idx.get(k) + idx = base_idx_get(k) if idx is not None: val = float(v) if val != 0.0: @@ -121,11 +163,15 @@ def load_data_efficiently(jsonl_path, base_schema, dll_schema, api_schema, num_l row_idx += 1 except Exception: continue - - y = y[:row_idx] - X = sp.csr_matrix((vals, (rows, cols)), shape=(row_idx, num_features), dtype=np.float32) + + y_arr = np.frombuffer(y_list, dtype=np.int8) + np_rows = np.frombuffer(rows, dtype=np.int32) + np_cols = np.frombuffer(cols, dtype=np.int32) + np_vals = np.frombuffer(vals, dtype=np.float32) - return X, pd.Series(y), final_features + X_local = sp.csr_matrix((np_vals, (np_rows, np_cols)), shape=(row_idx, num_features), dtype=np.float32) + + return X_local, y_arr #################################################################################################### @@ -233,6 +279,7 @@ def evaluate_and_save(model, X_test, y_test): #################################################################################################### if __name__ == "__main__": + multiprocessing.freeze_support() print("\n---------------- PE Malware Trainer v4.0 ----------------\n") if not os.path.exists(JSONL_PATH): diff --git a/Engine/Properties/train_user.py b/Engine/Properties/train_user.py index 73782b5..8c84fb0 100644 --- a/Engine/Properties/train_user.py +++ b/Engine/Properties/train_user.py @@ -1,4 +1,4 @@ -import os, sys, time, json, onnxmltools, zlib +import os, sys, time, json, zlib, array, onnxmltools, multiprocessing import numpy as np import pandas as pd import lightgbm as lgb @@ -35,7 +35,7 @@ LGBM_PARAMS = { 'objective': 'binary', 'metric': ['binary_logloss', 'auc'], 'boosting_type': 'gbdt', - 'num_leaves': 128, + 'num_leaves': 64, 'feature_fraction': 0.8, 'bagging_fraction': 1.0, 'bagging_freq': 5, @@ -55,14 +55,15 @@ def analyze_schema(jsonl_path): with open(jsonl_path, 'rb') as f: for line in f: - try: - data = JSON_LOADS(line) - base_data = data.get('Base') - if base_data: - base_keys.update(base_data.keys()) - num_lines += 1 - except Exception: - continue + num_lines += 1 + if not base_keys: + try: + data = JSON_LOADS(line) + base_data = data.get('Base') + if base_data: + base_keys.update(base_data.keys()) + except Exception: + pass base_schema = sorted([k for k in base_keys if k.lower().strip() not in EXCLUDE_FEATURES]) dll_schema = [f"DllHash_{i:03d}" for i in range(DLL_HASH_DIM)] @@ -73,28 +74,69 @@ def analyze_schema(jsonl_path): #################################################################################################### def load_data_efficiently(jsonl_path, base_schema, dll_schema, api_schema, num_lines): - print(f"[*] Stage 2: Building sparse feature matrix dynamically (Expected samples: {num_lines})...") + print(f"[*] Stage 2: Building sparse feature matrix dynamically via Multiprocessing...") final_features = base_schema + dll_schema + api_schema num_features = len(final_features) - - base_idx = {feat: i for i, feat in enumerate(base_schema)} offset_dll = len(base_schema) offset_api = offset_dll + DLL_HASH_DIM - rows, cols, vals = [], [], [] - y = np.zeros(num_lines, dtype=np.int8) + file_size = os.path.getsize(jsonl_path) + workers = min(os.cpu_count() or 4, 16) + chunk_size = file_size // workers + + tasks = [] + for i in range(workers): + start = i * chunk_size + end = file_size if i == workers - 1 else (i + 1) * chunk_size + tasks.append((jsonl_path, start, end, base_schema, num_features, offset_dll, offset_api)) + + with multiprocessing.Pool(workers) as pool: + results = pool.map(_chunk_worker, tasks) + + Xs = [r[0] for r in results if r[0].shape[0] > 0] + ys = [r[1] for r in results if r[1].shape[0] > 0] + + X_final = sp.vstack(Xs, format='csr') if Xs else sp.csr_matrix((0, num_features), dtype=np.float32) + y_final = np.concatenate(ys) if ys else np.array([], dtype=np.int8) + + return X_final, pd.Series(y_final), final_features - with open(jsonl_path, 'rb') as f: +def _chunk_worker(args): + filepath, start_byte, end_byte, base_schema, num_features, offset_dll, offset_api = args + + base_idx = {feat: i for i, feat in enumerate(base_schema)} + base_idx_get = base_idx.get + + rows = array.array('i') + cols = array.array('i') + vals = array.array('f') + y_list = array.array('b') + + with open(filepath, 'rb') as f: + if start_byte > 0: + f.seek(start_byte - 1) + if f.read(1) != b'\n': + f.seek(start_byte) + f.readline() + else: + f.seek(start_byte) + else: + f.seek(0) + row_idx = 0 - for line in f: + while f.tell() < end_byte: + line = f.readline() + if not line: + break + try: data = JSON_LOADS(line) - y[row_idx] = data.get('Label', 0) + y_list.append(data.get('Label', 0)) base_data = data.get('Base') if base_data: for k, v in base_data.items(): - idx = base_idx.get(k) + idx = base_idx_get(k) if idx is not None: val = float(v) if val != 0.0: @@ -121,11 +163,15 @@ def load_data_efficiently(jsonl_path, base_schema, dll_schema, api_schema, num_l row_idx += 1 except Exception: continue - - y = y[:row_idx] - X = sp.csr_matrix((vals, (rows, cols)), shape=(row_idx, num_features), dtype=np.float32) + + y_arr = np.frombuffer(y_list, dtype=np.int8) + np_rows = np.frombuffer(rows, dtype=np.int32) + np_cols = np.frombuffer(cols, dtype=np.int32) + np_vals = np.frombuffer(vals, dtype=np.float32) - return X, pd.Series(y), final_features + X_local = sp.csr_matrix((np_vals, (np_rows, np_cols)), shape=(row_idx, num_features), dtype=np.float32) + + return X_local, y_arr #################################################################################################### @@ -233,6 +279,7 @@ def evaluate_and_save(model, X_test, y_test): #################################################################################################### if __name__ == "__main__": + multiprocessing.freeze_support() print("\n---------------- PE Malware Trainer v4.0 ----------------\n") if not os.path.exists(JSONL_PATH):