87owo/PYAS
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add files via upload

Browse Source
This commit is contained in:
87owo authored and GitHub committed 2026-05-26 09:30:27 +08:00
1 parent 9f22fe6d97
commit ef5793fb15
4 files changed
+274 -166
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

Engine/Properties/convert.py
+73 -65
View File
@@ -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()
Engine/Properties/scanner.py
+62 -56
View File
@@ -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))
Engine/Properties/train_server.py
+69 -22
View File
@@ -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):
Engine/Properties/train_user.py
+70 -23
View File
@@ -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):