From 73d09f9b09c3c93b320c8da8121aeedb4c10decb Mon Sep 17 00:00:00 2001 From: 87owo <85057800+87owo@users.noreply.github.com> Date: Fri, 20 Feb 2026 13:49:18 +0800 Subject: [PATCH] Add files via upload --- Engine/Properties/convert.py | 545 ++++++++++++++++++----------------- Engine/Properties/train.py | 42 ++- 2 files changed, 313 insertions(+), 274 deletions(-) diff --git a/Engine/Properties/convert.py b/Engine/Properties/convert.py index 3ef2672..19b6d5e 100644 --- a/Engine/Properties/convert.py +++ b/Engine/Properties/convert.py @@ -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: diff --git a/Engine/Properties/train.py b/Engine/Properties/train.py index 5c24e93..988d2cd 100644 --- a/Engine/Properties/train.py +++ b/Engine/Properties/train.py @@ -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)