Compare commits

...
45 Commits
Author SHA1 Message Date
87owo c220035fd9 Add files via upload 2026-05-21 00:27:56 +08:00
87owo d7b98d3cde Add files via upload 2026-05-21 00:27:38 +08:00
87owo 0ab876633f Add files via upload 2026-05-21 00:27:01 +08:00
87owo cabe76ca54 Add files via upload 2026-05-21 00:26:14 +08:00
87owo c57224e258 Add files via upload 2026-05-21 00:25:35 +08:00
87owo b56101d2e6 Add files via upload 2026-05-18 21:31:43 +08:00
87owo 16211c28d1 Add files via upload 2026-05-17 16:48:13 +08:00
87owo 139d3e0c3d Add files via upload 2026-05-17 16:47:14 +08:00
87owo f0652fc5ec Add files via upload 2026-05-17 01:32:52 +08:00
87owo 55489cd634 Add files via upload 2026-05-17 01:31:49 +08:00
87owo 742887c0e4 Add files via upload 2026-05-15 22:34:23 +08:00
87owo c3e5cd2c0d Add files via upload 2026-05-15 22:33:13 +08:00
87owo 63e125627c Update README.md 2026-05-15 14:39:15 +08:00
87owo c0189e9d05 Add files via upload 2026-05-14 22:31:46 +08:00
87owo 133e3208ec Add files via upload 2026-05-14 20:02:53 +08:00
87owo 17c6fd6d0f Add files via upload 2026-05-13 18:37:07 +08:00
87owo 2aa7ea43b8 Add files via upload 2026-05-13 17:56:05 +08:00
87owo 6d7dbfac31 Add files via upload 2026-05-13 00:51:04 +08:00
87owo d4942c6dd6 Add files via upload 2026-05-13 00:49:54 +08:00
87owo 39ae3c24f6 Add files via upload 2026-05-12 22:41:11 +08:00
87owo a04d74a949 Add files via upload 2026-05-10 17:54:55 +08:00
87owo b93a7ebd94 Add files via upload 2026-05-10 17:45:08 +08:00
87owo cbc7b0340f Add files via upload 2026-05-10 17:38:28 +08:00
87owo ea7a197848 Add files via upload 2026-05-10 17:31:14 +08:00
87owo 82d4df889e Add files via upload 2026-05-10 13:06:24 +08:00
87owo df3a2a1010 Add files via upload 2026-05-10 13:06:05 +08:00
87owo 67cd6c2f59 Add files via upload 2026-05-10 02:30:07 +08:00
87owo b1788e7df9 Add files via upload 2026-05-10 02:28:27 +08:00
87owo 1cb8a81518 Add files via upload 2026-05-10 00:06:45 +08:00
87owo e1b61f2675 Update README.md 2026-05-09 22:53:01 +08:00
87owo 834d44856c Add files via upload 2026-05-09 22:49:48 +08:00
87owo c088eeb5bc Add files via upload 2026-05-09 22:47:09 +08:00
87owo 921510b7b7 Add files via upload 2026-05-09 22:27:36 +08:00
87owo 17a54f7a8a Add files via upload 2026-05-08 20:02:34 +08:00
87owo 43ff924efc Add files via upload 2026-05-08 18:24:36 +08:00
87owo f9bd4b6628 Add files via upload 2026-05-08 18:20:52 +08:00
87owo 702aa22139 Delete Interface/static/js/style.css 2026-05-08 18:20:37 +08:00
87owo 57d53f63e5 Add files via upload 2026-05-08 18:17:20 +08:00
87owo 835ada4a53 Add files via upload 2026-05-08 13:51:08 +08:00
87owo 7f204f3e1b Add files via upload 2026-05-08 10:53:11 +08:00
87owo 70e13b8414 Delete Engine/Properties/train_server_lfp.py 2026-05-07 14:13:19 +08:00
87owo df9bb8cbb8 Add files via upload 2026-05-07 14:12:56 +08:00
87owo cebdb32502 Add files via upload 2026-05-07 13:58:44 +08:00
87owo 58bd61343a Add files via upload 2026-05-06 14:27:54 +08:00
87owo 90faf1007e Add files via upload 2026-05-06 14:14:48 +08:00
18 changed files with 5578 additions and 5692 deletions

No files matched your search

File diff suppressed because it is too large. Load diff
+680
View File
@@ -0,0 +1,680 @@
import os, sys, math, json, pefile, datetime
import numpy as np
import onnxruntime as ort
import ctypes, ctypes.wintypes
import re, zlib
####################################################################################################
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,}')
@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
import numpy as np
arr = np.frombuffer(data, dtype=np.uint8)
counts = np.bincount(arr, minlength=256)
probs = counts[counts > 0] / len(arr)
return float(-np.sum(probs * np.log2(probs)))
@classmethod
def _extract_strings(cls, file_bytes):
res = {"StringCount": 0.0, "StringMeanLength": 0.0, "StringEntropy": 0.0}
for i in range(95):
res[f'StringHist_{i:02d}'] = 0.0
count = 0
total_len = 0
char_counts = [0] * 95
combined_data = bytearray()
for match in cls._STRING_PATTERN.finditer(file_bytes):
m = match.group()
count += 1
length = len(m)
total_len += length
combined_data.extend(m)
for b in m:
char_counts[b - 0x20] += 1
if count == 0:
return res
res["StringCount"] = float(count)
res["StringMeanLength"] = float(total_len) / count
res["StringEntropy"] = cls._calc_entropy(combined_data)
for i in range(95):
res[f'StringHist_{i:02d}'] = float(char_counts[i]) / total_len
return res
@classmethod
def _extract_histograms(cls, file_bytes):
import numpy as np
arr = np.frombuffer(file_bytes, dtype=np.uint8)
sz = len(arr)
res = {}
if sz == 0:
for i in range(256):
res[f'ByteHist_{i:02X}'] = 0.0
res[f'ByteEnt_{i:02X}'] = 0.0
return res
byte_counts = np.bincount(arr, minlength=256)
byte_hist = byte_counts / sz
for i, v in enumerate(byte_hist):
res[f'ByteHist_{i:02X}'] = float(v)
ent_hist = np.zeros(256, dtype=np.float64)
window = 2048
step = 1024
num_windows = (sz - window) // step + 1 if sz >= window else 1
for i in range(num_windows):
w = arr if sz < window else arr[i*step : i*step+window]
w_counts = np.bincount(w, minlength=256)
p = w_counts[w_counts > 0] / len(w)
entropy = -np.sum(p * np.log2(p))
ent_bin = int(entropy * 2.0)
if ent_bin > 15:
ent_bin = 15
byte_bins = w_counts.reshape(16, 16).sum(axis=1)
idx_start = ent_bin * 16
ent_hist[idx_start : idx_start+16] += byte_bins
ent_sum = np.sum(ent_hist)
if ent_sum > 0:
ent_hist /= ent_sum
for i, v in enumerate(ent_hist):
res[f'ByteEnt_{i:02X}'] = float(v)
return res
@classmethod
def _extract_overlay_features(cls, pe, fsize):
overlay_offset = pe.get_overlay_data_start_offset()
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
char_flags = [0x00000020, 0x00000040, 0x00000080, 0x02000000, 0x20000000, 0x40000000, 0x80000000]
for flag in char_flags:
base[f'Char_{flag:08X}_Count'] = 0.0
base[f'Char_{flag:08X}_MeanEntropy'] = 0.0
for i in range(50):
base[f'SectionHash_{i:02d}'] = 0.0
if hasattr(pe, 'sections'):
base['SectionCount'] = len(pe.sections)
entropies, raw_sizes, v_sizes = [], [], []
exec_sec, write_sec, read_sec, sec_exc = 0, 0, 0, 0
for section in pe.sections:
sec_name = section.Name.rstrip(b'\x00')
hash_val = zlib.crc32(sec_name) % 50
base[f'SectionHash_{hash_val:02d}'] += 1.0
try:
s_data = section.get_data()
s_entropy = cls._calc_entropy(s_data)
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 char_flags:
if section.Characteristics & flag:
base[f'Char_{flag:08X}_Count'] += 1.0
base[f'Char_{flag:08X}_MeanEntropy'] += 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'] = exec_sec
base['WritableSections'] = write_sec
base['ReadableSections'] = 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']
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'] = 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'] = 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
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
else:
base['IconCount'] = 0.0
base['DebugCount'] = len(pe.DIRECTORY_ENTRY_DEBUG) if hasattr(pe, 'DIRECTORY_ENTRY_DEBUG') else 0.0
base.update(cls._extract_strings(file_bytes))
base.update(cls._extract_histograms(file_bytes))
base.update(cls._extract_overlay_features(pe, fsize))
base.update(cls._extract_rich_header(pe))
base.update(cls._extract_ep_anomalies(pe))
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}")
+128 -98
View File
@@ -1,109 +1,138 @@
import os, sys, time, sqlite3, re, json
import pandas as pd
import os, sys, time, json, onnxmltools, zlib
import numpy as np
import pandas as pd
import lightgbm as lgb
import onnxmltools
import scipy.sparse as sp
from collections import defaultdict
from onnxmltools.convert.common.data_types import FloatTensorType
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, accuracy_score, roc_auc_score
from sklearn.metrics import accuracy_score, roc_auc_score
try:
import orjson
JSON_LOADS = orjson.loads
except ImportError:
JSON_LOADS = json.loads
####################################################################################################
DB_PATH = "pe_features.db"
JSONL_PATH = "pe_features.jsonl"
MODEL_FILE = "model.txt"
ONNX_FILE = "Pefile_General_S1.onnx"
FEATURE_FILE = "features.json"
TEST_SIZE = 0.1
TEST_SIZE = 0.0001
RANDOM_SEED = 42
DLL_HASH_DIM = 256
API_HASH_DIM = 1024
EXCLUDE_FEATURES = {
'label', 'filehash',
}
LGBM_PARAMS = {
'objective': 'binary',
'metric': ['binary_logloss', 'auc'],
'boosting_type': 'gbdt',
'num_leaves': 512,
'learning_rate': 0.05,
'num_leaves': 256,
'feature_fraction': 0.8,
'bagging_fraction': 0.8,
'bagging_fraction': 1.0,
'bagging_freq': 5,
'verbose': -1,
'seed': RANDOM_SEED,
'n_jobs': -1,
'max_depth': -1,
'min_data_in_leaf': 10,
'lambda_l1': 0.05,
'lambda_l2': 0.05,
'min_data_in_leaf': 50,
}
####################################################################################################
def clean_col_name(name):
return re.sub(r'[^A-Za-z0-9_]+', '', name)
def analyze_schema(jsonl_path):
print("[*] Stage 1: Scanning global schema...")
base_keys = set()
num_lines = 0
def get_raw_schema(db_path):
if not os.path.exists(db_path):
return None
conn = sqlite3.connect(db_path)
try:
cursor = conn.cursor()
cursor.execute("PRAGMA table_info(PeData)")
columns = [row[1] for row in cursor.fetchall()]
exclude = {'label', 'id', 'filename', 'filelength', 'rowid', 'probability', 'predictedlabel', 'score', 'filehash', 'timedatestamp'}
return [c for c in columns if c.lower().strip() not in exclude]
except Exception as e:
print(f"[-] Schema extraction error: {e}")
return None
finally:
conn.close()
with open(jsonl_path, 'rb') as f:
for line in f:
try:
data = JSON_LOADS(line)
base_data = data.get('Base')
if base_data:
base_keys.update(base_data.keys())
num_lines += 1
except Exception:
continue
####################################################################################################
def load_data(db_path, raw_cols):
print(f"[*] Connecting to database: {db_path}")
conn = sqlite3.connect(db_path)
req_cols = list(set(raw_cols + ['Label']))
sql = f"SELECT {', '.join(req_cols)} FROM PeData"
base_schema = sorted([k for k in base_keys if k.lower().strip() not in EXCLUDE_FEATURES])
dll_schema = [f"DllHash_{i:03d}" for i in range(DLL_HASH_DIM)]
api_schema = [f"ApiHash_{i:04d}" for i in range(API_HASH_DIM)]
try:
chunks = []
for chunk in pd.read_sql_query(sql, conn, chunksize=20000):
label_col = next((c for c in chunk.columns if c.lower().strip() == 'label'), None)
if not label_col:
print("[-] Critical Error: Label column not found.")
return None, None, None
cols_to_cast = [c for c in chunk.columns if c != label_col]
chunk[cols_to_cast] = chunk[cols_to_cast].astype(np.float32)
chunk[label_col] = chunk[label_col].astype(np.int8)
chunks.append(chunk)
if not chunks:
return None, None, None
df = pd.concat(chunks, ignore_index=True)
label_col = next((c for c in df.columns if c.lower().strip() == 'label'), None)
y = df[label_col].astype(int)
X = df.drop(columns=[label_col], errors='ignore')
valid_features = [c for c in raw_cols if c in X.columns]
X = X[valid_features].copy()
rename_map = {old: clean_col_name(old) for old in X.columns}
X = X.rename(columns=rename_map)
if X.columns.duplicated().any():
print("[-] Warning: Duplicate column names detected after cleaning. Deduplicating...")
X = X.loc[:, ~X.columns.duplicated()]
final_features = X.columns.tolist()
return X, y, final_features
finally:
conn.close()
return base_schema, dll_schema, api_schema, num_lines
####################################################################################################
def train_process(X, y):
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})...")
final_features = base_schema + dll_schema + api_schema
num_features = len(final_features)
base_idx = {feat: i for i, feat in enumerate(base_schema)}
offset_dll = len(base_schema)
offset_api = offset_dll + DLL_HASH_DIM
rows, cols, vals = [], [], []
y = np.zeros(num_lines, dtype=np.int8)
with open(jsonl_path, 'rb') as f:
row_idx = 0
for line in f:
try:
data = JSON_LOADS(line)
y[row_idx] = data.get('Label', 0)
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)
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(offset_dll + idx)
vals.append(1.0)
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(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)
return X, pd.Series(y), final_features
####################################################################################################
def _calculate_dynamic_lr(current_iter: int) -> float:
return max(0.01, 0.1 - (current_iter // 80) * 0.01)
def train_process(X, y, feature_names):
print(f"[*] Dataset shape: {X.shape}")
X_train, X_test, y_train, y_test = train_test_split(
@@ -117,11 +146,11 @@ def train_process(X, y):
weight_ratio_target = 0.005
final_pos_weight = base_weight * weight_ratio_target
print(f"[*] Balance Report: Safe={n_neg}, Malware={n_pos}")
print(f"[*] Weight Config: Base Balanced={base_weight:.4f}, Final Adjusted={final_pos_weight:.4f}")
print(f"[*] Sample distribution: Safe={n_neg}, Malware={n_pos}")
print(f"[*] Weight correction: Base={base_weight:.4f}, Final={final_pos_weight:.4f}")
train_data = lgb.Dataset(X_train, label=y_train)
valid_data = lgb.Dataset(X_test, label=y_test, reference=train_data)
train_data = lgb.Dataset(X_train, label=y_train, feature_name=feature_names)
valid_data = lgb.Dataset(X_test, label=y_test, reference=train_data, feature_name=feature_names)
params = LGBM_PARAMS.copy()
params['scale_pos_weight'] = final_pos_weight
@@ -133,14 +162,15 @@ def train_process(X, y):
params,
train_data,
valid_sets=[valid_data],
num_boost_round=1000,
num_boost_round=500,
callbacks=[
lgb.log_evaluation(period=100)
lgb.log_evaluation(period=50),
lgb.reset_parameter(learning_rate=_calculate_dynamic_lr)
]
)
elapsed = time.time() - start_time
print(f"\n[+] Training finished in {elapsed:.2f}s")
print(f"\n[+] Training complete, time elapsed {elapsed:.2f} seconds")
return model, X_test, y_test
@@ -156,12 +186,10 @@ def export_onnx_model(model, feature_names):
with open(FEATURE_FILE, 'w') as f:
json.dump(feature_names, f)
print(f"[+] ONNX export success. Features synced to {FEATURE_FILE}")
print(f"[+] ONNX export successful. Feature order saved to {FEATURE_FILE}")
except Exception as e:
print(f"[-] ONNX export failed: {e}")
####################################################################################################
def evaluate_and_save(model, X_test, y_test):
print("\n------------------- Evaluation Report -------------------\n")
y_prob = model.predict(X_test, num_iteration=model.best_iteration)
@@ -181,14 +209,14 @@ def evaluate_and_save(model, X_test, y_test):
detection_rate = (tp / total_malware * 100) if total_malware > 0 else 0.0
fpr = (fp / total_safe * 100) if total_safe > 0 else 0.0
print(f"Accuracy : {acc:.4f}")
print(f"ROC AUC : {auc:.4f}\n")
print(f"Accuracy : {acc:.6f}")
print(f"ROC AUC : {auc:.6f}\n")
print(f"True Positive (Class 1) : {detection_rate:.3f}% ({tp}/{total_malware})")
print(f"False Positive (Class 0) : {fpr:.3f}% ({fp}/{total_safe})")
print("\n---------------------------------------------------------\n")
print("[*] Top 20 Features (Importance by Gain):")
print("[*] Top 20 Features:")
feature_names = model.feature_name()
importance = model.feature_importance(importance_type='gain')
@@ -205,25 +233,27 @@ def evaluate_and_save(model, X_test, y_test):
####################################################################################################
if __name__ == "__main__":
print("\n---------------- PE Malware Trainer v3.1 ----------------\n")
print("\n---------------- PE Malware Trainer v4.0 ----------------\n")
if not os.path.exists(DB_PATH):
print(f"[-] Error: Database {DB_PATH} not found.")
if not os.path.exists(JSONL_PATH):
print(f"[-] Error: Feature file {JSONL_PATH} not found.")
sys.exit(1)
raw_schema = get_raw_schema(DB_PATH)
if not raw_schema:
print("[-] Error: Could not read schema.")
base_schema, dll_schema, api_schema, num_lines = analyze_schema(JSONL_PATH)
if num_lines == 0:
print("[-] Error: Feature file is empty.")
sys.exit(1)
X, y, _ = load_data(DB_PATH, raw_schema)
if X is None:
sys.exit(1)
print(f"[*] Detected Base features: {len(base_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)
try:
model, X_test, y_test = train_process(X, y)
model, X_test, y_test = train_process(X, y, final_features)
evaluate_and_save(model, X_test, y_test)
except KeyboardInterrupt:
print("\n\n[-] Training interrupted.")
print("\n\n[-] Training aborted.")
except Exception as e:
print(f"\n[-] Critical Error: {e}")
print(f"\n[-] System exception: {e}")
-216
View File
@@ -1,216 +0,0 @@
import os, sys, time, sqlite3, re, json
import pandas as pd
import numpy as np
import lightgbm as lgb
import onnxmltools
from onnxmltools.convert.common.data_types import FloatTensorType
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, accuracy_score, roc_auc_score
####################################################################################################
DB_PATH = "pe_features.db"
MODEL_FILE = "model.txt"
ONNX_FILE = "model.onnx"
FEATURE_FILE = "features.json"
TEST_SIZE = 0.0001
RANDOM_SEED = 42
LGBM_PARAMS = {
'objective': 'binary',
'metric': ['binary_logloss', 'auc'],
'boosting_type': 'gbdt',
'num_leaves': 512,
'learning_rate': 0.05,
'feature_fraction': 0.8,
'bagging_fraction': 0.8,
'bagging_freq': 5,
'verbose': -1,
'seed': RANDOM_SEED,
'n_jobs': -1,
'max_depth': -1,
'min_data_in_leaf': 10,
'lambda_l1': 0.1,
'lambda_l2': 0.1,
}
####################################################################################################
def clean_col_name(name):
return re.sub(r'[^A-Za-z0-9_]+', '', name)
def get_raw_schema(db_path):
if not os.path.exists(db_path):
return None
conn = sqlite3.connect(db_path)
try:
cursor = conn.cursor()
cursor.execute("PRAGMA table_info(PeData)")
columns = [row[1] for row in cursor.fetchall()]
exclude = {'label', 'id', 'filename', 'filelength', 'rowid', 'probability', 'predictedlabel', 'score', 'filehash', 'timedatestamp'}
return [c for c in columns if c.lower().strip() not in exclude]
except Exception as e:
print(f"[-] Schema extraction error: {e}")
return None
finally:
conn.close()
####################################################################################################
def load_data(db_path, raw_cols):
print(f"[*] Connecting to database: {db_path}")
conn = sqlite3.connect(db_path)
req_cols = list(set(raw_cols + ['Label']))
sql = f"SELECT {', '.join(req_cols)} FROM PeData"
try:
chunks = []
for chunk in pd.read_sql_query(sql, conn, chunksize=20000):
label_col = next((c for c in chunk.columns if c.lower().strip() == 'label'), None)
if not label_col:
print("[-] Critical Error: Label column not found.")
return None, None, None
cols_to_cast = [c for c in chunk.columns if c != label_col]
chunk[cols_to_cast] = chunk[cols_to_cast].astype(np.float32)
chunk[label_col] = chunk[label_col].astype(np.int8)
chunks.append(chunk)
if not chunks:
return None, None, None
df = pd.concat(chunks, ignore_index=True)
label_col = next((c for c in df.columns if c.lower().strip() == 'label'), None)
y = df[label_col].astype(int)
X = df.drop(columns=[label_col], errors='ignore')
valid_features = [c for c in raw_cols if c in X.columns]
X = X[valid_features].copy()
rename_map = {old: clean_col_name(old) for old in X.columns}
X = X.rename(columns=rename_map)
if X.columns.duplicated().any():
print("[-] Warning: Duplicate column names detected after cleaning. Deduplicating...")
X = X.loc[:, ~X.columns.duplicated()]
final_features = X.columns.tolist()
return X, y, final_features
finally:
conn.close()
####################################################################################################
def train_process(X, y):
print(f"[*] Dataset shape: {X.shape}")
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=TEST_SIZE, random_state=RANDOM_SEED, stratify=y
)
n_pos = (y_train == 1).sum()
n_neg = (y_train == 0).sum()
base_weight = n_neg / n_pos if n_pos > 0 else 1.0
weight_ratio_target = 0.001
final_pos_weight = base_weight * weight_ratio_target
print(f"[*] Balance Report: Safe={n_neg}, Malware={n_pos}")
print(f"[*] Weight Config: Base Balanced={base_weight:.4f}, Final Adjusted={final_pos_weight:.4f}")
train_data = lgb.Dataset(X_train, label=y_train)
valid_data = lgb.Dataset(X_test, label=y_test, reference=train_data)
params = LGBM_PARAMS.copy()
params['scale_pos_weight'] = final_pos_weight
print("\n[*] Starting LightGBM training...")
start_time = time.time()
model = lgb.train(
params,
train_data,
valid_sets=[valid_data],
num_boost_round=1000,
callbacks=[
lgb.log_evaluation(period=100)
]
)
elapsed = time.time() - start_time
print(f"\n[+] Training finished in {elapsed:.2f}s")
return model, X_test, y_test
####################################################################################################
def export_onnx_model(model, feature_names):
print(f"[*] Exporting ONNX model to {ONNX_FILE}...")
try:
initial_type = [('float_input', FloatTensorType([None, len(feature_names)]))]
onnx_model = onnxmltools.convert_lightgbm(model, initial_types=initial_type)
onnxmltools.utils.save_model(onnx_model, ONNX_FILE)
with open(FEATURE_FILE, 'w') as f:
json.dump(feature_names, f)
print(f"[+] ONNX export success. Features synced to {FEATURE_FILE}")
except Exception as e:
print(f"[-] ONNX export failed: {e}")
####################################################################################################
def evaluate_and_save(model, X_test, y_test):
print("\n------------------- Evaluation Report -------------------\n")
y_prob = model.predict(X_test, num_iteration=model.best_iteration)
y_pred = (y_prob > 0.5).astype(int)
acc = accuracy_score(y_test, y_pred)
auc = roc_auc_score(y_test, y_prob)
print(f"Accuracy : {acc:.4f}")
print(f"ROC AUC : {auc:.4f}")
print("\n" + classification_report(y_test, y_pred, zero_division=0))
print("---------------------------------------------------------\n")
print("[*] Top 20 Features (Importance by Gain):")
feature_names = model.feature_name()
importance = model.feature_importance(importance_type='gain')
feature_imp = pd.DataFrame(sorted(zip(importance, feature_names), reverse=True), columns=['Value', 'Feature'])
print(feature_imp.head(20).to_string(index=False))
print("\n---------------------------------------------------------\n")
model.save_model(MODEL_FILE)
print(f"[+] Native model saved to: {os.path.abspath(MODEL_FILE)}")
export_onnx_model(model, feature_names)
####################################################################################################
if __name__ == "__main__":
print("\n---------------- PE Malware Trainer v3.1 ----------------\n")
if not os.path.exists(DB_PATH):
print(f"[-] Error: Database {DB_PATH} not found.")
sys.exit(1)
raw_schema = get_raw_schema(DB_PATH)
if not raw_schema:
print("[-] Error: Could not read schema.")
sys.exit(1)
X, y, _ = load_data(DB_PATH, raw_schema)
if X is None:
sys.exit(1)
try:
model, X_test, y_test = train_process(X, y)
evaluate_and_save(model, X_test, y_test)
except KeyboardInterrupt:
print("\n\n[-] Training interrupted.")
except Exception as e:
print(f"\n[-] Critical Error: {e}")
+128 -98
View File
@@ -1,109 +1,138 @@
import os, sys, time, sqlite3, re, json
import pandas as pd
import os, sys, time, json, onnxmltools, zlib
import numpy as np
import pandas as pd
import lightgbm as lgb
import onnxmltools
import scipy.sparse as sp
from collections import defaultdict
from onnxmltools.convert.common.data_types import FloatTensorType
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report, accuracy_score, roc_auc_score
from sklearn.metrics import accuracy_score, roc_auc_score
try:
import orjson
JSON_LOADS = orjson.loads
except ImportError:
JSON_LOADS = json.loads
####################################################################################################
DB_PATH = "pe_features.db"
JSONL_PATH = "pe_features.jsonl"
MODEL_FILE = "model.txt"
ONNX_FILE = "Pefile_General_T1.onnx"
FEATURE_FILE = "features.json"
TEST_SIZE = 0.1
TEST_SIZE = 0.0001
RANDOM_SEED = 42
DLL_HASH_DIM = 256
API_HASH_DIM = 1024
EXCLUDE_FEATURES = {
'label', 'filehash',
}
LGBM_PARAMS = {
'objective': 'binary',
'metric': ['binary_logloss', 'auc'],
'boosting_type': 'gbdt',
'num_leaves': 64,
'learning_rate': 0.05,
'num_leaves': 128,
'feature_fraction': 0.8,
'bagging_fraction': 0.8,
'bagging_fraction': 1.0,
'bagging_freq': 5,
'verbose': -1,
'seed': RANDOM_SEED,
'n_jobs': -1,
'max_depth': -1,
'min_data_in_leaf': 10,
'lambda_l1': 0.05,
'lambda_l2': 0.05,
'min_data_in_leaf': 50,
}
####################################################################################################
def clean_col_name(name):
return re.sub(r'[^A-Za-z0-9_]+', '', name)
def analyze_schema(jsonl_path):
print("[*] Stage 1: Scanning global schema...")
base_keys = set()
num_lines = 0
def get_raw_schema(db_path):
if not os.path.exists(db_path):
return None
conn = sqlite3.connect(db_path)
try:
cursor = conn.cursor()
cursor.execute("PRAGMA table_info(PeData)")
columns = [row[1] for row in cursor.fetchall()]
exclude = {'label', 'id', 'filename', 'filelength', 'rowid', 'probability', 'predictedlabel', 'score', 'filehash', 'timedatestamp'}
return [c for c in columns if c.lower().strip() not in exclude]
except Exception as e:
print(f"[-] Schema extraction error: {e}")
return None
finally:
conn.close()
with open(jsonl_path, 'rb') as f:
for line in f:
try:
data = JSON_LOADS(line)
base_data = data.get('Base')
if base_data:
base_keys.update(base_data.keys())
num_lines += 1
except Exception:
continue
####################################################################################################
def load_data(db_path, raw_cols):
print(f"[*] Connecting to database: {db_path}")
conn = sqlite3.connect(db_path)
req_cols = list(set(raw_cols + ['Label']))
sql = f"SELECT {', '.join(req_cols)} FROM PeData"
base_schema = sorted([k for k in base_keys if k.lower().strip() not in EXCLUDE_FEATURES])
dll_schema = [f"DllHash_{i:03d}" for i in range(DLL_HASH_DIM)]
api_schema = [f"ApiHash_{i:04d}" for i in range(API_HASH_DIM)]
try:
chunks = []
for chunk in pd.read_sql_query(sql, conn, chunksize=20000):
label_col = next((c for c in chunk.columns if c.lower().strip() == 'label'), None)
if not label_col:
print("[-] Critical Error: Label column not found.")
return None, None, None
cols_to_cast = [c for c in chunk.columns if c != label_col]
chunk[cols_to_cast] = chunk[cols_to_cast].astype(np.float32)
chunk[label_col] = chunk[label_col].astype(np.int8)
chunks.append(chunk)
if not chunks:
return None, None, None
df = pd.concat(chunks, ignore_index=True)
label_col = next((c for c in df.columns if c.lower().strip() == 'label'), None)
y = df[label_col].astype(int)
X = df.drop(columns=[label_col], errors='ignore')
valid_features = [c for c in raw_cols if c in X.columns]
X = X[valid_features].copy()
rename_map = {old: clean_col_name(old) for old in X.columns}
X = X.rename(columns=rename_map)
if X.columns.duplicated().any():
print("[-] Warning: Duplicate column names detected after cleaning. Deduplicating...")
X = X.loc[:, ~X.columns.duplicated()]
final_features = X.columns.tolist()
return X, y, final_features
finally:
conn.close()
return base_schema, dll_schema, api_schema, num_lines
####################################################################################################
def train_process(X, y):
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})...")
final_features = base_schema + dll_schema + api_schema
num_features = len(final_features)
base_idx = {feat: i for i, feat in enumerate(base_schema)}
offset_dll = len(base_schema)
offset_api = offset_dll + DLL_HASH_DIM
rows, cols, vals = [], [], []
y = np.zeros(num_lines, dtype=np.int8)
with open(jsonl_path, 'rb') as f:
row_idx = 0
for line in f:
try:
data = JSON_LOADS(line)
y[row_idx] = data.get('Label', 0)
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)
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(offset_dll + idx)
vals.append(1.0)
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(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)
return X, pd.Series(y), final_features
####################################################################################################
def _calculate_dynamic_lr(current_iter: int) -> float:
return max(0.01, 0.1 - (current_iter // 80) * 0.01)
def train_process(X, y, feature_names):
print(f"[*] Dataset shape: {X.shape}")
X_train, X_test, y_train, y_test = train_test_split(
@@ -117,11 +146,11 @@ def train_process(X, y):
weight_ratio_target = 0.001
final_pos_weight = base_weight * weight_ratio_target
print(f"[*] Balance Report: Safe={n_neg}, Malware={n_pos}")
print(f"[*] Weight Config: Base Balanced={base_weight:.4f}, Final Adjusted={final_pos_weight:.4f}")
print(f"[*] Sample distribution: Safe={n_neg}, Malware={n_pos}")
print(f"[*] Weight correction: Base={base_weight:.4f}, Final={final_pos_weight:.4f}")
train_data = lgb.Dataset(X_train, label=y_train)
valid_data = lgb.Dataset(X_test, label=y_test, reference=train_data)
train_data = lgb.Dataset(X_train, label=y_train, feature_name=feature_names)
valid_data = lgb.Dataset(X_test, label=y_test, reference=train_data, feature_name=feature_names)
params = LGBM_PARAMS.copy()
params['scale_pos_weight'] = final_pos_weight
@@ -133,14 +162,15 @@ def train_process(X, y):
params,
train_data,
valid_sets=[valid_data],
num_boost_round=1000,
num_boost_round=500,
callbacks=[
lgb.log_evaluation(period=100)
lgb.log_evaluation(period=50),
lgb.reset_parameter(learning_rate=_calculate_dynamic_lr)
]
)
elapsed = time.time() - start_time
print(f"\n[+] Training finished in {elapsed:.2f}s")
print(f"\n[+] Training complete, time elapsed {elapsed:.2f} seconds")
return model, X_test, y_test
@@ -156,12 +186,10 @@ def export_onnx_model(model, feature_names):
with open(FEATURE_FILE, 'w') as f:
json.dump(feature_names, f)
print(f"[+] ONNX export success. Features synced to {FEATURE_FILE}")
print(f"[+] ONNX export successful. Feature order saved to {FEATURE_FILE}")
except Exception as e:
print(f"[-] ONNX export failed: {e}")
####################################################################################################
def evaluate_and_save(model, X_test, y_test):
print("\n------------------- Evaluation Report -------------------\n")
y_prob = model.predict(X_test, num_iteration=model.best_iteration)
@@ -181,14 +209,14 @@ def evaluate_and_save(model, X_test, y_test):
detection_rate = (tp / total_malware * 100) if total_malware > 0 else 0.0
fpr = (fp / total_safe * 100) if total_safe > 0 else 0.0
print(f"Accuracy : {acc:.4f}")
print(f"ROC AUC : {auc:.4f}\n")
print(f"Accuracy : {acc:.6f}")
print(f"ROC AUC : {auc:.6f}\n")
print(f"True Positive (Class 1) : {detection_rate:.3f}% ({tp}/{total_malware})")
print(f"False Positive (Class 0) : {fpr:.3f}% ({fp}/{total_safe})")
print("\n---------------------------------------------------------\n")
print("[*] Top 20 Features (Importance by Gain):")
print("[*] Top 20 Features:")
feature_names = model.feature_name()
importance = model.feature_importance(importance_type='gain')
@@ -205,25 +233,27 @@ def evaluate_and_save(model, X_test, y_test):
####################################################################################################
if __name__ == "__main__":
print("\n---------------- PE Malware Trainer v3.1 ----------------\n")
print("\n---------------- PE Malware Trainer v4.0 ----------------\n")
if not os.path.exists(DB_PATH):
print(f"[-] Error: Database {DB_PATH} not found.")
if not os.path.exists(JSONL_PATH):
print(f"[-] Error: Feature file {JSONL_PATH} not found.")
sys.exit(1)
raw_schema = get_raw_schema(DB_PATH)
if not raw_schema:
print("[-] Error: Could not read schema.")
base_schema, dll_schema, api_schema, num_lines = analyze_schema(JSONL_PATH)
if num_lines == 0:
print("[-] Error: Feature file is empty.")
sys.exit(1)
X, y, _ = load_data(DB_PATH, raw_schema)
if X is None:
sys.exit(1)
print(f"[*] Detected Base features: {len(base_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)
try:
model, X_test, y_test = train_process(X, y)
model, X_test, y_test = train_process(X, y, final_features)
evaluate_and_save(model, X_test, y_test)
except KeyboardInterrupt:
print("\n\n[-] Training interrupted.")
print("\n\n[-] Training aborted.")
except Exception as e:
print(f"\n[-] Critical Error: {e}")
print(f"\n[-] System exception: {e}")
+39 -1
View File
@@ -56,7 +56,7 @@ body {
}
.navigate-title {
font-size: 16px;
font-size: 17px;
font-weight: 600;
letter-spacing: -0.2px;
}
@@ -797,6 +797,44 @@ body.resizing-active * {
color: var(--text-primary);
}
.custom-context-menu {
position: absolute;
z-index: 10000;
background-color: var(--bg-panel);
border: 1px solid var(--border-color);
border-radius: 8px;
box-shadow: 0 4px 12px rgba(0, 0, 0, 0.1);
padding: 4px;
display: none;
min-width: 160px;
}
.custom-context-menu.show {
display: block;
}
.custom-context-menu-item {
display: flex;
align-items: center;
gap: 10px;
padding: 8px 12px;
border-radius: 6px;
cursor: pointer;
font-size: 13px;
color: var(--text-primary);
transition: background-color 0.2s;
}
.custom-context-menu-item:hover {
background-color: var(--bg-hover);
}
.custom-context-menu-item svg {
width: 16px;
height: 16px;
color: var(--text-secondary);
}
::-webkit-scrollbar { width: 6px; height: 6px; }
::-webkit-scrollbar-track { background: transparent; }
::-webkit-scrollbar-thumb { background: var(--border-color); border-radius: 10px; }
+1474 -1364
View File
File diff suppressed because it is too large. Load diff
File diff suppressed because it is too large. Load diff
+1 -1
View File
@@ -158,7 +158,7 @@
</div>
<div class="list-item" style="cursor: pointer;" data-target="repair_window">
<div class="item-info">
<h2 data-i18n>系統修復</h2><p data-i18n>此選項可以修復系統註冊項目</p>
<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>
+1331 -1885
View File
File diff suppressed because it is too large. Load diff
+479 -449
View File
File diff suppressed because it is too large. Load diff
+1 -1
View File
@@ -2,7 +2,7 @@ import pyinstaller_versionfile
pyinstaller_versionfile.create_versionfile(
output_file="versionfile.txt",
version='3.5.3.0',
version='3.5.6.0',
company_name="PYAS Security",
file_description="PYAS Security Antivirus",
internal_name="PYAS",
+20 -15
View File
@@ -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();
+10 -3
View File
@@ -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() {
+9 -3
View File
@@ -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() {
+2 -2
View File
@@ -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;
+14 -6
View File
@@ -2,7 +2,6 @@
"Rule_Registry_BlockList": [
"\\REGISTRY\\MACHINE\\BCD00000000",
"\\REGISTRY\\MACHINE\\BCD00000000\\*",
"*\\Image File Execution Options\\*",
"*\\Classes\\.exe",
"*\\Classes\\.bat",
@@ -78,12 +77,10 @@
"*\\FirmwareBootDevice",
"*\\BootDriverFlags",
"*\\SystemStartOptions",
"*\\CrashDumpEnabled",
"*\\DisableAutomaticRestartOnFailure",
"*\\SystemBiosVersion",
"*\\NoTheme",
"*\\SystemSetupInProgress",
"*\\CmdLine",
"*\\SetupType",
"*\\AutoReboot",
"*\\Userinit",
@@ -106,27 +103,38 @@
"\\Device\\HarddiskVolume*\\ProgramData\\Microsoft\\Windows Defender\\*",
"\\Device\\HarddiskVolume*\\ProgramData\\Microsoft\\Windows Defender Advanced Threat Protection\\*",
"*\\Sources\\SetupHost.exe",
"*\\Work\\*\\DismHost.exe",
"*\\Google\\Chrome\\Application\\chrome.exe",
"*\\Microsoft\\Edge\\Application\\msedge.exe",
"*\\Mozilla Firefox\\firefox.exe",
"*\\Thorium\\Application\\thorium.exe",
"*\\Steam\\steam.exe",
"*\\steam\\*\\steam.exe",
"*\\Battle.net\\*\\Battle.net.exe",
"*\\Epic Games\\Launcher\\Portal\\Binaries\\*\\EpicGamesLauncher.exe",
"*\\resources\\bin\\SDK\\DownloadSDKServer.exe",
"*\\ThunderEv.exe",
"*\\ThunderStart.exe",
"*\\resources\\bin\\SDK\\DownloadSDKServer.exe",
"*\\KuGou.exe",
"*\\XLGame.exe",
"*\\XLLiveUD.exe",
"*\\PC4399_WPFLauncher.exe",
"*\\bin\\APlayer.exe",
"*\\hrkill-*.exe",
"*\\bin\\HipsDaemon.exe",
"*\\360Safe.exe",
"*\\bilibili\\bilibili.exe",
"*\\Kingsoft\\WPS Office\\ksolaunch.exe",
"*\\Kingsoft\\WPS Office\\*\\wps.exe",
"*\\PC4399_WPFLauncher.exe",
"*\\bin\\APlayer.exe",
"*\\Neverness To Everness\\NTELauncher.exe",
"*\\Application\\AndrowsLauncher.exe",
],
"Rule_Process_ExploitableBlacklist": [
"*\\cmd.exe",
+2 -1
View File
@@ -26,6 +26,7 @@ pip install pandas
pip install scikit-learn
pip install lightgbm
pip install onnxmltools
pip install orjson
```
## File Information
@@ -131,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