Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
766b6e81d0 | ||
|
|
0349e01629 | ||
|
|
0bd9da2bbc | ||
|
|
f14c7c3f55 | ||
|
|
ef5793fb15 | ||
|
|
9f22fe6d97 | ||
|
|
3147d576ee | ||
|
|
2ceec2588c | ||
|
|
5b4d88b9c7 | ||
|
|
55043cca36 | ||
|
|
097ebe025f | ||
|
|
9f92307f64 | ||
|
|
c220035fd9 | ||
|
|
d7b98d3cde | ||
|
|
0ab876633f | ||
|
|
cabe76ca54 | ||
|
|
c57224e258 | ||
|
|
b56101d2e6 | ||
|
|
16211c28d1 | ||
|
|
139d3e0c3d | ||
|
|
f0652fc5ec | ||
|
|
55489cd634 | ||
|
|
742887c0e4 | ||
|
|
c3e5cd2c0d | ||
|
|
63e125627c | ||
|
|
c0189e9d05 | ||
|
|
133e3208ec | ||
|
|
17c6fd6d0f | ||
|
|
2aa7ea43b8 | ||
|
|
6d7dbfac31 | ||
|
|
d4942c6dd6 | ||
|
|
39ae3c24f6 |
No files matched your search
+315
-51
@@ -1,5 +1,8 @@
|
||||
import os, sys, math, time, datetime, pefile, hashlib
|
||||
import ctypes, ctypes.wintypes, multiprocessing
|
||||
import re, os, sys, math, time, zlib, datetime, hashlib
|
||||
import ctypes, ctypes.wintypes, pefile, multiprocessing
|
||||
import numpy as np
|
||||
|
||||
####################################################################################################
|
||||
|
||||
try:
|
||||
import orjson
|
||||
@@ -12,15 +15,10 @@ except ImportError:
|
||||
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
|
||||
TARGET_EXTENSIONS = {
|
||||
'.exe', '.dll', '.sys', '.ocx', '.scr', '.efi', '.acm', '.ax', '.cpl', '.drv', '.com', '.mui', '.pyd'
|
||||
}
|
||||
|
||||
####################################################################################################
|
||||
|
||||
@@ -110,6 +108,19 @@ 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):
|
||||
try:
|
||||
@@ -121,11 +132,243 @@ class FeatureExtractor:
|
||||
@staticmethod
|
||||
def _calc_entropy(data):
|
||||
if not data: return 0.0
|
||||
import numpy as np
|
||||
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}
|
||||
res.update(dict.fromkeys(cls._STRING_KEYS, 0.0))
|
||||
|
||||
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)
|
||||
|
||||
res["StringCount"] = float(count)
|
||||
res["StringMeanLength"] = float(total_len) / count
|
||||
res["StringEntropy"] = cls._calc_entropy(combined_data)
|
||||
|
||||
arr = np.frombuffer(combined_data, dtype=np.uint8)
|
||||
counts = np.bincount(arr, minlength=127)
|
||||
char_counts = counts[0x20:0x7F]
|
||||
|
||||
res.update(dict(zip(cls._STRING_KEYS, char_counts.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.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.astype(np.float64) / sz
|
||||
res.update(dict(zip(cls._BYTE_HIST_KEYS, byte_hist)))
|
||||
|
||||
ent_hist = np.zeros(256, dtype=np.float64)
|
||||
window = 2048
|
||||
step = 1024
|
||||
|
||||
if sz < window:
|
||||
w_counts = byte_counts
|
||||
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_windows = 10000
|
||||
|
||||
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
|
||||
|
||||
start_byte = w_start * step
|
||||
end_byte = start_byte + b_num * step + step
|
||||
b_arr = arr[start_byte:end_byte]
|
||||
|
||||
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:]
|
||||
|
||||
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())
|
||||
|
||||
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):
|
||||
@@ -134,9 +377,23 @@ class FeatureExtractor:
|
||||
|
||||
try:
|
||||
pe = pefile.PE(data=file_bytes, fast_load=True)
|
||||
base['TrustSigned'] = verify_signature(file_path)
|
||||
|
||||
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['FileEntropy'] = cls._calc_entropy(file_bytes)
|
||||
base['FileSize'] = fsize
|
||||
base['FileSize'] = float(fsize)
|
||||
|
||||
fh = pe.FILE_HEADER
|
||||
base['Machine'] = getattr(fh, 'Machine', 0)
|
||||
@@ -171,10 +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
|
||||
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)
|
||||
@@ -182,6 +438,10 @@ class FeatureExtractor:
|
||||
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)
|
||||
@@ -197,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
|
||||
@@ -211,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:
|
||||
@@ -238,16 +503,6 @@ class FeatureExtractor:
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
try:
|
||||
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']
|
||||
])
|
||||
except Exception: pass
|
||||
|
||||
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
|
||||
|
||||
@@ -260,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):
|
||||
@@ -271,20 +526,29 @@ 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))
|
||||
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)}
|
||||
@@ -297,17 +561,17 @@ 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()
|
||||
|
||||
def worker_init(shared_hashes):
|
||||
@@ -344,11 +608,11 @@ def _extract_worker(file_path, label):
|
||||
|
||||
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 if os.path.splitext(f)[1].lower() in TARGET_EXTENSIONS]
|
||||
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 valid PE files found.")
|
||||
print("[-] No files found.")
|
||||
return
|
||||
|
||||
print(f"[*] Found {total} files. Starting {MAX_WORKERS}-process extraction...\n")
|
||||
|
||||
@@ -0,0 +1,698 @@
|
||||
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_S1.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,}')
|
||||
|
||||
_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):
|
||||
try:
|
||||
f = float(val)
|
||||
return f if not (math.isinf(f) or math.isnan(f)) else 0.0
|
||||
except Exception:
|
||||
return 0.0
|
||||
|
||||
@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))
|
||||
|
||||
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)
|
||||
|
||||
res["StringCount"] = float(count)
|
||||
res["StringMeanLength"] = float(total_len) / count
|
||||
res["StringEntropy"] = cls._calc_entropy(combined_data)
|
||||
|
||||
arr = np.frombuffer(combined_data, dtype=np.uint8)
|
||||
counts = np.bincount(arr, minlength=127)
|
||||
char_counts = counts[0x20:0x7F]
|
||||
|
||||
res.update(dict(zip(cls._STRING_KEYS, char_counts.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.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.astype(np.float64) / sz
|
||||
res.update(dict(zip(cls._BYTE_HIST_KEYS, byte_hist)))
|
||||
|
||||
ent_hist = np.zeros(256, dtype=np.float64)
|
||||
window = 2048
|
||||
step = 1024
|
||||
|
||||
if sz < window:
|
||||
w_counts = byte_counts
|
||||
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_windows = 10000
|
||||
|
||||
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
|
||||
|
||||
start_byte = w_start * step
|
||||
end_byte = start_byte + b_num * step + step
|
||||
b_arr = arr[start_byte:end_byte]
|
||||
|
||||
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:]
|
||||
|
||||
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())
|
||||
|
||||
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
|
||||
try:
|
||||
fsize = os.path.getsize(file_path)
|
||||
if fsize == 0 or fsize > MAX_FILE_SIZE:
|
||||
return None
|
||||
|
||||
with open(file_path, "rb") as f:
|
||||
file_bytes = f.read()
|
||||
|
||||
base, dlls, apis = {}, 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['FileEntropy'] = cls._calc_entropy(file_bytes)
|
||||
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()
|
||||
|
||||
####################################################################################################
|
||||
|
||||
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 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}")
|
||||
with open(path, 'r') as f:
|
||||
return json.load(f)
|
||||
|
||||
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')) % 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})
|
||||
|
||||
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
|
||||
|
||||
print(f"\n[*] Scanning {len(files)} files...\n")
|
||||
print(f"{'RESULT':<10} | {'PROB':<8} | {'FILE'}")
|
||||
print("-" * 80)
|
||||
|
||||
safe_count = 0
|
||||
mal_count = 0
|
||||
|
||||
for fpath in files:
|
||||
try:
|
||||
data = FeatureExtractor.extract(fpath)
|
||||
if not data:
|
||||
print(f"{'ERROR':<10} | {'----':<8} | {fpath}")
|
||||
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
|
||||
|
||||
print(f"{label:<10} | {prob:.4f} | {fpath}")
|
||||
|
||||
except Exception as e:
|
||||
print(f"{'FAIL':<10} | {'----':<8} | {fpath} ({str(e)})")
|
||||
|
||||
print("-" * 80)
|
||||
print(f"[*] Summary: Safe={safe_count}, Malware={mal_count}, Total={len(files)}")
|
||||
|
||||
####################################################################################################
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("\n-------------------------- PE Malware Predictor v4.0 --------------------------\n")
|
||||
print(f"[*] © 2020-2026 87owo (PYAS Security)")
|
||||
print(f"[*] Loading model features and weights...")
|
||||
|
||||
base_dir = os.path.dirname(sys.executable) if getattr(sys, 'frozen', False) else os.path.dirname(os.path.abspath(__file__))
|
||||
model_path = os.path.join(base_dir, MODEL_FILE)
|
||||
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 len(sys.argv) > 1:
|
||||
for path in sys.argv[1:]:
|
||||
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}")
|
||||
@@ -1,4 +1,4 @@
|
||||
import os, sys, time, json, onnxmltools
|
||||
import os, sys, time, json, zlib, array, onnxmltools, multiprocessing
|
||||
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',
|
||||
}
|
||||
@@ -34,9 +43,7 @@ LGBM_PARAMS = {
|
||||
'seed': RANDOM_SEED,
|
||||
'n_jobs': -1,
|
||||
'max_depth': -1,
|
||||
'min_data_in_leaf': 20,
|
||||
'lambda_l1': 0.01,
|
||||
'lambda_l2': 0.01,
|
||||
'min_data_in_leaf': 50,
|
||||
}
|
||||
|
||||
####################################################################################################
|
||||
@@ -44,78 +51,127 @@ 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
|
||||
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 = sorted([f"Dll_{k}" for k, v in dll_counts.items() if v >= 100])
|
||||
api_schema = sorted([f"Api_{k}" for k, v in api_counts.items() if v >= 100])
|
||||
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
|
||||
|
||||
####################################################################################################
|
||||
|
||||
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)
|
||||
feature_idx = {feat: i for i, feat in enumerate(final_features)}
|
||||
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, 'r') 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)
|
||||
data = JSON_LOADS(line)
|
||||
y_list.append(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
|
||||
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
|
||||
|
||||
####################################################################################################
|
||||
|
||||
@@ -223,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):
|
||||
@@ -235,8 +292,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)
|
||||
|
||||
|
||||
+107
-50
@@ -1,4 +1,4 @@
|
||||
import os, sys, time, json, onnxmltools
|
||||
import os, sys, time, json, zlib, array, onnxmltools, multiprocessing
|
||||
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',
|
||||
}
|
||||
@@ -26,7 +35,7 @@ LGBM_PARAMS = {
|
||||
'objective': 'binary',
|
||||
'metric': ['binary_logloss', 'auc'],
|
||||
'boosting_type': 'gbdt',
|
||||
'num_leaves': 64,
|
||||
'num_leaves': 128,
|
||||
'feature_fraction': 0.8,
|
||||
'bagging_fraction': 1.0,
|
||||
'bagging_freq': 5,
|
||||
@@ -34,9 +43,7 @@ LGBM_PARAMS = {
|
||||
'seed': RANDOM_SEED,
|
||||
'n_jobs': -1,
|
||||
'max_depth': -1,
|
||||
'min_data_in_leaf': 20,
|
||||
'lambda_l1': 0.01,
|
||||
'lambda_l2': 0.01,
|
||||
'min_data_in_leaf': 50,
|
||||
}
|
||||
|
||||
####################################################################################################
|
||||
@@ -44,78 +51,127 @@ 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
|
||||
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 = sorted([f"Dll_{k}" for k, v in dll_counts.items() if v >= 100])
|
||||
api_schema = sorted([f"Api_{k}" for k, v in api_counts.items() if v >= 100])
|
||||
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
|
||||
|
||||
####################################################################################################
|
||||
|
||||
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)
|
||||
feature_idx = {feat: i for i, feat in enumerate(final_features)}
|
||||
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, 'r') 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)
|
||||
data = JSON_LOADS(line)
|
||||
y_list.append(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
|
||||
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
|
||||
|
||||
####################################################################################################
|
||||
|
||||
@@ -223,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):
|
||||
@@ -235,8 +292,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)
|
||||
|
||||
|
||||
+143
-11
@@ -49,6 +49,7 @@ window.AppI18n = {
|
||||
"scan_context_desc": "此選項可以啟用右鍵掃描功能",
|
||||
"theme_color": "顯色主題",
|
||||
"theme_color_desc": "此選項可以選擇主題顏色",
|
||||
"theme_system": "跟隨系統",
|
||||
"theme_white": "白色主題",
|
||||
"theme_red": "紅色主題",
|
||||
"theme_yellow": "黃色主題",
|
||||
@@ -134,7 +135,18 @@ window.AppI18n = {
|
||||
"res_net_block": "網路攔截",
|
||||
"res_drv_block": "驅動攔截",
|
||||
"res_sys_repair": "系統修復",
|
||||
"res_cfg_update": "設定更新"
|
||||
"res_cfg_update": "設定更新",
|
||||
"manage_startup": "啟動項管理",
|
||||
"manage_startup_desc": "此選項可以管理各應用啟動項",
|
||||
"col_status": "狀態",
|
||||
"status_enabled": "已啟用",
|
||||
"status_disabled": "已禁用",
|
||||
"btn_enable": "啟用",
|
||||
"btn_disable": "禁用",
|
||||
"search_item": "搜尋項目...",
|
||||
"type_reg": "登錄檔",
|
||||
"type_srv": "服務",
|
||||
"type_tsk": "排程任務",
|
||||
},
|
||||
"simplified_switch": {
|
||||
"app_name": "PYAS Security",
|
||||
@@ -185,6 +197,7 @@ window.AppI18n = {
|
||||
"scan_context_desc": "此选项可以启用右键扫描功能",
|
||||
"theme_color": "显色主题",
|
||||
"theme_color_desc": "此选项可以选择主题颜色",
|
||||
"theme_system": "跟随系统",
|
||||
"theme_white": "白色主题",
|
||||
"theme_red": "红色主题",
|
||||
"theme_yellow": "黄色主题",
|
||||
@@ -270,7 +283,18 @@ window.AppI18n = {
|
||||
"res_net_block": "网络拦截",
|
||||
"res_drv_block": "驱动拦截",
|
||||
"res_sys_repair": "系统修复",
|
||||
"res_cfg_update": "设置更新"
|
||||
"res_cfg_update": "设置更新",
|
||||
"manage_startup": "启动项管理",
|
||||
"manage_startup_desc": "此选项可以管理各应用启动项",
|
||||
"col_status": "状态",
|
||||
"status_enabled": "已启用",
|
||||
"status_disabled": "已禁用",
|
||||
"btn_enable": "启用",
|
||||
"btn_disable": "禁用",
|
||||
"search_item": "搜索项目...",
|
||||
"type_reg": "注册表",
|
||||
"type_srv": "服务",
|
||||
"type_tsk": "计划任务",
|
||||
},
|
||||
"english_switch": {
|
||||
"app_name": "PYAS Security",
|
||||
@@ -321,6 +345,7 @@ window.AppI18n = {
|
||||
"scan_context_desc": "Enable right-click scan functionality",
|
||||
"theme_color": "Color Theme",
|
||||
"theme_color_desc": "Select application color theme",
|
||||
"theme_system": "System Theme",
|
||||
"theme_white": "Light Mode",
|
||||
"theme_red": "Red Theme",
|
||||
"theme_yellow": "Yellow Theme",
|
||||
@@ -406,7 +431,18 @@ window.AppI18n = {
|
||||
"res_net_block": "Network Blocked",
|
||||
"res_drv_block": "Driver Blocked",
|
||||
"res_sys_repair": "System Repair",
|
||||
"res_cfg_update": "Config Update"
|
||||
"res_cfg_update": "Config Update",
|
||||
"manage_startup": "Startup Manager",
|
||||
"manage_startup_desc": "Manage startup applications and services",
|
||||
"col_status": "Status",
|
||||
"status_enabled": "Enabled",
|
||||
"status_disabled": "Disabled",
|
||||
"btn_enable": "Enable",
|
||||
"btn_disable": "Disable",
|
||||
"search_item": "Search items...",
|
||||
"type_reg": "Registry",
|
||||
"type_srv": "Service",
|
||||
"type_tsk": "Task",
|
||||
},
|
||||
"japanese_switch": {
|
||||
"app_name": "PYAS Security",
|
||||
@@ -457,6 +493,7 @@ window.AppI18n = {
|
||||
"scan_context_desc": "右クリックでスキャン機能を有効にします",
|
||||
"theme_color": "カラーテーマ",
|
||||
"theme_color_desc": "アプリケーションのテーマカラーを選択します",
|
||||
"theme_system": "システムに従う",
|
||||
"theme_white": "ライトモード",
|
||||
"theme_red": "レッドテーマ",
|
||||
"theme_yellow": "イエローテーマ",
|
||||
@@ -542,7 +579,18 @@ window.AppI18n = {
|
||||
"res_net_block": "ネットワークブロック",
|
||||
"res_drv_block": "ドライバーブロック",
|
||||
"res_sys_repair": "システム修復",
|
||||
"res_cfg_update": "設定の更新"
|
||||
"res_cfg_update": "設定の更新",
|
||||
"manage_startup": "起動項管理",
|
||||
"manage_startup_desc": "アプリケーションの起動項目を管理します",
|
||||
"col_status": "ステータス",
|
||||
"status_enabled": "有効",
|
||||
"status_disabled": "無効",
|
||||
"btn_enable": "有効化",
|
||||
"btn_disable": "無効化",
|
||||
"search_item": "項目を検索...",
|
||||
"type_reg": "レジストリ",
|
||||
"type_srv": "サービス",
|
||||
"type_tsk": "タスク",
|
||||
},
|
||||
"korean_switch": {
|
||||
"app_name": "PYAS Security",
|
||||
@@ -593,6 +641,7 @@ window.AppI18n = {
|
||||
"scan_context_desc": "우클릭 스캔 기능을 활성화합니다",
|
||||
"theme_color": "색상 테마",
|
||||
"theme_color_desc": "응용 프로그램 테마 색상을 선택합니다",
|
||||
"theme_system": "시스템 테마",
|
||||
"theme_white": "라이트 모드",
|
||||
"theme_red": "레드 테마",
|
||||
"theme_yellow": "옐로우 테마",
|
||||
@@ -678,7 +727,18 @@ window.AppI18n = {
|
||||
"res_net_block": "네트워크 차단",
|
||||
"res_drv_block": "드라이버 차단",
|
||||
"res_sys_repair": "시스템 복구",
|
||||
"res_cfg_update": "설정 업데이트"
|
||||
"res_cfg_update": "설정 업데이트",
|
||||
"manage_startup": "시작 항목 관리",
|
||||
"manage_startup_desc": "애플리케이션 시작 항목을 관리합니다",
|
||||
"col_status": "상태",
|
||||
"status_enabled": "사용",
|
||||
"status_disabled": "사용 안 함",
|
||||
"btn_enable": "사용",
|
||||
"btn_disable": "사용 안 함",
|
||||
"search_item": "항목 검색...",
|
||||
"type_reg": "레지스트리",
|
||||
"type_srv": "서비스",
|
||||
"type_tsk": "작업",
|
||||
},
|
||||
"french_switch": {
|
||||
"app_name": "PYAS Security",
|
||||
@@ -729,6 +789,7 @@ window.AppI18n = {
|
||||
"scan_context_desc": "Activer la fonctionnalité d'analyse par clic droit",
|
||||
"theme_color": "Thème de couleur",
|
||||
"theme_color_desc": "Sélectionnez le thème de couleur de l'application",
|
||||
"theme_system": "Thème du système",
|
||||
"theme_white": "Mode clair",
|
||||
"theme_red": "Thème rouge",
|
||||
"theme_yellow": "Thème jaune",
|
||||
@@ -814,7 +875,18 @@ window.AppI18n = {
|
||||
"res_net_block": "Réseau bloqué",
|
||||
"res_drv_block": "Pilote bloqué",
|
||||
"res_sys_repair": "Réparation du système",
|
||||
"res_cfg_update": "Mise à jour de la configuration"
|
||||
"res_cfg_update": "Mise à jour de la configuration",
|
||||
"manage_startup": "Gestionnaire de démarrage",
|
||||
"manage_startup_desc": "Gérer les applications et services de démarrage",
|
||||
"col_status": "Statut",
|
||||
"status_enabled": "Activé",
|
||||
"status_disabled": "Désactivé",
|
||||
"btn_enable": "Activer",
|
||||
"btn_disable": "Désactiver",
|
||||
"search_item": "Rechercher des éléments...",
|
||||
"type_reg": "Registre",
|
||||
"type_srv": "Service",
|
||||
"type_tsk": "Tâche",
|
||||
},
|
||||
"spanish_switch": {
|
||||
"app_name": "PYAS Security",
|
||||
@@ -865,6 +937,7 @@ window.AppI18n = {
|
||||
"scan_context_desc": "Habilitar la función de escaneo con el botón derecho",
|
||||
"theme_color": "Tema de color",
|
||||
"theme_color_desc": "Seleccione el tema de color de la aplicación",
|
||||
"theme_system": "Tema del sistema",
|
||||
"theme_white": "Modo claro",
|
||||
"theme_red": "Tema rojo",
|
||||
"theme_yellow": "Tema amarillo",
|
||||
@@ -950,7 +1023,18 @@ window.AppI18n = {
|
||||
"res_net_block": "Red bloqueada",
|
||||
"res_drv_block": "Controlador bloqueado",
|
||||
"res_sys_repair": "Reparación del sistema",
|
||||
"res_cfg_update": "Actualización de configuración"
|
||||
"res_cfg_update": "Actualización de configuración",
|
||||
"manage_startup": "Gestor de inicio",
|
||||
"manage_startup_desc": "Administre aplicaciones y servicios de inicio",
|
||||
"col_status": "Estado",
|
||||
"status_enabled": "Habilitado",
|
||||
"status_disabled": "Deshabilitado",
|
||||
"btn_enable": "Habilitar",
|
||||
"btn_disable": "Deshabilitar",
|
||||
"search_item": "Buscar elementos...",
|
||||
"type_reg": "Registro",
|
||||
"type_srv": "Servicio",
|
||||
"type_tsk": "Tarea",
|
||||
},
|
||||
"hindi_switch": {
|
||||
"app_name": "PYAS Security",
|
||||
@@ -1001,6 +1085,7 @@ window.AppI18n = {
|
||||
"scan_context_desc": "राइट-क्लिक स्कैन कार्यक्षमता सक्षम करें",
|
||||
"theme_color": "रंग थीम",
|
||||
"theme_color_desc": "एप्लिकेशन का रंग थीम चुनें",
|
||||
"theme_system": "सिस्टम थीम",
|
||||
"theme_white": "लाइट मोड",
|
||||
"theme_red": "लाल थीम",
|
||||
"theme_yellow": "पीला थीम",
|
||||
@@ -1086,7 +1171,18 @@ window.AppI18n = {
|
||||
"res_net_block": "नेटवर्क अवरुद्ध",
|
||||
"res_drv_block": "ड्राइवर अवरुद्ध",
|
||||
"res_sys_repair": "सिस्टम मरम्मत",
|
||||
"res_cfg_update": "कॉन्फ़िगरेशन अपडेट"
|
||||
"res_cfg_update": "कॉन्फ़िगरेशन अपडेट",
|
||||
"manage_startup": "स्टार्टअप प्रबंधक",
|
||||
"manage_startup_desc": "स्टार्टअप एप्लिकेशन और सेवाएं प्रबंधित करें",
|
||||
"col_status": "स्थिति",
|
||||
"status_enabled": "सक्षम",
|
||||
"status_disabled": "अक्षम",
|
||||
"btn_enable": "सक्षम करें",
|
||||
"btn_disable": "अक्षम करें",
|
||||
"search_item": "आइटम खोजें...",
|
||||
"type_reg": "रजिस्ट्री",
|
||||
"type_srv": "सेवा",
|
||||
"type_tsk": "कार्य",
|
||||
},
|
||||
"arabic_switch": {
|
||||
"app_name": "PYAS Security",
|
||||
@@ -1137,6 +1233,7 @@ window.AppI18n = {
|
||||
"scan_context_desc": "تمكين ميزة الفحص بزر الماوس الأيمن",
|
||||
"theme_color": "سمة اللون",
|
||||
"theme_color_desc": "تحديد سمة ألوان التطبيق",
|
||||
"theme_system": "سمة النظام",
|
||||
"theme_white": "الوضع الفاتح",
|
||||
"theme_red": "سمة حمراء",
|
||||
"theme_yellow": "سمة صفراء",
|
||||
@@ -1222,7 +1319,18 @@ window.AppI18n = {
|
||||
"res_net_block": "حظر الشبكة",
|
||||
"res_drv_block": "حظر برنامج التشغيل",
|
||||
"res_sys_repair": "إصلاح النظام",
|
||||
"res_cfg_update": "تحديث التكوين"
|
||||
"res_cfg_update": "تحديث التكوين",
|
||||
"manage_startup": "مدير بدء التشغيل",
|
||||
"manage_startup_desc": "إدارة تطبيقات وخدمات بدء التشغيل",
|
||||
"col_status": "الحالة",
|
||||
"status_enabled": "مُمكّن",
|
||||
"status_disabled": "مُعطّل",
|
||||
"btn_enable": "تمكين",
|
||||
"btn_disable": "تعطيل",
|
||||
"search_item": "البحث في العناصر...",
|
||||
"type_reg": "التسجيل",
|
||||
"type_srv": "خدمة",
|
||||
"type_tsk": "مهمة",
|
||||
},
|
||||
"russian_switch": {
|
||||
"app_name": "PYAS Security",
|
||||
@@ -1273,6 +1381,7 @@ window.AppI18n = {
|
||||
"scan_context_desc": "Включить функцию сканирования по правому клику",
|
||||
"theme_color": "Цветовая тема",
|
||||
"theme_color_desc": "Выберите цветовую тему приложения",
|
||||
"theme_system": "Системная тема",
|
||||
"theme_white": "Светлый режим",
|
||||
"theme_red": "Красная тема",
|
||||
"theme_yellow": "Желтая тема",
|
||||
@@ -1358,7 +1467,18 @@ window.AppI18n = {
|
||||
"res_net_block": "Сеть заблокирована",
|
||||
"res_drv_block": "Драйвер заблокирован",
|
||||
"res_sys_repair": "Восстановление системы",
|
||||
"res_cfg_update": "Обновление конфигурации"
|
||||
"res_cfg_update": "Обновление конфигурации",
|
||||
"manage_startup": "Диспетчер автозагрузки",
|
||||
"manage_startup_desc": "Управление приложениями и службами автозагрузки",
|
||||
"col_status": "Статус",
|
||||
"status_enabled": "Включено",
|
||||
"status_disabled": "Отключено",
|
||||
"btn_enable": "Включить",
|
||||
"btn_disable": "Отключить",
|
||||
"search_item": "Поиск элементов...",
|
||||
"type_reg": "Реестр",
|
||||
"type_srv": "Служба",
|
||||
"type_tsk": "Задача",
|
||||
},
|
||||
"slovenian_switch": {
|
||||
"app_name": "PYAS Security",
|
||||
@@ -1409,6 +1529,7 @@ window.AppI18n = {
|
||||
"scan_context_desc": "Omogoči funkcijo skeniranja z desnim klikom",
|
||||
"theme_color": "Barvna tema",
|
||||
"theme_color_desc": "Izberite barvno temo aplikacije",
|
||||
"theme_system": "Sistemska tema",
|
||||
"theme_white": "Svetel način",
|
||||
"theme_red": "Rdeča tema",
|
||||
"theme_yellow": "Rumena tema",
|
||||
@@ -1494,7 +1615,18 @@ window.AppI18n = {
|
||||
"res_net_block": "Omrežje blokirano",
|
||||
"res_drv_block": "Gonilnik blokiran",
|
||||
"res_sys_repair": "Popravilo sistema",
|
||||
"res_cfg_update": "Posodobitev konfiguracije"
|
||||
"res_cfg_update": "Posodobitev konfiguracije",
|
||||
"manage_startup": "Upravitelj zagona",
|
||||
"manage_startup_desc": "Upravljanje zagon aplikacij in storitev",
|
||||
"col_status": "Stanje",
|
||||
"status_enabled": "Omogočeno",
|
||||
"status_disabled": "Onemogočeno",
|
||||
"btn_enable": "Omogoči",
|
||||
"btn_disable": "Onemogoči",
|
||||
"search_item": "Išči elemente...",
|
||||
"type_reg": "Register",
|
||||
"type_srv": "Storitev",
|
||||
"type_tsk": "Opravilo",
|
||||
}
|
||||
},
|
||||
langMap: {
|
||||
|
||||
@@ -1,8 +1,15 @@
|
||||
['dragenter', 'dragover', 'dragleave', 'drop'].forEach(eventName => {
|
||||
document.addEventListener(eventName, e => {
|
||||
e.preventDefault();
|
||||
});
|
||||
});
|
||||
|
||||
document.addEventListener('DOMContentLoaded', () => {
|
||||
const { dict, langMap } = window.AppI18n;
|
||||
|
||||
const appState = {
|
||||
lang: "english_switch",
|
||||
theme: "system_switch",
|
||||
scanning: false,
|
||||
firstLaunch: false,
|
||||
taskmgrTimer: null,
|
||||
@@ -73,11 +80,12 @@ document.addEventListener('DOMContentLoaded', () => {
|
||||
const cols = {
|
||||
log: [{key: 'time_str', label: 'col_date', flex: 1.5}, {key: 'level', label: 'col_type', flex: 0.8}, {key: 'action', label: 'col_func', flex: 1.2, isI18n: true}, {key: 'source', label: 'col_path', flex: 3}],
|
||||
process: [{key: 'name', label: 'col_name', flex: 1.5}, {key: 'pid', label: 'col_pid', flex: 0.5}, {key: 'path', label: 'col_path', flex: 3}],
|
||||
startup: [{key: 'type', label: 'col_type', flex: 0.5, isI18n: true}, {key: 'name', label: 'col_name', flex: 1.5}, {key: 'status', label: 'col_status', flex: 0.5, isI18n: true}, {key: 'path', label: 'col_path', flex: 3}],
|
||||
virus: [{key: 'label', label: 'col_type', flex: 1, isI18n: true}, {key: 'path', label: 'col_path', flex: 3}],
|
||||
pathOnly: [{key: 'path', label: 'col_path', flex: 1}],
|
||||
junk: [{key: 'path', label: 'col_path', flex: 3}, {key: 'sizeStr', label: 'col_size', flex: 1}],
|
||||
popup: [{key: 'exe', label: 'col_prog', flex: 1}, {key: 'class', label: 'col_class', flex: 1}, {key: 'title', label: 'col_title', flex: 2}],
|
||||
repair: [{key: 'display', label: 'col_repair', flex: 1, isI18n: true}]
|
||||
repair: [{key: 'display', label: 'col_repair', flex: 1, isI18n: true}, {key: 'path', label: 'col_path', flex: 3}]
|
||||
};
|
||||
|
||||
const buildCustomSelectElement = (select) => {
|
||||
@@ -192,7 +200,12 @@ document.addEventListener('DOMContentLoaded', () => {
|
||||
};
|
||||
|
||||
const applyTheme = (theme) => {
|
||||
const themeVars = themes[theme];
|
||||
appState.theme = theme;
|
||||
const actualTheme = theme === "system_switch"
|
||||
? (window.matchMedia('(prefers-color-scheme: dark)').matches ? "black_switch" : "white_switch")
|
||||
: theme;
|
||||
|
||||
const themeVars = themes[actualTheme];
|
||||
if (themeVars) {
|
||||
for (const [key, value] of Object.entries(themeVars)) {
|
||||
document.documentElement.style.setProperty(key, value.startsWith("rgba") ? value : `rgb(${value})`);
|
||||
@@ -200,6 +213,10 @@ document.addEventListener('DOMContentLoaded', () => {
|
||||
}
|
||||
};
|
||||
|
||||
window.matchMedia('(prefers-color-scheme: dark)').addEventListener('change', () => {
|
||||
if (appState.theme === "system_switch") applyTheme("system_switch");
|
||||
});
|
||||
|
||||
const renderDataGrid = (containerSelector, dataList, columns, valKey, defaultChecked = false) => {
|
||||
const widget = document.querySelector(`${containerSelector} .manage-widget`) || document.querySelector(`${containerSelector} .virus-widget`);
|
||||
const listUl = document.querySelector(`${containerSelector} .manage-list`) || document.querySelector(`${containerSelector} .virus-list`);
|
||||
@@ -497,6 +514,46 @@ document.addEventListener('DOMContentLoaded', () => {
|
||||
}
|
||||
};
|
||||
|
||||
const fetchStartup = () => {
|
||||
if (!window.pywebview) return;
|
||||
const listUl = document.querySelector('#startup_window .manage-list');
|
||||
if (listUl && !listUl.children.length) {
|
||||
listUl.innerHTML = `<li class="manage-list-item" style="pointer-events: none; justify-content: center; height: 100%; border: none;"><div style="color: var(--text-secondary); padding: 40px; font-size: 15px;">${getMsg('msg_init')}</div></li>`;
|
||||
}
|
||||
window.pywebview.api.get_startup_list().then(items => {
|
||||
const mapped = items.map(item => ({ id: item.id, type: item.type, name: item.name, status: item.status === 'enabled' ? 'status_enabled' : 'status_disabled', path: item.path || 'path_unknown' }));
|
||||
renderDataGrid('#startup_window', mapped, cols.startup, 'id', false);
|
||||
});
|
||||
};
|
||||
|
||||
document.querySelector('#startup_window .primary-btn')?.addEventListener('click', () => {
|
||||
const checked = getCheckedValues('#startup_window');
|
||||
if (checked.length > 0 && window.pywebview) {
|
||||
const widget = document.querySelector('#startup_window .manage-widget');
|
||||
if (widget && widget.gridState && widget.gridState.lastDisplayData) {
|
||||
widget.gridState.lastDisplayData.forEach(item => {
|
||||
if (checked.includes(item.id)) item.status = 'status_enabled';
|
||||
});
|
||||
widget.renderData();
|
||||
}
|
||||
window.pywebview.api.manage_startup(checked, 'enable').then(fetchStartup);
|
||||
}
|
||||
});
|
||||
|
||||
document.querySelector('#startup_window .danger-btn')?.addEventListener('click', () => {
|
||||
const checked = getCheckedValues('#startup_window');
|
||||
if (checked.length > 0 && window.pywebview) {
|
||||
const widget = document.querySelector('#startup_window .manage-widget');
|
||||
if (widget && widget.gridState && widget.gridState.lastDisplayData) {
|
||||
widget.gridState.lastDisplayData.forEach(item => {
|
||||
if (checked.includes(item.id)) item.status = 'status_disabled';
|
||||
});
|
||||
widget.renderData();
|
||||
}
|
||||
window.pywebview.api.manage_startup(checked, 'disable').then(fetchStartup);
|
||||
}
|
||||
});
|
||||
|
||||
const refreshConfigLists = () => {
|
||||
if (!window.pywebview) return;
|
||||
window.pywebview.api.get_config().then(cfg => {
|
||||
@@ -517,7 +574,7 @@ document.addEventListener('DOMContentLoaded', () => {
|
||||
const input = document.querySelector('#junk_window .search-box input');
|
||||
if (input) input.value = '';
|
||||
} else if (oldActive.id === 'repair_window') {
|
||||
renderDataGrid('#repair_window', [], cols.repair, 'value', true);
|
||||
renderDataGrid('#repair_window', [], cols.repair, 'id', true);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -529,6 +586,8 @@ document.addEventListener('DOMContentLoaded', () => {
|
||||
if (appState.taskmgrActive) fetchProcs();
|
||||
else if (appState.taskmgrTimer) { clearTimeout(appState.taskmgrTimer); appState.taskmgrTimer = null; }
|
||||
|
||||
if (targetId === 'startup_window') fetchStartup();
|
||||
|
||||
if (['whitelist_window', 'quarantine_window', 'popup_window', 'custom_protect_window'].includes(targetId)) refreshConfigLists();
|
||||
|
||||
if (targetId === 'log_export_window' && window.pywebview) {
|
||||
@@ -624,7 +683,7 @@ document.addEventListener('DOMContentLoaded', () => {
|
||||
}
|
||||
};
|
||||
|
||||
window.triggerContextScan = (path) => {
|
||||
window.triggerContextScan = (targets) => {
|
||||
if (appState.scanning) return;
|
||||
switchPage('scan_window');
|
||||
if (window.pywebview) {
|
||||
@@ -642,7 +701,8 @@ document.addEventListener('DOMContentLoaded', () => {
|
||||
if (title) { title.setAttribute('data-i18n', 'status_scanning'); title.textContent = getMsg('status_scanning'); }
|
||||
if (text) { text.setAttribute('data-i18n', 'msg_init'); text.removeAttribute('data-dynamic-msg'); text.textContent = getMsg('msg_init'); }
|
||||
|
||||
window.pywebview.api.start_scan([path]);
|
||||
const paths = Array.isArray(targets) ? targets : [targets];
|
||||
window.pywebview.api.start_scan(paths);
|
||||
}
|
||||
};
|
||||
|
||||
@@ -701,7 +761,7 @@ document.addEventListener('DOMContentLoaded', () => {
|
||||
|
||||
document.querySelector('#repair_window .primary-btn')?.addEventListener('click', () => {
|
||||
if (!window.pywebview) return;
|
||||
window.pywebview.api.scan_system_repair().then(list => renderDataGrid('#repair_window', list, cols.repair, 'value', true));
|
||||
window.pywebview.api.scan_system_repair().then(list => renderDataGrid('#repair_window', list, cols.repair, 'id', true));
|
||||
});
|
||||
|
||||
document.querySelector('#repair_window .danger-btn')?.addEventListener('click', () => {
|
||||
@@ -709,12 +769,19 @@ document.addEventListener('DOMContentLoaded', () => {
|
||||
if (checked.length === 0 || !window.pywebview) return;
|
||||
window.pywebview.api.execute_system_repair(checked).then(() => {
|
||||
window.pywebview.api.show_alert(getMsg("title_prompt"), getMsg("msg_repaired"), "info");
|
||||
renderDataGrid('#repair_window', [], cols.repair, 'value');
|
||||
renderDataGrid('#repair_window', [], cols.repair, 'id', true);
|
||||
});
|
||||
});
|
||||
|
||||
['whitelist_window', 'quarantine_window', 'popup_window', 'custom_protect_window'].forEach(id => {
|
||||
const listKey = id === 'popup_window' ? 'block_list' : (id === 'custom_protect_window' ? 'custom_rule' : id.replace('_window', ''));
|
||||
const listKeyMap = {
|
||||
'popup_window': 'block_list',
|
||||
'custom_protect_window': 'custom_rule',
|
||||
'whitelist_window': 'white_list',
|
||||
'quarantine_window': 'quarantine'
|
||||
};
|
||||
const listKey = listKeyMap[id];
|
||||
|
||||
document.querySelector(`#${id} .modern-btn:not(.primary-btn)`)?.addEventListener('click', () => {
|
||||
const checked = getCheckedValues(`#${id}`);
|
||||
if (checked.length > 0 && window.pywebview) window.pywebview.api.remove_list_items(listKey, checked).then(refreshConfigLists);
|
||||
@@ -847,8 +914,9 @@ document.addEventListener('DOMContentLoaded', () => {
|
||||
document.querySelectorAll('select.modern-select').forEach(buildCustomSelectElement);
|
||||
renderDataGrid('#scan_window', appState.virusResults, cols.virus, 'path', true);
|
||||
renderDataGrid('#taskmgr_window', [], cols.process, 'pid', false);
|
||||
renderDataGrid('#startup_window', [], cols.startup, 'id', false);
|
||||
renderDataGrid('#junk_window', [], cols.junk, 'path', true);
|
||||
renderDataGrid('#repair_window', [], cols.repair, 'value', true);
|
||||
renderDataGrid('#repair_window', [], cols.repair, 'id', true);
|
||||
renderDataGrid('#whitelist_window', [], cols.pathOnly, 'path', false);
|
||||
renderDataGrid('#quarantine_window', [], cols.pathOnly, 'path', false);
|
||||
renderDataGrid('#custom_protect_window', [], cols.pathOnly, 'path', false);
|
||||
|
||||
@@ -150,6 +150,12 @@
|
||||
</div>
|
||||
<button class="icon-action-btn"><svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><polyline points="9 18 15 12 9 6"></polyline></svg></button>
|
||||
</div>
|
||||
<div class="list-item" style="cursor: pointer;" data-target="startup_window">
|
||||
<div class="item-info">
|
||||
<h2 data-i18n>啟動項管理</h2><p data-i18n>此選項可以管理各應用啟動項</p>
|
||||
</div>
|
||||
<button class="icon-action-btn"><svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><polyline points="9 18 15 12 9 6"></polyline></svg></button>
|
||||
</div>
|
||||
<div class="list-item" style="cursor: pointer;" data-target="junk_window">
|
||||
<div class="item-info">
|
||||
<h2 data-i18n>垃圾清理</h2><p data-i18n>此選項可以清理暫存垃圾檔案</p>
|
||||
@@ -214,6 +220,7 @@
|
||||
<h2 data-i18n>顯色主題</h2><p data-i18n>此選項可以選擇主題顏色</p>
|
||||
</div>
|
||||
<select id="theme_select" class="modern-select">
|
||||
<option value="system_switch" selected data-i18n>跟隨系統</option>
|
||||
<option value="white_switch" data-i18n>白色主題</option>
|
||||
<option value="red_switch" data-i18n>紅色主題</option>
|
||||
<option value="yellow_switch" data-i18n>黃色主題</option>
|
||||
@@ -320,6 +327,21 @@ Last updated: 2025-12-19
|
||||
<div class="manage-widget flex-grow"><ul class="manage-list"></ul></div>
|
||||
</section>
|
||||
|
||||
<section id="startup_window" class="page sub-page">
|
||||
<div class="sub-header-toolbar">
|
||||
<button class="icon-btn" data-target="tool_window"><svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><polyline points="15 18 9 12 15 6"></polyline></svg></button>
|
||||
<h2 data-i18n>啟動項管理</h2>
|
||||
<div class="spacer"></div>
|
||||
<div class="search-box">
|
||||
<svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><circle cx="11" cy="11" r="8"></circle><line x1="21" y1="21" x2="16.65" y2="16.65"></line></svg>
|
||||
<input type="text" data-i18n-placeholder="搜尋項目..." placeholder="搜尋項目...">
|
||||
</div>
|
||||
<button class="action-btn modern-btn primary-btn" data-i18n>啟用</button>
|
||||
<button class="action-btn modern-btn danger-btn" data-i18n>禁用</button>
|
||||
</div>
|
||||
<div class="manage-widget flex-grow"><ul class="manage-list"></ul></div>
|
||||
</section>
|
||||
|
||||
<section id="whitelist_window" class="page sub-page">
|
||||
<div class="sub-header-toolbar">
|
||||
<button class="icon-btn" data-target="tool_window"><svg viewBox="0 0 24 24" fill="none" stroke="currentColor" stroke-width="2" stroke-linecap="round" stroke-linejoin="round"><polyline points="15 18 9 12 15 6"></polyline></svg></button>
|
||||
|
||||
+353
-76
@@ -1,4 +1,4 @@
|
||||
import os, re, yara, time, math, json, numpy, base64, datetime, requests
|
||||
import os, re, yara, time, math, json, numpy, base64, datetime, requests, zlib
|
||||
import ctypes, ctypes.wintypes, hashlib, pefile, threading, onnxruntime
|
||||
|
||||
####################################################################################################
|
||||
@@ -154,7 +154,22 @@ class rule_scanner:
|
||||
rule_name = str(matches[0])
|
||||
label = rule_name.split("_")[0]
|
||||
level = rule_name.split("_")[-1]
|
||||
return f"{label}:WinPE/Unknown.A!rb", level
|
||||
return f"{label}:WinPE/Unknown.{level}!yr", level
|
||||
return False, False
|
||||
except Exception:
|
||||
return False, False
|
||||
|
||||
def yara_mem_scan(self, pid):
|
||||
try:
|
||||
if not self.rules:
|
||||
return False, False
|
||||
|
||||
matches = self.rules.match(pid=pid)
|
||||
if matches:
|
||||
rule_name = str(matches[0])
|
||||
label = rule_name.split("_")[0]
|
||||
level = rule_name.split("_")[-1]
|
||||
return f"{label}:WinPE/Unknown.{level}!ym", level
|
||||
return False, False
|
||||
except Exception:
|
||||
return False, False
|
||||
@@ -162,6 +177,19 @@ class rule_scanner:
|
||||
####################################################################################################
|
||||
|
||||
class pe_scanner:
|
||||
_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 = numpy.zeros(2049, dtype=numpy.float64)
|
||||
_C_LOG_C_TABLE[1:] = numpy.arange(1, 2049) * numpy.log2(numpy.arange(1, 2049))
|
||||
|
||||
def __init__(self):
|
||||
self.model = None
|
||||
self.input_name = None
|
||||
@@ -204,9 +232,234 @@ class pe_scanner:
|
||||
if not data:
|
||||
return 0.0
|
||||
arr = numpy.frombuffer(data, dtype=numpy.uint8)
|
||||
sz = len(arr)
|
||||
counts = numpy.bincount(arr, minlength=256)
|
||||
probs = counts[counts > 0] / len(arr)
|
||||
return float(-numpy.sum(probs * numpy.log2(probs)))
|
||||
counts = counts[counts > 0]
|
||||
return float(numpy.log2(sz) - numpy.sum(counts * numpy.log2(counts)) / sz)
|
||||
|
||||
def _extract_strings(self, file_bytes):
|
||||
res = {"StringCount": 0.0, "StringMeanLength": 0.0, "StringEntropy": 0.0}
|
||||
res.update(dict.fromkeys(self._STRING_KEYS, 0.0))
|
||||
|
||||
matches = self._STRING_PATTERN.findall(file_bytes)
|
||||
if not matches:
|
||||
return res
|
||||
|
||||
combined_data = b''.join(matches)
|
||||
count = len(matches)
|
||||
total_len = len(combined_data)
|
||||
|
||||
res["StringCount"] = float(count)
|
||||
res["StringMeanLength"] = float(total_len) / count
|
||||
res["StringEntropy"] = self._calc_entropy(combined_data)
|
||||
|
||||
arr = numpy.frombuffer(combined_data, dtype=numpy.uint8)
|
||||
counts = numpy.bincount(arr, minlength=127)
|
||||
char_counts = counts[0x20:0x7F]
|
||||
|
||||
res.update(dict(zip(self._STRING_KEYS, char_counts.astype(numpy.float64) / total_len)))
|
||||
return res
|
||||
|
||||
def _extract_histograms(self, file_bytes):
|
||||
arr = numpy.frombuffer(file_bytes, dtype=numpy.uint8)
|
||||
sz = len(arr)
|
||||
res = {}
|
||||
|
||||
if sz == 0:
|
||||
res.update(dict.fromkeys(self._BYTE_HIST_KEYS, 0.0))
|
||||
res.update(dict.fromkeys(self._BYTE_ENT_KEYS, 0.0))
|
||||
return res
|
||||
|
||||
byte_counts = numpy.bincount(arr, minlength=256)
|
||||
byte_hist = byte_counts.astype(numpy.float64) / sz
|
||||
res.update(dict(zip(self._BYTE_HIST_KEYS, byte_hist)))
|
||||
|
||||
ent_hist = numpy.zeros(256, dtype=numpy.float64)
|
||||
window = 2048
|
||||
step = 1024
|
||||
|
||||
if sz < window:
|
||||
w_counts = byte_counts
|
||||
p = w_counts[w_counts > 0] / float(sz)
|
||||
entropy = float(-numpy.sum(p * numpy.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_windows = 10000
|
||||
|
||||
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
|
||||
|
||||
start_byte = w_start * step
|
||||
end_byte = start_byte + b_num * step + step
|
||||
b_arr = arr[start_byte:end_byte]
|
||||
|
||||
arr_2d = b_arr.reshape(b_num + 1, step)
|
||||
offsets = (numpy.arange(b_num + 1, dtype=numpy.int32)[:, None] * 256)
|
||||
flat_idx = (arr_2d.astype(numpy.int32) + offsets).ravel()
|
||||
|
||||
chunk_counts = numpy.bincount(flat_idx, minlength=(b_num + 1) * 256).reshape(b_num + 1, 256)
|
||||
window_counts = chunk_counts[:-1] + chunk_counts[1:]
|
||||
|
||||
sum_c_log_c = numpy.sum(self._C_LOG_C_TABLE[window_counts], axis=1)
|
||||
entropies = 11.0 - sum_c_log_c / 2048.0
|
||||
|
||||
ent_bins = (entropies * 2.0).astype(numpy.int32)
|
||||
numpy.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 + numpy.arange(16)).ravel()
|
||||
numpy.add.at(ent_hist, flat_ent_bins, byte_bins.ravel())
|
||||
|
||||
ent_sum = numpy.sum(ent_hist)
|
||||
if ent_sum > 0:
|
||||
ent_hist /= ent_sum
|
||||
|
||||
res.update(dict(zip(self._BYTE_ENT_KEYS, ent_hist)))
|
||||
return res
|
||||
|
||||
def _extract_overlay_features(self, 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": self._calc_entropy(overlay_data)
|
||||
}
|
||||
|
||||
def _extract_rich_header(self, 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
|
||||
}
|
||||
|
||||
def _extract_ep_anomalies(self, 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
|
||||
}
|
||||
|
||||
def _extract_advanced_resources(self, 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(self._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
|
||||
|
||||
def _extract_load_config(self, 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
|
||||
|
||||
def _extract_security_directory(self, 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
|
||||
|
||||
def extract_features(self, file_path):
|
||||
pe = None
|
||||
@@ -223,6 +476,19 @@ class pe_scanner:
|
||||
apis = 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'] = 1.0 if self.signer.sign_verify(file_path) else 0.0
|
||||
base['FileEntropy'] = self._calc_entropy(file_bytes)
|
||||
base['FileSize'] = float(fsize)
|
||||
@@ -271,22 +537,22 @@ class pe_scanner:
|
||||
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
|
||||
res_flags_count = dict.fromkeys(self._CHAR_COUNT_KEYS.values(), 0.0)
|
||||
res_flags_ent = dict.fromkeys(self._CHAR_ENT_KEYS.values(), 0.0)
|
||||
res_sec_hash = dict.fromkeys(self._SECTION_HASH_KEYS, 0.0)
|
||||
|
||||
if hasattr(pe, 'sections'):
|
||||
base['SectionCount'] = len(pe.sections)
|
||||
entropies = []
|
||||
raw_sizes = []
|
||||
v_sizes = []
|
||||
exec_sec = 0
|
||||
write_sec = 0
|
||||
read_sec = 0
|
||||
sec_exc = 0
|
||||
exec_sec = write_sec = read_sec = sec_exc = 0
|
||||
|
||||
for section in pe.sections:
|
||||
sec_name = section.Name.rstrip(b'\x00')
|
||||
hash_val = zlib.crc32(sec_name) % 50
|
||||
res_sec_hash[self._SECTION_HASH_KEYS[hash_val]] += 1.0
|
||||
|
||||
try:
|
||||
s_data = section.get_data()
|
||||
s_entropy = self._calc_entropy(s_data)
|
||||
@@ -302,10 +568,10 @@ class pe_scanner:
|
||||
if section.Characteristics & 0x40000000: read_sec += 1
|
||||
if section.SizeOfRawData + section.PointerToRawData > fsize: sec_exc = 1
|
||||
|
||||
for flag in char_flags:
|
||||
for flag in self._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[self._CHAR_COUNT_KEYS[flag]] += 1.0
|
||||
res_flags_ent[self._CHAR_ENT_KEYS[flag]] += s_entropy
|
||||
|
||||
base['SectionMaxEntropy'] = max(entropies) if entropies else 0.0
|
||||
base['SectionMinEntropy'] = min(entropies) if entropies else 0.0
|
||||
@@ -321,9 +587,14 @@ class pe_scanner:
|
||||
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 self._CHAR_FLAGS:
|
||||
cnt_key = self._CHAR_COUNT_KEYS[flag]
|
||||
if res_flags_count[cnt_key] > 0:
|
||||
res_flags_ent[self._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:
|
||||
@@ -343,17 +614,6 @@ class pe_scanner:
|
||||
except Exception:
|
||||
continue
|
||||
|
||||
try:
|
||||
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']
|
||||
])
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
if hasattr(pe, 'DIRECTORY_ENTRY_TLS') and hasattr(pe.DIRECTORY_ENTRY_TLS, 'struct'):
|
||||
if getattr(pe.DIRECTORY_ENTRY_TLS.struct, 'AddressOfCallBacks', 0) != 0:
|
||||
base['HasTlsCallbacks'] = 1.0
|
||||
@@ -406,6 +666,15 @@ class pe_scanner:
|
||||
else:
|
||||
base['DebugCount'] = 0.0
|
||||
|
||||
base.update(self._extract_strings(file_bytes))
|
||||
base.update(self._extract_histograms(file_bytes))
|
||||
base.update(self._extract_overlay_features(pe, fsize))
|
||||
base.update(self._extract_rich_header(pe))
|
||||
base.update(self._extract_ep_anomalies(pe))
|
||||
base.update(self._extract_advanced_resources(pe))
|
||||
base.update(self._extract_load_config(pe))
|
||||
base.update(self._extract_security_directory(pe, file_bytes, fsize))
|
||||
|
||||
for k in base:
|
||||
base[k] = self._safe_float(base[k])
|
||||
|
||||
@@ -429,16 +698,29 @@ class pe_scanner:
|
||||
base = raw_data.get('Base', {})
|
||||
dlls = set(raw_data.get('DLLs', []))
|
||||
apis = set(raw_data.get('APIs', []))
|
||||
feat_map = {feat: i for i, feat in enumerate(self.feature_order)}
|
||||
|
||||
for i, feat in enumerate(self.feature_order):
|
||||
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 feat_map:
|
||||
vec[0, feat_map[k]] = v
|
||||
|
||||
for d in dlls:
|
||||
old_fn = f"Dll_{d}"
|
||||
if old_fn in feat_map:
|
||||
vec[0, feat_map[old_fn]] = 1.0
|
||||
h = zlib.crc32(d.encode('utf-8', 'ignore')) % 256
|
||||
new_fn = f"DllHash_{h:03d}"
|
||||
if new_fn in feat_map:
|
||||
vec[0, feat_map[new_fn]] += 1.0
|
||||
|
||||
for a in apis:
|
||||
old_fn = f"Api_{a}"
|
||||
if old_fn in feat_map:
|
||||
vec[0, feat_map[old_fn]] = 1.0
|
||||
h = zlib.crc32(a.encode('utf-8', 'ignore')) % 1024
|
||||
new_fn = f"ApiHash_{h:04d}"
|
||||
if new_fn in feat_map:
|
||||
vec[0, feat_map[new_fn]] += 1.0
|
||||
|
||||
outputs = self.model.run(None, {self.input_name: vec})
|
||||
result = outputs[1]
|
||||
@@ -514,7 +796,7 @@ class cloud_scanner:
|
||||
r = self._request("POST", f"/api/rescan/{sha256}")
|
||||
return r is not None and r.json().get('status') == 'success'
|
||||
|
||||
def upload_file(self, file_path, api_host, api_key, chunk_size=5242880, need_rescan=False):
|
||||
def upload_file(self, file_path, api_host, api_key, chunk_size=4194304, need_rescan=False, max_retries=3):
|
||||
try:
|
||||
self._init_session(api_host, api_key)
|
||||
sha256 = self.calc_hash(file_path)
|
||||
@@ -523,9 +805,7 @@ class cloud_scanner:
|
||||
|
||||
status_req = self._request("GET", f"/api/processing_status/{sha256}")
|
||||
if status_req:
|
||||
status_data = status_req.json()
|
||||
current_status = status_data.get('status')
|
||||
|
||||
current_status = status_req.json().get('status')
|
||||
if current_status == 'done':
|
||||
if need_rescan:
|
||||
self.rescan(sha256, api_host, api_key)
|
||||
@@ -537,40 +817,37 @@ class cloud_scanner:
|
||||
if file_size > 104857600:
|
||||
return False, sha256
|
||||
|
||||
if file_size > chunk_size:
|
||||
total_chunks = (file_size + chunk_size - 1) // chunk_size
|
||||
upload_id = os.urandom(16).hex()
|
||||
with open(file_path, 'rb') as f:
|
||||
for i in range(total_chunks):
|
||||
chunk_data = f.read(chunk_size)
|
||||
headers = {
|
||||
"X-Chunk-Index": str(i),
|
||||
"X-Total-Chunks": str(total_chunks),
|
||||
"X-Upload-ID": upload_id
|
||||
}
|
||||
|
||||
chunk_success = False
|
||||
for _ in range(3):
|
||||
r = self._request("POST", "/api/upload", files={'file': (os.path.basename(file_path), chunk_data)}, headers=headers)
|
||||
if r:
|
||||
chunk_success = True
|
||||
break
|
||||
time.sleep(1)
|
||||
|
||||
if not chunk_success:
|
||||
return False, sha256
|
||||
return True, sha256
|
||||
|
||||
total_chunks = max(1, (file_size + chunk_size - 1) // chunk_size)
|
||||
upload_id = os.urandom(16).hex()
|
||||
|
||||
with open(file_path, 'rb') as f:
|
||||
file_data = f.read()
|
||||
|
||||
for _ in range(3):
|
||||
r = self._request("POST", "/api/upload", files={'file': (os.path.basename(file_path), file_data)})
|
||||
if r:
|
||||
return True, sha256
|
||||
time.sleep(1)
|
||||
|
||||
return False, sha256
|
||||
for i in range(total_chunks):
|
||||
chunk_data = f.read(chunk_size)
|
||||
headers = {
|
||||
"X-Chunk-Index": str(i),
|
||||
"X-Total-Chunks": str(total_chunks),
|
||||
"X-Upload-ID": upload_id
|
||||
}
|
||||
|
||||
chunk_success = False
|
||||
for attempt in range(max_retries):
|
||||
r = self._request("POST", "/api/upload", files={'file': (os.path.basename(file_path), chunk_data)}, headers=headers)
|
||||
if r:
|
||||
if i == total_chunks - 1:
|
||||
try:
|
||||
resp = r.json()
|
||||
if 'url' in resp:
|
||||
sha256 = resp.get('url', '').split('/')[-1]
|
||||
except Exception:
|
||||
pass
|
||||
chunk_success = True
|
||||
break
|
||||
time.sleep(2 ** attempt)
|
||||
|
||||
if not chunk_success:
|
||||
return False, sha256
|
||||
|
||||
return True, sha256
|
||||
except Exception:
|
||||
return False, None
|
||||
|
||||
@@ -615,7 +892,7 @@ class cloud_scanner:
|
||||
sim_malicious_count += 1
|
||||
|
||||
if is_malicious and (valid_sim_count == 0 or sim_malicious_count == valid_sim_count):
|
||||
return f"General:WinPE/Unknown.{score}!cl"
|
||||
return f"General:WinPE/Malware.{score}!cl"
|
||||
except Exception:
|
||||
pass
|
||||
return False
|
||||
+1
-1
@@ -2,7 +2,7 @@ import pyinstaller_versionfile
|
||||
|
||||
pyinstaller_versionfile.create_versionfile(
|
||||
output_file="versionfile.txt",
|
||||
version='3.5.4.0',
|
||||
version='3.5.7.0',
|
||||
company_name="PYAS Security",
|
||||
file_description="PYAS Security Antivirus",
|
||||
internal_name="PYAS",
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
#include "DriverCommon.h"
|
||||
|
||||
DRIVER_DATA GlobalData;
|
||||
static BOOLEAN g_ImageNotifyRegistered = FALSE;
|
||||
|
||||
static NTSTATUS InstanceSetup(PCFLT_RELATED_OBJECTS FltObjects, FLT_INSTANCE_SETUP_FLAGS Flags, DEVICE_TYPE VolumeDeviceType, FLT_FILESYSTEM_TYPE VolumeFilesystemType) {
|
||||
UNREFERENCED_PARAMETER(FltObjects);
|
||||
@@ -27,9 +28,7 @@ static VOID InstanceTeardownComplete(PCFLT_RELATED_OBJECTS FltObjects, FLT_INSTA
|
||||
}
|
||||
|
||||
static NTSTATUS DriverUnload(FLT_FILTER_UNLOAD_FLAGS Flags) {
|
||||
if (!(Flags & FLTFL_FILTER_UNLOAD_MANDATORY)) {
|
||||
return STATUS_FLT_DO_NOT_DETACH;
|
||||
}
|
||||
UNREFERENCED_PARAMETER(Flags);
|
||||
|
||||
if (GlobalData.ServerPort) {
|
||||
FltCloseCommunicationPort(GlobalData.ServerPort);
|
||||
@@ -41,12 +40,14 @@ static NTSTATUS DriverUnload(FLT_FILTER_UNLOAD_FLAGS Flags) {
|
||||
GlobalData.FilterHandle = NULL;
|
||||
}
|
||||
|
||||
PsRemoveLoadImageNotifyRoutine(ImageLoadNotify);
|
||||
if (g_ImageNotifyRegistered) {
|
||||
PsRemoveLoadImageNotifyRoutine(ImageLoadNotify);
|
||||
g_ImageNotifyRegistered = FALSE;
|
||||
}
|
||||
|
||||
UninitializeRegistryProtection();
|
||||
UninitializeProcessProtection();
|
||||
|
||||
ExWaitForRundownProtectionRelease(&GlobalData.PortRundown);
|
||||
|
||||
UnloadRules();
|
||||
UninitializeRulesEngine();
|
||||
|
||||
@@ -61,12 +62,9 @@ static NTSTATUS PortMessage(PVOID PortCookie, PVOID InputBuffer, ULONG InputBuff
|
||||
if (InputBuffer && InputBufferLength >= sizeof(PYAS_USER_MESSAGE)) {
|
||||
PPYAS_USER_MESSAGE msg = (PPYAS_USER_MESSAGE)InputBuffer;
|
||||
if (msg->Command == 1 || msg->Command == 2) {
|
||||
|
||||
msg->Path[MAX_PATH_LEN - 1] = L'\0';
|
||||
|
||||
UNICODE_STRING us;
|
||||
RtlInitUnicodeString(&us, msg->Path);
|
||||
|
||||
if (us.Length > 0) {
|
||||
if (msg->Command == 1) {
|
||||
AddDynamicWhitelist(&us);
|
||||
@@ -178,22 +176,29 @@ extern "C" NTSTATUS DriverEntry(PDRIVER_OBJECT DriverObject, PUNICODE_STRING Reg
|
||||
}
|
||||
|
||||
if (!NT_SUCCESS(status)) {
|
||||
if (GlobalData.FilterHandle) {
|
||||
FltUnregisterFilter(GlobalData.FilterHandle);
|
||||
GlobalData.FilterHandle = NULL;
|
||||
}
|
||||
FltUnregisterFilter(GlobalData.FilterHandle);
|
||||
GlobalData.FilterHandle = NULL;
|
||||
UnloadRules();
|
||||
UninitializeRulesEngine();
|
||||
return status;
|
||||
}
|
||||
|
||||
InitializeProcessProtection();
|
||||
|
||||
InitializeRegistryProtection(DriverObject);
|
||||
PsSetLoadImageNotifyRoutine(ImageLoadNotify);
|
||||
|
||||
NTSTATUS imageNotifyStatus = PsSetLoadImageNotifyRoutine(ImageLoadNotify);
|
||||
if (NT_SUCCESS(imageNotifyStatus)) {
|
||||
g_ImageNotifyRegistered = TRUE;
|
||||
}
|
||||
|
||||
status = FltStartFiltering(GlobalData.FilterHandle);
|
||||
if (!NT_SUCCESS(status)) {
|
||||
PsRemoveLoadImageNotifyRoutine(ImageLoadNotify);
|
||||
if (g_ImageNotifyRegistered) {
|
||||
PsRemoveLoadImageNotifyRoutine(ImageLoadNotify);
|
||||
g_ImageNotifyRegistered = FALSE;
|
||||
}
|
||||
|
||||
UninitializeRegistryProtection();
|
||||
UninitializeProcessProtection();
|
||||
|
||||
|
||||
@@ -189,8 +189,7 @@ static OB_PREOP_CALLBACK_STATUS PreOpenThread(PVOID RegistrationContext, POB_PRE
|
||||
}
|
||||
|
||||
NTSTATUS InitializeProcessProtection() {
|
||||
static UNICODE_STRING Altitude = { 0 };
|
||||
RtlInitUnicodeString(&Altitude, L"320000");
|
||||
static UNICODE_STRING Altitude = RTL_CONSTANT_STRING(L"320000.PYAS.Ob");
|
||||
|
||||
ObCallbacks[0].ObjectType = PsProcessType;
|
||||
ObCallbacks[0].Operations = OB_OPERATION_HANDLE_CREATE | OB_OPERATION_HANDLE_DUPLICATE;
|
||||
@@ -208,7 +207,15 @@ NTSTATUS InitializeProcessProtection() {
|
||||
ObRegistration.RegistrationContext = NULL;
|
||||
ObRegistration.OperationRegistration = ObCallbacks;
|
||||
|
||||
return ObRegisterCallbacks(&ObRegistration, &ObRegistrationHandle);
|
||||
NTSTATUS status = ObRegisterCallbacks(&ObRegistration, &ObRegistrationHandle);
|
||||
|
||||
if (!NT_SUCCESS(status) && (status == STATUS_FLT_INSTANCE_ALTITUDE_COLLISION || status == STATUS_OBJECT_NAME_COLLISION)) {
|
||||
static UNICODE_STRING FallbackAltitude = RTL_CONSTANT_STRING(L"320000.PYAS.Ob.Fallback");
|
||||
ObRegistration.Altitude = FallbackAltitude;
|
||||
status = ObRegisterCallbacks(&ObRegistration, &ObRegistrationHandle);
|
||||
}
|
||||
|
||||
return status;
|
||||
}
|
||||
|
||||
VOID UninitializeProcessProtection() {
|
||||
|
||||
@@ -235,9 +235,15 @@ static NTSTATUS RegistryCallback(PVOID CallbackContext, PVOID Argument1, PVOID A
|
||||
}
|
||||
|
||||
NTSTATUS InitializeRegistryProtection(PDRIVER_OBJECT DriverObject) {
|
||||
UNICODE_STRING Altitude;
|
||||
RtlInitUnicodeString(&Altitude, L"320000");
|
||||
return CmRegisterCallbackEx(RegistryCallback, &Altitude, DriverObject, NULL, &Cookie, NULL);
|
||||
static UNICODE_STRING Altitude = RTL_CONSTANT_STRING(L"320000.PYAS.Cm");
|
||||
NTSTATUS status = CmRegisterCallbackEx(RegistryCallback, &Altitude, DriverObject, NULL, &Cookie, NULL);
|
||||
|
||||
if (!NT_SUCCESS(status)) {
|
||||
static UNICODE_STRING FallbackAltitude = RTL_CONSTANT_STRING(L"320000.PYAS.Cm.Fallback");
|
||||
status = CmRegisterCallbackEx(RegistryCallback, &FallbackAltitude, DriverObject, NULL, &Cookie, NULL);
|
||||
}
|
||||
|
||||
return status;
|
||||
}
|
||||
|
||||
VOID UninitializeRegistryProtection() {
|
||||
|
||||
@@ -525,7 +525,7 @@ BOOLEAN IsProcessTrusted(HANDLE ProcessId) {
|
||||
if (g_CacheInitialized) {
|
||||
KIRQL OldIrql;
|
||||
KeAcquireSpinLock(&TrustCacheLock, &OldIrql);
|
||||
ULONG Hash = ((ULONG)(ULONG_PTR)ProcessId) & (TRUST_CACHE_SIZE - 1);
|
||||
ULONG Hash = (((ULONG)(ULONG_PTR)ProcessId) >> 2) & (TRUST_CACHE_SIZE - 1);
|
||||
|
||||
if (TrustCache[Hash].Process == Process && TrustCache[Hash].ProcessCreateTime.QuadPart == createTime.QuadPart) {
|
||||
LARGE_INTEGER Now;
|
||||
@@ -580,7 +580,7 @@ update_cache:
|
||||
if (g_CacheInitialized) {
|
||||
KIRQL OldIrql;
|
||||
KeAcquireSpinLock(&TrustCacheLock, &OldIrql);
|
||||
ULONG Hash = ((ULONG)(ULONG_PTR)ProcessId) & (TRUST_CACHE_SIZE - 1);
|
||||
ULONG Hash = (((ULONG)(ULONG_PTR)ProcessId) >> 2) & (TRUST_CACHE_SIZE - 1);
|
||||
TrustCache[Hash].Process = Process;
|
||||
TrustCache[Hash].ProcessCreateTime = createTime;
|
||||
TrustCache[Hash].IsTrusted = isTrusted;
|
||||
|
||||
@@ -77,7 +77,6 @@
|
||||
"*\\FirmwareBootDevice",
|
||||
"*\\BootDriverFlags",
|
||||
"*\\SystemStartOptions",
|
||||
"*\\CrashDumpEnabled",
|
||||
"*\\DisableAutomaticRestartOnFailure",
|
||||
"*\\SystemBiosVersion",
|
||||
"*\\NoTheme",
|
||||
@@ -129,11 +128,15 @@
|
||||
"*\\360Safe.exe",
|
||||
"*\\bilibili\\bilibili.exe",
|
||||
|
||||
"*\\Kingsoft\\WPS Office\\ksolaunch.exe",
|
||||
"*\\Kingsoft\\WPS Office\\*\\wps.exe",
|
||||
|
||||
"*\\PC4399_WPFLauncher.exe",
|
||||
"*\\bin\\APlayer.exe",
|
||||
"*\\Kingsoft\\WPS Office\\ksolaunch.exe",
|
||||
"*\\Neverness To Everness\\NTELauncher.exe",
|
||||
"*\\Application\\AndrowsLauncher.exe",
|
||||
"*\\bin\\*\\codex.exe",
|
||||
"*\\bin\\codex.exe",
|
||||
],
|
||||
"Rule_Process_ExploitableBlacklist": [
|
||||
"*\\cmd.exe",
|
||||
|
||||
@@ -132,7 +132,7 @@ graph TD
|
||||
|
||||
The machine learning-based virus scanning engine training achieved extremely high accuracy.
|
||||
|
||||
<img width="1191" height="1323" alt="train" src="https://github.com/user-attachments/assets/2eea8426-1f82-44e3-9a38-f96f9257cf50" />
|
||||
<img width="895" height="1097" alt="train" src="https://github.com/user-attachments/assets/fa7cf0e2-01a1-4234-9358-718c2b8a2c9a" />
|
||||
|
||||
## Support System
|
||||
|
||||
|
||||
Reference in new issue
Block a user