Compare commits

...
79 Commits
Author SHA1 Message Date
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
87owo f8dbd12b69 Add files via upload 2026-05-06 12:54:54 +08:00
87owo cf6ba66dd2 Add files via upload 2026-05-04 23:29:36 +08:00
87owo 2080b5d10a Add files via upload 2026-05-04 23:29:14 +08:00
87owo dd180977d9 Add files via upload 2026-05-04 23:28:48 +08:00
87owo 221dcbaf1b Add files via upload 2026-05-04 23:28:13 +08:00
87owo 1a2098f4fb Add files via upload 2026-05-04 23:27:50 +08:00
87owo 73eab17110 Add files via upload 2026-05-03 11:34:05 +08:00
87owo 978a336358 Update README.md 2026-05-03 03:27:57 +08:00
87owo 6a0fd4e427 Update README.md 2026-05-03 03:23:58 +08:00
87owo 5fd4efd920 Add files via upload 2026-05-03 03:20:31 +08:00
87owo 43be2987f5 Add files via upload 2026-05-02 17:47:38 +08:00
87owo 06cc048fdc Add files via upload 2026-05-02 17:47:03 +08:00
87owo 5752471420 Add files via upload 2026-05-02 16:16:46 +08:00
87owo 3dfe12b90a Add files via upload 2026-05-02 15:48:47 +08:00
87owo ad1e3501b1 Add files via upload 2026-05-02 15:33:52 +08:00
87owo 1b46572e4c Add files via upload 2026-05-02 15:09:42 +08:00
87owo 9d8e42d34c Add files via upload 2026-05-02 15:09:28 +08:00
87owo 3f6e9dab3e Add files via upload 2026-05-02 15:09:06 +08:00
87owo ad85f11875 Add files via upload 2026-05-02 01:49:17 +08:00
87owo 9ae31834cc Add files via upload 2026-05-02 01:16:32 +08:00
87owo 49a4dc6942 Add files via upload 2026-05-02 01:16:05 +08:00
87owo a9f58cc30f Add files via upload 2026-05-02 01:15:51 +08:00
87owo bdb2198582 Add files via upload 2026-05-02 01:15:28 +08:00
87owo f32a77ca8b Add files via upload 2026-05-01 23:16:30 +08:00
87owo 121d3fb0ae Add files via upload 2026-05-01 23:15:59 +08:00
87owo 1dd72f7600 Add files via upload 2026-05-01 23:15:22 +08:00
87owo c7c168b007 Add files via upload 2026-05-01 23:14:45 +08:00
87owo 441b41c6af Add files via upload 2026-05-01 23:13:56 +08:00
87owo 49a907283a Add files via upload 2026-05-01 15:41:33 +08:00
87owo 0ffc52f9ba Add files via upload 2026-05-01 13:10:33 +08:00
87owo 4e96ec49e7 Update README.md 2026-04-29 21:04:53 +08:00
87owo dfb4646f51 Delete Embed directory 2026-04-29 21:04:30 +08:00
87owo a8e82c38c9 Add files via upload 2026-04-29 18:09:24 +08:00
87owo dba23289fd Add files via upload 2026-04-29 17:41:00 +08:00
87owo c2c1eb98d9 Add files via upload 2026-04-29 16:19:44 +08:00
87owo 364176912d Add files via upload 2026-04-28 14:38:35 +08:00
87owo 202245baad Add files via upload 2026-04-28 14:10:07 +08:00
87owo c3fa99640c Add files via upload 2026-04-27 15:03:53 +08:00
87owo a7c34b5a6a Add files via upload 2026-04-27 10:19:47 +08:00
87owo 05de092642 Update README.md 2026-04-26 22:35:28 +08:00
87owo 1bfa40c841 Add files via upload 2026-04-26 21:13:42 +08:00
87owo c325435ae8 Add files via upload 2026-04-26 15:41:47 +08:00
87owo b978bd41c3 Add files via upload 2026-04-26 14:31:23 +08:00
87owo f03d8733b0 Add files via upload 2026-04-26 14:30:49 +08:00
87owo a624b3b759 Update README.md 2026-04-26 14:29:57 +08:00
87owo d683398ba9 Add files via upload 2026-04-26 14:28:24 +08:00
87owo 42f692b82d Add files via upload 2026-04-26 14:27:33 +08:00
87owo 855fe9115e Delete Engine/Pattern directory 2026-04-26 14:26:30 +08:00
87owo d6bd383822 Add files via upload 2026-04-26 02:13:58 +08:00
87owo b1b77b73ac Add files via upload 2026-04-25 21:45:46 +08:00
87owo 8d87adcc50 Add files via upload 2026-04-25 21:45:32 +08:00
87owo 49bf6c54e7 Delete Interface/main.py 2026-04-25 20:53:22 +08:00
87owo 7e92e0b1d9 Add files via upload 2026-04-25 20:52:55 +08:00
87owo f0805d320c Add files via upload 2026-04-25 19:19:36 +08:00
30 changed files with 5298 additions and 5969 deletions

No files matched your search

-16
View File
@@ -1,16 +0,0 @@
FROM python:3.10-slim
RUN apt-get update && apt-get install -y \
build-essential \
libfuzzy-dev \
git \
&& rm -rf /var/lib/apt/lists/*
WORKDIR /app
COPY requirements.txt .
RUN pip install --no-cache-dir -r requirements.txt
COPY . .
CMD ["python", "-u", "predict_onnx_linux.py"]
-211
View File
@@ -1,211 +0,0 @@
import os, time, uuid, logging, requests, hashlib
from pathlib import Path
from typing import Optional, List, Dict, Any
####################################################################################################
class PYAS_Client:
def __init__(self, api_key: str, host: str, timeout: int = 30):
self.base_url = host.rstrip('/')
self.timeout = timeout
self.logger = logging.getLogger("PYAS_API")
self.session = requests.Session()
self.session.headers.update({"X-API-Key": api_key, "User-Agent": "PYAS-API/1.1"})
self.logger.info(f"Initialized PYAS API Client targeting {self.base_url}")
def _request(self, method: str, endpoint: str, **kwargs) -> Optional[requests.Response]:
self.logger.debug(f"Request: {method} {endpoint}")
try:
r = self.session.request(method, f"{self.base_url}{endpoint}", timeout=self.timeout, **kwargs)
if r.status_code == 200:
return r
error_map = {401: "API Key invalid", 402: "Insufficient points", 404: "Not found", 500: "Server Error"}
self.logger.error(f"[{endpoint}] Failed: {error_map.get(r.status_code, f'HTTP {r.status_code}')}")
except Exception as e:
self.logger.error(f"Connection error: {e}")
return None
####################################################################################################
def _calculate_sha256(self, path: Path) -> str:
self.logger.debug(f"Calculating SHA256 for {path}")
sha256_hash = hashlib.sha256()
with open(path, "rb") as f:
for byte_block in iter(lambda: f.read(1048576), b""):
sha256_hash.update(byte_block)
return sha256_hash.hexdigest()
def upload_file(self, file_path: str, chunk_size: int = 5242880) -> Optional[str]:
path = Path(file_path)
if not path.exists():
self.logger.error(f"File not found: {file_path}")
return None
file_size = path.stat().st_size
self.logger.info(f"Initiating upload for {path.name} ({file_size} bytes)")
if file_size > 536870912:
self.logger.error(f"File size exceeds 512MB limit: {file_size} bytes")
return None
local_sha256 = self._calculate_sha256(path)
status_req = self._request("GET", f"/api/processing_status/{local_sha256}")
if status_req and status_req.json().get('status') in ['done', 'processing', 'queued']:
self.logger.info(f"File already exists or processing on server: {local_sha256}")
return local_sha256
if file_size > chunk_size:
return self._upload_chunked(path, chunk_size, file_size)
return self._upload_standard(path)
####################################################################################################
def _upload_standard(self, path: Path, max_retries: int = 3) -> Optional[str]:
self.logger.info(f"Starting standard upload for {path.name}")
with open(path, 'rb') as f:
file_data = f.read()
for attempt in range(max_retries):
r = self._request("POST", "/api/upload", files={'file': (path.name, file_data)})
if r and 'url' in r.json():
sha256 = r.json().get('url', '').split('/')[-1]
self.logger.info(f"Standard upload successful: {sha256}")
return sha256
self.logger.warning(f"Standard upload retry {attempt + 1}/{max_retries}")
time.sleep(1)
self.logger.error("Standard upload failed after retries")
return None
def _upload_chunked(self, path: Path, chunk_size: int, file_size: int, max_retries: int = 3) -> Optional[str]:
total_chunks = (file_size + chunk_size - 1) // chunk_size
upload_id = uuid.uuid4().hex
self.logger.info(f"Starting chunked upload for {path.name} (Chunks: {total_chunks}, ID: {upload_id})")
with open(path, 'rb') as f:
for i in range(total_chunks):
chunk_data = f.read(chunk_size)
headers = {"X-Chunk-Index": str(i), "X-Total-Chunks": str(total_chunks), "X-Upload-ID": upload_id}
chunk_success = False
for attempt in range(max_retries):
r = self._request("POST", "/api/upload", files={'file': (path.name, chunk_data)}, headers=headers)
if r:
chunk_success = True
self.logger.debug(f"Uploaded chunk {i + 1}/{total_chunks}")
if i == total_chunks - 1 and 'url' in r.json():
sha256 = r.json().get('url', '').split('/')[-1]
self.logger.info(f"Chunked upload successful: {sha256}")
return sha256
break
self.logger.warning(f"Chunk upload retry {attempt + 1}/{max_retries} for chunk {i}")
time.sleep(1)
if not chunk_success:
self.logger.error(f"Failed to upload chunk {i + 1}/{total_chunks}")
return None
return None
####################################################################################################
def wait_for_analysis(self, sha256: str, interval: int = 5, max_retries: int = 10) -> bool:
self.logger.info(f"Waiting for analysis completion: {sha256}")
for attempt in range(max_retries):
r = self._request("GET", f"/api/processing_status/{sha256}")
if r:
status = r.json().get('status')
self.logger.debug(f"Analysis status: {status}")
if status == 'done':
self.logger.info(f"Analysis completed: {sha256}")
return True
if status in ['missing', 'error', 'failed']:
self.logger.error(f"Analysis failed with status: {status}")
return False
time.sleep(interval)
self.logger.error(f"Analysis polling timed out for {sha256}")
return False
def get_report(self, sha256: str) -> Optional[Dict[str, Any]]:
self.logger.info(f"Fetching report for {sha256}")
r = self._request("GET", f"/api/report/{sha256}")
if r:
return r.json()
self.logger.error(f"Failed to fetch report for {sha256}")
return None
####################################################################################################
def rescan(self, sha256: str) -> bool:
self.logger.info(f"Triggering rescan for {sha256}")
r = self._request("POST", f"/api/rescan/{sha256}")
success = bool(r and r.json().get('status') == 'success')
if success:
self.logger.info(f"Rescan triggered successfully for {sha256}")
else:
self.logger.error(f"Rescan failed for {sha256}")
return success
def search_files(self, query: str, limit: int = 50, page: int = 0, snapshot: Optional[str] = None) -> List[Dict[str, Any]]:
self.logger.info(f"Searching files (Query: '{query}', Limit: {limit}, Page: {page})")
payload = {'query': query, 'page': page}
if snapshot:
payload['snapshot'] = snapshot
r = self._request("POST", "/api/search_more", data=payload)
results = r.json().get('results', [])[:limit] if r else []
self.logger.info(f"Search returned {len(results)} results")
return results
def download_sample(self, sha256: str, save_dir: str = "downloads") -> Optional[str]:
self.logger.info(f"Downloading sample {sha256}")
os.makedirs(save_dir, exist_ok=True)
path = os.path.join(save_dir, f"{sha256}.zip")
r = self._request("GET", f"/api/download/{sha256}", stream=True)
if r:
with open(path, 'wb') as f:
for chunk in r.iter_content(8192):
if chunk:
f.write(chunk)
self.logger.info(f"Sample downloaded successfully to {path}")
return path
self.logger.error(f"Failed to download sample {sha256}")
return None
####################################################################################################
if __name__ == "__main__":
logging.basicConfig(level=logging.DEBUG, format='[%(asctime)s] %(levelname)s [%(name)s] - %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
API_KEY = ""
HOST = "https://pyas-security.com/"
client = PYAS_Client(API_KEY, HOST)
target_file = r"test.exe"
sha = client.upload_file(target_file)
if sha and client.wait_for_analysis(sha):
report = client.get_report(sha)
if report:
meta = report.get('data', {}).get('metadata', {})
logging.info(f"Result: {meta.get('label')} ({meta.get('score')}%)")
for item in client.search_files('date>=2020; ext=exe | dll', limit=3):
client.download_sample(item['sha256'])
-1
View File
@@ -1 +0,0 @@
["FileEntropy", "SectionMeanEntropy", "IsExe", "IsDll", "IsDriver", "Is64Bit", "Machine", "Magic", "CheckSum", "ImageBase", "SizeOfImage", "SizeOfHeaders", "Characteristics", "DllCharacteristics", "Subsystem", "LoaderFlags", "MajorLinkerVersion", "MinorLinkerVersion", "SizeOfCode", "SizeOfInitializedData", "SizeOfUninitializedData", "AddressOfEntryPoint", "BaseOfCode", "SectionAlignment", "FileAlignment", "MajorOperatingSystemVersion", "MinorOperatingSystemVersion", "MajorImageVersion", "MinorImageVersion", "MajorSubsystemVersion", "MinorSubsystemVersion", "SizeOfStackReserve", "SizeOfStackCommit", "SizeOfHeapReserve", "SizeOfHeapCommit", "NumberOfSections", "NumberOfRvaAndSizes", "PointerToSymbolTable", "NumberOfSymbols", "SizeOfOptionalHeader", "TextSectionMaxEntropy", "TextSectionMeanEntropy", "TextSizeRatio", "DataSectionMaxEntropy", "DataSectionMeanEntropy", "DataSizeRatio", "RsrcSectionMaxEntropy", "RsrcSectionMeanEntropy", "RsrcSizeRatio", "SectionCount", "ExecutableSections", "WritableSections", "ReadableSections", "SectionException", "Char_00000020_Count", "Char_00000020_MeanEntropy", "Char_00000040_Count", "Char_00000040_MeanEntropy", "Char_00000080_Count", "Char_00000080_MeanEntropy", "Char_02000000_Count", "Char_02000000_MeanEntropy", "Char_20000000_Count", "Char_20000000_MeanEntropy", "Char_40000000_Count", "Char_40000000_MeanEntropy", "Char_80000000_Count", "Char_80000000_MeanEntropy", "IconCount", "ApiCount", "ExportCount", "DebugCount", "ExceptionCount", "ImportCount", "ImportFunctionCount", "FileDescriptionLength", "FileVersionLength", "ProductNameLength", "ProductVersionLength", "CompanyNameLength", "LegalCopyrightLength", "CommentsLength", "InternalNameLength", "LegalTrademarksLength", "SpecialBuildLength", "PrivateBuildLength", "TrustSigned", "IsDebug", "IsPatched", "IsPrivateBuild", "IsPreRelease", "IsSpecialBuild", "IsAdmin", "IsInstall", "HasTlsCallbacks", "HasInvalidTimestamp", "HasRelocationDirectory", "HasPacked", "FileTimeException", "Cat_AntiDebug", "Cat_Audio", "Cat_COM", "Cat_Camera", "Cat_Clipboard", "Cat_DataObfuscation", "Cat_DotNet", "Cat_Encryption", "Cat_FileIO", "Cat_Graphics", "Cat_Injection", "Cat_Input", "Cat_Keylogging", "Cat_Memory", "Cat_MultiThreading", "Cat_Native", "Cat_Network", "Cat_Privileges", "Cat_ProcessControl", "Cat_Registry", "Cat_Resource", "Cat_ScreenCapture", "Cat_Services", "Cat_Synchronization", "Cat_WindowControl", "Dll_advapi32_dll", "Dll_bcrypt_dll", "Dll_comctl32_dll", "Dll_crypt32_dll", "Dll_dbghelp_dll", "Dll_gdi32_dll", "Dll_hal_dll", "Dll_imagehlp_dll", "Dll_iphlpapi_dll", "Dll_kernel32_dll", "Dll_mscoree_dll", "Dll_netapi32_dll", "Dll_ntdll_dll", "Dll_ole32_dll", "Dll_oleaut32_dll", "Dll_psapi_dll", "Dll_shell32_dll", "Dll_shlwapi_dll", "Dll_urlmon_dll", "Dll_user32_dll", "Dll_vbscript_dll", "Dll_version_dll", "Dll_winhttp_dll", "Dll_wininet_dll", "Dll_winmm_dll", "Dll_ws2_32_dll", "Dll_wsock32_dll", "Dll_wtsapi32_dll", "Api_AdjustTokenPrivileges", "Api_Base64Decode", "Api_BitBlt", "Api_CallNextHookEx", "Api_CallWindowProcA", "Api_CapCreateCaptureWindowA", "Api_CapCreateCaptureWindowW", "Api_CharLowerA", "Api_CharUpperA", "Api_CheckRemoteDebuggerPresent", "Api_CloseClipboard", "Api_CloseHandle", "Api_CoCreateInstance", "Api_CoInitialize", "Api_CoUninitialize", "Api_CompareStringA", "Api_ConnectNamedPipe", "Api_ControlService", "Api_CopyFileA", "Api_CopyFileW", "Api_CreateCompatibleBitmap", "Api_CreateCompatibleDC", "Api_CreateDCA", "Api_CreateDCW", "Api_CreateEventA", "Api_CreateEventW", "Api_CreateFileA", "Api_CreateFileMappingA", "Api_CreateFileMappingW", "Api_CreateFileW", "Api_CreateMutexA", "Api_CreateMutexW", "Api_CreateProcessA", "Api_CreateProcessW", "Api_CreateRemoteThread", "Api_CreateServiceA", "Api_CreateServiceW", "Api_CreateThread", "Api_CreateThreadpoolWork", "Api_CreateWindowExA", "Api_CryptAcquireContextA", "Api_CryptAcquireContextW", "Api_CryptCreateHash", "Api_CryptDecodeObject", "Api_CryptDecrypt", "Api_CryptDeriveKey", "Api_CryptDestroyHash", "Api_CryptDestroyKey", "Api_CryptEncrypt", "Api_CryptExportKey", "Api_CryptGenKey", "Api_CryptHashData", "Api_CryptImportKey", "Api_CryptReleaseContext", "Api_D3D11CreateDevice", "Api_DecodePointer", "Api_DefWindowProcA", "Api_DeleteCriticalSection", "Api_DeleteDC", "Api_DeleteFileA", "Api_DeleteFileW", "Api_DeleteObject", "Api_DeleteService", "Api_DestroyWindow", "Api_DeviceIoControl", "Api_Direct3DCreate9", "Api_DispatchMessageA", "Api_DnsQuery_A", "Api_DnsQuery_W", "Api_EncodePointer", "Api_EnterCriticalSection", "Api_EnumChildWindows", "Api_EnumWindows", "Api_ExitProcess", "Api_ExitThread", "Api_FindClose", "Api_FindFirstFileA", "Api_FindNextFileA", "Api_FindResourceA", "Api_FindWindowA", "Api_FindWindowW", "Api_FlushFileBuffers", "Api_FreeEnvironmentStringsW", "Api_FreeLibrary", "Api_FreeResource", "Api_GdipSaveImageToFile", "Api_GetACP", "Api_GetActiveWindow", "Api_GetAsyncKeyState", "Api_GetCapture", "Api_GetClientRect", "Api_GetClipboardData", "Api_GetCommandLineA", "Api_Line truncated
BIN
View File
Binary file not shown.
-611
View File
@@ -1,611 +0,0 @@
import os, sys, time, re, math, pefile, json, datetime, mmap
import numpy as np
import onnxruntime as ort
from signify.authenticode.signed_file.pe import SignedPEFile
from signify.authenticode.verification_result import AuthenticodeVerificationResult
####################################################################################################
MODEL_FILE = "model.onnx"
FEATURE_FILE = "features.json"
MAX_FILE_SIZE = 4 * 1024 * 1024 * 1024
####################################################################################################
class WinTrust:
@staticmethod
def verify(file_path):
if not file_path or not os.path.exists(file_path) or os.path.getsize(file_path) == 0:
return 0
try:
with open(file_path, 'rb') as f:
with mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ) as mm:
pe = SignedPEFile(mm)
if not pe.signatures:
return 0
status, _ = pe.explain_verify()
if status == AuthenticodeVerificationResult.OK:
return 1
except Exception:
pass
return 0
####################################################################################################
class FeatureExtractor:
API_CAT_MAPPING = {
'ProcessControl': {
'CreateProcessA', 'CreateProcessW', 'WinExec', 'ShellExecuteA', 'ShellExecuteW',
'ShellExecuteExW', 'ExitProcess', 'TerminateProcess', 'OpenProcess', 'GetExitCodeProcess',
'SetThreadPriority', 'GetThreadPriority', 'GetCurrentProcess', 'GetCurrentProcessId',
'GetModuleFileNameA', 'GetCommandLineA', 'GetModuleFileNameW', 'GetStartupInfoW'
},
'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',
'VirtualFree', 'VirtualQuery'
},
'Synchronization': {
'WaitForSingleObject', 'WaitForSingleObjectEx', 'WaitForMultipleObjects', 'WaitForMultipleObjectsEx',
'CreateMutexA', 'CreateMutexW', 'OpenMutexW', 'ReleaseMutex', 'CreateEventA', 'CreateEventW',
'OpenEventW', 'SetEvent', 'ResetEvent', 'EnterCriticalSection', 'LeaveCriticalSection',
'InitializeCriticalSection', 'DeleteCriticalSection', 'Sleep', 'SleepEx',
'InitializeCriticalSectionAndSpinCount', 'InterlockedDecrement', 'InterlockedIncrement',
'InterlockedExchange', 'InterlockedCompareExchange', 'SetTimer', 'KillTimer'
},
'MultiThreading': {
'CreateThread', 'ResumeThread', 'SuspendThread', 'ExitThread', 'TerminateThread',
'GetCurrentThread', 'GetCurrentThreadId', 'TlsAlloc', 'TlsSetValue', 'TlsGetValue',
'CreateThreadpoolWork', 'SubmitThreadpoolWork', 'TlsFree'
},
'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',
'GetStringTypeW', 'LCMapStringW', 'IsValidCodePage', 'DecodePointer', 'EncodePointer'
},
'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',
'CloseHandle', 'GetFileType', 'SetStdHandle', 'SetFilePointerEx', 'FindClose',
'SetHandleCount', 'SetEndOfFile', 'GetFullPathNameA', 'GetFileAttributesA'
},
'Registry': {
'RegOpenKeyExA', 'RegOpenKeyExW', 'RegSetValueExA', 'RegSetValueExW',
'RegCreateKeyExA', 'RegCreateKeyExW', 'RegDeleteKeyA', 'RegDeleteKeyW',
'RegEnumValueA', 'RegEnumValueW', 'RegQueryValueExA', 'RegQueryValueExW',
'RegDeleteValueA', 'RegDeleteValueW', 'RegCloseKey'
},
'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',
'GetLastError', 'GetACP', 'RaiseException', 'SetLastError', 'GetEnvironmentStringsW',
'FreeEnvironmentStringsW', 'GetLocaleInfoA', 'lstrlenA', 'GetVersion', 'GetSystemInfo',
'CompareStringA'
},
'DotNet': {
'_CorExeMain', '_CorDllMain'
},
'AntiDebug': {
'IsDebuggerPresent', 'CheckRemoteDebuggerPresent', 'OutputDebugStringA', 'OutputDebugStringW',
'GetTickCount', 'GetTickCount64', 'QueryPerformanceCounter', 'FindWindowA', 'FindWindowW',
'EnumWindows', 'EnumChildWindows', 'GetWindowRect', 'GetClientRect', 'SetWindowDisplayAffinity',
'UnhandledExceptionFilter', 'SetUnhandledExceptionFilter', 'IsProcessorFeaturePresent',
'SetErrorMode', 'GetTimeZoneInformation', 'GetDriveTypeW', 'GetDiskFreeSpaceA'
},
'Keylogging': {
'GetAsyncKeyState', 'GetKeyState', 'GetKeyboardState', 'MapVirtualKeyA', 'MapVirtualKeyW',
'ToAscii', 'ToAsciiEx', 'ToUnicode', 'ToUnicodeEx', 'GetKeyNameTextA', 'GetForegroundWindow',
'KeybdEvent', 'SendInput', 'MapVirtualKeyExA', 'MapVirtualKeyExW', 'CallNextHookEx'
},
'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'
},
'Memory': {
'HeapAlloc', 'HeapFree', 'HeapSize', 'HeapCreate', 'GetProcessHeap',
'LocalAlloc', 'LocalFree', 'GlobalAlloc', 'GlobalFree', 'GlobalUnlock',
'HeapDestroy', 'GlobalLock', 'SysFreeString', 'SysAllocStringLen'
},
'Resource': {
'LoadResource', 'SizeofResource', 'LockResource', 'FreeResource', 'FindResourceA'
},
'WindowControl': {
'ShowWindow', 'DestroyWindow', 'TranslateMessage', 'DispatchMessageA', 'GetWindow',
'PeekMessageA', 'GetWindowLongA', 'CallWindowProcA', 'SetWindowLongA', 'GetWindowTextA',
'ScreenToClient', 'GetActiveWindow', 'CreateWindowExA', 'DefWindowProcA', 'GetMessageA',
'RegisterClassA', 'UnregisterClassA', 'MessageBoxA'
},
'COM': {
'CoCreateInstance', 'CoUninitialize', 'CoInitialize'
}
}
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", "SectionMeanEntropy", "IsExe", "IsDll", "IsDriver", "Is64Bit", "Machine", "Magic",
"CheckSum", "ImageBase", "SizeOfImage", "SizeOfHeaders",
"Characteristics", "DllCharacteristics", "Subsystem", "LoaderFlags",
"MajorLinkerVersion", "MinorLinkerVersion", "SizeOfCode",
"SizeOfInitializedData", "SizeOfUninitializedData", "AddressOfEntryPoint",
"BaseOfCode", "SectionAlignment", "FileAlignment",
"MajorOperatingSystemVersion", "MinorOperatingSystemVersion",
"MajorImageVersion", "MinorImageVersion",
"MajorSubsystemVersion", "MinorSubsystemVersion",
"SizeOfStackReserve", "SizeOfStackCommit", "SizeOfHeapReserve", "SizeOfHeapCommit",
"NumberOfSections", "NumberOfRvaAndSizes", "PointerToSymbolTable", "NumberOfSymbols",
"SizeOfOptionalHeader",
"TextSectionMaxEntropy", "TextSectionMeanEntropy", "TextSizeRatio",
"DataSectionMaxEntropy", "DataSectionMeanEntropy", "DataSizeRatio",
"RsrcSectionMaxEntropy", "RsrcSectionMeanEntropy", "RsrcSizeRatio",
"SectionCount", "ExecutableSections",
"WritableSections", "ReadableSections", "SectionException",
"Char_00000020_Count", "Char_00000020_MeanEntropy",
"Char_00000040_Count", "Char_00000040_MeanEntropy",
"Char_00000080_Count", "Char_00000080_MeanEntropy",
"Char_02000000_Count", "Char_02000000_MeanEntropy",
"Char_20000000_Count", "Char_20000000_MeanEntropy",
"Char_40000000_Count", "Char_40000000_MeanEntropy",
"Char_80000000_Count", "Char_80000000_MeanEntropy",
"IconCount", "ApiCount", "ExportCount", "DebugCount", "ExceptionCount",
"ImportCount", "ImportFunctionCount",
"FileDescriptionLength", "FileVersionLength", "ProductNameLength",
"ProductVersionLength", "CompanyNameLength", "LegalCopyrightLength",
"CommentsLength", "InternalNameLength", "LegalTrademarksLength",
"SpecialBuildLength", "PrivateBuildLength",
"TrustSigned", "IsDebug", "IsPatched", "IsPrivateBuild", "IsPreRelease",
"IsSpecialBuild", "IsAdmin", "IsInstall", "HasTlsCallbacks",
"HasInvalidTimestamp", "HasRelocationDirectory", "HasPacked", "FileTimeException"
]
CAT_SCHEMA = [f"Cat_{k}" for k in sorted(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
_API_TO_CATS = {}
for _cat, _apis in API_CAT_MAPPING.items():
for _api in _apis:
if _api not in _API_TO_CATS:
_API_TO_CATS[_api] = []
_API_TO_CATS[_api].append(_cat)
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',
}
INSTALLER_SIGS = [b"Nullsoft.NSIS", b"Inno Setup", b"7-Zip.7zip", b"InstallShield"]
####################################################################################################
@staticmethod
def _safe_float(val):
try:
f = float(val)
if math.isinf(f) or math.isnan(f): return 0.0
return f
except Exception:
return 0.0
@staticmethod
def _calc_entropy(data):
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)))
@staticmethod
def _safe_div(n, d):
return float(n) / d if d > 0 else 0.0
####################################################################################################
@classmethod
def extract(cls, file_path):
fts = {k: 0.0 for k in cls.SCHEMA}
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()
pe = pefile.PE(data=file_bytes, fast_load=True)
fts['TrustSigned'] = WinTrust.verify(file_path)
fh = pe.FILE_HEADER
fts['Machine'] = fh.Machine
fts['NumberOfSections'] = fh.NumberOfSections
fts['TimeDateStamp'] = fh.TimeDateStamp
fts['PointerToSymbolTable'] = fh.PointerToSymbolTable
fts['NumberOfSymbols'] = fh.NumberOfSymbols
fts['SizeOfOptionalHeader'] = fh.SizeOfOptionalHeader
fts['Characteristics'] = fh.Characteristics
curr_ts = datetime.datetime.utcnow().timestamp()
fts['HasInvalidTimestamp'] = 1 if (fh.TimeDateStamp < 631152000 or fh.TimeDateStamp > curr_ts + 2592000) else 0
fts['FileTimeException'] = 1 if fh.TimeDateStamp == 0 else 0
if hasattr(pe, 'OPTIONAL_HEADER'):
op = pe.OPTIONAL_HEADER
fts['Magic'] = getattr(op, 'Magic', 0)
fts['MajorLinkerVersion'] = getattr(op, 'MajorLinkerVersion', 0)
fts['MinorLinkerVersion'] = getattr(op, 'MinorLinkerVersion', 0)
fts['SizeOfCode'] = getattr(op, 'SizeOfCode', 0)
fts['SizeOfInitializedData'] = getattr(op, 'SizeOfInitializedData', 0)
fts['SizeOfUninitializedData'] = getattr(op, 'SizeOfUninitializedData', 0)
fts['AddressOfEntryPoint'] = getattr(op, 'AddressOfEntryPoint', 0)
fts['BaseOfCode'] = getattr(op, 'BaseOfCode', 0)
fts['ImageBase'] = getattr(op, 'ImageBase', 0)
fts['SectionAlignment'] = getattr(op, 'SectionAlignment', 0)
fts['FileAlignment'] = getattr(op, 'FileAlignment', 0)
fts['MajorOperatingSystemVersion'] = getattr(op, 'MajorOperatingSystemVersion', 0)
fts['MinorOperatingSystemVersion'] = getattr(op, 'MinorOperatingSystemVersion', 0)
fts['MajorImageVersion'] = getattr(op, 'MajorImageVersion', 0)
fts['MinorImageVersion'] = getattr(op, 'MinorImageVersion', 0)
fts['MajorSubsystemVersion'] = getattr(op, 'MajorSubsystemVersion', 0)
fts['MinorSubsystemVersion'] = getattr(op, 'MinorSubsystemVersion', 0)
fts['SizeOfImage'] = getattr(op, 'SizeOfImage', 0)
fts['SizeOfHeaders'] = getattr(op, 'SizeOfHeaders', 0)
fts['CheckSum'] = getattr(op, 'CheckSum', 0)
fts['Subsystem'] = getattr(op, 'Subsystem', 0)
fts['DllCharacteristics'] = getattr(op, 'DllCharacteristics', 0)
fts['SizeOfStackReserve'] = getattr(op, 'SizeOfStackReserve', 0)
fts['SizeOfStackCommit'] = getattr(op, 'SizeOfStackCommit', 0)
fts['SizeOfHeapReserve'] = getattr(op, 'SizeOfHeapReserve', 0)
fts['SizeOfHeapCommit'] = getattr(op, 'SizeOfHeapCommit', 0)
fts['LoaderFlags'] = getattr(op, 'LoaderFlags', 0)
fts['NumberOfRvaAndSizes'] = getattr(op, 'NumberOfRvaAndSizes', 0)
fts['Is64Bit'] = 1 if fh.Machine in (0x8664, 0xAA64, 0x0200) else 0
fts['IsExe'] = 1 if pe.is_exe() else 0
fts['IsDll'] = 1 if pe.is_dll() else 0
fts['IsDriver'] = 1 if fts['Subsystem'] == 1 else 0
if hasattr(op, 'DATA_DIRECTORY'):
reloc_idx = pefile.DIRECTORY_ENTRY.get('IMAGE_DIRECTORY_ENTRY_BASERELOC', 5)
if len(op.DATA_DIRECTORY) > reloc_idx and op.DATA_DIRECTORY[reloc_idx].Size > 0:
fts['HasRelocationDirectory'] = 1
try:
fts['FileEntropy'] = cls._calc_entropy(file_bytes)
except Exception:
pass
sec_entropies = []
text_entropies = []
data_entropies = []
rsrc_entropies = []
char_flags = [0x00000020, 0x00000040, 0x00000080, 0x02000000, 0x20000000, 0x40000000, 0x80000000]
if hasattr(pe, 'sections'):
fts['SectionCount'] = len(pe.sections)
for section in pe.sections:
try:
name = section.Name.decode('ascii', 'ignore').strip('\x00')
except Exception:
name = ""
try:
s_data = section.get_data()
s_entropy = cls._calc_entropy(s_data)
except Exception:
s_data = b""
s_entropy = 0.0
sec_entropies.append(s_entropy)
if any(p in name.lower() for p in cls.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
for flag in char_flags:
if section.Characteristics & flag:
fts[f'Char_{flag:08X}_Count'] += 1
fts[f'Char_{flag:08X}_MeanEntropy'] += s_entropy
ratio = cls._safe_div(section.SizeOfRawData, fsize)
if name.startswith('.text'):
text_entropies.append(s_entropy)
fts['TextSizeRatio'] += ratio
elif name.startswith('.data'):
data_entropies.append(s_entropy)
fts['DataSizeRatio'] += ratio
elif name.startswith('.rsrc'):
rsrc_entropies.append(s_entropy)
fts['RsrcSizeRatio'] += ratio
if b"requireAdministrator" in s_data: fts['IsAdmin'] = 1
for sig in cls.INSTALLER_SIGS:
if sig in s_data:
fts['IsInstall'] = 1
break
fts['SectionMeanEntropy'] = cls._safe_div(sum(sec_entropies), len(sec_entropies))
fts['TextSectionMaxEntropy'] = max(text_entropies) if text_entropies else 0.0
fts['TextSectionMeanEntropy'] = cls._safe_div(sum(text_entropies), len(text_entropies))
fts['DataSectionMaxEntropy'] = max(data_entropies) if data_entropies else 0.0
fts['DataSectionMeanEntropy'] = cls._safe_div(sum(data_entropies), len(data_entropies))
fts['RsrcSectionMaxEntropy'] = max(rsrc_entropies) if rsrc_entropies else 0.0
fts['RsrcSectionMeanEntropy'] = cls._safe_div(sum(rsrc_entropies), len(rsrc_entropies))
for flag in char_flags:
fts[f'Char_{flag:08X}_MeanEntropy'] = cls._safe_div(fts[f'Char_{flag:08X}_MeanEntropy'], fts[f'Char_{flag:08X}_Count'])
try:
pe.parse_data_directories(directories=[
pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_IMPORT'],
pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_EXPORT'],
pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_RESOURCE'],
pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_DEBUG'],
pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_TLS']
])
except Exception:
pass
if hasattr(pe, 'DIRECTORY_ENTRY_TLS') and hasattr(pe.DIRECTORY_ENTRY_TLS, 'struct'):
if getattr(pe.DIRECTORY_ENTRY_TLS.struct, 'AddressOfCallBacks', 0) != 0:
fts['HasTlsCallbacks'] = 1
if hasattr(pe, 'VS_FIXEDFILEINFO') and len(pe.VS_FIXEDFILEINFO) > 0:
flags = getattr(pe.VS_FIXEDFILEINFO[0], 'FileFlags', 0)
fts['IsDebug'] = 1 if flags & 0x1 else 0
fts['IsPreRelease'] = 1 if flags & 0x2 else 0
fts['IsPatched'] = 1 if flags & 0x4 else 0
fts['IsPrivateBuild'] = 1 if flags & 0x8 else 0
fts['IsSpecialBuild'] = 1 if flags & 0x20 else 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)
v = val.decode('utf-8', 'ignore') if isinstance(val, bytes) else str(val)
length = len(v)
if k == 'FileDescription': fts['FileDescriptionLength'] = length
elif k == 'FileVersion': fts['FileVersionLength'] = length
elif k == 'ProductName': fts['ProductNameLength'] = length
elif k == 'ProductVersion': fts['ProductVersionLength'] = length
elif k == 'CompanyName': fts['CompanyNameLength'] = length
elif k == 'LegalCopyright': fts['LegalCopyrightLength'] = length
elif k == 'Comments': fts['CommentsLength'] = length
elif k == 'InternalName': fts['InternalNameLength'] = length
elif k == 'LegalTrademarks': fts['LegalTrademarksLength'] = length
elif k == 'SpecialBuild': fts['SpecialBuildLength'] = length
elif k == 'PrivateBuild': fts['PrivateBuildLength'] = length
except Exception:
continue
if hasattr(pe, 'DIRECTORY_ENTRY_IMPORT'):
fts['ImportCount'] = len(pe.DIRECTORY_ENTRY_IMPORT)
for entry in pe.DIRECTORY_ENTRY_IMPORT:
if getattr(entry, 'dll', None):
try:
dll_name = entry.dll.decode('ascii', 'ignore').lower()
if dll_name in cls.TARGET_DLLS:
fts[f"Dll_{dll_name.replace('.', '_')}"] = 1.0
except Exception:
pass
imports = getattr(entry, 'imports', [])
for imp in imports:
fts['ImportFunctionCount'] += 1
if not getattr(imp, 'name', None): continue
try:
func_name = imp.name.decode('ascii', 'ignore')
if func_name in cls._ALL_APIS:
fts[f'Api_{func_name}'] += 1
fts['ApiCount'] += 1
for cat in cls._API_TO_CATS.get(func_name, []):
fts[f'Cat_{cat}'] += 1
except Exception:
continue
if hasattr(pe, 'DIRECTORY_ENTRY_EXPORT') and hasattr(pe.DIRECTORY_ENTRY_EXPORT, 'symbols'):
fts['ExportCount'] = len(pe.DIRECTORY_ENTRY_EXPORT.symbols)
if hasattr(pe, 'DIRECTORY_ENTRY_RESOURCE'):
for entry in pe.DIRECTORY_ENTRY_RESOURCE.entries:
if getattr(entry, 'id', None) == 3 and hasattr(entry, 'directory'):
fts['IconCount'] += len(getattr(entry.directory, 'entries', []))
if hasattr(pe, 'OPTIONAL_HEADER') and hasattr(pe.OPTIONAL_HEADER, 'DATA_DIRECTORY') and len(pe.OPTIONAL_HEADER.DATA_DIRECTORY) > 3:
exc_dir = pe.OPTIONAL_HEADER.DATA_DIRECTORY[3]
if exc_dir.Size > 0:
fts['ExceptionCount'] = exc_dir.Size // 12 if fts['Machine'] in (0x8664, 0xAA64) else 0
if hasattr(pe, 'DIRECTORY_ENTRY_DEBUG'):
fts['DebugCount'] = len(pe.DIRECTORY_ENTRY_DEBUG)
for k in fts:
fts[k] = cls._safe_float(fts[k])
return fts
except Exception:
return None
finally:
if pe:
pe.close()
####################################################################################################
class ModelPredictor:
def __init__(self, model_path, feature_path):
self.sess = None
self.input_name = None
self.feature_order = []
self.model_features = self._load_features(feature_path)
self._load_model(model_path)
self._build_input_map()
def _load_features(self, path):
if os.path.exists(path):
try:
with open(path, 'r') as f:
return json.load(f)
except Exception:
pass
return [re.sub(r'[^A-Za-z0-9_]+', '', f) for f in FeatureExtractor.SCHEMA]
def _load_model(self, model_path):
if not os.path.exists(model_path):
raise FileNotFoundError(f"Model missing: {model_path}")
self.sess = ort.InferenceSession(model_path, providers=['CPUExecutionProvider'])
model_inputs = self.sess.get_inputs()
self.input_name = model_inputs[0].name
expected_dim = model_inputs[0].shape[1]
current_dim = len(self.model_features)
if current_dim != expected_dim:
raise ValueError(f"Dimension Mismatch: Expected {expected_dim}, got {current_dim}.")
def _build_input_map(self):
local_map = {re.sub(r'[^A-Za-z0-9_]+', '', raw_name): raw_name for raw_name in FeatureExtractor.SCHEMA}
self.feature_order = [local_map.get(req_feat) for req_feat in self.model_features]
def predict(self, raw_data):
vec_data = [raw_data.get(key, 0.0) if key else 0.0 for key in self.feature_order]
vec = np.array(vec_data, dtype=np.float32).reshape(1, -1)
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 get_threat_name(prob_float):
prob_pct = min(int(prob_float * 100), 99)
if prob_pct >= 80:
return f"General:WinPE/Malware.{prob_pct}!ml"
elif prob_pct >= 50:
return f"General:WinPE/Suspicious.{prob_pct}!ml"
else:
return "Undetected"
def scan_target(target, predictor):
if not os.path.exists(target):
print(json.dumps({"result": "error", "prob": 0.0, "file": target, "details": "Path not found"}))
sys.stdout.flush()
return
try:
data = FeatureExtractor.extract(target)
if not data:
print(json.dumps({"result": "error", "prob": 0.0, "file": target, "details": "Extraction failed"}))
sys.stdout.flush()
return
prob = predictor.predict(data)
label = get_threat_name(prob)
print(json.dumps({"result": label, "prob": round(prob, 4), "file": target}))
sys.stdout.flush()
except Exception as e:
print(json.dumps({"result": "error", "prob": 0.0, "file": target, "details": str(e)}))
sys.stdout.flush()
####################################################################################################
if __name__ == "__main__":
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(json.dumps({"result": "fatal", "details": str(e)}))
sys.stdout.flush()
sys.exit(1)
if len(sys.argv) > 1:
for path in sys.argv[1:]:
scan_target(path.strip(), predictor)
else:
for line in sys.stdin:
target = line.strip()
if not target or target.lower() in ['q', 'exit']:
break
scan_target(target, predictor)
-4
View File
@@ -1,4 +0,0 @@
onnxruntime
numpy
pefile
signify[full]
-55
View File
@@ -1,55 +0,0 @@
from PIL import Image
import numpy, os
####################################################################################################
def preprocess_image(file_data, size, channels=1):
width, height = size
wah = int(numpy.ceil(numpy.sqrt(len(file_data) / channels)))
arr = numpy.frombuffer(file_data, dtype=numpy.uint8)
img = numpy.zeros(wah*wah*channels, dtype=numpy.uint8)
img[:len(file_data)] = arr
image = Image.fromarray(img.reshape((wah, wah)), 'L')
return image.resize((width, height), Image.Resampling.NEAREST)
####################################################################################################
def get_type(file_path):
suffix = {".com", ".dll", ".drv", ".exe", ".ocx", ".scr", ".sys", ".mui", ".cpl"}
try:
file_content = None
if not os.path.splitext(file_path)[-1].lower() in suffix:
return file_content
with open(file_path, 'rb') as f:
file_content = f.read()
except:
pass
return file_content
####################################################################################################
def file_to_images(input_dir, output_dir):
for root, dirs, files in os.walk(input_dir):
for file in files:
try:
file_path = os.path.join(root, file)
file_content = get_type(file_path)
if file_content:
print(file_path)
relative_path = os.path.relpath(root, input_dir)
output_path_dir = os.path.join(output_dir, relative_path)
if not os.path.exists(output_path_dir):
os.makedirs(output_path_dir)
image = preprocess_image(file_content, (224, 224))
safe_name = os.path.basename(file_path).replace('/', '_').replace('\\', '_')
image.save(os.path.join(output_path_dir, f"{safe_name}.png"))
except Exception as e:
print(e)
####################################################################################################
input_directory = input('Enter the folder path: ')
output_directory = "./Image_Output"
file_to_images(input_directory, output_directory)
input('Conversion Complete')
-34
View File
@@ -1,34 +0,0 @@
import os, hashlib
####################################################################################################
def calc_sha256(path):
with open(path, "rb") as f:
return hashlib.sha256(f.read()).digest()
def dedup_realtime(root):
seen = set()
deleted = 0
for dirpath, _, files in os.walk(root):
for name in files:
p = os.path.join(dirpath, name)
try:
h = calc_sha256(p)
except Exception:
continue
if h in seen:
try:
os.remove(p)
print(p)
deleted += 1
except Exception:
pass
else:
seen.add(h)
print(f"deleted {deleted}")
####################################################################################################
if __name__ == "__main__":
dedup_realtime(input("Input file path: "))
input('Deduplicate Complete')
-165
View File
@@ -1,165 +0,0 @@
import tensorflow as tf
import os, tf2onnx, numpy, random
from PIL import Image
from tensorflow.keras import *
from collections import Counter
####################################################################################################
Image.MAX_IMAGE_PIXELS = None
tf.keras.mixed_precision.set_global_policy('mixed_float16')
####################################################################################################
def load_dataset(data_dir, image_size=(224, 224), val_split=0.0001, batch_size=256, color_mode="grayscale"):
channels = 1 if color_mode == "grayscale" else 3
file_paths, labels, class_indices = get_file_list(data_dir)
combined = list(zip(file_paths, labels))
random.shuffle(combined)
file_paths[:], labels[:] = zip(*combined)
num_val = int(val_split * len(file_paths))
train_files = file_paths[num_val:]
train_labels = labels[num_val:]
val_files = file_paths[:num_val]
val_labels = labels[:num_val]
train_ds = tf.data.Dataset.from_tensor_slices((train_files, train_labels))
train_ds = train_ds.map(lambda f, l: parse_fn(f, l, image_size, channels), num_parallel_calls=tf.data.AUTOTUNE)
train_ds = train_ds.batch(batch_size).prefetch(tf.data.AUTOTUNE)
val_ds = tf.data.Dataset.from_tensor_slices((val_files, val_labels))
val_ds = val_ds.map(lambda f, l: parse_fn(f, l, image_size, channels), num_parallel_calls=tf.data.AUTOTUNE)
val_ds = val_ds.batch(batch_size).prefetch(tf.data.AUTOTUNE)
train_ds.class_indices = class_indices
return train_ds, val_ds
####################################################################################################
def parse_fn(filename, label, image_size, channels):
img = tf.io.read_file(filename)
img = tf.image.decode_image(img, channels=channels, expand_animations=False)
img = tf.cond(tf.reduce_all(tf.equal(tf.shape(img)[:2], image_size)), lambda: tf.cast(img, tf.float32), lambda: tf.image.resize(img, image_size))
img = tf.cast(img, tf.float32) / 255.0
label = tf.cast(label, tf.float32)
return img, label
def get_file_list(directory):
file_paths, labels, class_indices = [], [], {}
for i, cls in enumerate(sorted([d for d in os.listdir(directory) if os.path.isdir(os.path.join(directory, d))])):
class_indices[cls] = i
cls_dir = os.path.join(directory, cls)
for root, _, files in os.walk(cls_dir):
for file in files:
if file.lower().endswith(('.jpg', '.jpeg', '.png', '.bmp', '.gif')):
file_paths.append(os.path.join(root, file))
labels.append(i)
return file_paths, labels, class_indices
####################################################################################################
def multi_conv(x, strides):
convs = []
for k in [(1, 7), (1, 1), (3, 3), (5, 5)]:
c = layers.DepthwiseConv2D(kernel_size=k, strides=strides, padding='same', activation='gelu')(x)
convs.append(c)
x = layers.Concatenate()(convs)
x = layers.BatchNormalization()(x)
return layers.Activation('gelu')(x)
def gated_conv(x, strides):
c = layers.DepthwiseConv2D(3, strides=strides, padding='same', activation='gelu')(x)
g = layers.DepthwiseConv2D(3, strides=strides, padding='same', activation='sigmoid')(x)
return layers.Multiply()([c, g])
def coordinate_attention(x, reduction):
input_shape = tf.shape(x)
h, w = input_shape[1], input_shape[2]
filters = x.shape[-1]
mip = max(8, filters // reduction)
x_h = layers.Lambda(lambda t: tf.reduce_mean(t, axis=2, keepdims=True))(x)
x_w = layers.Lambda(lambda t: tf.reduce_mean(t, axis=1, keepdims=True))(x)
x_w_t = layers.Permute((2, 1, 3))(x_w)
x_cat = layers.Concatenate(axis=1)([x_h, x_w_t])
x_cat = layers.Conv2D(mip, 1, padding='valid', use_bias=False)(x_cat)
x_cat = layers.BatchNormalization()(x_cat)
x_cat = layers.Activation('gelu')(x_cat)
x_cat_h, x_cat_w = layers.Lambda(lambda t: tf.split(t[0], [t[1], t[2]], axis=1))([x_cat, h, w])
x_h_sigmoid = layers.Conv2D(filters, 1, padding='valid', activation='sigmoid')(x_cat_h)
x_cat_w = layers.Permute((2, 1, 3))(x_cat_w)
x_w_sigmoid = layers.Conv2D(filters, 1, padding='valid', activation='sigmoid')(x_cat_w)
return layers.Multiply()([x, x_h_sigmoid, x_w_sigmoid])
def build_model(input_shape):
i = layers.Input(shape=input_shape)
x = multi_conv(i, strides=1)
x = gated_conv(x, strides=2)
x = multi_conv(x, strides=2)
x = multi_conv(x, strides=2)
x = multi_conv(x, strides=2)
x = multi_conv(x, strides=2)
x = coordinate_attention(x, reduction=4)
x = layers.GlobalMaxPooling2D()(x)
x = layers.Dense(x.shape[-1] // 2, activation='gelu', dtype='float32')(x)
o = layers.Dense(1, activation='sigmoid', dtype='float32')(x)
return models.Model(i, o)
####################################################################################################
def calculate_lr_cosine_sq(epoch, total_epochs=30, lr_start=5e-4, lr_end=5e-7):
cos_val = 0.5 * (1 + numpy.cos(epoch / total_epochs * numpy.pi))
return lr_end + (lr_start - lr_end) * (cos_val ** 2)
def lr_scheduler(epoch):
return calculate_lr_cosine_sq(epoch)
def binary_focal_loss(gamma=2.0, alpha=0.25):
def focal_loss(y_true, y_pred):
y_true = tf.cast(y_true, tf.float32)
y_pred = tf.clip_by_value(y_pred, tf.keras.backend.epsilon(), 1.0 - tf.keras.backend.epsilon())
p_t = tf.where(tf.equal(y_true, 1), y_pred, 1 - y_pred)
alpha_t = tf.where(tf.equal(y_true, 1), alpha, 1 - alpha)
loss = -alpha_t * tf.math.pow(1.0 - p_t, gamma) * tf.math.log(p_t)
return tf.reduce_mean(loss)
return focal_loss
class CustomModelCheckpoint(callbacks.Callback):
def on_epoch_end(self, epoch, logs=False):
models.save_model(self.model, f"PYAS_Model_Epoch_{epoch + 1}.h5")
spec = (tf.TensorSpec((None,) + self.model.input.shape[1:], tf.float32),)
model_proto, _ = tf2onnx.convert.from_keras(self.model, input_signature=spec, opset=17)
with open(f"PYAS_Model_Epoch_{epoch + 1}.onnx", "wb") as f:
f.write(model_proto.SerializeToString())
####################################################################################################
train_ds, val_ds = load_dataset(r'.\Image_Full_Pefile')
for imgs, _ in train_ds.take(1):
train_ds.image_shape = imgs.shape[1:]
break
try:
model = models.load_model('PYAS_Model_Epoch_0.h5', compile=False)
print("Load model from disk")
except:
print("Creating a new model")
model = build_model(train_ds.image_shape)
print(f"Total parameters: {model.count_params()}")
model.compile(optimizer='adam', loss=binary_focal_loss(gamma=2.0, alpha=0.25),
metrics=[metrics.BinaryAccuracy(name='accuracy'), metrics.AUC(name='auc'),
metrics.Precision(name='precision'), metrics.Recall(name='recall'),
metrics.FalsePositives(name='fp'), metrics.TrueNegatives(name='tn')])
model.fit(train_ds, epochs=30, callbacks=[CustomModelCheckpoint(), callbacks.LearningRateScheduler(lr_scheduler)], validation_data=val_ds, class_weight={0: 10.0, 1: 1.0})
input('Training Complete')
+287 -567
View File
@@ -1,12 +1,23 @@
import os, sys, time, ctypes, sqlite3, datetime, pefile, math, hashlib
import numpy as np
import ctypes.wintypes
import os, sys, math, time, datetime, pefile, hashlib
import ctypes, ctypes.wintypes, multiprocessing
try:
import orjson
JSON_DUMPS = lambda x: orjson.dumps(x).decode('utf-8')
JSON_LOADS = orjson.loads
HAS_ORJSON = True
except ImportError:
import json
JSON_DUMPS = json.dumps
JSON_LOADS = json.loads
HAS_ORJSON = False
####################################################################################################
DB_PATH = "pe_features.db"
JSONL_PATH = "pe_features.jsonl"
BATCH_SIZE = 1000
MAX_FILE_SIZE = 4 * 1024 * 1024 * 1024
MAX_WORKERS = os.cpu_count() - 2
TARGET_EXTENSIONS = {
'.exe', '.dll', '.sys', '.ocx', '.scr', '.efi', '.acm', '.ax', '.cpl', '.drv', '.com', '.mui', '.pyd'
}
@@ -60,24 +71,23 @@ V2_GUID = GUID(0x00AAC56B, 0xCD44, 0x11D0, (ctypes.c_ubyte * 8)(0x8C, 0xC2, 0x00
####################################################################################################
def verify_signature(file_path):
if os.name != 'nt':
return 0
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
abs_path = os.path.abspath(file_path)
fi = WINTRUST_FILE_INFO()
fi.cbStruct = ctypes.sizeof(WINTRUST_FILE_INFO)
fi.pcwszFilePath = abs_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
@@ -94,628 +104,338 @@ def verify_signature(file_path):
td.pSignatureSettings = None
return 1 if wintrust.WinVerifyTrust(None, ctypes.byref(V2_GUID), ctypes.byref(td)) == 0 else 0
except Exception as e:
print(e)
except Exception:
return 0
####################################################################################################
class Features:
API_CAT_MAPPING = {
'ProcessControl': {
'CreateProcessA', 'CreateProcessW', 'WinExec', 'ShellExecuteA', 'ShellExecuteW',
'ShellExecuteExW', 'ExitProcess', 'TerminateProcess', 'OpenProcess', 'GetExitCodeProcess',
'SetThreadPriority', 'GetThreadPriority', 'GetCurrentProcess', 'GetCurrentProcessId',
'GetModuleFileNameA', 'GetCommandLineA', 'GetModuleFileNameW', 'GetStartupInfoW'
},
'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',
'VirtualFree', 'VirtualQuery'
},
'Synchronization': {
'WaitForSingleObject', 'WaitForSingleObjectEx', 'WaitForMultipleObjects', 'WaitForMultipleObjectsEx',
'CreateMutexA', 'CreateMutexW', 'OpenMutexW', 'ReleaseMutex', 'CreateEventA', 'CreateEventW',
'OpenEventW', 'SetEvent', 'ResetEvent', 'EnterCriticalSection', 'LeaveCriticalSection',
'InitializeCriticalSection', 'DeleteCriticalSection', 'Sleep', 'SleepEx',
'InitializeCriticalSectionAndSpinCount', 'InterlockedDecrement', 'InterlockedIncrement',
'InterlockedExchange', 'InterlockedCompareExchange', 'SetTimer', 'KillTimer'
},
'MultiThreading': {
'CreateThread', 'ResumeThread', 'SuspendThread', 'ExitThread', 'TerminateThread',
'GetCurrentThread', 'GetCurrentThreadId', 'TlsAlloc', 'TlsSetValue', 'TlsGetValue',
'CreateThreadpoolWork', 'SubmitThreadpoolWork', 'TlsFree'
},
'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',
'GetStringTypeW', 'LCMapStringW', 'IsValidCodePage', 'DecodePointer', 'EncodePointer'
},
'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',
'CloseHandle', 'GetFileType', 'SetStdHandle', 'SetFilePointerEx', 'FindClose',
'SetHandleCount', 'SetEndOfFile', 'GetFullPathNameA', 'GetFileAttributesA'
},
'Registry': {
'RegOpenKeyExA', 'RegOpenKeyExW', 'RegSetValueExA', 'RegSetValueExW',
'RegCreateKeyExA', 'RegCreateKeyExW', 'RegDeleteKeyA', 'RegDeleteKeyW',
'RegEnumValueA', 'RegEnumValueW', 'RegQueryValueExA', 'RegQueryValueExW',
'RegDeleteValueA', 'RegDeleteValueW', 'RegCloseKey'
},
'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',
'GetLastError', 'GetACP', 'RaiseException', 'SetLastError', 'GetEnvironmentStringsW',
'FreeEnvironmentStringsW', 'GetLocaleInfoA', 'lstrlenA', 'GetVersion', 'GetSystemInfo',
'CompareStringA'
},
'DotNet': {
'_CorExeMain', '_CorDllMain'
},
'AntiDebug': {
'IsDebuggerPresent', 'CheckRemoteDebuggerPresent', 'OutputDebugStringA', 'OutputDebugStringW',
'GetTickCount', 'GetTickCount64', 'QueryPerformanceCounter', 'FindWindowA', 'FindWindowW',
'EnumWindows', 'EnumChildWindows', 'GetWindowRect', 'GetClientRect', 'SetWindowDisplayAffinity',
'UnhandledExceptionFilter', 'SetUnhandledExceptionFilter', 'IsProcessorFeaturePresent',
'SetErrorMode', 'GetTimeZoneInformation', 'GetDriveTypeW', 'GetDiskFreeSpaceA'
},
'Keylogging': {
'GetAsyncKeyState', 'GetKeyState', 'GetKeyboardState', 'MapVirtualKeyA', 'MapVirtualKeyW',
'ToAscii', 'ToAsciiEx', 'ToUnicode', 'ToUnicodeEx', 'GetKeyNameTextA', 'GetForegroundWindow',
'KeybdEvent', 'SendInput', 'MapVirtualKeyExA', 'MapVirtualKeyExW', 'CallNextHookEx'
},
'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'
},
'Memory': {
'HeapAlloc', 'HeapFree', 'HeapSize', 'HeapCreate', 'GetProcessHeap',
'LocalAlloc', 'LocalFree', 'GlobalAlloc', 'GlobalFree', 'GlobalUnlock',
'HeapDestroy', 'GlobalLock', 'SysFreeString', 'SysAllocStringLen'
},
'Resource': {
'LoadResource', 'SizeofResource', 'LockResource', 'FreeResource', 'FindResourceA'
},
'WindowControl': {
'ShowWindow', 'DestroyWindow', 'TranslateMessage', 'DispatchMessageA', 'GetWindow',
'PeekMessageA', 'GetWindowLongA', 'CallWindowProcA', 'SetWindowLongA', 'GetWindowTextA',
'ScreenToClient', 'GetActiveWindow', 'CreateWindowExA', 'DefWindowProcA', 'GetMessageA',
'RegisterClassA', 'UnregisterClassA', 'MessageBoxA'
},
'COM': {
'CoCreateInstance', 'CoUninitialize', 'CoInitialize'
}
}
class FeatureExtractor:
@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
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'
}
@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)))
BASE_SCHEMA = [
"FileEntropy", "SectionMeanEntropy", "IsExe", "IsDll", "IsDriver", "Is64Bit", "Machine", "Magic",
"TimeDateStamp", "CheckSum", "ImageBase", "SizeOfImage", "SizeOfHeaders",
"Characteristics", "DllCharacteristics", "Subsystem", "LoaderFlags",
"MajorLinkerVersion", "MinorLinkerVersion", "SizeOfCode",
"SizeOfInitializedData", "SizeOfUninitializedData", "AddressOfEntryPoint",
"BaseOfCode", "SectionAlignment", "FileAlignment",
"MajorOperatingSystemVersion", "MinorOperatingSystemVersion",
"MajorImageVersion", "MinorImageVersion",
"MajorSubsystemVersion", "MinorSubsystemVersion",
"SizeOfStackReserve", "SizeOfStackCommit", "SizeOfHeapReserve", "SizeOfHeapCommit",
"NumberOfSections", "NumberOfRvaAndSizes", "PointerToSymbolTable", "NumberOfSymbols",
"SizeOfOptionalHeader",
"TextSectionMaxEntropy", "TextSectionMeanEntropy", "TextSizeRatio",
"DataSectionMaxEntropy", "DataSectionMeanEntropy", "DataSizeRatio",
"RsrcSectionMaxEntropy", "RsrcSectionMeanEntropy", "RsrcSizeRatio",
"SectionCount", "ExecutableSections",
"WritableSections", "ReadableSections", "SectionException",
"Char_00000020_Count", "Char_00000020_MeanEntropy",
"Char_00000040_Count", "Char_00000040_MeanEntropy",
"Char_00000080_Count", "Char_00000080_MeanEntropy",
"Char_02000000_Count", "Char_02000000_MeanEntropy",
"Char_20000000_Count", "Char_20000000_MeanEntropy",
"Char_40000000_Count", "Char_40000000_MeanEntropy",
"Char_80000000_Count", "Char_80000000_MeanEntropy",
"IconCount", "ApiCount", "ExportCount", "DebugCount", "ExceptionCount",
"ImportCount", "ImportFunctionCount",
"FileDescriptionLength", "FileVersionLength", "ProductNameLength",
"ProductVersionLength", "CompanyNameLength", "LegalCopyrightLength",
"CommentsLength", "InternalNameLength", "LegalTrademarksLength",
"SpecialBuildLength", "PrivateBuildLength",
"TrustSigned", "IsDebug", "IsPatched", "IsPrivateBuild", "IsPreRelease",
"IsSpecialBuild", "IsAdmin", "IsInstall", "HasTlsCallbacks",
"HasInvalidTimestamp", "HasRelocationDirectory", "HasPacked", "FileTimeException"
]
CAT_SCHEMA = [f"Cat_{k}" for k in sorted(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
_API_TO_CATS = {}
for _cat, _apis in API_CAT_MAPPING.items():
for _api in _apis:
if _api not in _API_TO_CATS:
_API_TO_CATS[_api] = []
_API_TO_CATS[_api].append(_cat)
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',
}
INSTALLER_SIGS = [b"Nullsoft.NSIS", b"Inno Setup", b"7-Zip.7zip", b"InstallShield"]
####################################################################################################
def safe_float(val):
try:
f = float(val)
if math.isinf(f) or math.isnan(f): return 0.0
return f
except:
return 0.0
####################################################################################################
def calc_entropy(data):
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 float(n) / d if d > 0 else 0.0
def process_single_file_content(file_bytes, file_path):
fsize = len(file_bytes)
if fsize == 0: return None
fts = {k: 0.0 for k in Features.SCHEMA}
pe = None
try:
pe = pefile.PE(data=file_bytes, fast_load=True)
fts['TrustSigned'] = verify_signature(file_path)
fh = pe.FILE_HEADER
fts['Machine'] = fh.Machine
fts['NumberOfSections'] = fh.NumberOfSections
fts['TimeDateStamp'] = fh.TimeDateStamp
fts['PointerToSymbolTable'] = fh.PointerToSymbolTable
fts['NumberOfSymbols'] = fh.NumberOfSymbols
fts['SizeOfOptionalHeader'] = fh.SizeOfOptionalHeader
fts['Characteristics'] = fh.Characteristics
curr_ts = datetime.datetime.utcnow().timestamp()
fts['HasInvalidTimestamp'] = 1.0 if (fh.TimeDateStamp < 631152000 or fh.TimeDateStamp > curr_ts + 2592000) else 0.0
fts['FileTimeException'] = 1.0 if fh.TimeDateStamp == 0 else 0.0
if hasattr(pe, 'OPTIONAL_HEADER'):
op = pe.OPTIONAL_HEADER
fts['Magic'] = getattr(op, 'Magic', 0)
fts['MajorLinkerVersion'] = getattr(op, 'MajorLinkerVersion', 0)
fts['MinorLinkerVersion'] = getattr(op, 'MinorLinkerVersion', 0)
fts['SizeOfCode'] = getattr(op, 'SizeOfCode', 0)
fts['SizeOfInitializedData'] = getattr(op, 'SizeOfInitializedData', 0)
fts['SizeOfUninitializedData'] = getattr(op, 'SizeOfUninitializedData', 0)
fts['AddressOfEntryPoint'] = getattr(op, 'AddressOfEntryPoint', 0)
fts['BaseOfCode'] = getattr(op, 'BaseOfCode', 0)
fts['ImageBase'] = getattr(op, 'ImageBase', 0)
fts['SectionAlignment'] = getattr(op, 'SectionAlignment', 0)
fts['FileAlignment'] = getattr(op, 'FileAlignment', 0)
fts['MajorOperatingSystemVersion'] = getattr(op, 'MajorOperatingSystemVersion', 0)
fts['MinorOperatingSystemVersion'] = getattr(op, 'MinorOperatingSystemVersion', 0)
fts['MajorImageVersion'] = getattr(op, 'MajorImageVersion', 0)
fts['MinorImageVersion'] = getattr(op, 'MinorImageVersion', 0)
fts['MajorSubsystemVersion'] = getattr(op, 'MajorSubsystemVersion', 0)
fts['MinorSubsystemVersion'] = getattr(op, 'MinorSubsystemVersion', 0)
fts['SizeOfImage'] = getattr(op, 'SizeOfImage', 0)
fts['SizeOfHeaders'] = getattr(op, 'SizeOfHeaders', 0)
fts['CheckSum'] = getattr(op, 'CheckSum', 0)
fts['Subsystem'] = getattr(op, 'Subsystem', 0)
fts['DllCharacteristics'] = getattr(op, 'DllCharacteristics', 0)
fts['SizeOfStackReserve'] = getattr(op, 'SizeOfStackReserve', 0)
fts['SizeOfStackCommit'] = getattr(op, 'SizeOfStackCommit', 0)
fts['SizeOfHeapReserve'] = getattr(op, 'SizeOfHeapReserve', 0)
fts['SizeOfHeapCommit'] = getattr(op, 'SizeOfHeapCommit', 0)
fts['LoaderFlags'] = getattr(op, 'LoaderFlags', 0)
fts['NumberOfRvaAndSizes'] = getattr(op, 'NumberOfRvaAndSizes', 0)
fts['Is64Bit'] = 1.0 if fh.Machine in (0x8664, 0xAA64, 0x0200) else 0.0
fts['IsExe'] = 1.0 if pe.is_exe() else 0.0
fts['IsDll'] = 1.0 if pe.is_dll() else 0.0
fts['IsDriver'] = 1.0 if fts['Subsystem'] == 1 else 0.0
if hasattr(op, 'DATA_DIRECTORY'):
reloc_idx = pefile.DIRECTORY_ENTRY.get('IMAGE_DIRECTORY_ENTRY_BASERELOC', 5)
if len(op.DATA_DIRECTORY) > reloc_idx and op.DATA_DIRECTORY[reloc_idx].Size > 0:
fts['HasRelocationDirectory'] = 1.0
@classmethod
def extract(cls, file_bytes, file_path, fsize):
if fsize == 0: return None
base, dlls, apis, pe = {}, set(), set(), None
try:
fts['FileEntropy'] = calc_entropy(file_bytes)
except Exception:
pass
pe = pefile.PE(data=file_bytes, fast_load=True)
base['TrustSigned'] = verify_signature(file_path)
base['FileEntropy'] = cls._calc_entropy(file_bytes)
base['FileSize'] = 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)
sec_entropies = []
text_entropies = []
data_entropies = []
rsrc_entropies = []
char_flags = [0x00000020, 0x00000040, 0x00000080, 0x02000000, 0x20000000, 0x40000000, 0x80000000]
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
if hasattr(pe, 'sections'):
fts['SectionCount'] = len(pe.sections)
for section in pe.sections:
try:
name = section.Name.decode('ascii', 'ignore').strip('\x00')
except Exception:
name = ""
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
try:
s_data = section.get_data()
s_entropy = calc_entropy(s_data)
except Exception:
s_data = b""
s_entropy = 0.0
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
sec_entropies.append(s_entropy)
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
if any(p in name.lower() for p in Features.PACKERS): fts['HasPacked'] = 1.0
if section.Characteristics & 0x20000000: fts['ExecutableSections'] += 1.0
if section.Characteristics & 0x80000000: fts['WritableSections'] += 1.0
if section.Characteristics & 0x40000000: fts['ReadableSections'] += 1.0
if section.SizeOfRawData + section.PointerToRawData > fsize: fts['SectionException'] = 1.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:
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 section.Characteristics & flag:
fts[f'Char_{flag:08X}_Count'] += 1.0
fts[f'Char_{flag:08X}_MeanEntropy'] += s_entropy
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']
ratio = safe_div(section.SizeOfRawData, fsize)
if name.startswith('.text'):
text_entropies.append(s_entropy)
fts['TextSizeRatio'] += ratio
elif name.startswith('.data'):
data_entropies.append(s_entropy)
fts['DataSizeRatio'] += ratio
elif name.startswith('.rsrc'):
rsrc_entropies.append(s_entropy)
fts['RsrcSizeRatio'] += ratio
if b"requireAdministrator" in s_data: fts['IsAdmin'] = 1.0
for sig in Features.INSTALLER_SIGS:
if sig in s_data:
fts['IsInstall'] = 1.0
break
fts['SectionMeanEntropy'] = safe_div(sum(sec_entropies), len(sec_entropies))
fts['TextSectionMaxEntropy'] = max(text_entropies) if text_entropies else 0.0
fts['TextSectionMeanEntropy'] = safe_div(sum(text_entropies), len(text_entropies))
fts['DataSectionMaxEntropy'] = max(data_entropies) if data_entropies else 0.0
fts['DataSectionMeanEntropy'] = safe_div(sum(data_entropies), len(data_entropies))
fts['RsrcSectionMaxEntropy'] = max(rsrc_entropies) if rsrc_entropies else 0.0
fts['RsrcSectionMeanEntropy'] = safe_div(sum(rsrc_entropies), len(rsrc_entropies))
for flag in char_flags:
fts[f'Char_{flag:08X}_MeanEntropy'] = safe_div(fts[f'Char_{flag:08X}_MeanEntropy'], fts[f'Char_{flag:08X}_Count'])
try:
pe.parse_data_directories(directories=[
pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_IMPORT'],
pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_EXPORT'],
pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_RESOURCE'],
pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_DEBUG'],
pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_TLS']
])
except Exception:
pass
if hasattr(pe, 'DIRECTORY_ENTRY_TLS') and hasattr(pe.DIRECTORY_ENTRY_TLS, 'struct'):
if getattr(pe.DIRECTORY_ENTRY_TLS.struct, 'AddressOfCallBacks', 0) != 0:
fts['HasTlsCallbacks'] = 1.0
if hasattr(pe, 'VS_FIXEDFILEINFO') and len(pe.VS_FIXEDFILEINFO) > 0:
flags = getattr(pe.VS_FIXEDFILEINFO[0], 'FileFlags', 0)
fts['IsDebug'] = 1.0 if flags & 0x1 else 0.0
fts['IsPreRelease'] = 1.0 if flags & 0x2 else 0.0
fts['IsPatched'] = 1.0 if flags & 0x4 else 0.0
fts['IsPrivateBuild'] = 1.0 if flags & 0x8 else 0.0
fts['IsSpecialBuild'] = 1.0 if flags & 0x20 else 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)
v = val.decode('utf-8', 'ignore') if isinstance(val, bytes) else str(val)
length = len(v)
if k == 'FileDescription': fts['FileDescriptionLength'] = length
elif k == 'FileVersion': fts['FileVersionLength'] = length
elif k == 'ProductName': fts['ProductNameLength'] = length
elif k == 'ProductVersion': fts['ProductVersionLength'] = length
elif k == 'CompanyName': fts['CompanyNameLength'] = length
elif k == 'LegalCopyright': fts['LegalCopyrightLength'] = length
elif k == 'Comments': fts['CommentsLength'] = length
elif k == 'InternalName': fts['InternalNameLength'] = length
elif k == 'LegalTrademarks': fts['LegalTrademarksLength'] = length
elif k == 'SpecialBuild': fts['SpecialBuildLength'] = length
elif k == 'PrivateBuild': fts['PrivateBuildLength'] = length
except Exception:
continue
if hasattr(pe, 'DIRECTORY_ENTRY_IMPORT'):
fts['ImportCount'] = len(pe.DIRECTORY_ENTRY_IMPORT)
for entry in pe.DIRECTORY_ENTRY_IMPORT:
if getattr(entry, 'dll', None):
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 Exception:
pass
string_keys = ['FileDescription', 'FileVersion', 'ProductName', 'ProductVersion', 'CompanyName', 'LegalCopyright', 'Comments', 'InternalName', 'LegalTrademarks', 'SpecialBuild', 'PrivateBuild']
for key in string_keys:
base[f'{key}Length'] = 0.0
imports = getattr(entry, 'imports', [])
for imp in imports:
fts['ImportFunctionCount'] += 1.0
if not getattr(imp, 'name', None): continue
try:
func_name = imp.name.decode('ascii', 'ignore')
if func_name in Features._ALL_APIS:
fts[f'Api_{func_name}'] += 1.0
fts['ApiCount'] += 1.0
for cat in Features._API_TO_CATS.get(func_name, []):
fts[f'Cat_{cat}'] += 1.0
except Exception:
continue
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_EXPORT') and hasattr(pe.DIRECTORY_ENTRY_EXPORT, 'symbols'):
fts['ExportCount'] = len(pe.DIRECTORY_ENTRY_EXPORT.symbols)
if hasattr(pe, 'DIRECTORY_ENTRY_RESOURCE'):
for entry in pe.DIRECTORY_ENTRY_RESOURCE.entries:
if getattr(entry, 'id', None) == 3 and hasattr(entry, 'directory'):
fts['IconCount'] += len(getattr(entry.directory, 'entries', []))
if hasattr(pe, 'OPTIONAL_HEADER') and hasattr(pe.OPTIONAL_HEADER, 'DATA_DIRECTORY') and len(pe.OPTIONAL_HEADER.DATA_DIRECTORY) > 3:
exc_dir = pe.OPTIONAL_HEADER.DATA_DIRECTORY[3]
if exc_dir.Size > 0:
fts['ExceptionCount'] = exc_dir.Size // 12 if fts['Machine'] in (0x8664, 0xAA64) else 0.0
try:
pe.parse_data_directories(directories=[
pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_IMPORT'],
pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_EXPORT'],
pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_RESOURCE'],
pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_DEBUG'],
pefile.DIRECTORY_ENTRY['IMAGE_DIRECTORY_ENTRY_TLS']
])
except Exception: pass
if hasattr(pe, 'DIRECTORY_ENTRY_TLS') and hasattr(pe.DIRECTORY_ENTRY_TLS, 'struct'):
base['HasTlsCallbacks'] = 1.0 if getattr(pe.DIRECTORY_ENTRY_TLS.struct, 'AddressOfCallBacks', 0) != 0 else 0.0
if hasattr(pe, 'DIRECTORY_ENTRY_DEBUG'):
fts['DebugCount'] = len(pe.DIRECTORY_ENTRY_DEBUG)
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
for k in fts:
fts[k] = safe_float(fts[k])
return fts
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
except pefile.PEFormatError:
return None
except Exception:
return None
finally:
if pe:
pe.close()
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
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()
####################################################################################################
def init_db():
conn = sqlite3.connect(DB_PATH)
conn.execute("PRAGMA synchronous = OFF")
conn.execute("PRAGMA journal_mode = WAL")
conn.execute("PRAGMA cache_size = 10000")
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 load_existing_hashes():
hashes = set()
if os.path.exists(JSONL_PATH):
with open(JSONL_PATH, 'r', encoding='utf-8') as f:
for line in f:
try:
data = JSON_LOADS(line)
if 'FileHash' in data: hashes.add(data['FileHash'])
except Exception: pass
return hashes
def load_existing_hashes(conn):
print("[*] Loading existing hashes into memory...")
####################################################################################################
GLOBAL_HASHES = set()
def worker_init(shared_hashes):
global GLOBAL_HASHES
GLOBAL_HASHES = shared_hashes
def _extract_worker_wrapper(args):
return _extract_worker(*args)
def _extract_worker(file_path, label):
try:
cursor = conn.execute("SELECT FileHash FROM PeData")
return {row[0] for row in cursor.fetchall()}
except:
return set()
fsize = os.path.getsize(file_path)
if fsize == 0 or fsize > MAX_FILE_SIZE:
return 'error', None
with open(file_path, "rb") as f:
file_bytes = f.read()
sha256 = hashlib.sha256(file_bytes).hexdigest()
if sha256 in GLOBAL_HASHES:
return 'skipped', None
res = FeatureExtractor.extract(file_bytes, file_path, fsize)
if res:
final_res = {"Label": label, "FileHash": sha256, **res}
return 'success', (sha256, JSON_DUMPS(final_res) + "\n")
return 'error', None
except Exception:
return 'error', None
####################################################################################################
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 = []
if os.path.isfile(path):
files.append(path)
else:
for r, _, fs in os.walk(path):
for f in fs:
if os.path.splitext(f)[1].lower() in TARGET_EXTENSIONS:
files.append(os.path.join(r, f))
def scan_and_save(path, label, existing_hashes):
print(f"\n[*] Scanning path: {path} (Label={label})")
files = [path] if os.path.isfile(path) else [os.path.join(r, f) for r, _, fs in os.walk(path) for f in fs if os.path.splitext(f)[1].lower() in TARGET_EXTENSIONS]
total = len(files)
if total == 0:
print("[-] No valid PE files found.")
return
print(f"[*] Found {total} files. Processing...\n")
print(f"[*] Found {total} files. Starting {MAX_WORKERS}-process extraction...\n")
if HAS_ORJSON:
print("[+] orjson acceleration enabled.")
batch = []
count = 0
errors = 0
skipped = 0
count, errors, skipped, processed_files = 0, 0, 0, 0
start_time = time.time()
batch_buffer = []
placeholders = ",".join(["?"] * (len(Features.SCHEMA) + 2))
sql = f"INSERT INTO PeData VALUES ({placeholders})"
pool = multiprocessing.Pool(processes=MAX_WORKERS, initializer=worker_init, initargs=(existing_hashes,))
tasks = [(fpath, label) for fpath in files]
try:
for i, file_path in enumerate(files):
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:
for status, data in pool.imap_unordered(_extract_worker_wrapper, tasks):
processed_files += 1
if status == 'error':
errors += 1
if len(batch) >= BATCH_SIZE:
try:
conn.executemany(sql, batch)
conn.commit()
batch = []
except Exception as batch_err:
print(f"\n[-] Batch Error: {batch_err}")
batch = []
elif status == 'skipped':
skipped += 1
elif status == 'success':
sha256, json_str = data
if sha256 not in existing_hashes:
existing_hashes.add(sha256)
batch_buffer.append(json_str)
count += 1
else:
skipped += 1
if len(batch_buffer) >= BATCH_SIZE:
with open(JSONL_PATH, 'a', encoding='utf-8') as f_out:
f_out.writelines(batch_buffer)
batch_buffer.clear()
elapsed = time.time() - start_time
rate = (i + 1) / elapsed if elapsed > 0 else 0
sys.stdout.write(f"\r[{i+1}/{total}] New: {count} | Skip: {skipped} | Err: {errors} | {rate:.1f} files/s")
rate = processed_files / elapsed if elapsed > 0 else 0
sys.stdout.write(f"\r[{processed_files}/{total}] Added: {count} | Skipped: {skipped} | Errors: {errors} | Speed: {rate:.1f} files/s")
sys.stdout.flush()
pool.close()
pool.join()
except KeyboardInterrupt:
print("\n\n[-] Interrupted by user.")
print("\n\n[-] User aborted. Force terminating all child processes...")
pool.terminate()
pool.join()
except Exception as e:
print(f"\n[-] Unexpected Error: {e}")
print(f"\n[-] Critical error: {e}")
pool.terminate()
pool.join()
finally:
if batch:
if batch_buffer:
try:
conn.executemany(sql, batch)
conn.commit()
except Exception:
pass
with open(JSONL_PATH, 'a', encoding='utf-8') as f_out:
f_out.writelines(batch_buffer)
except Exception as save_err:
print(f"\n[-] Failed to save remaining buffer: {save_err}")
print(f"\n\n[+] Scan finished. Total New: {count}, Total Skipped: {skipped}")
print(f"\n\n[+] Scan complete. Total added: {count}, Total skipped: {skipped}")
####################################################################################################
if __name__ == "__main__":
print("\n---------------- PE Dataset Builder v3.2 ----------------\n")
multiprocessing.freeze_support()
if os.path.exists(DB_PATH):
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()
print("\n---------------- PE Dataset Builder v4.0 ----------------\n")
existing_hashes = load_existing_hashes(conn)
print(f"[+] Loaded {len(existing_hashes)} hashes from database.\n")
existing_hashes = load_existing_hashes()
print(f"[+] Found {len(existing_hashes)} existing hashes.\n")
try:
while True:
print("-" * 60)
raw_path = input("\n[*] Enter Path (File or Folder): ").strip()
target_path = raw_path.strip('"').strip("'")
print("-" * 57)
target_path = input("\n[*] Enter path (file or directory): ").strip().strip('"').strip("'")
if target_path.lower() in ['exit', 'q']: break
if not target_path or not os.path.exists(target_path):
print("[-] Invalid path.")
continue
while True:
label_input = input("[*] Enter Label (0=Safe, 1=Malware): ").strip()
if label_input.lower() in ['exit', 'q']:
conn.close()
sys.exit()
if label_input in ['0', '1']:
label = int(label_input)
break
label_input = input("[*] Enter label (0=Safe, 1=Malware): ").strip()
if label_input.lower() in ['exit', 'q']: sys.exit()
if label_input in ['0', '1']: break
print("[-] Please enter 0 or 1.")
scan_and_save(conn, target_path, label, existing_hashes)
scan_and_save(target_path, int(label_input), existing_hashes)
except KeyboardInterrupt:
print("\n\n[-] Interrupted.")
finally:
conn.close()
print(f"\n[+] Database saved to: {os.path.abspath(DB_PATH)}")
print("\n\n[-] Aborted.")
+134 -101
View File
@@ -1,109 +1,128 @@
import os, sys, time, sqlite3, re, json
import pandas as pd
import os, sys, time, json, onnxmltools
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
####################################################################################################
DB_PATH = "pe_features.db"
JSONL_PATH = "pe_features.jsonl"
MODEL_FILE = "model.txt"
ONNX_FILE = "model.onnx"
ONNX_FILE = "Pefile_General_S1.onnx"
FEATURE_FILE = "features.json"
TEST_SIZE = 0.0001
RANDOM_SEED = 42
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.1,
'lambda_l2': 0.1,
'min_data_in_leaf': 20,
'lambda_l1': 0.01,
'lambda_l2': 0.01,
}
####################################################################################################
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()
dll_counts = defaultdict(int)
api_counts = defaultdict(int)
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, 'r') as f:
for line in f:
try:
data = json.loads(line)
base_keys.update(data.get('Base', {}).keys())
for d in data.get('DLLs', []):
dll_counts[d] += 1
for a in data.get('APIs', []):
api_counts[a] += 1
num_lines += 1
except Exception:
continue
####################################################################################################
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 = sorted([f"Dll_{k}" for k, v in dll_counts.items() if v >= 100])
api_schema = sorted([f"Api_{k}" for k, v in api_counts.items() if v >= 100])
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)
feature_idx = {feat: i for i, feat in enumerate(final_features)}
rows, cols, vals = [], [], []
y = np.zeros(num_lines, dtype=np.int8)
with open(jsonl_path, 'r') as f:
row_idx = 0
for line in f:
try:
data = json.loads(line)
y[row_idx] = data.get('Label', 0)
for k, v in data.get('Base', {}).items():
idx = feature_idx.get(k)
if idx is not None:
val = float(v)
if val != 0.0:
rows.append(row_idx)
cols.append(idx)
vals.append(val)
for d in data.get('DLLs', []):
idx = feature_idx.get(f"Dll_{d}")
if idx is not None:
rows.append(row_idx)
cols.append(idx)
vals.append(1.0)
for a in data.get('APIs', []):
idx = feature_idx.get(f"Api_{a}")
if idx is not None:
rows.append(row_idx)
cols.append(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(
@@ -114,14 +133,14 @@ 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_target = 0.1
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 +152,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 +176,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)
@@ -170,12 +188,25 @@ def evaluate_and_save(model, X_test, y_test):
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))
tp = np.sum((y_test == 1) & (y_pred == 1))
fn = np.sum((y_test == 1) & (y_pred == 0))
tn = np.sum((y_test == 0) & (y_pred == 0))
fp = np.sum((y_test == 0) & (y_pred == 1))
total_malware = tp + fn
total_safe = tn + fp
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:.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")
print("[*] Top 20 Features (Importance by Gain):")
print("\n---------------------------------------------------------\n")
print("[*] Top 20 Features:")
feature_names = model.feature_name()
importance = model.feature_importance(importance_type='gain')
@@ -192,25 +223,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"[*] Filtered DLL features: {len(dll_schema)}")
print(f"[*] Filtered 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}")
+132 -99
View File
@@ -1,109 +1,128 @@
import os, sys, time, sqlite3, re, json
import pandas as pd
import os, sys, time, json, onnxmltools
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
####################################################################################################
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.0001
RANDOM_SEED = 42
EXCLUDE_FEATURES = {
'label', 'filehash',
}
LGBM_PARAMS = {
'objective': 'binary',
'metric': ['binary_logloss', 'auc'],
'boosting_type': 'gbdt',
'num_leaves': 63,
'learning_rate': 0.05,
'num_leaves': 64,
'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.1,
'lambda_l2': 0.1,
'min_data_in_leaf': 20,
'lambda_l1': 0.01,
'lambda_l2': 0.01,
}
####################################################################################################
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()
dll_counts = defaultdict(int)
api_counts = defaultdict(int)
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, 'r') as f:
for line in f:
try:
data = json.loads(line)
base_keys.update(data.get('Base', {}).keys())
for d in data.get('DLLs', []):
dll_counts[d] += 1
for a in data.get('APIs', []):
api_counts[a] += 1
num_lines += 1
except Exception:
continue
####################################################################################################
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 = sorted([f"Dll_{k}" for k, v in dll_counts.items() if v >= 100])
api_schema = sorted([f"Api_{k}" for k, v in api_counts.items() if v >= 100])
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)
feature_idx = {feat: i for i, feat in enumerate(final_features)}
rows, cols, vals = [], [], []
y = np.zeros(num_lines, dtype=np.int8)
with open(jsonl_path, 'r') as f:
row_idx = 0
for line in f:
try:
data = json.loads(line)
y[row_idx] = data.get('Label', 0)
for k, v in data.get('Base', {}).items():
idx = feature_idx.get(k)
if idx is not None:
val = float(v)
if val != 0.0:
rows.append(row_idx)
cols.append(idx)
vals.append(val)
for d in data.get('DLLs', []):
idx = feature_idx.get(f"Dll_{d}")
if idx is not None:
rows.append(row_idx)
cols.append(idx)
vals.append(1.0)
for a in data.get('APIs', []):
idx = feature_idx.get(f"Api_{a}")
if idx is not None:
rows.append(row_idx)
cols.append(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 +136,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 +152,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 +176,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)
@@ -170,12 +188,25 @@ def evaluate_and_save(model, X_test, y_test):
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))
tp = np.sum((y_test == 1) & (y_pred == 1))
fn = np.sum((y_test == 1) & (y_pred == 0))
tn = np.sum((y_test == 0) & (y_pred == 0))
fp = np.sum((y_test == 0) & (y_pred == 1))
total_malware = tp + fn
total_safe = tn + fp
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:.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")
print("[*] Top 20 Features (Importance by Gain):")
print("\n---------------------------------------------------------\n")
print("[*] Top 20 Features:")
feature_names = model.feature_name()
importance = model.feature_importance(importance_type='gain')
@@ -192,25 +223,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"[*] Filtered DLL features: {len(dll_schema)}")
print(f"[*] Filtered 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}")
+187 -12
View File
@@ -56,7 +56,7 @@ body {
}
.navigate-title {
font-size: 16px;
font-size: 17px;
font-weight: 600;
letter-spacing: -0.2px;
}
@@ -292,17 +292,14 @@ body {
border: 1px solid var(--border-color);
border-radius: 12px;
background-color: var(--bg-panel);
overflow-y: auto;
overflow-x: auto;
}
.virus-list {
list-style: none;
height: 100%;
overflow-y: auto;
overflow-x: auto;
padding: 12px;
}
.virus-list .manage-list-item {
padding: 0;
margin: 0;
min-width: max-content;
}
@@ -498,6 +495,7 @@ body {
height: 16px;
accent-color: var(--accent-color);
cursor: pointer;
flex-shrink: 0;
}
.manage-list-item-content {
@@ -510,13 +508,16 @@ body {
.custom-select-wrapper {
position: relative;
min-width: 120px;
min-width: 100px;
display: inline-grid;
vertical-align: middle;
user-select: none;
}
select.hidden + .custom-select-wrapper {
display: none !important;
}
.custom-select-trigger {
grid-area: 1 / 1;
display: flex;
justify-content: space-between;
align-items: center;
@@ -530,7 +531,24 @@ select.hidden + .custom-select-wrapper {
cursor: pointer;
transition: border-color 0.2s, box-shadow 0.2s;
gap: 12px;
box-sizing: border-box;
}
.custom-select-phantom {
grid-area: 1 / 1;
visibility: hidden;
pointer-events: none;
display: grid;
padding: 8px 16px;
font-size: 14px;
font-weight: 500;
box-sizing: border-box;
}
.custom-select-phantom span {
grid-area: 1 / 1;
white-space: nowrap;
padding-right: 28px;
}
.custom-select-wrapper.open .custom-select-trigger,
.custom-select-trigger:hover {
border-color: var(--accent-color);
@@ -550,7 +568,10 @@ select.hidden + .custom-select-wrapper {
position: absolute;
top: calc(100% + 8px);
left: 0;
right: 0;
min-width: 100%;
width: max-content;
white-space: nowrap;
background-color: var(--bg-panel);
border: 1px solid var(--border-color);
border-radius: 10px;
@@ -558,12 +579,10 @@ select.hidden + .custom-select-wrapper {
z-index: 100;
list-style: none;
margin: 0;
max-height: 0;
padding: 0 6px;
border-width: 0;
overflow: hidden;
opacity: 0;
visibility: hidden;
transform: translateY(-10px);
@@ -660,6 +679,162 @@ select.hidden + .custom-select-wrapper {
opacity: 1 !important;
}
.manage-list-header {
display: flex;
align-items: center;
gap: 12px;
padding: 12px 16px;
border-bottom: 1px solid var(--border-color);
background-color: var(--bg-panel);
font-weight: 600;
font-size: 13px;
color: var(--text-secondary);
position: sticky;
top: 0;
z-index: 10;
}
.manage-list-header input[type="checkbox"] {
width: 16px;
height: 16px;
accent-color: var(--accent-color);
cursor: pointer;
flex-shrink: 0;
}
.header-cols-container, .manage-list-item-content {
display: flex;
flex-grow: 1;
min-width: 0;
gap: 16px;
align-items: center;
}
.col-header {
position: relative;
cursor: pointer;
display: flex;
align-items: center;
padding: 0 8px;
border-right: 1px solid var(--border-color);
background-color: var(--bg-panel);
user-select: none;
-webkit-user-select: none;
}
.col-header:last-child {
border-right: none;
}
.header-text {
overflow: hidden;
text-overflow: ellipsis;
white-space: nowrap;
flex-grow: 1;
}
.manage-list-header, .col-header, .col-resizer {
user-select: none;
-webkit-user-select: none;
}
.col-resizer {
position: absolute;
right: -5px;
top: 0;
bottom: 0;
width: 10px;
cursor: col-resize;
z-index: 10;
background-color: transparent !important;
}
.col-resizer::after {
content: "";
position: absolute;
left: 50%;
top: 0;
bottom: 0;
width: 1px;
background-color: transparent;
transition: background-color 0.2s;
}
.col-resizer:hover::after,
.col-resizer.resizing::after {
background-color: var(--text-secondary) !important;
opacity: 0.6;
}
body.resizing-active,
body.resizing-active * {
user-select: none !important;
-webkit-user-select: none !important;
cursor: col-resize !important;
}
.col-resizer:hover {
background-color: var(--accent-color);
opacity: 0.5;
}
.col-header:hover {
color: var(--text-primary);
}
.sort-icon {
font-size: 10px;
margin-left: 4px;
color: var(--accent-color);
width: 12px;
}
.row-col {
overflow: hidden;
text-overflow: ellipsis;
white-space: nowrap;
font-size: 14px;
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; }
File diff suppressed because it is too large. Load diff
+781 -1810
View File
File diff suppressed because it is too large. Load diff
+26 -20
View File
@@ -5,6 +5,7 @@
<meta name="viewport" content="width=device-width, initial-scale=1.0">
<title>PYAS Security</title>
<link rel="stylesheet" href="/static/css/style.css">
<script defer src="/static/js/i18n.js"></script>
<script defer src="/static/js/main.js"></script>
</head>
<body>
@@ -59,13 +60,13 @@
<div class="quick-actions">
<div class="list-item" style="cursor: pointer;" id="quick_scan_btn" data-action="scan">
<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>
<div class="list-item" style="cursor: pointer;" data-target="log_export_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>
@@ -80,16 +81,15 @@
<div class="progress-widget">
<div class="header-group">
<h2 class="section-title" data-i18n>病毒掃描</h2>
<p class="section-desc" id="progress_text" data-i18n>此選項可以選擇掃描方式</p>
<p class="section-desc" id="progress_text" data-i18n>此選項可以選擇進階掃描方式</p>
</div>
<div class="action-group modern-action-group">
<button class="action-btn modern-btn danger-btn hidden" id="solve_btn" data-i18n>刪除</button>
<select id="scan_method_select" class="modern-select">
<option value="none" selected disabled data-i18n>選擇</option>
<option value="smart" data-i18n>智能</option>
<option value="file" data-i18n>檔案</option>
<option value="path" data-i18n>路徑</option>
<option value="full" data-i18n>全盤</option>
<option value="smart" data-i18n>智能掃描</option>
<option value="file" data-i18n>檔案掃描</option>
<option value="path" data-i18n>路徑掃描</option>
<option value="full" data-i18n>全盤掃描</option>
</select>
<button class="action-btn modern-btn hidden" id="stop_btn" data-i18n>停止</button>
</div>
@@ -105,37 +105,37 @@
<div class="list-widget">
<div class="list-item">
<div class="item-info">
<h2 data-i18n>進程防護</h2><p data-i18n>此選項可以攔截進程病毒</p>
<h2 data-i18n>進程防護</h2><p data-i18n>此選項可以終止惡意進程運行</p>
</div>
<label class="toggle-switch"><input type="checkbox" checked><span class="slider"></span></label>
</div>
<div class="list-item">
<div class="item-info">
<h2 data-i18n>檔案防護</h2><p data-i18n>此選項可以攔截檔案變更</p>
<h2 data-i18n>檔案防護</h2><p data-i18n>此選項可以隔離惡意檔案建立</p>
</div>
<label class="toggle-switch"><input type="checkbox" checked><span class="slider"></span></label>
</div>
<div class="list-item">
<div class="item-info">
<h2 data-i18n>系統防護</h2><p data-i18n>此選項可以攔截系統項目</p>
<h2 data-i18n>系統防護</h2><p data-i18n>此選項可以實時修復系統異常</p>
</div>
<label class="toggle-switch"><input type="checkbox" checked><span class="slider"></span></label>
</div>
<div class="list-item">
<div class="item-info">
<h2 data-i18n>驅動防護</h2><p data-i18n>此選項可以增強自身防護</p>
<h2 data-i18n>驅動防護</h2><p data-i18n>此選項可以實時攔截惡意行為</p>
</div>
<label class="toggle-switch"><input type="checkbox" checked><span class="slider"></span></label>
</div>
<div class="list-item">
<div class="item-info">
<h2 data-i18n>網路防護</h2><p data-i18n>此選項可以攔截網路通訊</p>
<h2 data-i18n>網路防護</h2><p data-i18n>此選項可以攔截惡意網路連接</p>
</div>
<label class="toggle-switch"><input type="checkbox" checked><span class="slider"></span></label>
</div>
<div class="list-item" style="cursor: pointer;" data-target="custom_protect_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>
@@ -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>
@@ -170,13 +170,13 @@
</div>
<div class="list-item" style="cursor: pointer;" data-target="whitelist_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>
<div class="list-item" style="cursor: pointer;" data-target="quarantine_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>
@@ -187,22 +187,28 @@
<div class="list-widget">
<div class="list-item">
<div class="item-info">
<h2 data-i18n>增強模式</h2><p data-i18n>此選項可以增強掃描引擎</p>
<h2 data-i18n>增強模式</h2><p data-i18n>此選項可以增強靜態掃描引擎</p>
</div>
<label class="toggle-switch"><input type="checkbox"><span class="slider"></span></label>
</div>
<div class="list-item">
<div class="item-info">
<h2 data-i18n>擴展引擎</h2><p data-i18n>此選項可以使用擴展引擎</p>
<h2 data-i18n>擴展引擎</h2><p data-i18n>此選項可以使用擴展規則引擎</p>
</div>
<label class="toggle-switch"><input type="checkbox"><span class="slider"></span></label>
</div>
<div class="list-item">
<div class="item-info">
<h2 data-i18n>雲端服務</h2><p data-i18n>此選項可以上報雲端引擎</p>
<h2 data-i18n>雲端服務</h2><p data-i18n>此選項可以自動上報可疑檔案</p>
</div>
<label class="toggle-switch"><input type="checkbox" checked><span class="slider"></span></label>
</div>
<div class="list-item">
<div class="item-info">
<h2 data-i18n>選單掃描</h2><p data-i18n>此選項可以啟用右鍵掃描功能</p>
</div>
<label class="toggle-switch"><input type="checkbox"><span class="slider"></span></label>
</div>
<div class="list-item">
<div class="item-info">
<h2 data-i18n>顯色主題</h2><p data-i18n>此選項可以選擇主題顏色</p>
+1730 -1353
View File
File diff suppressed because it is too large. Load diff
+202 -521
View File
@@ -1,8 +1,5 @@
import os, yara, time, numpy, base64, requests, pefile, math, json, re, datetime
import ctypes, ctypes.wintypes, hashlib, onnxruntime
from PIL import Image
from io import BytesIO
import os, re, yara, time, math, json, numpy, base64, datetime, requests
import ctypes, ctypes.wintypes, hashlib, pefile, threading, onnxruntime
####################################################################################################
@@ -164,317 +161,7 @@ class rule_scanner:
####################################################################################################
class cnn_scanner:
def __init__(self):
self.models = []
self.labels = ["White:WinPE/Unknown", "General:WinPE/Unknown"]
self.detect_set = {"General:WinPE/Unknown"}
self.resize = (224, 224)
def load_path(self, path, callback=None):
for root, _, files in os.walk(path):
for file in files:
full_path = os.path.join(root, file)
if callback:
callback(full_path)
self.load_file(full_path)
def load_file(self, file):
if file.lower().endswith('.onnx'):
try:
session = onnxruntime.InferenceSession(file, providers=['CPUExecutionProvider'])
self.models.append(session)
except Exception:
return False
####################################################################################################
def model_scan(self, file_path, full_output=False):
if not self.models:
return (False, False) if not full_output else ([], str(file_path), None)
data = self.get_data(file_path)
if not data:
return (False, False) if not full_output else ([], str(file_path), None)
image = self.preprocess_image(data, self.resize)
if not image:
return (False, False) if not full_output else ([], str(file_path), None)
arr = numpy.asarray(image).astype('float32') / 255.0
arr = numpy.expand_dims(arr, axis=0)
if arr.ndim == 3:
arr = numpy.expand_dims(arr, axis=-1)
results = []
best_malicious = None
for model in self.models:
input_meta = model.get_inputs()[0]
input_name = input_meta.name
input_shape = input_meta.shape
curr_arr = arr.copy()
if len(input_shape) == 4:
if input_shape[1] in (1, 3) and input_shape[3] not in (1, 3):
curr_arr = curr_arr.transpose(0, 3, 1, 2)
try:
probs = model.run(None, {input_name: curr_arr})[0]
pred_prob = float(probs[0][0]) if hasattr(probs[0], '__len__') else float(probs[0])
except Exception:
continue
is_malicious = pred_prob > 0.5
idx = 1 if is_malicious else 0
label = self.labels[idx]
conf = pred_prob if is_malicious else (1.0 - pred_prob)
conf = round(conf * 100, 2)
if full_output:
results.append(("Whole File", label, conf, self.pil_to_base64(image)))
if label in self.detect_set:
if best_malicious is None or conf > best_malicious[1]:
best_malicious = (f"{label}.{int(conf)}!dl", int(conf))
if full_output:
results.sort(key=lambda x: (x[1] not in self.detect_set, -x[2]))
malicious_count = sum(1 for _, lbl, _, _ in results if lbl in self.detect_set)
total = len(results)
return results, str(file_path), f"{malicious_count}/{total}"
return best_malicious if best_malicious else (False, False)
####################################################################################################
def pil_to_base64(self, img):
buf = BytesIO()
img.save(buf, format='PNG')
return base64.b64encode(buf.getvalue()).decode()
def preprocess_image(self, data, size, channels=1):
width, height = size
wah = int(numpy.ceil(numpy.sqrt(len(data) / channels)))
arr = numpy.frombuffer(data, dtype=numpy.uint8)
imgbuf = numpy.zeros(wah * wah * channels, dtype=numpy.uint8)
imgbuf[:len(arr)] = arr
image = Image.fromarray(imgbuf.reshape((wah, wah)), 'L')
return image.resize((width, height), Image.Resampling.NEAREST)
def get_data(self, file_path):
try:
with open(file_path, 'rb') as f:
return f.read()
except Exception:
return None
####################################################################################################
class pe_scanner:
API_CAT_MAPPING = {
'ProcessControl': {
'CreateProcessA', 'CreateProcessW', 'WinExec', 'ShellExecuteA', 'ShellExecuteW',
'ShellExecuteExW', 'ExitProcess', 'TerminateProcess', 'OpenProcess', 'GetExitCodeProcess',
'SetThreadPriority', 'GetThreadPriority', 'GetCurrentProcess', 'GetCurrentProcessId',
'GetModuleFileNameA', 'GetCommandLineA', 'GetModuleFileNameW', 'GetStartupInfoW'
},
'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',
'VirtualFree', 'VirtualQuery'
},
'Synchronization': {
'WaitForSingleObject', 'WaitForSingleObjectEx', 'WaitForMultipleObjects', 'WaitForMultipleObjectsEx',
'CreateMutexA', 'CreateMutexW', 'OpenMutexW', 'ReleaseMutex', 'CreateEventA', 'CreateEventW',
'OpenEventW', 'SetEvent', 'ResetEvent', 'EnterCriticalSection', 'LeaveCriticalSection',
'InitializeCriticalSection', 'DeleteCriticalSection', 'Sleep', 'SleepEx',
'InitializeCriticalSectionAndSpinCount', 'InterlockedDecrement', 'InterlockedIncrement',
'InterlockedExchange', 'InterlockedCompareExchange', 'SetTimer', 'KillTimer'
},
'MultiThreading': {
'CreateThread', 'ResumeThread', 'SuspendThread', 'ExitThread', 'TerminateThread',
'GetCurrentThread', 'GetCurrentThreadId', 'TlsAlloc', 'TlsSetValue', 'TlsGetValue',
'CreateThreadpoolWork', 'SubmitThreadpoolWork', 'TlsFree'
},
'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',
'GetStringTypeW', 'LCMapStringW', 'IsValidCodePage', 'DecodePointer', 'EncodePointer'
},
'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',
'CloseHandle', 'GetFileType', 'SetStdHandle', 'SetFilePointerEx', 'FindClose',
'SetHandleCount', 'SetEndOfFile', 'GetFullPathNameA', 'GetFileAttributesA'
},
'Registry': {
'RegOpenKeyExA', 'RegOpenKeyExW', 'RegSetValueExA', 'RegSetValueExW',
'RegCreateKeyExA', 'RegCreateKeyExW', 'RegDeleteKeyA', 'RegDeleteKeyW',
'RegEnumValueA', 'RegEnumValueW', 'RegQueryValueExA', 'RegQueryValueExW',
'RegDeleteValueA', 'RegDeleteValueW', 'RegCloseKey'
},
'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',
'GetLastError', 'GetACP', 'RaiseException', 'SetLastError', 'GetEnvironmentStringsW',
'FreeEnvironmentStringsW', 'GetLocaleInfoA', 'lstrlenA', 'GetVersion', 'GetSystemInfo',
'CompareStringA'
},
'DotNet': {
'_CorExeMain', '_CorDllMain'
},
'AntiDebug': {
'IsDebuggerPresent', 'CheckRemoteDebuggerPresent', 'OutputDebugStringA', 'OutputDebugStringW',
'GetTickCount', 'GetTickCount64', 'QueryPerformanceCounter', 'FindWindowA', 'FindWindowW',
'EnumWindows', 'EnumChildWindows', 'GetWindowRect', 'GetClientRect', 'SetWindowDisplayAffinity',
'UnhandledExceptionFilter', 'SetUnhandledExceptionFilter', 'IsProcessorFeaturePresent',
'SetErrorMode', 'GetTimeZoneInformation', 'GetDriveTypeW', 'GetDiskFreeSpaceA'
},
'Keylogging': {
'GetAsyncKeyState', 'GetKeyState', 'GetKeyboardState', 'MapVirtualKeyA', 'MapVirtualKeyW',
'ToAscii', 'ToAsciiEx', 'ToUnicode', 'ToUnicodeEx', 'GetKeyNameTextA', 'GetForegroundWindow',
'KeybdEvent', 'SendInput', 'MapVirtualKeyExA', 'MapVirtualKeyExW', 'CallNextHookEx'
},
'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'
},
'Memory': {
'HeapAlloc', 'HeapFree', 'HeapSize', 'HeapCreate', 'GetProcessHeap',
'LocalAlloc', 'LocalFree', 'GlobalAlloc', 'GlobalFree', 'GlobalUnlock',
'HeapDestroy', 'GlobalLock', 'SysFreeString', 'SysAllocStringLen'
},
'Resource': {
'LoadResource', 'SizeofResource', 'LockResource', 'FreeResource', 'FindResourceA'
},
'WindowControl': {
'ShowWindow', 'DestroyWindow', 'TranslateMessage', 'DispatchMessageA', 'GetWindow',
'PeekMessageA', 'GetWindowLongA', 'CallWindowProcA', 'SetWindowLongA', 'GetWindowTextA',
'ScreenToClient', 'GetActiveWindow', 'CreateWindowExA', 'DefWindowProcA', 'GetMessageA',
'RegisterClassA', 'UnregisterClassA', 'MessageBoxA'
},
'COM': {
'CoCreateInstance', 'CoUninitialize', 'CoInitialize'
}
}
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", "SectionMeanEntropy", "IsExe", "IsDll", "IsDriver", "Is64Bit",
"Machine", "Magic", "CheckSum", "ImageBase", "SizeOfImage", "SizeOfHeaders",
"Characteristics", "DllCharacteristics", "Subsystem", "LoaderFlags",
"MajorLinkerVersion", "MinorLinkerVersion", "SizeOfCode",
"SizeOfInitializedData", "SizeOfUninitializedData", "AddressOfEntryPoint",
"BaseOfCode", "SectionAlignment", "FileAlignment",
"MajorOperatingSystemVersion", "MinorOperatingSystemVersion",
"MajorImageVersion", "MinorImageVersion",
"MajorSubsystemVersion", "MinorSubsystemVersion",
"SizeOfStackReserve", "SizeOfStackCommit", "SizeOfHeapReserve", "SizeOfHeapCommit",
"NumberOfSections", "NumberOfRvaAndSizes", "PointerToSymbolTable", "NumberOfSymbols",
"SizeOfOptionalHeader",
"TextSectionMaxEntropy", "TextSectionMeanEntropy", "TextSizeRatio",
"DataSectionMaxEntropy", "DataSectionMeanEntropy", "DataSizeRatio",
"RsrcSectionMaxEntropy", "RsrcSectionMeanEntropy", "RsrcSizeRatio",
"SectionCount", "ExecutableSections",
"WritableSections", "ReadableSections", "SectionException",
"Char_00000020_Count", "Char_00000020_MeanEntropy",
"Char_00000040_Count", "Char_00000040_MeanEntropy",
"Char_00000080_Count", "Char_00000080_MeanEntropy",
"Char_02000000_Count", "Char_02000000_MeanEntropy",
"Char_20000000_Count", "Char_20000000_MeanEntropy",
"Char_40000000_Count", "Char_40000000_MeanEntropy",
"Char_80000000_Count", "Char_80000000_MeanEntropy",
"IconCount", "ApiCount", "ExportCount", "DebugCount", "ExceptionCount",
"ImportCount", "ImportFunctionCount",
"FileDescriptionLength", "FileVersionLength", "ProductNameLength",
"ProductVersionLength", "CompanyNameLength", "LegalCopyrightLength",
"CommentsLength", "InternalNameLength", "LegalTrademarksLength",
"SpecialBuildLength", "PrivateBuildLength",
"TrustSigned", "IsDebug", "IsPatched", "IsPrivateBuild", "IsPreRelease",
"IsSpecialBuild", "IsAdmin", "IsInstall", "HasTlsCallbacks",
"HasInvalidTimestamp", "HasRelocationDirectory", "HasPacked", "FileTimeException"
]
CAT_SCHEMA = [f"Cat_{k}" for k in sorted(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
_API_TO_CATS = {}
for _cat, _apis in API_CAT_MAPPING.items():
for _api in _apis:
if _api not in _API_TO_CATS:
_API_TO_CATS[_api] = []
_API_TO_CATS[_api].append(_cat)
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',
}
INSTALLER_SIGS = [b"Nullsoft.NSIS", b"Inno Setup", b"7-Zip.7zip", b"InstallShield"]
####################################################################################################
def __init__(self):
self.model = None
self.input_name = None
@@ -498,168 +185,163 @@ class pe_scanner:
self.model = onnxruntime.InferenceSession(model_path, providers=['CPUExecutionProvider'])
self.input_name = self.model.get_inputs()[0].name
feat_path = os.path.join(os.path.dirname(model_path), "features.json")
model_features = []
if os.path.exists(feat_path):
with open(feat_path, 'r') as f:
model_features = json.load(f)
else:
model_features = [re.sub(r'[^A-Za-z0-9_]+', '', f) for f in self.SCHEMA]
local_map = {re.sub(r'[^A-Za-z0-9_]+', '', raw_name): raw_name for raw_name in self.SCHEMA}
self.feature_order = [local_map.get(req_feat) for req_feat in model_features]
with open(feat_path, 'r', encoding='utf-8') as f:
self.feature_order = json.load(f)
except Exception:
pass
####################################################################################################
def _safe_float(self, 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 Exception:
return 0.0
def _safe_div(self, n, d):
return float(n) / d if d > 0 else 0.0
def _calc_entropy(self, data):
if not data: return 0.0
if not data:
return 0.0
arr = numpy.frombuffer(data, dtype=numpy.uint8)
counts = numpy.bincount(arr, minlength=256)
probs = counts[counts > 0] / len(arr)
return float(-numpy.sum(probs * numpy.log2(probs)))
####################################################################################################
def extract_features(self, file_path):
fts = {k: 0.0 for k in self.SCHEMA}
pe = None
try:
fsize = os.path.getsize(file_path)
if fsize == 0 or fsize > 512 * 1024 * 1024:
if fsize == 0 or fsize > 536870912:
return None
with open(file_path, "rb") as f:
file_bytes = f.read()
base = {}
dlls = set()
apis = set()
pe = pefile.PE(data=file_bytes, fast_load=True)
fts['TrustSigned'] = 1 if self.signer.sign_verify(file_path) else 0
base['TrustSigned'] = 1.0 if self.signer.sign_verify(file_path) else 0.0
base['FileEntropy'] = self._calc_entropy(file_bytes)
base['FileSize'] = float(fsize)
fh = pe.FILE_HEADER
fts['Machine'] = fh.Machine
fts['NumberOfSections'] = fh.NumberOfSections
fts['TimeDateStamp'] = fh.TimeDateStamp
fts['PointerToSymbolTable'] = fh.PointerToSymbolTable
fts['NumberOfSymbols'] = fh.NumberOfSymbols
fts['SizeOfOptionalHeader'] = fh.SizeOfOptionalHeader
fts['Characteristics'] = fh.Characteristics
curr_ts = datetime.datetime.now(datetime.timezone.utc).timestamp()
fts['HasInvalidTimestamp'] = 1 if (fh.TimeDateStamp < 631152000 or fh.TimeDateStamp > curr_ts + 2592000) else 0
fts['FileTimeException'] = 1 if fh.TimeDateStamp == 0 else 0
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
fts['Magic'] = getattr(op, 'Magic', 0)
fts['MajorLinkerVersion'] = getattr(op, 'MajorLinkerVersion', 0)
fts['MinorLinkerVersion'] = getattr(op, 'MinorLinkerVersion', 0)
fts['SizeOfCode'] = getattr(op, 'SizeOfCode', 0)
fts['SizeOfInitializedData'] = getattr(op, 'SizeOfInitializedData', 0)
fts['SizeOfUninitializedData'] = getattr(op, 'SizeOfUninitializedData', 0)
fts['AddressOfEntryPoint'] = getattr(op, 'AddressOfEntryPoint', 0)
fts['BaseOfCode'] = getattr(op, 'BaseOfCode', 0)
fts['ImageBase'] = getattr(op, 'ImageBase', 0)
fts['SectionAlignment'] = getattr(op, 'SectionAlignment', 0)
fts['FileAlignment'] = getattr(op, 'FileAlignment', 0)
fts['MajorOperatingSystemVersion'] = getattr(op, 'MajorOperatingSystemVersion', 0)
fts['MinorOperatingSystemVersion'] = getattr(op, 'MinorOperatingSystemVersion', 0)
fts['MajorImageVersion'] = getattr(op, 'MajorImageVersion', 0)
fts['MinorImageVersion'] = getattr(op, 'MinorImageVersion', 0)
fts['MajorSubsystemVersion'] = getattr(op, 'MajorSubsystemVersion', 0)
fts['MinorSubsystemVersion'] = getattr(op, 'MinorSubsystemVersion', 0)
fts['SizeOfImage'] = getattr(op, 'SizeOfImage', 0)
fts['SizeOfHeaders'] = getattr(op, 'SizeOfHeaders', 0)
fts['CheckSum'] = getattr(op, 'CheckSum', 0)
fts['Subsystem'] = getattr(op, 'Subsystem', 0)
fts['DllCharacteristics'] = getattr(op, 'DllCharacteristics', 0)
fts['SizeOfStackReserve'] = getattr(op, 'SizeOfStackReserve', 0)
fts['SizeOfStackCommit'] = getattr(op, 'SizeOfStackCommit', 0)
fts['SizeOfHeapReserve'] = getattr(op, 'SizeOfHeapReserve', 0)
fts['SizeOfHeapCommit'] = getattr(op, 'SizeOfHeapCommit', 0)
fts['LoaderFlags'] = getattr(op, 'LoaderFlags', 0)
fts['NumberOfRvaAndSizes'] = getattr(op, 'NumberOfRvaAndSizes', 0)
fts['Is64Bit'] = 1 if fh.Machine in (0x8664, 0xAA64, 0x0200) else 0
fts['IsExe'] = 1 if pe.is_exe() else 0
fts['IsDll'] = 1 if pe.is_dll() else 0
fts['IsDriver'] = 1 if fts['Subsystem'] == 1 else 0
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_field in fields:
base[f_field] = getattr(op, f_field, 0)
base['IsDriver'] = 1.0 if base.get('Subsystem') == 1 else 0.0
if hasattr(op, 'DATA_DIRECTORY'):
reloc_idx = pefile.DIRECTORY_ENTRY.get('IMAGE_DIRECTORY_ENTRY_BASERELOC', 5)
if len(op.DATA_DIRECTORY) > reloc_idx and op.DATA_DIRECTORY[reloc_idx].Size > 0:
fts['HasRelocationDirectory'] = 1
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
try:
fts['FileEntropy'] = self._calc_entropy(file_bytes)
except Exception:
pass
sec_entropies = []
text_entropies = []
data_entropies = []
rsrc_entropies = []
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
if hasattr(pe, 'sections'):
fts['SectionCount'] = len(pe.sections)
base['SectionCount'] = len(pe.sections)
entropies = []
raw_sizes = []
v_sizes = []
exec_sec = 0
write_sec = 0
read_sec = 0
sec_exc = 0
for section in pe.sections:
try:
name = section.Name.decode('ascii', 'ignore').strip('\x00')
except Exception:
name = ""
try:
try:
s_data = section.get_data()
s_entropy = self._calc_entropy(s_data)
except Exception:
s_data = b""
except Exception:
s_entropy = 0.0
sec_entropies.append(s_entropy)
if any(p in name.lower() for p in self.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
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:
fts[f'Char_{flag:08X}_Count'] += 1
fts[f'Char_{flag:08X}_MeanEntropy'] += s_entropy
base[f'Char_{flag:08X}_Count'] += 1.0
base[f'Char_{flag:08X}_MeanEntropy'] += s_entropy
ratio = self._safe_div(section.SizeOfRawData, fsize)
if name.startswith('.text'):
text_entropies.append(s_entropy)
fts['TextSizeRatio'] += ratio
elif name.startswith('.data'):
data_entropies.append(s_entropy)
fts['DataSizeRatio'] += ratio
elif name.startswith('.rsrc'):
rsrc_entropies.append(s_entropy)
fts['RsrcSizeRatio'] += ratio
if b"requireAdministrator" in s_data: fts['IsAdmin'] = 1
for sig in self.INSTALLER_SIGS:
if sig in s_data:
fts['IsInstall'] = 1
break
base['SectionMaxEntropy'] = max(entropies) if entropies else 0.0
base['SectionMinEntropy'] = min(entropies) if entropies else 0.0
base['SectionMeanEntropy'] = sum(entropies) / len(entropies) if entropies else 0.0
base['SectionMaxRawSize'] = max(raw_sizes) if raw_sizes else 0.0
base['SectionMinRawSize'] = min(raw_sizes) if raw_sizes else 0.0
base['SectionMeanRawSize'] = sum(raw_sizes) / len(raw_sizes) if raw_sizes else 0.0
base['SectionMaxVSize'] = max(v_sizes) if v_sizes else 0.0
base['SectionMinVSize'] = min(v_sizes) if v_sizes else 0.0
base['SectionMeanVSize'] = sum(v_sizes) / len(v_sizes) if v_sizes else 0.0
base['ExecutableSections'] = float(exec_sec)
base['WritableSections'] = float(write_sec)
base['ReadableSections'] = float(read_sec)
base['SectionException'] = float(sec_exc)
fts['SectionMeanEntropy'] = self._safe_div(sum(sec_entropies), len(sec_entropies))
fts['TextSectionMaxEntropy'] = max(text_entropies) if text_entropies else 0.0
fts['TextSectionMeanEntropy'] = self._safe_div(sum(text_entropies), len(text_entropies))
fts['DataSectionMaxEntropy'] = max(data_entropies) if data_entropies else 0.0
fts['DataSectionMeanEntropy'] = self._safe_div(sum(data_entropies), len(data_entropies))
fts['RsrcSectionMaxEntropy'] = max(rsrc_entropies) if rsrc_entropies else 0.0
fts['RsrcSectionMeanEntropy'] = self._safe_div(sum(rsrc_entropies), len(rsrc_entropies))
for flag in char_flags:
if base[f'Char_{flag:08X}_Count'] > 0:
base[f'Char_{flag:08X}_MeanEntropy'] = base[f'Char_{flag:08X}_MeanEntropy'] / base[f'Char_{flag:08X}_Count']
for flag in char_flags:
fts[f'Char_{flag:08X}_MeanEntropy'] = self._safe_div(fts[f'Char_{flag:08X}_MeanEntropy'], fts[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
try:
pe.parse_data_directories(directories=[
@@ -674,110 +356,90 @@ class pe_scanner:
if hasattr(pe, 'DIRECTORY_ENTRY_TLS') and hasattr(pe.DIRECTORY_ENTRY_TLS, 'struct'):
if getattr(pe.DIRECTORY_ENTRY_TLS.struct, 'AddressOfCallBacks', 0) != 0:
fts['HasTlsCallbacks'] = 1
base['HasTlsCallbacks'] = 1.0
if hasattr(pe, 'VS_FIXEDFILEINFO') and len(pe.VS_FIXEDFILEINFO) > 0:
flags = getattr(pe.VS_FIXEDFILEINFO[0], 'FileFlags', 0)
fts['IsDebug'] = 1 if flags & 0x1 else 0
fts['IsPreRelease'] = 1 if flags & 0x2 else 0
fts['IsPatched'] = 1 if flags & 0x4 else 0
fts['IsPrivateBuild'] = 1 if flags & 0x8 else 0
fts['IsSpecialBuild'] = 1 if flags & 0x20 else 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)
v = val.decode('utf-8', 'ignore') if isinstance(val, bytes) else str(val)
length = len(v)
if k == 'FileDescription': fts['FileDescriptionLength'] = length
elif k == 'FileVersion': fts['FileVersionLength'] = length
elif k == 'ProductName': fts['ProductNameLength'] = length
elif k == 'ProductVersion': fts['ProductVersionLength'] = length
elif k == 'CompanyName': fts['CompanyNameLength'] = length
elif k == 'LegalCopyright': fts['LegalCopyrightLength'] = length
elif k == 'Comments': fts['CommentsLength'] = length
elif k == 'InternalName': fts['InternalNameLength'] = length
elif k == 'LegalTrademarks': fts['LegalTrademarksLength'] = length
elif k == 'SpecialBuild': fts['SpecialBuildLength'] = length
elif k == 'PrivateBuild': fts['PrivateBuildLength'] = length
except Exception:
continue
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'):
fts['ImportCount'] = len(pe.DIRECTORY_ENTRY_IMPORT)
base['ImportCount'] = float(len(pe.DIRECTORY_ENTRY_IMPORT))
func_count = 0
for entry in pe.DIRECTORY_ENTRY_IMPORT:
if getattr(entry, 'dll', None):
try:
dll_name = entry.dll.decode('ascii', 'ignore').lower()
if dll_name in self.TARGET_DLLS:
fts[f"Dll_{dll_name.replace('.', '_')}"] = 1.0
except Exception:
dlls.add(entry.dll.decode('ascii', 'ignore').lower())
except Exception:
pass
imports = getattr(entry, 'imports', [])
for imp in imports:
fts['ImportFunctionCount'] += 1
if not getattr(imp, 'name', None): continue
try:
func_name = imp.name.decode('ascii', 'ignore')
if func_name in self._ALL_APIS:
fts[f'Api_{func_name}'] += 1
fts['ApiCount'] += 1
for cat in self._API_TO_CATS.get(func_name, []):
fts[f'Cat_{cat}'] += 1
except Exception:
continue
for imp in getattr(entry, 'imports', []):
func_count += 1
if getattr(imp, 'name', None):
try:
apis.add(imp.name.decode('ascii', 'ignore'))
except Exception:
pass
base['ImportFunctionCount'] = float(func_count)
else:
base['ImportCount'] = 0.0
base['ImportFunctionCount'] = 0.0
if hasattr(pe, 'DIRECTORY_ENTRY_EXPORT') and hasattr(pe.DIRECTORY_ENTRY_EXPORT, 'symbols'):
fts['ExportCount'] = len(pe.DIRECTORY_ENTRY_EXPORT.symbols)
base['ExportCount'] = float(len(pe.DIRECTORY_ENTRY_EXPORT.symbols))
else:
base['ExportCount'] = 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'):
fts['IconCount'] += len(getattr(entry.directory, 'entries', []))
if hasattr(pe, 'OPTIONAL_HEADER') and hasattr(pe.OPTIONAL_HEADER, 'DATA_DIRECTORY') and len(pe.OPTIONAL_HEADER.DATA_DIRECTORY) > 3:
exc_dir = pe.OPTIONAL_HEADER.DATA_DIRECTORY[3]
if exc_dir.Size > 0:
fts['ExceptionCount'] = exc_dir.Size // 12 if fts['Machine'] in (0x8664, 0xAA64) else 0
if hasattr(pe, 'DIRECTORY_ENTRY_DEBUG'):
fts['DebugCount'] = len(pe.DIRECTORY_ENTRY_DEBUG)
icon_count += len(getattr(entry.directory, 'entries', []))
base['IconCount'] = float(icon_count)
else:
base['IconCount'] = 0.0
for k in fts:
fts[k] = self._safe_float(fts[k])
return fts
if hasattr(pe, 'DIRECTORY_ENTRY_DEBUG'):
base['DebugCount'] = float(len(pe.DIRECTORY_ENTRY_DEBUG))
else:
base['DebugCount'] = 0.0
for k in base:
base[k] = self._safe_float(base[k])
return {"Base": base, "DLLs": list(dlls), "APIs": list(apis)}
except Exception:
return None
finally:
if pe:
if pe:
pe.close()
####################################################################################################
def pe_scan(self, file_path):
if not self.model:
def pe_scan(self, file_path, enhanced_mode=False):
if not self.model or not self.feature_order:
return False, False
try:
data = self.extract_features(file_path)
if not data:
raw_data = self.extract_features(file_path)
if not raw_data:
return False, False
vec_data = []
for key in self.feature_order:
if key:
vec_data.append(data.get(key, 0.0))
vec = numpy.zeros((1, len(self.feature_order)), dtype=numpy.float32)
base = raw_data.get('Base', {})
dlls = set(raw_data.get('DLLs', []))
apis = set(raw_data.get('APIs', []))
for i, feat in enumerate(self.feature_order):
if feat.startswith('Dll_'):
if feat[4:] in dlls:
vec[0, i] = 1.0
elif feat.startswith('Api_'):
if feat[4:] in apis:
vec[0, i] = 1.0
else:
vec_data.append(0.0)
vec[0, i] = base.get(feat, 0.0)
vec = numpy.array(vec_data, dtype=numpy.float32).reshape(1, -1)
outputs = self.model.run(None, {self.input_name: vec})
result = outputs[1]
prob = 0.0
@@ -786,14 +448,17 @@ class pe_scanner:
prob_dict = result[0]
if hasattr(prob_dict, 'get'):
prob = float(prob_dict.get(1, prob_dict.get('1', 0.0)))
elif isinstance(result, numpy.ndarray):
if result.ndim == 2 and result.shape[1] > 1:
prob = float(result[0][1])
score = int(prob * 100)
if prob > 0.5:
return f"General:WinPE/Unknown.{score}!ml", score
if score >= 80:
return f"General:WinPE/Malware.{score}!ml", score
elif enhanced_mode and score >= 50:
return f"General:WinPE/Suspicious.{score}!ml", score
return False, False
except Exception:
return False, False
@@ -806,17 +471,26 @@ class cloud_scanner:
self.api_host = None
self.api_key = None
self.timeout = 30
self.lock = threading.RLock()
def _init_session(self, api_host, api_key):
if self.session is None or self.api_host != api_host or self.api_key != api_key:
self.api_host = api_host.rstrip('/')
self.api_key = api_key
self.session = requests.Session()
self.session.headers.update({"X-API-Key": api_key, "User-Agent": "PYAS-Engine/1.2"})
with self.lock:
if self.session is None or self.api_host != api_host or self.api_key != api_key:
self.api_host = api_host.rstrip('/')
self.api_key = api_key
self.session = requests.Session()
self.session.headers.update({"X-API-Key": api_key, "User-Agent": "PYAS-Engine/1.1"})
def _request(self, method, endpoint, **kwargs):
with self.lock:
current_session = self.session
current_host = self.api_host
if not current_session:
return None
try:
r = self.session.request(method, f"{self.api_host}{endpoint}", timeout=self.timeout, **kwargs)
r = current_session.request(method, f"{current_host}{endpoint}", timeout=self.timeout, **kwargs)
if r.status_code == 200:
return r
except Exception:
@@ -840,7 +514,7 @@ class cloud_scanner:
r = self._request("POST", f"/api/rescan/{sha256}")
return r is not None and r.json().get('status') == 'success'
def upload_file(self, file_path, api_host, api_key, chunk_size=5242880):
def upload_file(self, file_path, api_host, api_key, chunk_size=5242880, need_rescan=False):
try:
self._init_session(api_host, api_key)
sha256 = self.calc_hash(file_path)
@@ -848,9 +522,16 @@ class cloud_scanner:
return False, None
status_req = self._request("GET", f"/api/processing_status/{sha256}")
if status_req and status_req.json().get('status') != 'missing':
self.rescan(sha256, api_host, api_key)
return True, sha256
if status_req:
status_data = status_req.json()
current_status = status_data.get('status')
if current_status == 'done':
if need_rescan:
self.rescan(sha256, api_host, api_key)
return True, sha256
elif current_status in ('queued', 'processing'):
return True, sha256
file_size = os.path.getsize(file_path)
if file_size > 104857600:
+1 -1
View File
@@ -2,7 +2,7 @@ import pyinstaller_versionfile
pyinstaller_versionfile.create_versionfile(
output_file="versionfile.txt",
version='3.5.0.0',
version='3.5.4.0',
company_name="PYAS Security",
file_description="PYAS Security Antivirus",
internal_name="PYAS",
+2 -2
View File
@@ -114,8 +114,8 @@ typedef struct _DRIVER_DATA {
PEPROCESS PyasProcess;
ULONG PyasPid;
LARGE_INTEGER Cookie;
FAST_MUTEX PortMutex;
FAST_MUTEX TrackerMutex;
KSPIN_LOCK PortMutex;
KSPIN_LOCK TrackerMutex;
EX_RUNDOWN_REF PortRundown;
} DRIVER_DATA, * PDRIVER_DATA;
+12 -6
View File
@@ -90,7 +90,8 @@ static NTSTATUS PortConnect(PFLT_PORT ClientPort, PVOID ServerPortCookie, PVOID
UNREFERENCED_PARAMETER(SizeOfContext);
UNREFERENCED_PARAMETER(ConnectionPortCookie);
ExAcquireFastMutex(&GlobalData.PortMutex);
KIRQL OldIrql;
KeAcquireSpinLock(&GlobalData.PortMutex, &OldIrql);
GlobalData.ClientPort = ClientPort;
GlobalData.PyasPid = (ULONG)(ULONG_PTR)PsGetCurrentProcessId();
@@ -98,7 +99,7 @@ static NTSTATUS PortConnect(PFLT_PORT ClientPort, PVOID ServerPortCookie, PVOID
ExReInitializeRundownProtection(&GlobalData.PortRundown);
ExReleaseFastMutex(&GlobalData.PortMutex);
KeReleaseSpinLock(&GlobalData.PortMutex, OldIrql);
return STATUS_SUCCESS;
}
@@ -106,9 +107,14 @@ static NTSTATUS PortConnect(PFLT_PORT ClientPort, PVOID ServerPortCookie, PVOID
static VOID PortDisconnect(PVOID ConnectionCookie) {
UNREFERENCED_PARAMETER(ConnectionCookie);
ExAcquireFastMutex(&GlobalData.PortMutex);
KIRQL OldIrql;
KeAcquireSpinLock(&GlobalData.PortMutex, &OldIrql);
GlobalData.ClientPort = NULL;
ExReleaseFastMutex(&GlobalData.PortMutex);
GlobalData.PyasPid = 0;
GlobalData.PyasProcess = NULL;
KeReleaseSpinLock(&GlobalData.PortMutex, OldIrql);
ExWaitForRundownProtectionRelease(&GlobalData.PortRundown);
}
@@ -147,8 +153,8 @@ extern "C" NTSTATUS DriverEntry(PDRIVER_OBJECT DriverObject, PUNICODE_STRING Reg
RtlZeroMemory(&GlobalData, sizeof(GlobalData));
GlobalData.DriverObject = DriverObject;
ExInitializeFastMutex(&GlobalData.PortMutex);
ExInitializeFastMutex(&GlobalData.TrackerMutex);
KeInitializeSpinLock(&GlobalData.PortMutex);
KeInitializeSpinLock(&GlobalData.TrackerMutex);
ExInitializeRundownProtection(&GlobalData.PortRundown);
+1 -1
View File
@@ -3,7 +3,7 @@ Signature = "$WINDOWS NT$"
Class = ActivityMonitor
ClassGuid = {b86dff51-a31e-4bac-b3cf-e8cfe75c9fc2}
Provider = %ManufacturerName%
DriverVer = 04/12/2026,3.5.0.0
DriverVer = 04/26/2026,3.5.1.0
CatalogFile = PYAS_Driver.cat
PnpLockdown = 1
+5 -6
View File
@@ -199,6 +199,7 @@ FLT_PREOP_CALLBACK_STATUS ProtectFile_PreWrite(PFLT_CALLBACK_DATA Data, PCFLT_RE
UNREFERENCED_PARAMETER(FltObjects);
if (Data->RequestorMode == KernelMode) return FLT_PREOP_SUCCESS_NO_CALLBACK;
if (KeGetCurrentIrql() > APC_LEVEL) return FLT_PREOP_SUCCESS_NO_CALLBACK;
if (Data->Iopb->IrpFlags & (IRP_PAGING_IO | IRP_SYNCHRONOUS_PAGING_IO | IRP_NOCACHE)) {
return FLT_PREOP_SUCCESS_NO_CALLBACK;
@@ -235,12 +236,10 @@ FLT_PREOP_CALLBACK_STATUS ProtectFile_PreWrite(PFLT_CALLBACK_DATA Data, PCFLT_RE
SystemBuffer = MmGetSystemAddressForMdlSafe(Mdl, NormalPagePriority | MdlMappingNoExecute);
}
if (SystemBuffer) {
if (CheckRansomActivity(Pid, &nameInfo->Name, SystemBuffer, WriteLength, TRUE)) {
SendMessageToUser(5001, (ULONG)(ULONG_PTR)Pid, nameInfo->Name.Buffer, nameInfo->Name.Length);
Data->IoStatus.Status = STATUS_ACCESS_DENIED;
callbackStatus = FLT_PREOP_COMPLETE;
}
if (CheckRansomActivity(Pid, &nameInfo->Name, SystemBuffer, WriteLength, TRUE)) {
SendMessageToUser(5001, (ULONG)(ULONG_PTR)Pid, nameInfo->Name.Buffer, nameInfo->Name.Length);
Data->IoStatus.Status = STATUS_ACCESS_DENIED;
callbackStatus = FLT_PREOP_COMPLETE;
}
}
}
+38 -4
View File
@@ -34,6 +34,36 @@ VOID ImageLoadNotify(PUNICODE_STRING FullImageName, HANDLE ProcessId, PIMAGE_INF
}
}
static BOOLEAN IsCriticalSystemProcess(HANDLE ProcessId) {
if (ProcessId == (HANDLE)4) return TRUE;
if (KeGetCurrentIrql() != PASSIVE_LEVEL) return FALSE;
NTSTATUS status;
BOOLEAN isCritical = FALSE;
PEPROCESS Process = NULL;
PUNICODE_STRING imageFileName = NULL;
status = PsLookupProcessByProcessId(ProcessId, &Process);
if (!NT_SUCCESS(status)) return FALSE;
status = SeLocateProcessImageName(Process, &imageFileName);
if (NT_SUCCESS(status) && imageFileName && imageFileName->Buffer) {
if (WildcardMatch(L"*\\Windows\\System32\\csrss.exe", imageFileName->Buffer, imageFileName->Length) ||
WildcardMatch(L"*\\Windows\\System32\\lsass.exe", imageFileName->Buffer, imageFileName->Length) ||
WildcardMatch(L"*\\Windows\\System32\\services.exe", imageFileName->Buffer, imageFileName->Length) ||
WildcardMatch(L"*\\Windows\\System32\\wininit.exe", imageFileName->Buffer, imageFileName->Length) ||
WildcardMatch(L"*\\Windows\\System32\\smss.exe", imageFileName->Buffer, imageFileName->Length)) {
isCritical = TRUE;
}
ExFreePool(imageFileName);
}
ObDereferenceObject(Process);
return isCritical;
}
static BOOLEAN IsBlacklistedAdminTool(HANDLE ProcessId) {
if (KeGetCurrentIrql() != PASSIVE_LEVEL) return FALSE;
@@ -79,6 +109,8 @@ static OB_PREOP_CALLBACK_STATUS PreOpenProcess(PVOID RegistrationContext, POB_PR
if (SourcePid == (HANDLE)4) return OB_PREOP_SUCCESS;
if (IsTargetProtected(TargetPid)) {
if (IsCriticalSystemProcess(SourcePid)) return OB_PREOP_SUCCESS;
BOOLEAN bIsTrusted = IsProcessTrusted(SourcePid);
if (bIsTrusted) {
@@ -99,10 +131,10 @@ static OB_PREOP_CALLBACK_STATUS PreOpenProcess(PVOID RegistrationContext, POB_PR
PROCESS_SET_INFORMATION |
PROCESS_SET_QUOTA;
OperationInformation->Parameters->CreateHandleInformation.DesiredAccess &= ~DenyMask;
OperationInformation->Parameters->CreateHandleInformation.DesiredAccess &= (ACCESS_MASK)(~DenyMask);
if (OperationInformation->Operation == OB_OPERATION_HANDLE_DUPLICATE) {
OperationInformation->Parameters->DuplicateHandleInformation.DesiredAccess &= ~DenyMask;
OperationInformation->Parameters->DuplicateHandleInformation.DesiredAccess &= (ACCESS_MASK)(~DenyMask);
}
}
@@ -128,6 +160,8 @@ static OB_PREOP_CALLBACK_STATUS PreOpenThread(PVOID RegistrationContext, POB_PRE
if (SourcePid == (HANDLE)4) return OB_PREOP_SUCCESS;
if (IsTargetProtected(TargetPid)) {
if (IsCriticalSystemProcess(SourcePid)) return OB_PREOP_SUCCESS;
BOOLEAN bIsTrusted = IsProcessTrusted(SourcePid);
if (bIsTrusted) {
@@ -144,10 +178,10 @@ static OB_PREOP_CALLBACK_STATUS PreOpenThread(PVOID RegistrationContext, POB_PRE
THREAD_SET_INFORMATION |
THREAD_SET_THREAD_TOKEN;
OperationInformation->Parameters->CreateHandleInformation.DesiredAccess &= ~DenyMask;
OperationInformation->Parameters->CreateHandleInformation.DesiredAccess &= (ACCESS_MASK)(~DenyMask);
if (OperationInformation->Operation == OB_OPERATION_HANDLE_DUPLICATE) {
OperationInformation->Parameters->DuplicateHandleInformation.DesiredAccess &= ~DenyMask;
OperationInformation->Parameters->DuplicateHandleInformation.DesiredAccess &= (ACCESS_MASK)(~DenyMask);
}
}
+40 -36
View File
@@ -149,30 +149,32 @@ static NTSTATUS RegistryCallback(PVOID CallbackContext, PVOID Argument1, PVOID A
(Info->ValueName ? Info->ValueName->Length : 0) +
sizeof(WCHAR);
PWCHAR Buffer = (PWCHAR)PyasAllocate(FullSize);
if (Buffer) {
UNICODE_STRING FullPath;
FullPath.Buffer = Buffer;
FullPath.Length = 0;
FullPath.MaximumLength = (USHORT)FullSize;
if (FullSize <= 0xFFFF) {
PWCHAR Buffer = (PWCHAR)PyasAllocate(FullSize);
if (Buffer) {
UNICODE_STRING FullPath;
FullPath.Buffer = Buffer;
FullPath.Length = 0;
FullPath.MaximumLength = (USHORT)FullSize;
RtlCopyUnicodeString(&FullPath, KeyName);
RtlCopyUnicodeString(&FullPath, KeyName);
if (Info->ValueName && Info->ValueName->Length > 0) {
if (FullPath.Length > 0 && FullPath.Buffer[(FullPath.Length / sizeof(WCHAR)) - 1] != L'\\') {
RtlAppendUnicodeToString(&FullPath, L"\\");
if (Info->ValueName && Info->ValueName->Length > 0) {
if (FullPath.Length > 0 && FullPath.Buffer[(FullPath.Length / sizeof(WCHAR)) - 1] != L'\\') {
RtlAppendUnicodeToString(&FullPath, L"\\");
}
RtlAppendUnicodeStringToString(&FullPath, Info->ValueName);
}
RtlAppendUnicodeStringToString(&FullPath, Info->ValueName);
}
if (CheckRegistryRule(&FullPath)) {
HANDLE Pid = PsGetCurrentProcessId();
if (!IsProcessTrusted(Pid)) {
SendMessageToUser(3001, (ULONG)(ULONG_PTR)Pid, FullPath.Buffer, FullPath.Length);
status = STATUS_ACCESS_DENIED;
if (CheckRegistryRule(&FullPath)) {
HANDLE Pid = PsGetCurrentProcessId();
if (!IsProcessTrusted(Pid)) {
SendMessageToUser(3001, (ULONG)(ULONG_PTR)Pid, FullPath.Buffer, FullPath.Length);
status = STATUS_ACCESS_DENIED;
}
}
PyasFree(Buffer);
}
PyasFree(Buffer);
}
}
}
@@ -194,30 +196,32 @@ static NTSTATUS RegistryCallback(PVOID CallbackContext, PVOID Argument1, PVOID A
(Info->ValueName ? Info->ValueName->Length : 0) +
sizeof(WCHAR);
PWCHAR Buffer = (PWCHAR)PyasAllocate(FullSize);
if (Buffer) {
UNICODE_STRING FullPath;
FullPath.Buffer = Buffer;
FullPath.Length = 0;
FullPath.MaximumLength = (USHORT)FullSize;
if (FullSize <= 0xFFFF) {
PWCHAR Buffer = (PWCHAR)PyasAllocate(FullSize);
if (Buffer) {
UNICODE_STRING FullPath;
FullPath.Buffer = Buffer;
FullPath.Length = 0;
FullPath.MaximumLength = (USHORT)FullSize;
RtlCopyUnicodeString(&FullPath, KeyName);
RtlCopyUnicodeString(&FullPath, KeyName);
if (Info->ValueName && Info->ValueName->Length > 0) {
if (FullPath.Length > 0 && FullPath.Buffer[(FullPath.Length / sizeof(WCHAR)) - 1] != L'\\') {
RtlAppendUnicodeToString(&FullPath, L"\\");
if (Info->ValueName && Info->ValueName->Length > 0) {
if (FullPath.Length > 0 && FullPath.Buffer[(FullPath.Length / sizeof(WCHAR)) - 1] != L'\\') {
RtlAppendUnicodeToString(&FullPath, L"\\");
}
RtlAppendUnicodeStringToString(&FullPath, Info->ValueName);
}
RtlAppendUnicodeStringToString(&FullPath, Info->ValueName);
}
if (CheckRegistryRule(&FullPath)) {
HANDLE Pid = PsGetCurrentProcessId();
if (!IsProcessTrusted(Pid)) {
SendMessageToUser(3001, (ULONG)(ULONG_PTR)Pid, FullPath.Buffer, FullPath.Length);
status = STATUS_ACCESS_DENIED;
if (CheckRegistryRule(&FullPath)) {
HANDLE Pid = PsGetCurrentProcessId();
if (!IsProcessTrusted(Pid)) {
SendMessageToUser(3001, (ULONG)(ULONG_PTR)Pid, FullPath.Buffer, FullPath.Length);
status = STATUS_ACCESS_DENIED;
}
}
PyasFree(Buffer);
}
PyasFree(Buffer);
}
}
}
+178 -96
View File
@@ -10,19 +10,21 @@ constexpr auto TRUST_CACHE_TTL_SEC = 300;
typedef struct _RANSOM_TRACKER {
PEPROCESS Process;
LARGE_INTEGER ProcessCreateTime;
ULONG ActivityCount;
LARGE_INTEGER LastActivityTime;
} RANSOM_TRACKER, * PRANSOM_TRACKER;
typedef struct _TRUST_CACHE_ENTRY {
PEPROCESS Process;
LARGE_INTEGER ProcessCreateTime;
BOOLEAN IsTrusted;
LARGE_INTEGER CacheTime;
} TRUST_CACHE_ENTRY, * PTRUST_CACHE_ENTRY;
static RANSOM_TRACKER RansomTrackers[MAX_TRACKERS];
static TRUST_CACHE_ENTRY TrustCache[TRUST_CACHE_SIZE];
static FAST_MUTEX TrustCacheLock;
static KSPIN_LOCK TrustCacheLock;
static BOOLEAN g_CacheInitialized = FALSE;
static ERESOURCE RuleLock;
@@ -107,7 +109,7 @@ static BOOLEAN HasSuffix(PCUNICODE_STRING String, PCWSTR Suffix) {
}
BOOLEAN WildcardMatch(PCWSTR Pattern, PCWSTR String, USHORT StringLengthBytes) {
if (!Pattern || !String) return FALSE;
if (Pattern == NULL || String == NULL) return FALSE;
USHORT StringLenChars = StringLengthBytes / sizeof(WCHAR);
PCWSTR mp = NULL;
@@ -123,7 +125,7 @@ BOOLEAN WildcardMatch(PCWSTR Pattern, PCWSTR String, USHORT StringLengthBytes) {
Pattern++;
String++;
}
else if (mp) {
else if (mp != NULL) {
Pattern = mp;
String = cp++;
}
@@ -134,7 +136,8 @@ BOOLEAN WildcardMatch(PCWSTR Pattern, PCWSTR String, USHORT StringLengthBytes) {
while (*Pattern == L'*') {
Pattern++;
}
return !*Pattern;
return (*Pattern == L'\0') ? TRUE : FALSE;
}
static ULONG ProcessJsonUnescape(PWCHAR Buffer, ULONG LengthChars) {
@@ -202,17 +205,23 @@ static VOID ParseAndLoadRules(PCHAR JsonContent, ULONG ContentLength, PCSTR KeyN
while (Ptr < End) {
SkipWhitespace(&Ptr, End);
if (Ptr >= End) break;
if (*Ptr == ']') {
Ptr++;
if (Ptr >= End || *Ptr == ']') {
if (Ptr < End) Ptr++;
return;
}
if (*Ptr == '"') {
PCHAR StartQuote = ++Ptr;
BOOLEAN Escaped = FALSE;
while (Ptr < End) {
if (*Ptr == '"' && *(Ptr - 1) != '\\') break;
if (*Ptr == '"' && Escaped == FALSE) break;
if (*Ptr == '\\') {
Escaped = (Escaped == FALSE) ? TRUE : FALSE;
}
else {
Escaped = FALSE;
}
Ptr++;
}
@@ -242,9 +251,6 @@ static VOID ParseAndLoadRules(PCHAR JsonContent, ULONG ContentLength, PCSTR KeyN
Ptr++;
}
}
else if (*Ptr == ',') {
Ptr++;
}
else {
Ptr++;
}
@@ -259,7 +265,7 @@ static VOID ParseAndLoadRules(PCHAR JsonContent, ULONG ContentLength, PCSTR KeyN
VOID InitializeRulesEngine() {
ExInitializeResourceLite(&RuleLock);
ExInitializeFastMutex(&TrustCacheLock);
KeInitializeSpinLock(&TrustCacheLock);
RtlZeroMemory(TrustCache, sizeof(TrustCache));
g_CacheInitialized = TRUE;
}
@@ -298,25 +304,33 @@ static VOID RemoveRule(PRULE_NODE* Head, PUNICODE_STRING RuleStr) {
}
VOID AddDynamicWhitelist(PUNICODE_STRING RuleStr) {
KeEnterCriticalRegion();
ExAcquireResourceExclusiveLite(&RuleLock, TRUE);
AddRule(&g_ProcessTrustedPaths, RuleStr);
AddRule(&g_FileExceptionPaths, RuleStr);
ExAcquireFastMutex(&TrustCacheLock);
KIRQL OldIrql;
KeAcquireSpinLock(&TrustCacheLock, &OldIrql);
RtlZeroMemory(TrustCache, sizeof(TrustCache));
ExReleaseFastMutex(&TrustCacheLock);
KeReleaseSpinLock(&TrustCacheLock, OldIrql);
ExReleaseResourceLite(&RuleLock);
KeLeaveCriticalRegion();
}
VOID RemoveDynamicWhitelist(PUNICODE_STRING RuleStr) {
KeEnterCriticalRegion();
ExAcquireResourceExclusiveLite(&RuleLock, TRUE);
RemoveRule(&g_ProcessTrustedPaths, RuleStr);
RemoveRule(&g_FileExceptionPaths, RuleStr);
ExAcquireFastMutex(&TrustCacheLock);
KIRQL OldIrql;
KeAcquireSpinLock(&TrustCacheLock, &OldIrql);
RtlZeroMemory(TrustCache, sizeof(TrustCache));
ExReleaseFastMutex(&TrustCacheLock);
KeReleaseSpinLock(&TrustCacheLock, OldIrql);
ExReleaseResourceLite(&RuleLock);
KeLeaveCriticalRegion();
}
NTSTATUS LoadRulesFromDisk(PUNICODE_STRING RegistryPath) {
@@ -359,58 +373,58 @@ NTSTATUS LoadRulesFromDisk(PUNICODE_STRING RegistryPath) {
}
if (Info->Type == REG_EXPAND_SZ || Info->Type == REG_SZ) {
PathBufferSize = ResultLength + 1024;
PathBufferSize = Info->DataLength + 1024;
PathBuffer = (PWCHAR)PyasAllocate(PathBufferSize);
if (PathBuffer) {
RtlZeroMemory(PathBuffer, PathBufferSize);
PWCHAR Src = (PWCHAR)Info->Data;
if (Info->DataLength > 0) {
RtlCopyMemory(PathBuffer, Info->Data, Info->DataLength);
}
if (NT_SUCCESS(RtlStringCbCopyW(PathBuffer, PathBufferSize, Src))) {
PWCHAR LastSlash = NULL;
PWCHAR Current = PathBuffer;
while (*Current) {
if (*Current == L'\\') LastSlash = Current;
Current++;
}
PWCHAR LastSlash = NULL;
PWCHAR Current = PathBuffer;
while (*Current) {
if (*Current == L'\\') LastSlash = Current;
Current++;
}
if (LastSlash) {
*LastSlash = L'\0';
if (LastSlash) {
*LastSlash = L'\0';
SIZE_T CurrentPathLen = wcslen(PathBuffer);
const WCHAR FilterSuffix[] = L"\\Filter";
SIZE_T FilterLen = (sizeof(FilterSuffix) / sizeof(WCHAR)) - 1;
SIZE_T CurrentPathLen = wcslen(PathBuffer);
const WCHAR FilterSuffix[] = L"\\Filter";
SIZE_T FilterLen = (sizeof(FilterSuffix) / sizeof(WCHAR)) - 1;
if (CurrentPathLen >= FilterLen) {
PWCHAR SuffixStart = PathBuffer + CurrentPathLen - FilterLen;
BOOLEAN Match = TRUE;
for (SIZE_T i = 0; i < FilterLen; i++) {
if (RtlDowncaseUnicodeChar(SuffixStart[i]) != RtlDowncaseUnicodeChar(FilterSuffix[i])) {
Match = FALSE;
break;
}
}
if (Match) {
*SuffixStart = L'\0';
if (CurrentPathLen >= FilterLen) {
PWCHAR SuffixStart = PathBuffer + CurrentPathLen - FilterLen;
BOOLEAN Match = TRUE;
for (SIZE_T i = 0; i < FilterLen; i++) {
if (RtlDowncaseUnicodeChar(SuffixStart[i]) != RtlDowncaseUnicodeChar(FilterSuffix[i])) {
Match = FALSE;
break;
}
}
RtlStringCbCatW(PathBuffer, PathBufferSize, L"\\Rules\\Rules_Driver_P1.json");
if (wcsncmp(PathBuffer, L"\\??\\", 4) != 0 &&
wcsncmp(PathBuffer, L"\\SystemRoot", 11) != 0 &&
wcsncmp(PathBuffer, L"\\DosDevices\\", 12) != 0) {
PWCHAR TmpBuffer = (PWCHAR)PyasAllocate(PathBufferSize + 16);
if (TmpBuffer) {
RtlStringCbCopyW(TmpBuffer, PathBufferSize + 16, L"\\??\\");
RtlStringCbCatW(TmpBuffer, PathBufferSize + 16, PathBuffer);
PyasFree(PathBuffer);
PathBuffer = TmpBuffer;
}
if (Match) {
*SuffixStart = L'\0';
}
RtlInitUnicodeString(&FinalPath, PathBuffer);
}
RtlStringCbCatW(PathBuffer, PathBufferSize, L"\\Rules\\Rules_Driver_P1.json");
if (wcsncmp(PathBuffer, L"\\??\\", 4) != 0 &&
wcsncmp(PathBuffer, L"\\SystemRoot", 11) != 0 &&
wcsncmp(PathBuffer, L"\\DosDevices\\", 12) != 0) {
PWCHAR TmpBuffer = (PWCHAR)PyasAllocate(PathBufferSize + 16);
if (TmpBuffer) {
RtlStringCbCopyW(TmpBuffer, PathBufferSize + 16, L"\\??\\");
RtlStringCbCatW(TmpBuffer, PathBufferSize + 16, PathBuffer);
PyasFree(PathBuffer);
PathBuffer = TmpBuffer;
}
}
RtlInitUnicodeString(&FinalPath, PathBuffer);
}
}
}
@@ -435,6 +449,7 @@ NTSTATUS LoadRulesFromDisk(PUNICODE_STRING RegistryPath) {
if (NT_SUCCESS(status)) {
((PCHAR)FileBuffer)[FileInfo.EndOfFile.LowPart] = '\0';
KeEnterCriticalRegion();
ExAcquireResourceExclusiveLite(&RuleLock, TRUE);
ParseAndLoadRules((PCHAR)FileBuffer, FileInfo.EndOfFile.LowPart, "Rule_Registry_BlockList", &g_RegistryBlockList);
ParseAndLoadRules((PCHAR)FileBuffer, FileInfo.EndOfFile.LowPart, "Rule_Registry_TrustedList", &g_RegistryTrustedList);
@@ -444,6 +459,7 @@ NTSTATUS LoadRulesFromDisk(PUNICODE_STRING RegistryPath) {
ParseAndLoadRules((PCHAR)FileBuffer, FileInfo.EndOfFile.LowPart, "Rule_File_ExceptionPaths", &g_FileExceptionPaths);
ParseAndLoadRules((PCHAR)FileBuffer, FileInfo.EndOfFile.LowPart, "Rule_File_RansomwareExtensions", &g_FileRansomExts);
ExReleaseResourceLite(&RuleLock);
KeLeaveCriticalRegion();
}
PyasFree(FileBuffer);
}
@@ -456,6 +472,7 @@ NTSTATUS LoadRulesFromDisk(PUNICODE_STRING RegistryPath) {
}
VOID UnloadRules() {
KeEnterCriticalRegion();
ExAcquireResourceExclusiveLite(&RuleLock, TRUE);
FreeList(&g_RegistryBlockList);
FreeList(&g_RegistryTrustedList);
@@ -465,9 +482,14 @@ VOID UnloadRules() {
FreeList(&g_FileExceptionPaths);
FreeList(&g_FileRansomExts);
ExReleaseResourceLite(&RuleLock);
KeLeaveCriticalRegion();
}
NTSTATUS GetProcessImageName(HANDLE ProcessId, PUNICODE_STRING* ImageName) {
if (KeGetCurrentIrql() != PASSIVE_LEVEL) {
return STATUS_UNSUCCESSFUL;
}
NTSTATUS status;
PEPROCESS Process = NULL;
@@ -492,33 +514,42 @@ BOOLEAN IsProcessTrusted(HANDLE ProcessId) {
if ((ULONG)(ULONG_PTR)ProcessId == GlobalData.PyasPid) return TRUE;
if (ProcessId == (HANDLE)4) return TRUE;
if (KeGetCurrentIrql() != PASSIVE_LEVEL) return FALSE;
PEPROCESS Process = NULL;
if (!NT_SUCCESS(PsLookupProcessByProcessId(ProcessId, &Process))) {
return FALSE;
}
LARGE_INTEGER createTime;
createTime.QuadPart = PsGetProcessCreateTimeQuadPart(Process);
if (g_CacheInitialized) {
ExAcquireFastMutex(&TrustCacheLock);
KIRQL OldIrql;
KeAcquireSpinLock(&TrustCacheLock, &OldIrql);
ULONG Hash = ((ULONG)(ULONG_PTR)ProcessId) & (TRUST_CACHE_SIZE - 1);
if (TrustCache[Hash].Process == Process) {
if (TrustCache[Hash].Process == Process && TrustCache[Hash].ProcessCreateTime.QuadPart == createTime.QuadPart) {
LARGE_INTEGER Now;
KeQuerySystemTime(&Now);
if ((Now.QuadPart - TrustCache[Hash].CacheTime.QuadPart) < (TRUST_CACHE_TTL_SEC * 10000000LL)) {
BOOLEAN cachedResult = TrustCache[Hash].IsTrusted;
ExReleaseFastMutex(&TrustCacheLock);
KeReleaseSpinLock(&TrustCacheLock, OldIrql);
ObDereferenceObject(Process);
return cachedResult;
}
}
ExReleaseFastMutex(&TrustCacheLock);
KeReleaseSpinLock(&TrustCacheLock, OldIrql);
}
if (KeGetCurrentIrql() != PASSIVE_LEVEL) {
ObDereferenceObject(Process);
return FALSE;
}
PUNICODE_STRING imageFileName = NULL;
NTSTATUS status = GetProcessImageName(ProcessId, &imageFileName);
BOOLEAN isTrusted = FALSE;
KeEnterCriticalRegion();
ExAcquireResourceSharedLite(&RuleLock, TRUE);
if (NT_SUCCESS(status) && imageFileName && imageFileName->Buffer) {
@@ -547,15 +578,18 @@ BOOLEAN IsProcessTrusted(HANDLE ProcessId) {
update_cache:
if (g_CacheInitialized) {
ExAcquireFastMutex(&TrustCacheLock);
KIRQL OldIrql;
KeAcquireSpinLock(&TrustCacheLock, &OldIrql);
ULONG Hash = ((ULONG)(ULONG_PTR)ProcessId) & (TRUST_CACHE_SIZE - 1);
TrustCache[Hash].Process = Process;
TrustCache[Hash].ProcessCreateTime = createTime;
TrustCache[Hash].IsTrusted = isTrusted;
KeQuerySystemTime(&TrustCache[Hash].CacheTime);
ExReleaseFastMutex(&TrustCacheLock);
KeReleaseSpinLock(&TrustCacheLock, OldIrql);
}
ExReleaseResourceLite(&RuleLock);
KeLeaveCriticalRegion();
if (imageFileName) ExFreePool(imageFileName);
ObDereferenceObject(Process);
@@ -575,12 +609,14 @@ BOOLEAN CheckRegistryRule(PCUNICODE_STRING KeyName) {
return FALSE;
}
KeEnterCriticalRegion();
ExAcquireResourceSharedLite(&RuleLock, TRUE);
PRULE_NODE AllowNode = g_RegistryTrustedList;
while (AllowNode) {
if (WildcardMatch(AllowNode->Pattern.Buffer, KeyName->Buffer, KeyName->Length)) {
ExReleaseResourceLite(&RuleLock);
KeLeaveCriticalRegion();
return FALSE;
}
AllowNode = AllowNode->Next;
@@ -596,18 +632,21 @@ BOOLEAN CheckRegistryRule(PCUNICODE_STRING KeyName) {
Node = Node->Next;
}
ExReleaseResourceLite(&RuleLock);
KeLeaveCriticalRegion();
return Match;
}
BOOLEAN CheckFileExtensionRule(PCUNICODE_STRING FileName) {
if (!FileName || !FileName->Buffer) return FALSE;
KeEnterCriticalRegion();
ExAcquireResourceSharedLite(&RuleLock, TRUE);
PRULE_NODE ExNode = g_FileExceptionPaths;
while (ExNode) {
if (WildcardMatch(ExNode->Pattern.Buffer, FileName->Buffer, FileName->Length)) {
ExReleaseResourceLite(&RuleLock);
KeLeaveCriticalRegion();
return FALSE;
}
ExNode = ExNode->Next;
@@ -623,6 +662,7 @@ BOOLEAN CheckFileExtensionRule(PCUNICODE_STRING FileName) {
Node = Node->Next;
}
ExReleaseResourceLite(&RuleLock);
KeLeaveCriticalRegion();
return Match;
}
@@ -633,12 +673,14 @@ BOOLEAN CheckProtectedPathRule(PCUNICODE_STRING FileName) {
return FALSE;
}
KeEnterCriticalRegion();
ExAcquireResourceSharedLite(&RuleLock, TRUE);
PRULE_NODE ExNode = g_FileExceptionPaths;
while (ExNode) {
if (WildcardMatch(ExNode->Pattern.Buffer, FileName->Buffer, FileName->Length)) {
ExReleaseResourceLite(&RuleLock);
KeLeaveCriticalRegion();
return FALSE;
}
ExNode = ExNode->Next;
@@ -654,13 +696,16 @@ BOOLEAN CheckProtectedPathRule(PCUNICODE_STRING FileName) {
Node = Node->Next;
}
ExReleaseResourceLite(&RuleLock);
KeLeaveCriticalRegion();
return Match;
}
static BOOLEAN IsHighEntropy(PVOID Buffer, ULONG Length) {
if (!Buffer || Length < 256) return FALSE;
ULONG Histogram[256] = { 0 };
PUSHORT Histogram = (PUSHORT)PyasAllocate(256 * sizeof(USHORT));
if (!Histogram) return FALSE;
PUCHAR Ptr = (PUCHAR)Buffer;
ULONG ScanLen = (Length > 1024) ? 1024 : Length;
@@ -670,16 +715,20 @@ static BOOLEAN IsHighEntropy(PVOID Buffer, ULONG Length) {
}
}
__except (EXCEPTION_EXECUTE_HANDLER) {
PyasFree(Histogram);
return FALSE;
}
ULONG MaxFreq = 0;
USHORT MaxFreq = 0;
for (int i = 0; i < 256; i++) {
if (Histogram[i] > MaxFreq) MaxFreq = Histogram[i];
}
ULONG ExpectedAvg = ScanLen / 256;
return (MaxFreq < (ExpectedAvg + HIGH_ENTROPY_THRESHOLD));
BOOLEAN Result = (MaxFreq < (ExpectedAvg + HIGH_ENTROPY_THRESHOLD));
PyasFree(Histogram);
return Result;
}
static BOOLEAN IsNaturallyCompressed(PCUNICODE_STRING FileName) {
@@ -707,19 +756,33 @@ static BOOLEAN IsNoisyRansomPath(PCUNICODE_STRING FileName) {
return FALSE;
}
static BOOLEAN IsUserDirectoryPath(PCUNICODE_STRING FileName) {
if (!FileName || !FileName->Buffer) return FALSE;
return WildcardMatch(L"*\\Users\\*", FileName->Buffer, FileName->Length);
}
static BOOLEAN IsExplorerProcess(HANDLE ProcessId) {
if (KeGetCurrentIrql() != PASSIVE_LEVEL) return FALSE;
if (KeGetCurrentIrql() > APC_LEVEL) return FALSE;
PUNICODE_STRING imageFileName = NULL;
NTSTATUS status = GetProcessImageName(ProcessId, &imageFileName);
BOOLEAN isExplorer = FALSE;
if (NT_SUCCESS(status) && imageFileName && imageFileName->Buffer) {
if (HasSuffix(imageFileName, L"\\explorer.exe")) isExplorer = TRUE;
if (NT_SUCCESS(status) && imageFileName) {
if (imageFileName->Buffer && HasSuffix(imageFileName, L"\\explorer.exe")) {
isExplorer = TRUE;
}
ExFreePool(imageFileName);
}
return isExplorer;
}
BOOLEAN CheckRansomActivity(HANDLE ProcessId, PUNICODE_STRING FileName, PVOID WriteBuffer, ULONG WriteLength, BOOLEAN IsWrite) {
if (!IsUserDirectoryPath(FileName)) {
return FALSE;
}
if (IsNoisyRansomPath(FileName)) {
return FALSE;
}
@@ -746,9 +809,12 @@ BOOLEAN CheckRansomActivity(HANDLE ProcessId, PUNICODE_STRING FileName, PVOID Wr
LARGE_INTEGER Now = { 0 };
KeQuerySystemTime(&Now);
LARGE_INTEGER createTime;
createTime.QuadPart = PsGetProcessCreateTimeQuadPart(Process);
BOOLEAN Result = FALSE;
ExAcquireFastMutex(&GlobalData.TrackerMutex);
KIRQL OldIrql;
KeAcquireSpinLock(&GlobalData.TrackerMutex, &OldIrql);
PRANSOM_TRACKER Tracker = NULL;
PRANSOM_TRACKER CandidateSlot = NULL;
@@ -761,7 +827,7 @@ BOOLEAN CheckRansomActivity(HANDLE ProcessId, PUNICODE_STRING FileName, PVOID Wr
LruSlot = Current;
}
if (Current->Process == Process) {
if (Current->Process == Process && Current->ProcessCreateTime.QuadPart == createTime.QuadPart) {
Tracker = Current;
break;
}
@@ -789,6 +855,7 @@ BOOLEAN CheckRansomActivity(HANDLE ProcessId, PUNICODE_STRING FileName, PVOID Wr
}
Tracker->Process = Process;
Tracker->ProcessCreateTime = createTime;
Tracker->ActivityCount = 0;
Tracker->LastActivityTime = Now;
}
@@ -808,7 +875,8 @@ BOOLEAN CheckRansomActivity(HANDLE ProcessId, PUNICODE_STRING FileName, PVOID Wr
if (Tracker->ActivityCount >= RANSOM_COUNT_THRESHOLD) {
Result = TRUE;
}
ExReleaseFastMutex(&GlobalData.TrackerMutex);
KeReleaseSpinLock(&GlobalData.TrackerMutex, OldIrql);
ObDereferenceObject(Process);
return Result;
}
@@ -816,34 +884,48 @@ BOOLEAN CheckRansomActivity(HANDLE ProcessId, PUNICODE_STRING FileName, PVOID Wr
NTSTATUS SendMessageToUser(ULONG Code, ULONG Pid, PWCHAR Path, USHORT PathSize) {
if (KeGetCurrentIrql() > APC_LEVEL) return STATUS_UNSUCCESSFUL;
PYAS_MESSAGE msg;
RtlZeroMemory(&msg, sizeof(msg));
msg.MessageCode = Code;
msg.ProcessId = Pid;
if (Path && PathSize > 0) {
size_t MaxSize = sizeof(msg.Path) - sizeof(WCHAR);
size_t BytesToCopy = PathSize;
if (BytesToCopy > MaxSize) BytesToCopy = MaxSize;
__try {
RtlCopyMemory(msg.Path, Path, BytesToCopy);
msg.Path[BytesToCopy / sizeof(WCHAR)] = L'\0';
}
__except (EXCEPTION_EXECUTE_HANDLER) {
return STATUS_ACCESS_VIOLATION;
}
}
if (!ExAcquireRundownProtection(&GlobalData.PortRundown)) {
return STATUS_PORT_DISCONNECTED;
}
NTSTATUS status = STATUS_PORT_DISCONNECTED;
if (GlobalData.ClientPort) {
LARGE_INTEGER timeout;
timeout.QuadPart = -(5 * 10000);
status = FltSendMessage(GlobalData.FilterHandle, &GlobalData.ClientPort, &msg, sizeof(msg), NULL, NULL, &timeout);
PPYAS_MESSAGE msg = (PPYAS_MESSAGE)PyasAllocate(sizeof(PYAS_MESSAGE));
if (msg) {
msg->MessageCode = Code;
msg->ProcessId = Pid;
if (Path && PathSize > 0) {
size_t MaxSize = sizeof(msg->Path) - sizeof(WCHAR);
size_t BytesToCopy = PathSize > MaxSize ? MaxSize : PathSize;
__try {
RtlCopyMemory(msg->Path, Path, BytesToCopy);
msg->Path[BytesToCopy / sizeof(WCHAR)] = L'\0';
}
__except (EXCEPTION_EXECUTE_HANDLER) {
PyasFree(msg);
ExReleaseRundownProtection(&GlobalData.PortRundown);
return STATUS_ACCESS_VIOLATION;
}
}
LARGE_INTEGER timeout;
if (KeGetCurrentIrql() == PASSIVE_LEVEL) {
timeout.QuadPart = -(5 * 10000);
}
else {
timeout.QuadPart = 0;
}
status = FltSendMessage(GlobalData.FilterHandle, &GlobalData.ClientPort, msg, sizeof(PYAS_MESSAGE), NULL, NULL, &timeout);
PyasFree(msg);
}
else {
status = STATUS_INSUFFICIENT_RESOURCES;
}
}
ExReleaseRundownProtection(&GlobalData.PortRundown);
+4 -4
View File
@@ -9,11 +9,11 @@
#define VER_ORIGINALFILENAME_STR "PYAS_Driver.sys"
#define VER_PRODUCTNAME_STR "PYAS"
#define VER_FILEVERSION 3,5.0,0
#define VER_FILEVERSION_STR "3,5.0,0"
#define VER_FILEVERSION 3,5,1,0
#define VER_FILEVERSION_STR "3,5,1,0"
#define VER_PRODUCTVERSION 3,5.0,0
#define VER_PRODUCTVERSION_STR "3,5.0,0"
#define VER_PRODUCTVERSION 3,5,1,0
#define VER_PRODUCTVERSION_STR "3,5,1,0"
#define VER_COMPANYNAME_STR "PYAS Security"
+16 -6
View File
@@ -2,7 +2,6 @@
"Rule_Registry_BlockList": [
"\\REGISTRY\\MACHINE\\BCD00000000",
"\\REGISTRY\\MACHINE\\BCD00000000\\*",
"*\\Image File Execution Options\\*",
"*\\Classes\\.exe",
"*\\Classes\\.bat",
@@ -47,7 +46,6 @@
"*\\DisableTaskMgr",
"*\\DisableRegistryTools",
"*\\DisableChangePassword",
"*\\Wallpaper",
"*\\NoComponents",
"*\\NoStartMenuMorePrograms",
"*\\NoActiveDesktop",
@@ -84,7 +82,6 @@
"*\\SystemBiosVersion",
"*\\NoTheme",
"*\\SystemSetupInProgress",
"*\\CmdLine",
"*\\SetupType",
"*\\AutoReboot",
"*\\Userinit",
@@ -107,24 +104,36 @@
"\\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",
"*\\hrkill-*.exe",
"*\\bin\\HipsDaemon.exe",
"*\\360Safe.exe",
"*\\bilibili\\bilibili.exe",
"*\\PC4399_WPFLauncher.exe",
"*\\bin\\APlayer.exe",
"*\\Kingsoft\\WPS Office\\ksolaunch.exe",
"*\\Neverness To Everness\\NTELauncher.exe",
"*\\Application\\AndrowsLauncher.exe",
],
"Rule_Process_ExploitableBlacklist": [
"*\\cmd.exe",
@@ -163,9 +172,9 @@
],
"Rule_File_ProtectedPaths": [
"*\\PYAS.exe",
"*\\PYAS_Driver.sys",
"*\\ProgramData\\PYAS\\*.json",
"*\\Plugins\\Filter\\PYAS_Driver.sys",
"*\\Plugins\\Rules\\Rules_Driver_P1.json",
"*\\ProgramData\\PYAS\\*.json",
"*\\Engine\\Pattern\\*.onnx",
"*\\Engine\\Properties\\*.onnx",
@@ -203,6 +212,7 @@
"*\\ventoy\\Boot\\*",
"*\\AdminRegionConfig.*",
"*\\nt_qq\\nt_*\\*",
"*\\cache\\*",
],
"Rule_File_RansomwareExtensions": [
".zip", ".7z", ".rar", ".tar", ".gz",
+9 -11
View File
@@ -1,6 +1,6 @@
# PYAS
# PYAS Security
Antivirus software written in Python and C++ that blocks threats through deep learning and behavioral monitoring!
Antivirus software written in Python and C++ that blocks threats through Machine Learning and behavioral monitoring!
<img width="2440" height="1600" alt="PYAS_All" src="https://github.com/user-attachments/assets/1991aad3-64cc-4266-ac67-dab70c891ce7" />
@@ -24,10 +24,9 @@ Non-essential requirements installation, only used for model training or other f
```
pip install pandas
pip install scikit-learn
pip install tensorflow
pip install tf2onnx
pip install lightgbm
pip install onnxmltools
pip install orjson
```
## File Information
@@ -36,15 +35,9 @@ The following lists the storage locations of all relevant code and other related
```
PYAS/
├── Embed/ # Deploy an embedded command line antivirus engine
│ ├── Dockerfile
│ └── ...
│
├── Engine/ # Complete code for Yara signatures and AI model training
│ ├── Heuristic/
│ │ └── ... # Yara Rules
│ ├── Pattern/
│ │ └── ... # CNN Model
│ └── Properties/
│ └── ... # LightGBM Model
│
@@ -84,7 +77,6 @@ graph TD
subgraph Engines ["Scanning Engines (PYAS_Engine.py)"]
PEScanner["PE/ML Scanner (pe_scanner)"]:::userMode
CNNScanner["CNN Scanner (cnn_scanner)"]:::userMode
RuleScanner["YARA Scanner (rule_scanner)"]:::userMode
CloudScanner["Cloud API (cloud_scanner)"]:::userMode
SignScanner["Signature Verify (sign_scanner)"]:::userMode
@@ -136,6 +128,12 @@ graph TD
BootProtect -->|IOCTL_DISK_FORMAT_TRACKS...| Disk
```
## Machine Learning
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" />
## Support System
Please ensure your computer supports Microsoft Visual C++ 2015-2022 Redistributable and Edge Webview2.