Add files via upload
This commit is contained in:
87owo authored and GitHub committed
1 parent
ddcbe1cd51
commit
533e7dfd15
6 files changed
+2399
No files matched your search
File diff suppressed because it is too large.
Load diff
@@ -0,0 +1,55 @@
|
||||
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')
|
||||
@@ -0,0 +1,34 @@
|
||||
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')
|
||||
@@ -0,0 +1,163 @@
|
||||
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.00001, 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]
|
||||
num_classes = len(class_indices)
|
||||
|
||||
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_classes), 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_classes), 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, num_classes):
|
||||
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.one_hot(label, depth=num_classes)
|
||||
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, num_classes):
|
||||
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(num_classes, activation='softmax', 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 categorical_focal_loss(gamma=2.0, alpha=0.25):
|
||||
def focal_loss(y_true, y_pred):
|
||||
y_pred = tf.clip_by_value(y_pred, tf.keras.backend.epsilon(), 1.0 - tf.keras.backend.epsilon())
|
||||
cross_entropy = -y_true * tf.math.log(y_pred)
|
||||
loss = alpha * tf.math.pow(1.0 - y_pred, gamma) * cross_entropy
|
||||
return tf.reduce_sum(loss, axis=-1)
|
||||
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,) + train_ds.image_shape, 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')
|
||||
print("Load model from disk")
|
||||
except:
|
||||
print("Creating a new model")
|
||||
model = build_model(train_ds.image_shape, len(train_ds.class_indices))
|
||||
|
||||
print(f"Total parameters: {model.count_params()}")
|
||||
model.compile(optimizer='adam', loss=categorical_focal_loss(gamma=2.0, alpha=0.25),
|
||||
metrics=['accuracy', 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)
|
||||
input('Training Complete')
|
||||
@@ -0,0 +1,513 @@
|
||||
import os, sys, time, ctypes, sqlite3, datetime, pefile, math
|
||||
import numpy as np
|
||||
import ctypes.wintypes
|
||||
|
||||
####################################################################################################
|
||||
|
||||
DB_PATH = "pe_features.db"
|
||||
BATCH_SIZE = 100
|
||||
MAX_FILE_SIZE = 256 * 1024 * 1024
|
||||
TARGET_EXTENSIONS = {
|
||||
'.exe', '.dll', '.sys', '.ocx', '.scr', '.efi', '.acm', '.ax', '.cpl', '.drv', '.com'
|
||||
}
|
||||
|
||||
####################################################################################################
|
||||
|
||||
class Features:
|
||||
SCHEMA = [
|
||||
"FileEntropy", "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",
|
||||
|
||||
"TextSection", "TextSizeRatio", "DataSection", "DataSizeRatio",
|
||||
"RsrcSection", "RsrcSizeRatio", "SectionCount", "ExecutableSections",
|
||||
"WritableSections", "ReadableSections", "SectionException",
|
||||
"IconCount", "ApiCount", "ExportCount", "DebugCount", "ExceptionCount",
|
||||
"FileDescriptionLength", "FileVersionLength", "ProductNameLength",
|
||||
"ProductVersionLength", "CompanyNameLength", "LegalCopyrightLength",
|
||||
"CommentsLength", "InternalNameLength", "LegalTrademarksLength",
|
||||
"SpecialBuildLength", "PrivateBuildLength",
|
||||
|
||||
"TrustSigned", "IsDebug", "IsPatched", "IsPrivateBuild", "IsPreRelease",
|
||||
"IsSpecialBuild", "IsAdmin", "IsInstall", "HasTlsCallbacks",
|
||||
"HasInvalidTimestamp", "HasRelocationDirectory", "HasPacked", "FileTimeException",
|
||||
|
||||
"_CorExeMain", "_CorDllMain",
|
||||
|
||||
"GetDC", "CreateDCW", "CreateDCA", "BitBlt", "StretchBlt", "CreateCompatibleDC",
|
||||
"CreateCompatibleBitmap", "SelectObject", "DeleteDC", "DeleteObject", "GetDeviceCaps",
|
||||
"GetSystemMetrics",
|
||||
|
||||
"SetWindowsHookExA", "SetWindowsHookExW", "UnhookWindowsHookEx",
|
||||
"GetAsyncKeyState", "GetKeyState", "GetKeyboardState",
|
||||
"MapVirtualKeyA", "MapVirtualKeyW", "MapVirtualKeyExA", "MapVirtualKeyExW",
|
||||
"ToAscii", "ToAsciiEx", "ToUnicode", "ToUnicodeEx", "KeybdEvent", "SendInput",
|
||||
|
||||
"GetCursorPos", "SetCursorPos", "MouseEvent", "GetDoubleClickTime",
|
||||
"GetCapture", "SetCapture",
|
||||
|
||||
"WaveInOpen", "WaveInClose", "WaveInStart", "WaveInStop",
|
||||
"CapCreateCaptureWindowA", "CapCreateCaptureWindowW",
|
||||
"OpenClipboard", "CloseClipboard", "GetClipboardData", "SetClipboardData",
|
||||
|
||||
"GetDesktopWindow", "GetForegroundWindow", "GetWindowDC", "GetWindowRect",
|
||||
"GetClientRect", "PrintWindow", "SetWindowDisplayAffinity", "FindWindowA", "FindWindowW",
|
||||
"EnumWindows", "EnumChildWindows",
|
||||
|
||||
"Direct3DCreate9", "D3D11CreateDevice", "GdipSaveImageToFile",
|
||||
|
||||
"CreateProcessA", "CreateProcessW", "CreateThread", "ResumeThread", "SuspendThread",
|
||||
"OpenProcess", "TerminateProcess", "GetCurrentProcess", "GetCurrentProcessId",
|
||||
"ExitProcess", "WinExec", "ShellExecuteA", "ShellExecuteW",
|
||||
|
||||
"CreateRemoteThread", "QueueUserAPC", "VirtualAlloc", "VirtualAllocEx",
|
||||
"VirtualProtect", "VirtualProtectEx", "WriteProcessMemory", "ReadProcessMemory",
|
||||
"NtUnmapViewOfSection", "SetThreadContext", "GetThreadContext", "Wow64SetThreadContext",
|
||||
"ZwUnmapViewOfSection", "CreateFileMappingA", "CreateFileMappingW", "MapViewOfFile",
|
||||
"UnmapViewOfSection",
|
||||
|
||||
"LoadLibraryA", "LoadLibraryW", "LoadLibraryExA", "LoadLibraryExW",
|
||||
"GetProcAddress", "GetModuleHandleA", "GetModuleHandleW", "FreeLibrary",
|
||||
|
||||
"CreateFileA", "CreateFileW", "WriteFile", "ReadFile", "DeleteFileA", "DeleteFileW",
|
||||
"CopyFileA", "CopyFileW", "MoveFileA", "MoveFileW", "FindFirstFileA", "FindNextFileA",
|
||||
"GetTempPathA", "GetTempPathW", "GetTempFileNameA", "GetTempFileNameW",
|
||||
"SetFileAttributesA", "SetFileAttributesW", "DeviceIoControl", "SetFileTime",
|
||||
|
||||
"RegOpenKeyExA", "RegOpenKeyExW", "RegSetValueExA", "RegSetValueExW",
|
||||
"RegCreateKeyExA", "RegCreateKeyExW", "RegDeleteKeyA", "RegDeleteKeyW",
|
||||
"RegEnumValueA", "RegEnumValueW", "RegQueryValueExA", "RegQueryValueExW",
|
||||
"RegDeleteValueA", "RegDeleteValueW",
|
||||
|
||||
"OpenSCManagerA", "OpenSCManagerW", "CreateServiceA", "CreateServiceW",
|
||||
"StartServiceA", "StartServiceW", "ControlService", "DeleteService",
|
||||
|
||||
"AdjustTokenPrivileges", "LookupPrivilegeValueA", "LookupPrivilegeValueW",
|
||||
"OpenProcessToken", "NetUserAdd", "NetLocalGroupAddMembers",
|
||||
|
||||
"Socket", "Connect", "Send", "Recv", "WSAStartup", "gethostbyname", "getaddrinfo",
|
||||
"WSAIoctl", "WSASocketA", "WSASocketW",
|
||||
"URLDownloadToFileA", "URLDownloadToFileW",
|
||||
"InternetOpenA", "InternetOpenW", "InternetConnectA", "InternetConnectW",
|
||||
"InternetOpenUrlA", "InternetOpenUrlW",
|
||||
"HttpOpenRequestA", "HttpOpenRequestW", "HttpSendRequestA", "HttpSendRequestW",
|
||||
"InternetReadFile", "DnsQuery_A", "DnsQuery_W",
|
||||
|
||||
"IsDebuggerPresent", "CheckRemoteDebuggerPresent", "OutputDebugStringA", "OutputDebugStringW",
|
||||
"GetTickCount", "QueryPerformanceCounter", "Sleep", "GetSystemTimeAsFileTime",
|
||||
"GetLocalTime", "GlobalMemoryStatus", "GetVersionExA", "GetVersionExW",
|
||||
"GetComputerNameA", "GetComputerNameW", "GetUserNameA", "GetUserNameW",
|
||||
|
||||
"CryptAcquireContextA", "CryptAcquireContextW", "CryptCreateHash", "CryptHashData",
|
||||
"CryptDeriveKey", "CryptEncrypt", "CryptDecrypt", "CryptDestroyKey", "CryptDestroyHash",
|
||||
"CryptReleaseContext", "CryptGenKey", "CryptImportKey", "CryptExportKey",
|
||||
"RtlDecompressBuffer", "ConnectNamedPipe", "PeekNamedPipe"
|
||||
]
|
||||
|
||||
SCHEMA_SET = set(SCHEMA)
|
||||
PACKERS = {
|
||||
'upx', 'aspack', 'asprotect', 'pecompact', 'upack', 'fsg', 'mew',
|
||||
'mpress', 'ezip', 'pklt', 'shrink', 'petite', 'telock',
|
||||
|
||||
'themida', 'winlicense', 'tmd', 'vmp', 'enigma', 'obsidium',
|
||||
'pelock', 'exestealth', 'yoda', 'armadillo', 'zprotect',
|
||||
'sforce', 'starforce', 'qihoo', 'wisevec', 'megastop',
|
||||
}
|
||||
INSTALLER_SIGS = [b"Nullsoft.NSIS", b"Inno Setup", b"7-Zip.7zip", b"InstallShield"]
|
||||
|
||||
####################################################################################################
|
||||
|
||||
def verify_signature(file_path):
|
||||
if os.name != 'nt':
|
||||
return 0
|
||||
try:
|
||||
wintrust = ctypes.windll.wintrust
|
||||
WINTRUST_ACTION_GENERIC_VERIFY_V2 = ctypes.c_char_p(b'\x6b\xc5\xaa\x00\x44\xcd\xd0\x11\x8c\xc2\x00\xc0\x4f\xc2\x95\xee')
|
||||
|
||||
class WINTRUST_FILE_INFO(ctypes.Structure):
|
||||
_fields_ = [
|
||||
("cbStruct", ctypes.wintypes.DWORD),
|
||||
("pcwszFilePath", ctypes.wintypes.LPCWSTR),
|
||||
("hFile", ctypes.wintypes.HANDLE),
|
||||
("pgKnownSubject", ctypes.wintypes.LPVOID)
|
||||
]
|
||||
|
||||
class WINTRUST_DATA(ctypes.Structure):
|
||||
_fields_ = [
|
||||
("cbStruct", ctypes.wintypes.DWORD),
|
||||
("pPolicyCallbackData", ctypes.wintypes.LPVOID),
|
||||
("pSIPClientData", ctypes.wintypes.LPVOID),
|
||||
("dwUIChoice", ctypes.wintypes.DWORD),
|
||||
("fdwRevocationChecks", ctypes.wintypes.DWORD),
|
||||
("dwUnionChoice", ctypes.wintypes.DWORD),
|
||||
("pFile", ctypes.POINTER(WINTRUST_FILE_INFO)),
|
||||
("dwStateAction", ctypes.wintypes.DWORD),
|
||||
("hWVTStateData", ctypes.wintypes.HANDLE),
|
||||
("pwszURLReference", ctypes.wintypes.LPCWSTR),
|
||||
("dwProvFlags", ctypes.wintypes.DWORD),
|
||||
("dwUIContext", ctypes.wintypes.DWORD),
|
||||
("pSignatureSettings", ctypes.wintypes.LPVOID)
|
||||
]
|
||||
|
||||
fi = WINTRUST_FILE_INFO(ctypes.sizeof(WINTRUST_FILE_INFO), file_path, None, None)
|
||||
td = WINTRUST_DATA(ctypes.sizeof(WINTRUST_DATA), None, None, 2, 1, 1, ctypes.pointer(fi), 0, None, None, 0x00000080, 0, None)
|
||||
action = ctypes.wintypes.GUID.from_buffer_copy(WINTRUST_ACTION_GENERIC_VERIFY_V2)
|
||||
return 1 if wintrust.WinVerifyTrust(None, ctypes.byref(action), ctypes.byref(td)) == 0 else 0
|
||||
except:
|
||||
return 0
|
||||
|
||||
####################################################################################################
|
||||
|
||||
def safe_float(val):
|
||||
try:
|
||||
f = float(val)
|
||||
if math.isinf(f) or math.isnan(f):
|
||||
return 0.0
|
||||
return f
|
||||
except:
|
||||
return 0.0
|
||||
|
||||
def process_single_file(file_path):
|
||||
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 n / d if d > 0 else 0.0
|
||||
|
||||
try:
|
||||
fsize = os.path.getsize(file_path)
|
||||
if fsize == 0 or fsize > MAX_FILE_SIZE:
|
||||
return None
|
||||
except:
|
||||
return None
|
||||
|
||||
fts = {k: 0 for k in Features.SCHEMA}
|
||||
pe = None
|
||||
try:
|
||||
pe = pefile.PE(file_path, 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 if (fh.TimeDateStamp < 631152000 or fh.TimeDateStamp > curr_ts + 2592000) else 0
|
||||
|
||||
if hasattr(pe, 'OPTIONAL_HEADER'):
|
||||
op = pe.OPTIONAL_HEADER
|
||||
fts['Magic'] = op.Magic
|
||||
fts['MajorLinkerVersion'] = op.MajorLinkerVersion
|
||||
fts['MinorLinkerVersion'] = op.MinorLinkerVersion
|
||||
fts['SizeOfCode'] = op.SizeOfCode
|
||||
fts['SizeOfInitializedData'] = op.SizeOfInitializedData
|
||||
fts['SizeOfUninitializedData'] = op.SizeOfUninitializedData
|
||||
fts['AddressOfEntryPoint'] = op.AddressOfEntryPoint
|
||||
fts['BaseOfCode'] = op.BaseOfCode
|
||||
fts['ImageBase'] = op.ImageBase
|
||||
fts['SectionAlignment'] = op.SectionAlignment
|
||||
fts['FileAlignment'] = op.FileAlignment
|
||||
fts['MajorOperatingSystemVersion'] = op.MajorOperatingSystemVersion
|
||||
fts['MinorOperatingSystemVersion'] = op.MinorOperatingSystemVersion
|
||||
fts['MajorImageVersion'] = op.MajorImageVersion
|
||||
fts['MinorImageVersion'] = op.MinorImageVersion
|
||||
fts['MajorSubsystemVersion'] = op.MajorSubsystemVersion
|
||||
fts['MinorSubsystemVersion'] = op.MinorSubsystemVersion
|
||||
fts['SizeOfImage'] = op.SizeOfImage
|
||||
fts['SizeOfHeaders'] = op.SizeOfHeaders
|
||||
fts['CheckSum'] = op.CheckSum
|
||||
fts['Subsystem'] = op.Subsystem
|
||||
fts['DllCharacteristics'] = op.DllCharacteristics
|
||||
fts['SizeOfStackReserve'] = op.SizeOfStackReserve
|
||||
fts['SizeOfStackCommit'] = op.SizeOfStackCommit
|
||||
fts['SizeOfHeapReserve'] = op.SizeOfHeapReserve
|
||||
fts['SizeOfHeapCommit'] = op.SizeOfHeapCommit
|
||||
fts['LoaderFlags'] = op.LoaderFlags
|
||||
fts['NumberOfRvaAndSizes'] = op.NumberOfRvaAndSizes
|
||||
|
||||
fts['Is64Bit'] = 1 if fh.Machine == 0x8664 else 0
|
||||
fts['IsExe'] = 1 if pe.is_exe() else 0
|
||||
fts['IsDll'] = 1 if pe.is_dll() else 0
|
||||
fts['IsDriver'] = 1 if pe.is_driver() else 0
|
||||
|
||||
if len(op.DATA_DIRECTORY) > 9 and op.DATA_DIRECTORY[9].VirtualAddress > 0:
|
||||
fts['HasTlsCallbacks'] = 1
|
||||
if len(op.DATA_DIRECTORY) > 5 and op.DATA_DIRECTORY[5].Size > 0:
|
||||
fts['HasRelocationDirectory'] = 1
|
||||
|
||||
pe.parse_sections(pe.FILE_HEADER.NumberOfSections)
|
||||
|
||||
file_data = pe.get_memory_mapped_image()
|
||||
fts['FileEntropy'] = _calc_entropy(file_data)
|
||||
|
||||
fts['SectionCount'] = len(pe.sections)
|
||||
for section in pe.sections:
|
||||
try:
|
||||
name = section.Name.decode('ascii', 'ignore').strip('\x00')
|
||||
except:
|
||||
name = ""
|
||||
s_data = section.get_data()
|
||||
s_entropy = _calc_entropy(s_data)
|
||||
|
||||
if any(p in name.lower() for p in Features.PACKERS):
|
||||
fts['HasPacked'] = 1
|
||||
if section.Characteristics & 0x20000000:
|
||||
fts['ExecutableSections'] += 1
|
||||
if section.Characteristics & 0x80000000:
|
||||
fts['WritableSections'] += 1
|
||||
if section.Characteristics & 0x40000000:
|
||||
fts['ReadableSections'] += 1
|
||||
if section.SizeOfRawData + section.PointerToRawData > fsize:
|
||||
fts['SectionException'] = 1
|
||||
|
||||
ratio = _safe_div(section.SizeOfRawData, fsize)
|
||||
if name.startswith('.text'):
|
||||
fts['TextSection'] = s_entropy
|
||||
fts['TextSizeRatio'] = ratio
|
||||
elif name.startswith('.data'):
|
||||
fts['DataSection'] = s_entropy
|
||||
fts['DataSizeRatio'] = ratio
|
||||
elif name.startswith('.rsrc'):
|
||||
fts['RsrcSection'] = s_entropy
|
||||
fts['RsrcSizeRatio'] = ratio
|
||||
if b"requireAdministrator" in s_data:
|
||||
fts['IsAdmin'] = 1
|
||||
|
||||
for sig in Features.INSTALLER_SIGS:
|
||||
if sig in s_data:
|
||||
fts['IsInstall'] = 1
|
||||
break
|
||||
|
||||
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']
|
||||
])
|
||||
|
||||
if hasattr(pe, 'DIRECTORY_ENTRY_IMPORT'):
|
||||
for entry in pe.DIRECTORY_ENTRY_IMPORT:
|
||||
for imp in entry.imports:
|
||||
fts['ApiCount'] += 1
|
||||
if not imp or not hasattr(imp, 'name') or not imp.name:
|
||||
continue
|
||||
try:
|
||||
name = imp.name.decode('ascii', 'ignore')
|
||||
if name in Features.SCHEMA_SET:
|
||||
fts[name] = 1
|
||||
except:
|
||||
pass
|
||||
|
||||
if hasattr(pe, 'DIRECTORY_ENTRY_EXPORT'):
|
||||
fts['ExportCount'] = len(pe.DIRECTORY_ENTRY_EXPORT.symbols)
|
||||
|
||||
if hasattr(pe, 'DIRECTORY_ENTRY_RESOURCE'):
|
||||
for entry in pe.DIRECTORY_ENTRY_RESOURCE.entries:
|
||||
if entry.id == 3 and hasattr(entry, 'directory'):
|
||||
fts['IconCount'] = len(entry.directory.entries)
|
||||
break
|
||||
|
||||
if hasattr(pe, 'DIRECTORY_ENTRY_DEBUG'):
|
||||
fts['DebugCount'] = len(pe.DIRECTORY_ENTRY_DEBUG)
|
||||
|
||||
for k in fts:
|
||||
fts[k] = safe_float(fts[k])
|
||||
|
||||
return fts
|
||||
|
||||
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 = ", ".join([f"{c} REAL" for c in Features.SCHEMA])
|
||||
conn.execute(f"CREATE TABLE IF NOT EXISTS PeData ({cols}, Label INTEGER)")
|
||||
conn.commit()
|
||||
return conn
|
||||
|
||||
def prune_db_duplicates(conn):
|
||||
try:
|
||||
print("\n[*] Post-processing: Checking for and removing duplicate entries...")
|
||||
all_cols = ", ".join(Features.SCHEMA + ["Label"])
|
||||
|
||||
sql_count_before = "SELECT COUNT(*) FROM PeData"
|
||||
before = conn.execute(sql_count_before).fetchone()[0]
|
||||
|
||||
sql_dedup = f"""
|
||||
DELETE FROM PeData
|
||||
WHERE rowid NOT IN (
|
||||
SELECT MIN(rowid)
|
||||
FROM PeData
|
||||
GROUP BY {all_cols}
|
||||
)
|
||||
"""
|
||||
conn.execute(sql_dedup)
|
||||
conn.commit()
|
||||
|
||||
after = conn.execute(sql_count_before).fetchone()[0]
|
||||
removed = before - after
|
||||
|
||||
if removed > 0:
|
||||
print(f"[+] Deduplication complete. Removed {removed} duplicate rows.")
|
||||
print(" (Duplicates occur if the same file is scanned multiple times)")
|
||||
else:
|
||||
print("[+] Deduplication complete. No duplicates found.")
|
||||
|
||||
except Exception as e:
|
||||
print(f"[-] Deduplication Warning: {e}")
|
||||
|
||||
####################################################################################################
|
||||
|
||||
def scan_and_save(conn, path, label):
|
||||
print(f"\n[*] Scanning: {path} (Label={label})")
|
||||
print(f"[*] Config: Max File Size = {MAX_FILE_SIZE/1024/1024:.0f} MB")
|
||||
|
||||
files = []
|
||||
print("[*] Walking directory tree...")
|
||||
if os.path.isfile(path):
|
||||
files.append(path)
|
||||
else:
|
||||
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))
|
||||
|
||||
total = len(files)
|
||||
if total == 0:
|
||||
print("[-] No valid PE files found in this path.")
|
||||
return
|
||||
|
||||
print(f"[*] Found {total} files. Processing...")
|
||||
|
||||
batch = []
|
||||
count = 0
|
||||
errors = 0
|
||||
start_time = time.time()
|
||||
|
||||
placeholders = ",".join(["?"] * (len(Features.SCHEMA) + 1))
|
||||
sql = f"INSERT INTO PeData VALUES ({placeholders})"
|
||||
|
||||
try:
|
||||
for i, file_path in enumerate(files):
|
||||
res = process_single_file(file_path)
|
||||
|
||||
if res:
|
||||
vals = [res.get(c, 0.0) for c in Features.SCHEMA]
|
||||
vals.append(float(label))
|
||||
batch.append(tuple(vals))
|
||||
count += 1
|
||||
else:
|
||||
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}. Retrying strictly...")
|
||||
for row in batch:
|
||||
try:
|
||||
conn.execute(sql, row)
|
||||
except Exception:
|
||||
errors += 1
|
||||
count -= 1
|
||||
conn.commit()
|
||||
batch = []
|
||||
|
||||
elapsed = time.time() - start_time
|
||||
rate = (i + 1) / elapsed if elapsed > 0 else 0
|
||||
sys.stdout.write(f"\r[{i+1}/{total}] OK: {count} | Skip/Err: {errors} | {rate:.1f} files/s")
|
||||
sys.stdout.flush()
|
||||
|
||||
except KeyboardInterrupt:
|
||||
print("\n\n[-] Interrupted by user.")
|
||||
except Exception as e:
|
||||
print(f"\n[-] Unexpected Error: {e}")
|
||||
finally:
|
||||
if batch:
|
||||
print(f"\n[*] Saving remaining {len(batch)} items...")
|
||||
try:
|
||||
conn.executemany(sql, batch)
|
||||
conn.commit()
|
||||
except Exception as final_err:
|
||||
print(f"[-] Final Batch Error: {final_err}. Retrying strictly...")
|
||||
for row in batch:
|
||||
try:
|
||||
conn.execute(sql, row)
|
||||
except Exception:
|
||||
pass
|
||||
conn.commit()
|
||||
|
||||
print(f"\n[+] Scan finished/stopped. Total added: {count}")
|
||||
prune_db_duplicates(conn)
|
||||
|
||||
####################################################################################################
|
||||
|
||||
if __name__ == "__main__":
|
||||
print("\n---------------- PE Dataset Builder v2.6 ----------------\n")
|
||||
print("Type 'exit' or 'q' at any prompt to quit.")
|
||||
|
||||
if os.path.exists(DB_PATH):
|
||||
print(f"[-] Note: If SCHEMA changed, please delete {DB_PATH} before running.")
|
||||
|
||||
conn = init_db()
|
||||
try:
|
||||
while True:
|
||||
print("\n---------------------------------------------------------\n")
|
||||
raw_path = input("[*] Enter Path (File or Folder): ").strip()
|
||||
target_path = raw_path.strip('"').strip("'")
|
||||
|
||||
if target_path.lower() in ['exit', 'q']:
|
||||
break
|
||||
|
||||
if not target_path or not os.path.exists(target_path):
|
||||
print("[-] Invalid path. Please try again.")
|
||||
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
|
||||
else:
|
||||
print("[-] Please enter 0 or 1.")
|
||||
|
||||
scan_and_save(conn, target_path, label)
|
||||
except KeyboardInterrupt:
|
||||
print("\n\n[-] Interrupted.")
|
||||
finally:
|
||||
conn.close()
|
||||
print(f"\n[+] Database saved to: {os.path.abspath(DB_PATH)}")
|
||||
@@ -0,0 +1,195 @@
|
||||
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 = "pe_malware_lgbm.txt"
|
||||
ONNX_FILE = "pe_malware.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': 63,
|
||||
'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': 20,
|
||||
'lambda_l1': 0.1,
|
||||
'lambda_l2': 0.1
|
||||
}
|
||||
|
||||
####################################################################################################
|
||||
|
||||
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_lower = {'label', 'id', 'filename', 'filelength', 'rowid', 'probability', 'predictedlabel', 'score'}
|
||||
return [c for c in columns if c.lower().strip() not in exclude_lower]
|
||||
except:
|
||||
return None
|
||||
finally:
|
||||
conn.close()
|
||||
|
||||
def load_and_clean_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:
|
||||
df = pd.read_sql_query(sql, conn)
|
||||
if df.empty: return None, None, None
|
||||
|
||||
label_col = next((c for c in df.columns if c.lower().strip() == 'label'), None)
|
||||
if not label_col:
|
||||
print("[-] Critical Error: Label column not found.")
|
||||
return None, None, 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()
|
||||
|
||||
X = X.rename(columns=lambda x: re.sub('[^A-Za-z0-9_]+', '', x))
|
||||
|
||||
if X.isnull().values.any():
|
||||
X = X.fillna(X.median(numeric_only=True))
|
||||
|
||||
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_base = 1.0
|
||||
weight_ratio_target = 0.1
|
||||
final_pos_weight = base_weight * (weight_ratio_target / weight_ratio_base)
|
||||
|
||||
print(f"[*] Balance Report: Safe={n_neg}, Malware={n_pos}")
|
||||
print(f"[*] Weight Config: Base Balanced={base_weight:.4f}, Final Adjusted (Favor Safe)={final_pos_weight:.4f}")
|
||||
|
||||
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 (Full Force, No Early Stopping)...")
|
||||
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.0 ----------------\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_and_clean_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}")
|
||||
Reference in new issue
Block a user