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

Add files via upload

Browse Source
This commit is contained in:
87owo authored and GitHub committed 2026-02-20 13:49:18 +08:00
1 parent 4e7c6efa2a
commit 73d09f9b09
2 files changed
+313 -274
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

Engine/Properties/convert.py
+287 -258
View File
@@ -1,20 +1,175 @@
import os, sys, time, ctypes, sqlite3, datetime, pefile, math
import os, sys, time, ctypes, sqlite3, datetime, pefile, math, hashlib
import numpy as np
import ctypes.wintypes
####################################################################################################
DB_PATH = "pe_features.db"
BATCH_SIZE = 100
MAX_FILE_SIZE = 256 * 1024 * 1024
BATCH_SIZE = 1000
MAX_FILE_SIZE = 256 * 1024 * 1024
TARGET_EXTENSIONS = {
'.exe', '.dll', '.sys', '.ocx', '.scr', '.efi', '.acm', '.ax', '.cpl', '.drv', '.com'
'.exe', '.dll', '.sys', '.ocx', '.scr', '.efi', '.acm', '.ax', '.cpl', '.drv', '.com', '.mui', '.pyd'
}
####################################################################################################
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(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),
("pFile", ctypes.POINTER(WINTRUST_FILE_INFO)),
("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, (0x8C, 0xC2, 0x00, 0xC0, 0x4F, 0xC2, 0x95, 0xEE))
def verify_signature(file_path):
if os.name != 'nt': return 0
try:
fi = WINTRUST_FILE_INFO(ctypes.sizeof(WINTRUST_FILE_INFO), file_path, None, None)
td = WINTRUST_DATA(ctypes.sizeof(WINTRUST_DATA), None, None, 2, 1, 1, ctypes.pointer(fi), 0, None, None, 0x80, 0, None)
return 1 if ctypes.windll.wintrust.WinVerifyTrust(None, ctypes.byref(V2_GUID), ctypes.byref(td)) == 0 else 0
except:
return 0
####################################################################################################
class Features:
SCHEMA = [
API_CAT_MAPPING = {
'ProcessControl': {
'CreateProcessA', 'CreateProcessW', 'WinExec', 'ShellExecuteA', 'ShellExecuteW',
'ShellExecuteExW', 'ExitProcess', 'TerminateProcess', 'OpenProcess', 'GetExitCodeProcess',
'SetThreadPriority', 'GetThreadPriority', 'GetCurrentProcess', 'GetCurrentProcessId'
},
'Injection': {
'VirtualAlloc', 'VirtualAllocEx', 'VirtualProtect', 'VirtualProtectEx', 'WriteProcessMemory',
'ReadProcessMemory', 'CreateRemoteThread', 'QueueUserAPC', 'SetThreadContext', 'GetThreadContext',
'Wow64SetThreadContext', 'NtUnmapViewOfSection', 'ZwUnmapViewOfSection', 'RtlCreateUserThread',
'SetWindowsHookExA', 'SetWindowsHookExW', 'UnhookWindowsHookEx', 'LoadLibraryA', 'LoadLibraryW',
'LoadLibraryExA', 'LoadLibraryExW', 'GetProcAddress', 'GetModuleHandleA', 'GetModuleHandleW',
'FreeLibrary', 'CreateFileMappingA', 'CreateFileMappingW', 'MapViewOfFile', 'UnmapViewOfSection'
},
'Synchronization': {
'WaitForSingleObject', 'WaitForSingleObjectEx', 'WaitForMultipleObjects', 'WaitForMultipleObjectsEx',
'CreateMutexA', 'CreateMutexW', 'OpenMutexW', 'ReleaseMutex', 'CreateEventA', 'CreateEventW',
'OpenEventW', 'SetEvent', 'ResetEvent', 'EnterCriticalSection', 'LeaveCriticalSection',
'InitializeCriticalSection', 'DeleteCriticalSection', 'Sleep', 'SleepEx'
},
'MultiThreading': {
'CreateThread', 'ResumeThread', 'SuspendThread', 'ExitThread', 'TerminateThread',
'GetCurrentThread', 'GetCurrentThreadId', 'TlsAlloc', 'TlsSetValue', 'TlsGetValue',
'CreateThreadpoolWork', 'SubmitThreadpoolWork'
},
'Network': {
'socket', 'connect', 'send', 'recv', 'bind', 'listen', 'accept', 'gethostbyname', 'getaddrinfo',
'WSAStartup', 'WSACleanup', 'WSAIoctl', 'WSASocketA', 'WSASocketW', 'closesocket', 'htons',
'InternetOpenA', 'InternetOpenW', 'InternetConnectA', 'InternetConnectW', 'InternetOpenUrlA',
'InternetOpenUrlW', 'HttpOpenRequestA', 'HttpOpenRequestW', 'HttpSendRequestA', 'HttpSendRequestW',
'InternetReadFile', 'URLDownloadToFileA', 'URLDownloadToFileW', 'DnsQuery_A', 'DnsQuery_W'
},
'Encryption': {
'CryptAcquireContextA', 'CryptAcquireContextW', 'CryptCreateHash', 'CryptHashData',
'CryptDeriveKey', 'CryptEncrypt', 'CryptDecrypt', 'CryptDestroyKey', 'CryptDestroyHash',
'CryptReleaseContext', 'CryptGenKey', 'CryptImportKey', 'CryptExportKey'
},
'DataObfuscation': {
'RtlDecompressBuffer', 'MultiByteToWideChar', 'WideCharToMultiByte', 'Base64Decode',
'CryptDecodeObject', 'IsDBCSLeadByte', 'CharUpperA', 'CharLowerA'
},
'FileIO': {
'CreateFileA', 'CreateFileW', 'WriteFile', 'ReadFile', 'DeleteFileA', 'DeleteFileW',
'CopyFileA', 'CopyFileW', 'MoveFileA', 'MoveFileW', 'FindFirstFileA', 'FindNextFileA',
'GetTempPathA', 'GetTempPathW', 'GetTempFileNameA', 'GetTempFileNameW', 'SetFileAttributesA',
'SetFileAttributesW', 'DeviceIoControl', 'SetFileTime', 'GetFileSize', 'GetFileSizeEx',
'SetFilePointer', 'FlushFileBuffers', 'ConnectNamedPipe', 'PeekNamedPipe'
},
'Registry': {
'RegOpenKeyExA', 'RegOpenKeyExW', 'RegSetValueExA', 'RegSetValueExW',
'RegCreateKeyExA', 'RegCreateKeyExW', 'RegDeleteKeyA', 'RegDeleteKeyW',
'RegEnumValueA', 'RegEnumValueW', 'RegQueryValueExA', 'RegQueryValueExW',
'RegDeleteValueA', 'RegDeleteValueW'
},
'Services': {
'OpenSCManagerA', 'OpenSCManagerW', 'CreateServiceA', 'CreateServiceW',
'StartServiceA', 'StartServiceW', 'ControlService', 'DeleteService'
},
'Privileges': {
'AdjustTokenPrivileges', 'LookupPrivilegeValueA', 'LookupPrivilegeValueW',
'OpenProcessToken', 'NetUserAdd', 'NetLocalGroupAddMembers', 'IsAdmin',
'GetUserNameA', 'GetUserNameW'
},
'Native': {
'RtlUnwind', 'RtlVirtualUnwind', 'RtlCaptureContext', 'RtlLookupFunctionEntry',
'NtClose', 'NtQueryInformationProcess', 'RtlAllocateHeap', 'RtlFreeHeap',
'GetSystemTimeAsFileTime', 'GetLocalTime', 'GlobalMemoryStatus',
'GetVersionExA', 'GetVersionExW', 'GetComputerNameA', 'GetComputerNameW'
},
'DotNet': {
'_CorExeMain', '_CorDllMain'
},
'AntiDebug': {
'IsDebuggerPresent', 'CheckRemoteDebuggerPresent', 'OutputDebugStringA', 'OutputDebugStringW',
'GetTickCount', 'GetTickCount64', 'QueryPerformanceCounter', 'FindWindowA', 'FindWindowW',
'EnumWindows', 'EnumChildWindows', 'GetWindowRect', 'GetClientRect', 'SetWindowDisplayAffinity'
},
'Keylogging': {
'GetAsyncKeyState', 'GetKeyState', 'GetKeyboardState', 'MapVirtualKeyA', 'MapVirtualKeyW',
'ToAscii', 'ToAsciiEx', 'ToUnicode', 'ToUnicodeEx', 'GetKeyNameTextA', 'GetForegroundWindow',
'KeybdEvent', 'SendInput', 'MapVirtualKeyExA', 'MapVirtualKeyExW'
},
'Input': {
'GetCursorPos', 'SetCursorPos', 'MouseEvent', 'GetDoubleClickTime',
'GetCapture', 'SetCapture'
},
'ScreenCapture': {
'BitBlt', 'StretchBlt', 'GetDC', 'GetWindowDC', 'CreateCompatibleDC', 'CreateCompatibleBitmap',
'GdipSaveImageToFile', 'PrintWindow', 'GetDesktopWindow', 'CreateDCW', 'CreateDCA',
'SelectObject', 'DeleteDC', 'DeleteObject', 'GetDeviceCaps', 'GetSystemMetrics'
},
'Graphics': {
'Direct3DCreate9', 'D3D11CreateDevice'
},
'Audio': {
'WaveInOpen', 'WaveInClose', 'WaveInStart', 'WaveInStop'
},
'Clipboard': {
'OpenClipboard', 'CloseClipboard', 'GetClipboardData', 'SetClipboardData'
},
'Camera': {
'CapCreateCaptureWindowA', 'CapCreateCaptureWindowW'
}
}
TARGET_DLLS = {
'kernel32.dll', 'user32.dll', 'gdi32.dll', 'advapi32.dll', 'shell32.dll', 'ntdll.dll', 'hal.dll',
'ws2_32.dll', 'wsock32.dll', 'wininet.dll', 'winhttp.dll', 'urlmon.dll', 'crypt32.dll', 'bcrypt.dll',
'psapi.dll', 'dbghelp.dll', 'imagehlp.dll', 'shlwapi.dll', 'ole32.dll', 'oleaut32.dll', 'comctl32.dll',
'mscoree.dll', 'vbscript.dll', 'netapi32.dll', 'iphlpapi.dll', 'wtsapi32.dll', 'version.dll', 'winmm.dll'
}
BASE_SCHEMA = [
"FileEntropy", "IsExe", "IsDll", "IsDriver", "Is64Bit", "Machine", "Magic",
"TimeDateStamp", "CheckSum", "ImageBase", "SizeOfImage", "SizeOfHeaders",
"Characteristics", "DllCharacteristics", "Subsystem", "LoaderFlags",
@@ -32,6 +187,8 @@ class Features:
"RsrcSection", "RsrcSizeRatio", "SectionCount", "ExecutableSections",
"WritableSections", "ReadableSections", "SectionException",
"IconCount", "ApiCount", "ExportCount", "DebugCount", "ExceptionCount",
"ImportCount", "ImportFunctionCount",
"FileDescriptionLength", "FileVersionLength", "ProductNameLength",
"ProductVersionLength", "CompanyNameLength", "LegalCopyrightLength",
"CommentsLength", "InternalNameLength", "LegalTrademarksLength",
@@ -39,85 +196,20 @@ class Features:
"TrustSigned", "IsDebug", "IsPatched", "IsPrivateBuild", "IsPreRelease",
"IsSpecialBuild", "IsAdmin", "IsInstall", "HasTlsCallbacks",
"HasInvalidTimestamp", "HasRelocationDirectory", "HasPacked", "FileTimeException",
"_CorExeMain", "_CorDllMain",
"GetDC", "CreateDCW", "CreateDCA", "BitBlt", "StretchBlt", "CreateCompatibleDC",
"CreateCompatibleBitmap", "SelectObject", "DeleteDC", "DeleteObject", "GetDeviceCaps",
"GetSystemMetrics",
"SetWindowsHookExA", "SetWindowsHookExW", "UnhookWindowsHookEx",
"GetAsyncKeyState", "GetKeyState", "GetKeyboardState",
"MapVirtualKeyA", "MapVirtualKeyW", "MapVirtualKeyExA", "MapVirtualKeyExW",
"ToAscii", "ToAsciiEx", "ToUnicode", "ToUnicodeEx", "KeybdEvent", "SendInput",
"GetCursorPos", "SetCursorPos", "MouseEvent", "GetDoubleClickTime",
"GetCapture", "SetCapture",
"WaveInOpen", "WaveInClose", "WaveInStart", "WaveInStop",
"CapCreateCaptureWindowA", "CapCreateCaptureWindowW",
"OpenClipboard", "CloseClipboard", "GetClipboardData", "SetClipboardData",
"GetDesktopWindow", "GetForegroundWindow", "GetWindowDC", "GetWindowRect",
"GetClientRect", "PrintWindow", "SetWindowDisplayAffinity", "FindWindowA", "FindWindowW",
"EnumWindows", "EnumChildWindows",
"Direct3DCreate9", "D3D11CreateDevice", "GdipSaveImageToFile",
"CreateProcessA", "CreateProcessW", "CreateThread", "ResumeThread", "SuspendThread",
"OpenProcess", "TerminateProcess", "GetCurrentProcess", "GetCurrentProcessId",
"ExitProcess", "WinExec", "ShellExecuteA", "ShellExecuteW",
"CreateRemoteThread", "QueueUserAPC", "VirtualAlloc", "VirtualAllocEx",
"VirtualProtect", "VirtualProtectEx", "WriteProcessMemory", "ReadProcessMemory",
"NtUnmapViewOfSection", "SetThreadContext", "GetThreadContext", "Wow64SetThreadContext",
"ZwUnmapViewOfSection", "CreateFileMappingA", "CreateFileMappingW", "MapViewOfFile",
"UnmapViewOfSection",
"LoadLibraryA", "LoadLibraryW", "LoadLibraryExA", "LoadLibraryExW",
"GetProcAddress", "GetModuleHandleA", "GetModuleHandleW", "FreeLibrary",
"CreateFileA", "CreateFileW", "WriteFile", "ReadFile", "DeleteFileA", "DeleteFileW",
"CopyFileA", "CopyFileW", "MoveFileA", "MoveFileW", "FindFirstFileA", "FindNextFileA",
"GetTempPathA", "GetTempPathW", "GetTempFileNameA", "GetTempFileNameW",
"SetFileAttributesA", "SetFileAttributesW", "DeviceIoControl", "SetFileTime",
"RegOpenKeyExA", "RegOpenKeyExW", "RegSetValueExA", "RegSetValueExW",
"RegCreateKeyExA", "RegCreateKeyExW", "RegDeleteKeyA", "RegDeleteKeyW",
"RegEnumValueA", "RegEnumValueW", "RegQueryValueExA", "RegQueryValueExW",
"RegDeleteValueA", "RegDeleteValueW",
"OpenSCManagerA", "OpenSCManagerW", "CreateServiceA", "CreateServiceW",
"StartServiceA", "StartServiceW", "ControlService", "DeleteService",
"AdjustTokenPrivileges", "LookupPrivilegeValueA", "LookupPrivilegeValueW",
"OpenProcessToken", "NetUserAdd", "NetLocalGroupAddMembers",
"Socket", "Connect", "Send", "Recv", "WSAStartup", "gethostbyname", "getaddrinfo",
"WSAIoctl", "WSASocketA", "WSASocketW",
"URLDownloadToFileA", "URLDownloadToFileW",
"InternetOpenA", "InternetOpenW", "InternetConnectA", "InternetConnectW",
"InternetOpenUrlA", "InternetOpenUrlW",
"HttpOpenRequestA", "HttpOpenRequestW", "HttpSendRequestA", "HttpSendRequestW",
"InternetReadFile", "DnsQuery_A", "DnsQuery_W",
"IsDebuggerPresent", "CheckRemoteDebuggerPresent", "OutputDebugStringA", "OutputDebugStringW",
"GetTickCount", "QueryPerformanceCounter", "Sleep", "GetSystemTimeAsFileTime",
"GetLocalTime", "GlobalMemoryStatus", "GetVersionExA", "GetVersionExW",
"GetComputerNameA", "GetComputerNameW", "GetUserNameA", "GetUserNameW",
"CryptAcquireContextA", "CryptAcquireContextW", "CryptCreateHash", "CryptHashData",
"CryptDeriveKey", "CryptEncrypt", "CryptDecrypt", "CryptDestroyKey", "CryptDestroyHash",
"CryptReleaseContext", "CryptGenKey", "CryptImportKey", "CryptExportKey",
"RtlDecompressBuffer", "ConnectNamedPipe", "PeekNamedPipe"
"HasInvalidTimestamp", "HasRelocationDirectory", "HasPacked", "FileTimeException"
]
CAT_SCHEMA = [f"Cat_{k}" for k in API_CAT_MAPPING.keys()]
DLL_SCHEMA = [f"Dll_{d.replace('.', '_')}" for d in sorted(TARGET_DLLS)]
_ALL_APIS = {api for apis in API_CAT_MAPPING.values() for api in apis}
API_SCHEMA = [f"Api_{api}" for api in sorted(_ALL_APIS)]
SCHEMA = BASE_SCHEMA + CAT_SCHEMA + DLL_SCHEMA + API_SCHEMA
SCHEMA_SET = set(SCHEMA)
PACKERS = {
'upx', 'aspack', 'asprotect', 'pecompact', 'upack', 'fsg', 'mew',
'mpress', 'ezip', 'pklt', 'shrink', 'petite', 'telock',
'themida', 'winlicense', 'tmd', 'vmp', 'enigma', 'obsidium',
'pelock', 'exestealth', 'yoda', 'armadillo', 'zprotect',
'sforce', 'starforce', 'qihoo', 'wisevec', 'megastop',
@@ -126,79 +218,33 @@ class Features:
####################################################################################################
def verify_signature(file_path):
if os.name != 'nt':
return 0
try:
wintrust = ctypes.windll.wintrust
WINTRUST_ACTION_GENERIC_VERIFY_V2 = ctypes.c_char_p(b'\x6b\xc5\xaa\x00\x44\xcd\xd0\x11\x8c\xc2\x00\xc0\x4f\xc2\x95\xee')
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(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),
("pFile", ctypes.POINTER(WINTRUST_FILE_INFO)),
("dwStateAction", ctypes.wintypes.DWORD),
("hWVTStateData", ctypes.wintypes.HANDLE),
("pwszURLReference", ctypes.wintypes.LPCWSTR),
("dwProvFlags", ctypes.wintypes.DWORD),
("dwUIContext", ctypes.wintypes.DWORD),
("pSignatureSettings", ctypes.wintypes.LPVOID)
]
fi = WINTRUST_FILE_INFO(ctypes.sizeof(WINTRUST_FILE_INFO), file_path, None, None)
td = WINTRUST_DATA(ctypes.sizeof(WINTRUST_DATA), None, None, 2, 1, 1, ctypes.pointer(fi), 0, None, None, 0x00000080, 0, None)
action = ctypes.wintypes.GUID.from_buffer_copy(WINTRUST_ACTION_GENERIC_VERIFY_V2)
return 1 if wintrust.WinVerifyTrust(None, ctypes.byref(action), ctypes.byref(td)) == 0 else 0
except:
return 0
####################################################################################################
def safe_float(val):
try:
f = float(val)
if math.isinf(f) or math.isnan(f):
return 0.0
if math.isinf(f) or math.isnan(f): return 0.0
return f
except:
return 0.0
def process_single_file(file_path):
####################################################################################################
def process_single_file_content(file_bytes, file_path):
def _calc_entropy(data):
if not data:
return 0.0
if not data: return 0.0
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)))
def _safe_div(n, d):
return n / d if d > 0 else 0.0
def _safe_div(n, d): return n / d if d > 0 else 0.0
try:
fsize = os.path.getsize(file_path)
if fsize == 0 or fsize > MAX_FILE_SIZE:
return None
except:
return None
fsize = len(file_bytes)
if fsize == 0: return None
fts = {k: 0 for k in Features.SCHEMA}
fts = {k: 0.0 for k in Features.SCHEMA}
pe = None
try:
pe = pefile.PE(file_path, fast_load=True)
pe = pefile.PE(data=file_bytes, fast_load=True)
fts['TrustSigned'] = verify_signature(file_path)
@@ -244,41 +290,36 @@ def process_single_file(file_path):
fts['SizeOfHeapCommit'] = op.SizeOfHeapCommit
fts['LoaderFlags'] = op.LoaderFlags
fts['NumberOfRvaAndSizes'] = op.NumberOfRvaAndSizes
fts['Is64Bit'] = 1 if fh.Machine == 0x8664 else 0
fts['IsExe'] = 1 if pe.is_exe() else 0
fts['IsDll'] = 1 if pe.is_dll() else 0
fts['IsDriver'] = 1 if pe.is_driver() else 0
if len(op.DATA_DIRECTORY) > 9 and op.DATA_DIRECTORY[9].VirtualAddress > 0:
fts['HasTlsCallbacks'] = 1
if len(op.DATA_DIRECTORY) > 5 and op.DATA_DIRECTORY[5].Size > 0:
fts['HasRelocationDirectory'] = 1
if len(op.DATA_DIRECTORY) > 9 and op.DATA_DIRECTORY[9].VirtualAddress > 0: fts['HasTlsCallbacks'] = 1
if len(op.DATA_DIRECTORY) > 5 and op.DATA_DIRECTORY[5].Size > 0: fts['HasRelocationDirectory'] = 1
pe.parse_sections(pe.FILE_HEADER.NumberOfSections)
file_data = pe.get_memory_mapped_image()
fts['FileEntropy'] = _calc_entropy(file_data)
try:
fts['FileEntropy'] = _calc_entropy(file_bytes)
except:
fts['FileEntropy'] = 0.0
fts['SectionCount'] = len(pe.sections)
for section in pe.sections:
try:
name = section.Name.decode('ascii', 'ignore').strip('\x00')
except:
name = ""
s_data = section.get_data()
s_entropy = _calc_entropy(s_data)
try: name = section.Name.decode('ascii', 'ignore').strip('\x00')
except: name = ""
try:
s_data = section.get_data()
s_entropy = _calc_entropy(s_data)
except:
s_data = b""
s_entropy = 0.0
if any(p in name.lower() for p in Features.PACKERS):
fts['HasPacked'] = 1
if section.Characteristics & 0x20000000:
fts['ExecutableSections'] += 1
if section.Characteristics & 0x80000000:
fts['WritableSections'] += 1
if section.Characteristics & 0x40000000:
fts['ReadableSections'] += 1
if section.SizeOfRawData + section.PointerToRawData > fsize:
fts['SectionException'] = 1
if any(p in name.lower() for p in Features.PACKERS): fts['HasPacked'] = 1
if section.Characteristics & 0x20000000: fts['ExecutableSections'] += 1
if section.Characteristics & 0x80000000: fts['WritableSections'] += 1
if section.Characteristics & 0x40000000: fts['ReadableSections'] += 1
if section.SizeOfRawData + section.PointerToRawData > fsize: fts['SectionException'] = 1
ratio = _safe_div(section.SizeOfRawData, fsize)
if name.startswith('.text'):
@@ -290,9 +331,7 @@ def process_single_file(file_path):
elif name.startswith('.rsrc'):
fts['RsrcSection'] = s_entropy
fts['RsrcSizeRatio'] = ratio
if b"requireAdministrator" in s_data:
fts['IsAdmin'] = 1
if b"requireAdministrator" in s_data: fts['IsAdmin'] = 1
for sig in Features.INSTALLER_SIGS:
if sig in s_data:
fts['IsInstall'] = 1
@@ -306,40 +345,40 @@ def process_single_file(file_path):
])
if hasattr(pe, 'DIRECTORY_ENTRY_IMPORT'):
fts['ImportCount'] = len(pe.DIRECTORY_ENTRY_IMPORT)
for entry in pe.DIRECTORY_ENTRY_IMPORT:
try:
dll_name = entry.dll.decode('ascii', 'ignore').lower()
if dll_name in Features.TARGET_DLLS: fts[f"Dll_{dll_name.replace('.', '_')}"] = 1.0
except: pass
for imp in entry.imports:
fts['ApiCount'] += 1
if not imp or not hasattr(imp, 'name') or not imp.name:
continue
fts['ImportFunctionCount'] += 1
if not imp or not hasattr(imp, 'name') or not imp.name: continue
try:
name = imp.name.decode('ascii', 'ignore')
if name in Features.SCHEMA_SET:
fts[name] = 1
except:
pass
func_name = imp.name.decode('ascii', 'ignore')
if func_name in Features._ALL_APIS:
fts[f'Api_{func_name}'] += 1
for cat, api_set in Features.API_CAT_MAPPING.items():
if func_name in api_set:
fts[f'Cat_{cat}'] += 1
break
except: pass
if hasattr(pe, 'DIRECTORY_ENTRY_EXPORT'):
fts['ExportCount'] = len(pe.DIRECTORY_ENTRY_EXPORT.symbols)
if hasattr(pe, 'DIRECTORY_ENTRY_EXPORT'): fts['ExportCount'] = len(pe.DIRECTORY_ENTRY_EXPORT.symbols)
if hasattr(pe, 'DIRECTORY_ENTRY_RESOURCE'):
for entry in pe.DIRECTORY_ENTRY_RESOURCE.entries:
if entry.id == 3 and hasattr(entry, 'directory'):
fts['IconCount'] = len(entry.directory.entries)
break
if hasattr(pe, 'DIRECTORY_ENTRY_DEBUG'):
fts['DebugCount'] = len(pe.DIRECTORY_ENTRY_DEBUG)
for k in fts:
fts[k] = safe_float(fts[k])
if hasattr(pe, 'DIRECTORY_ENTRY_DEBUG'): fts['DebugCount'] = len(pe.DIRECTORY_ENTRY_DEBUG)
for k in fts: fts[k] = safe_float(fts[k])
return fts
except Exception:
except Exception as e:
return None
finally:
if pe:
pe.close()
if pe: pe.close()
####################################################################################################
@@ -348,50 +387,28 @@ def init_db():
conn.execute("PRAGMA synchronous = OFF")
conn.execute("PRAGMA journal_mode = WAL")
conn.execute("PRAGMA cache_size = 10000")
cols = ", ".join([f"{c} REAL" for c in Features.SCHEMA])
cols = "FileHash TEXT, " + ", ".join([f"{c} REAL" for c in Features.SCHEMA])
conn.execute(f"CREATE TABLE IF NOT EXISTS PeData ({cols}, Label INTEGER)")
conn.execute(f"CREATE UNIQUE INDEX IF NOT EXISTS idx_pedata_hash ON PeData (FileHash)")
conn.commit()
return conn
def prune_db_duplicates(conn):
def load_existing_hashes(conn):
print("[*] Loading existing hashes into memory...")
try:
print("\n[*] Post-processing: Checking for and removing duplicate entries...")
all_cols = ", ".join(Features.SCHEMA + ["Label"])
sql_count_before = "SELECT COUNT(*) FROM PeData"
before = conn.execute(sql_count_before).fetchone()[0]
sql_dedup = f"""
DELETE FROM PeData
WHERE rowid NOT IN (
SELECT MIN(rowid)
FROM PeData
GROUP BY {all_cols}
)
"""
conn.execute(sql_dedup)
conn.commit()
after = conn.execute(sql_count_before).fetchone()[0]
removed = before - after
if removed > 0:
print(f"[+] Deduplication complete. Removed {removed} duplicate rows.")
print(" (Duplicates occur if the same file is scanned multiple times)")
else:
print("[+] Deduplication complete. No duplicates found.")
except Exception as e:
print(f"[-] Deduplication Warning: {e}")
cursor = conn.execute("SELECT FileHash FROM PeData")
return {row[0] for row in cursor.fetchall()}
except:
return set()
####################################################################################################
def scan_and_save(conn, path, label):
def scan_and_save(conn, path, label, existing_hashes):
print(f"\n[*] Scanning: {path} (Label={label})")
print(f"[*] Config: Max File Size = {MAX_FILE_SIZE/1024/1024:.0f} MB")
files = []
print("[*] Walking directory tree...")
if os.path.isfile(path):
files.append(path)
else:
@@ -402,29 +419,49 @@ def scan_and_save(conn, path, label):
total = len(files)
if total == 0:
print("[-] No valid PE files found in this path.")
print("[-] No valid PE files found.")
return
print(f"[*] Found {total} files. Processing...")
print(f"[*] Found {total} files. Processing...\n")
batch = []
count = 0
errors = 0
skipped = 0
start_time = time.time()
placeholders = ",".join(["?"] * (len(Features.SCHEMA) + 1))
placeholders = ",".join(["?"] * (len(Features.SCHEMA) + 2))
sql = f"INSERT INTO PeData VALUES ({placeholders})"
try:
for i, file_path in enumerate(files):
res = process_single_file(file_path)
if res:
vals = [res.get(c, 0.0) for c in Features.SCHEMA]
vals.append(float(label))
batch.append(tuple(vals))
count += 1
else:
try:
fsize = os.path.getsize(file_path)
if fsize == 0 or fsize > MAX_FILE_SIZE:
errors += 1
continue
with open(file_path, "rb") as f:
file_bytes = f.read()
sha256 = hashlib.sha256(file_bytes).hexdigest()
if sha256 in existing_hashes:
skipped += 1
else:
res = process_single_file_content(file_bytes, file_path)
if res:
vals = [sha256]
vals.extend([res.get(c, 0.0) for c in Features.SCHEMA])
vals.append(float(label))
batch.append(tuple(vals))
existing_hashes.add(sha256)
count += 1
else:
errors += 1
except Exception:
errors += 1
if len(batch) >= BATCH_SIZE:
@@ -433,19 +470,12 @@ def scan_and_save(conn, path, label):
conn.commit()
batch = []
except Exception as batch_err:
print(f"\n[-] Batch Error: {batch_err}. Retrying strictly...")
for row in batch:
try:
conn.execute(sql, row)
except Exception:
errors += 1
count -= 1
conn.commit()
print(f"\n[-] Batch Error: {batch_err}")
batch = []
elapsed = time.time() - start_time
rate = (i + 1) / elapsed if elapsed > 0 else 0
sys.stdout.write(f"\r[{i+1}/{total}] OK: {count} | Skip/Err: {errors} | {rate:.1f} files/s")
sys.stdout.write(f"\r[{i+1}/{total}] New: {count} | Skip: {skipped} | Err: {errors} | {rate:.1f} files/s")
sys.stdout.flush()
except KeyboardInterrupt:
@@ -454,43 +484,43 @@ def scan_and_save(conn, path, label):
print(f"\n[-] Unexpected Error: {e}")
finally:
if batch:
print(f"\n[*] Saving remaining {len(batch)} items...")
try:
conn.executemany(sql, batch)
conn.commit()
except Exception as final_err:
print(f"[-] Final Batch Error: {final_err}. Retrying strictly...")
for row in batch:
try:
conn.execute(sql, row)
except Exception:
pass
conn.commit()
except Exception:
pass
print(f"\n[+] Scan finished/stopped. Total added: {count}")
prune_db_duplicates(conn)
print(f"\n\n[+] Scan finished. Total New: {count}, Total Skipped: {skipped}")
####################################################################################################
if __name__ == "__main__":
print("\n---------------- PE Dataset Builder v2.6 ----------------\n")
print("Type 'exit' or 'q' at any prompt to quit.")
print("\n---------------- PE Dataset Builder v3.1 ----------------\n")
if os.path.exists(DB_PATH):
print(f"[-] Note: If SCHEMA changed, please delete {DB_PATH} before running.")
try:
conn_check = sqlite3.connect(DB_PATH)
cursor = conn_check.execute("PRAGMA table_info(PeData)")
cols = [row[1] for row in cursor.fetchall()]
if 'FileHash' not in cols or len(cols) != len(Features.SCHEMA) + 2:
print(f"[-] WARNING: DB Schema mismatch. Please delete {DB_PATH} to rebuild.")
conn_check.close()
except: pass
conn = init_db()
existing_hashes = load_existing_hashes(conn)
print(f"[+] Loaded {len(existing_hashes)} hashes from database.\n")
try:
while True:
print("\n---------------------------------------------------------\n")
raw_path = input("[*] Enter Path (File or Folder): ").strip()
print("-" * 60)
raw_path = input("\n[*] Enter Path (File or Folder): ").strip()
target_path = raw_path.strip('"').strip("'")
if target_path.lower() in ['exit', 'q']:
break
if target_path.lower() in ['exit', 'q']: break
if not target_path or not os.path.exists(target_path):
print("[-] Invalid path. Please try again.")
print("[-] Invalid path.")
continue
while True:
@@ -498,14 +528,13 @@ if __name__ == "__main__":
if label_input.lower() in ['exit', 'q']:
conn.close()
sys.exit()
if label_input in ['0', '1']:
label = int(label_input)
break
else:
print("[-] Please enter 0 or 1.")
print("[-] Please enter 0 or 1.")
scan_and_save(conn, target_path, label)
scan_and_save(conn, target_path, label, existing_hashes)
except KeyboardInterrupt:
print("\n\n[-] Interrupted.")
finally:
Engine/Properties/train.py
+26 -16
View File
@@ -3,6 +3,7 @@ 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
@@ -10,8 +11,8 @@ from sklearn.metrics import classification_report, accuracy_score, roc_auc_score
####################################################################################################
DB_PATH = "pe_features.db"
MODEL_FILE = "pe_malware_lgbm.txt"
ONNX_FILE = "pe_malware.onnx"
MODEL_FILE = "model.txt"
ONNX_FILE = "model.onnx"
FEATURE_FILE = "features.json"
TEST_SIZE = 0.0001
RANDOM_SEED = 42
@@ -31,11 +32,14 @@ LGBM_PARAMS = {
'max_depth': -1,
'min_data_in_leaf': 20,
'lambda_l1': 0.1,
'lambda_l2': 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
@@ -44,14 +48,16 @@ def get_raw_schema(db_path):
cursor = conn.cursor()
cursor.execute("PRAGMA table_info(PeData)")
columns = [row[1] for row in cursor.fetchall()]
exclude_lower = {'label', 'id', 'filename', 'filelength', 'rowid', 'probability', 'predictedlabel', 'score'}
return [c for c in columns if c.lower().strip() not in exclude_lower]
exclude = {'label', 'id', 'filename', 'filelength', 'rowid', 'probability', 'predictedlabel', 'score', 'filehash'}
return [c for c in columns if c.lower().strip() not in exclude]
except:
return None
finally:
conn.close()
def load_and_clean_data(db_path, raw_cols):
####################################################################################################
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']))
@@ -72,10 +78,12 @@ def load_and_clean_data(db_path, raw_cols):
valid_features = [c for c in raw_cols if c in X.columns]
X = X[valid_features].copy()
X = X.rename(columns=lambda x: re.sub('[^A-Za-z0-9_]+', '', x))
rename_map = {old: clean_col_name(old) for old in X.columns}
X = X.rename(columns=rename_map)
if X.isnull().values.any():
X = X.fillna(X.median(numeric_only=True))
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
@@ -95,13 +103,11 @@ def train_process(X, y):
n_neg = (y_train == 0).sum()
base_weight = n_neg / n_pos if n_pos > 0 else 1.0
weight_ratio_base = 1.0
weight_ratio_target = 0.1
final_pos_weight = base_weight * (weight_ratio_target / weight_ratio_base)
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 (Favor Safe)={final_pos_weight:.4f}")
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)
@@ -109,7 +115,7 @@ def train_process(X, y):
params = LGBM_PARAMS.copy()
params['scale_pos_weight'] = final_pos_weight
print("\n[*] Starting LightGBM training (Full Force, No Early Stopping)...")
print("\n[*] Starting LightGBM training...")
start_time = time.time()
model = lgb.train(
@@ -127,6 +133,8 @@ def train_process(X, y):
return model, X_test, y_test
####################################################################################################
def export_onnx_model(model, feature_names):
print(f"[*] Exporting ONNX model to {ONNX_FILE}...")
try:
@@ -141,6 +149,8 @@ def export_onnx_model(model, feature_names):
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)
@@ -171,7 +181,7 @@ def evaluate_and_save(model, X_test, y_test):
####################################################################################################
if __name__ == "__main__":
print("\n---------------- PE Malware Trainer v3.0 ----------------\n")
print("\n---------------- PE Malware Trainer v3.1 ----------------\n")
if not os.path.exists(DB_PATH):
print(f"[-] Error: Database {DB_PATH} not found.")
@@ -182,7 +192,7 @@ if __name__ == "__main__":
print("[-] Error: Could not read schema.")
sys.exit(1)
X, y, _ = load_and_clean_data(DB_PATH, raw_schema)
X, y, _ = load_data(DB_PATH, raw_schema)
if X is None:
sys.exit(1)