diff --git a/Engine/Properties/convert.py b/Engine/Properties/convert.py index 65f3174..ef472f3 100644 --- a/Engine/Properties/convert.py +++ b/Engine/Properties/convert.py @@ -1,5 +1,8 @@ import os, sys, math, time, datetime, pefile, hashlib import ctypes, ctypes.wintypes, multiprocessing +import re, zlib + +#################################################################################################### try: import orjson @@ -12,8 +15,6 @@ except ImportError: JSON_LOADS = json.loads HAS_ORJSON = False -#################################################################################################### - JSONL_PATH = "pe_features.jsonl" BATCH_SIZE = 1000 MAX_FILE_SIZE = 4 * 1024 * 1024 * 1024 @@ -107,6 +108,8 @@ def verify_signature(file_path): #################################################################################################### class FeatureExtractor: + _STRING_PATTERN = re.compile(b'[\x20-\x7E]{5,}') + @staticmethod def _safe_float(val): try: @@ -124,6 +127,86 @@ class FeatureExtractor: probs = counts[counts > 0] / len(arr) return float(-np.sum(probs * np.log2(probs))) + @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 + + count = 0 + total_len = 0 + char_counts = [0] * 95 + combined_data = bytearray() + + for match in cls._STRING_PATTERN.finditer(file_bytes): + m = match.group() + count += 1 + length = len(m) + total_len += length + combined_data.extend(m) + for b in m: + char_counts[b - 0x20] += 1 + + if count == 0: + return res + + res["StringCount"] = float(count) + res["StringMeanLength"] = float(total_len) / count + res["StringEntropy"] = cls._calc_entropy(combined_data) + + for i in range(95): + res[f'StringHist_{i:02d}'] = float(char_counts[i]) / total_len + + return res + + @classmethod + def _extract_histograms(cls, file_bytes): + import numpy as np + arr = np.frombuffer(file_bytes, dtype=np.uint8) + sz = len(arr) + res = {} + + if sz == 0: + for i in range(256): + res[f'ByteHist_{i:02X}'] = 0.0 + res[f'ByteEnt_{i:02X}'] = 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) + + ent_hist = np.zeros(256, dtype=np.float64) + window = 2048 + step = 1024 + num_windows = (sz - window) // step + 1 if sz >= window else 1 + + for i in range(num_windows): + w = arr if sz < window else arr[i*step : i*step+window] + w_counts = np.bincount(w, minlength=256) + + p = w_counts[w_counts > 0] / len(w) + entropy = -np.sum(p * np.log2(p)) + + ent_bin = int(entropy * 2.0) + if ent_bin > 15: + ent_bin = 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 + + ent_sum = np.sum(ent_hist) + if ent_sum > 0: + ent_hist /= ent_sum + + for i, v in enumerate(ent_hist): + res[f'ByteEnt_{i:02X}'] = float(v) + + return res + @classmethod def _extract_overlay_features(cls, pe, fsize): overlay_offset = pe.get_overlay_data_start_offset() @@ -332,12 +415,19 @@ class FeatureExtractor: 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 + 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 + base[f'SectionHash_{hash_val:02d}'] += 1.0 + try: s_data = section.get_data() s_entropy = cls._calc_entropy(s_data) @@ -432,6 +522,8 @@ class FeatureExtractor: base['DebugCount'] = 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)) @@ -459,8 +551,6 @@ def load_existing_hashes(): except Exception: pass return hashes -#################################################################################################### - GLOBAL_HASHES = set() def worker_init(shared_hashes): diff --git a/Engine/Properties/scanner.py b/Engine/Properties/scanner.py index 771e2d1..4fb2fe2 100644 --- a/Engine/Properties/scanner.py +++ b/Engine/Properties/scanner.py @@ -2,6 +2,7 @@ import os, sys, math, json, pefile, datetime import numpy as np import onnxruntime as ort import ctypes, ctypes.wintypes +import re, zlib #################################################################################################### @@ -98,6 +99,8 @@ class WinTrust: #################################################################################################### class FeatureExtractor: + _STRING_PATTERN = re.compile(b'[\x20-\x7E]{5,}') + @staticmethod def _safe_float(val): try: @@ -115,6 +118,86 @@ class FeatureExtractor: probs = counts[counts > 0] / len(arr) return float(-np.sum(probs * np.log2(probs))) + @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 + + count = 0 + total_len = 0 + char_counts = [0] * 95 + combined_data = bytearray() + + for match in cls._STRING_PATTERN.finditer(file_bytes): + m = match.group() + count += 1 + length = len(m) + total_len += length + combined_data.extend(m) + for b in m: + char_counts[b - 0x20] += 1 + + if count == 0: + return res + + res["StringCount"] = float(count) + res["StringMeanLength"] = float(total_len) / count + res["StringEntropy"] = cls._calc_entropy(combined_data) + + for i in range(95): + res[f'StringHist_{i:02d}'] = float(char_counts[i]) / total_len + + return res + + @classmethod + def _extract_histograms(cls, file_bytes): + import numpy as np + arr = np.frombuffer(file_bytes, dtype=np.uint8) + sz = len(arr) + res = {} + + if sz == 0: + for i in range(256): + res[f'ByteHist_{i:02X}'] = 0.0 + res[f'ByteEnt_{i:02X}'] = 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) + + ent_hist = np.zeros(256, dtype=np.float64) + window = 2048 + step = 1024 + num_windows = (sz - window) // step + 1 if sz >= window else 1 + + for i in range(num_windows): + w = arr if sz < window else arr[i*step : i*step+window] + w_counts = np.bincount(w, minlength=256) + + p = w_counts[w_counts > 0] / len(w) + entropy = -np.sum(p * np.log2(p)) + + ent_bin = int(entropy * 2.0) + if ent_bin > 15: + ent_bin = 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 + + ent_sum = np.sum(ent_hist) + if ent_sum > 0: + ent_hist /= ent_sum + + for i, v in enumerate(ent_hist): + res[f'ByteEnt_{i:02X}'] = float(v) + + return res + @classmethod def _extract_overlay_features(cls, pe, fsize): overlay_offset = pe.get_overlay_data_start_offset() @@ -329,12 +412,19 @@ class FeatureExtractor: 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 + 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 + base[f'SectionHash_{hash_val:02d}'] += 1.0 + try: s_data = section.get_data() s_entropy = cls._calc_entropy(s_data) @@ -429,6 +519,8 @@ class FeatureExtractor: base['DebugCount'] = 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)) @@ -456,6 +548,8 @@ class ModelPredictor: if current_dim != expected_dim: raise ValueError(f"Dimension mismatch! Model expects {expected_dim} features, but feature file lists {current_dim}.") + self.feat_map = {feat: i for i, feat in enumerate(self.model_features)} + def _load_features(self, path): if not os.path.exists(path): raise FileNotFoundError(f"Feature file not found: {path}") @@ -465,18 +559,24 @@ class ModelPredictor: def predict(self, raw_data): vec = np.zeros((1, len(self.model_features)), dtype=np.float32) base = raw_data.get('Base', {}) - dlls = set(raw_data.get('DLLs', [])) - apis = set(raw_data.get('APIs', [])) + dlls = raw_data.get('DLLs', []) + apis = raw_data.get('APIs', []) - for i, feat in enumerate(self.model_features): - if feat.startswith('Dll_'): - if feat[4:] in dlls: - vec[0, i] = 1.0 - elif feat.startswith('Api_'): - if feat[4:] in apis: - vec[0, i] = 1.0 - else: - vec[0, i] = base.get(feat, 0.0) + 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')) % 256 + feat_name = f"DllHash_{h:03d}" + 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')) % 1024 + feat_name = f"ApiHash_{h:04d}" + 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}) diff --git a/Engine/Properties/train_server.py b/Engine/Properties/train_server.py index 70c08ea..3edfc16 100644 --- a/Engine/Properties/train_server.py +++ b/Engine/Properties/train_server.py @@ -1,4 +1,4 @@ -import os, sys, time, json, onnxmltools +import os, sys, time, json, onnxmltools, zlib import numpy as np import pandas as pd import lightgbm as lgb @@ -9,6 +9,12 @@ from onnxmltools.convert.common.data_types import FloatTensorType from sklearn.model_selection import train_test_split from sklearn.metrics import accuracy_score, roc_auc_score +try: + import orjson + JSON_LOADS = orjson.loads +except ImportError: + JSON_LOADS = json.loads + #################################################################################################### JSONL_PATH = "pe_features.jsonl" @@ -18,6 +24,9 @@ FEATURE_FILE = "features.json" TEST_SIZE = 0.0001 RANDOM_SEED = 42 +DLL_HASH_DIM = 256 +API_HASH_DIM = 1024 + EXCLUDE_FEATURES = { 'label', 'filehash', } @@ -42,26 +51,22 @@ LGBM_PARAMS = { def analyze_schema(jsonl_path): print("[*] Stage 1: Scanning global schema...") base_keys = set() - dll_counts = defaultdict(int) - api_counts = defaultdict(int) num_lines = 0 - with open(jsonl_path, 'r') as f: + with open(jsonl_path, 'rb') as f: for line in f: try: - data = json.loads(line) - base_keys.update(data.get('Base', {}).keys()) - for d in data.get('DLLs', []): - dll_counts[d] += 1 - for a in data.get('APIs', []): - api_counts[a] += 1 + data = JSON_LOADS(line) + base_data = data.get('Base') + if base_data: + base_keys.update(base_data.keys()) num_lines += 1 except Exception: continue base_schema = sorted([k for k in base_keys if k.lower().strip() not in EXCLUDE_FEATURES]) - dll_schema = sorted([f"Dll_{k}" for k, v in dll_counts.items() if v >= 1000]) - api_schema = sorted([f"Api_{k}" for k, v in api_counts.items() if v >= 1000]) + dll_schema = [f"DllHash_{i:03d}" for i in range(DLL_HASH_DIM)] + api_schema = [f"ApiHash_{i:04d}" for i in range(API_HASH_DIM)] return base_schema, dll_schema, api_schema, num_lines @@ -71,39 +76,46 @@ def load_data_efficiently(jsonl_path, base_schema, dll_schema, api_schema, num_l print(f"[*] Stage 2: Building sparse feature matrix dynamically (Expected samples: {num_lines})...") final_features = base_schema + dll_schema + api_schema num_features = len(final_features) - feature_idx = {feat: i for i, feat in enumerate(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) - with open(jsonl_path, 'r') as f: + with open(jsonl_path, 'rb') as f: row_idx = 0 for line in f: try: - data = json.loads(line) + data = JSON_LOADS(line) y[row_idx] = data.get('Label', 0) - for k, v in data.get('Base', {}).items(): - idx = feature_idx.get(k) - if idx is not None: - val = float(v) - if val != 0.0: - rows.append(row_idx) - cols.append(idx) - vals.append(val) + base_data = data.get('Base') + if base_data: + for k, v in base_data.items(): + idx = base_idx.get(k) + if idx is not None: + val = float(v) + if val != 0.0: + rows.append(row_idx) + cols.append(idx) + vals.append(val) - for d in data.get('DLLs', []): - idx = feature_idx.get(f"Dll_{d}") - if idx is not None: + dll_data = data.get('DLLs') + if dll_data: + for d in dll_data: + idx = zlib.crc32(d.encode('utf-8', 'ignore')) % DLL_HASH_DIM rows.append(row_idx) - cols.append(idx) + cols.append(offset_dll + idx) vals.append(1.0) - for a in data.get('APIs', []): - idx = feature_idx.get(f"Api_{a}") - if idx is not None: + api_data = data.get('APIs') + if api_data: + for a in api_data: + idx = zlib.crc32(a.encode('utf-8', 'ignore')) % API_HASH_DIM rows.append(row_idx) - cols.append(idx) + cols.append(offset_api + idx) vals.append(1.0) row_idx += 1 @@ -233,8 +245,8 @@ if __name__ == "__main__": sys.exit(1) print(f"[*] Detected Base features: {len(base_schema)}") - print(f"[*] Filtered DLL features: {len(dll_schema)}") - print(f"[*] Filtered API features: {len(api_schema)}") + print(f"[*] Hashing DLL features: {len(dll_schema)}") + print(f"[*] Hashing API features: {len(api_schema)}") X, y, final_features = load_data_efficiently(JSONL_PATH, base_schema, dll_schema, api_schema, num_lines) diff --git a/Engine/Properties/train_user.py b/Engine/Properties/train_user.py index be9db1b..73782b5 100644 --- a/Engine/Properties/train_user.py +++ b/Engine/Properties/train_user.py @@ -1,4 +1,4 @@ -import os, sys, time, json, onnxmltools +import os, sys, time, json, onnxmltools, zlib import numpy as np import pandas as pd import lightgbm as lgb @@ -9,6 +9,12 @@ from onnxmltools.convert.common.data_types import FloatTensorType from sklearn.model_selection import train_test_split from sklearn.metrics import accuracy_score, roc_auc_score +try: + import orjson + JSON_LOADS = orjson.loads +except ImportError: + JSON_LOADS = json.loads + #################################################################################################### JSONL_PATH = "pe_features.jsonl" @@ -18,6 +24,9 @@ FEATURE_FILE = "features.json" TEST_SIZE = 0.0001 RANDOM_SEED = 42 +DLL_HASH_DIM = 256 +API_HASH_DIM = 1024 + EXCLUDE_FEATURES = { 'label', 'filehash', } @@ -42,26 +51,22 @@ LGBM_PARAMS = { def analyze_schema(jsonl_path): print("[*] Stage 1: Scanning global schema...") base_keys = set() - dll_counts = defaultdict(int) - api_counts = defaultdict(int) num_lines = 0 - with open(jsonl_path, 'r') as f: + with open(jsonl_path, 'rb') as f: for line in f: try: - data = json.loads(line) - base_keys.update(data.get('Base', {}).keys()) - for d in data.get('DLLs', []): - dll_counts[d] += 1 - for a in data.get('APIs', []): - api_counts[a] += 1 + data = JSON_LOADS(line) + base_data = data.get('Base') + if base_data: + base_keys.update(base_data.keys()) num_lines += 1 except Exception: continue base_schema = sorted([k for k in base_keys if k.lower().strip() not in EXCLUDE_FEATURES]) - dll_schema = sorted([f"Dll_{k}" for k, v in dll_counts.items() if v >= 1000]) - api_schema = sorted([f"Api_{k}" for k, v in api_counts.items() if v >= 1000]) + dll_schema = [f"DllHash_{i:03d}" for i in range(DLL_HASH_DIM)] + api_schema = [f"ApiHash_{i:04d}" for i in range(API_HASH_DIM)] return base_schema, dll_schema, api_schema, num_lines @@ -71,39 +76,46 @@ def load_data_efficiently(jsonl_path, base_schema, dll_schema, api_schema, num_l print(f"[*] Stage 2: Building sparse feature matrix dynamically (Expected samples: {num_lines})...") final_features = base_schema + dll_schema + api_schema num_features = len(final_features) - feature_idx = {feat: i for i, feat in enumerate(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) - with open(jsonl_path, 'r') as f: + with open(jsonl_path, 'rb') as f: row_idx = 0 for line in f: try: - data = json.loads(line) + data = JSON_LOADS(line) y[row_idx] = data.get('Label', 0) - for k, v in data.get('Base', {}).items(): - idx = feature_idx.get(k) - if idx is not None: - val = float(v) - if val != 0.0: - rows.append(row_idx) - cols.append(idx) - vals.append(val) + base_data = data.get('Base') + if base_data: + for k, v in base_data.items(): + idx = base_idx.get(k) + if idx is not None: + val = float(v) + if val != 0.0: + rows.append(row_idx) + cols.append(idx) + vals.append(val) - for d in data.get('DLLs', []): - idx = feature_idx.get(f"Dll_{d}") - if idx is not None: + dll_data = data.get('DLLs') + if dll_data: + for d in dll_data: + idx = zlib.crc32(d.encode('utf-8', 'ignore')) % DLL_HASH_DIM rows.append(row_idx) - cols.append(idx) + cols.append(offset_dll + idx) vals.append(1.0) - for a in data.get('APIs', []): - idx = feature_idx.get(f"Api_{a}") - if idx is not None: + api_data = data.get('APIs') + if api_data: + for a in api_data: + idx = zlib.crc32(a.encode('utf-8', 'ignore')) % API_HASH_DIM rows.append(row_idx) - cols.append(idx) + cols.append(offset_api + idx) vals.append(1.0) row_idx += 1 @@ -233,8 +245,8 @@ if __name__ == "__main__": sys.exit(1) print(f"[*] Detected Base features: {len(base_schema)}") - print(f"[*] Filtered DLL features: {len(dll_schema)}") - print(f"[*] Filtered API features: {len(api_schema)}") + print(f"[*] Hashing DLL features: {len(dll_schema)}") + print(f"[*] Hashing API features: {len(api_schema)}") X, y, final_features = load_data_efficiently(JSONL_PATH, base_schema, dll_schema, api_schema, num_lines)