Scribd
Upload
10 views9 pages

Malware Detection with Deep Learning

The document outlines a machine learning workflow for malware detection using various data sources, including API calls, PE headers, and images. It details the process of loading data, feature extraction, preprocessing, model building with a custom BiLSTM architecture, and training the model with resampling techniques for class imbalance. Finally, it includes evaluation metrics for both validation and test sets to assess model performance.

Uploaded by

raktouche
Copyright
© All Rights Reserved
We take content rights seriously. If you suspect this is your content, claim it here.
Available Formats
Download as TXT, PDF, TXT or read online on Scribd
10 views9 pages

Malware Detection with Deep Learning

The document outlines a machine learning workflow for malware detection using various data sources, including API calls, PE headers, and images. It details the process of loading data, feature extraction, preprocessing, model building with a custom BiLSTM architecture, and training the model with resampling techniques for class imbalance. Finally, it includes evaluation metrics for both validation and test sets to assess model performance.

Uploaded by

raktouche
Copyright
© All Rights Reserved
We take content rights seriously. If you suspect this is your content, claim it here.
Available Formats
Download as TXT, PDF, TXT or read online on Scribd

!

pip install -U scikit-learn imbalanced-learn

import numpy as np
import pandas as pd
import time
from functools import reduce
import [Link] as plt
from collections import defaultdict
from itertools import product

# ML & CV imports
from [Link] import StandardScaler, LabelEncoder
from sklearn.model_selection import train_test_split, KFold, StratifiedKFold
from [Link] import (confusion_matrix, ConfusionMatrixDisplay,
roc_curve, auc, classification_report,
accuracy_score)
try:
from imblearn.over_sampling import SMOTE, ADASYN
except ImportError:
print("Please install imbalanced-learn: pip install imbalanced-learn")
exit()

# TensorFlow imports
import tensorflow as tf
from [Link] import Model
from [Link] import (Input, Dense, Dropout, Conv2D, Conv1D,
MaxPooling2D,
Flatten, Embedding, Bidirectional,
LSTM,GlobalAveragePooling2D,
Add, Reshape, BatchNormalization,
GlobalMaxPooling1D, concatenate)
from [Link] import EarlyStopping
from [Link] import Tokenizer
from [Link] import pad_sequences

from [Link] import AUC

# --- Configuration ---


RANDOM_STATE = 42
MAX_NUM_WORDS = 15000
MAX_SEQ_LEN = 100
IMG_FEATURES = 1024
LABEL_COL = 'malware'

# --- 1. Load & Align Data ---


print("--- 1. Loading and Aligning Data ---")
# Assume paths dictionary and loading logic is correct as provided before
# ... (loading code using paths, dfs, common_hashes) ...
paths = {
'api': '/kaggle/input/olivera-dataset/API
Calls/dynamic_api_call_sequence_per_malware_100_0_306.csv',
'hdr': '/kaggle/input/olivera-dataset/PE section
header/pe_section_headers.csv',
'img': '/kaggle/input/olivera-dataset/Raw PE as image/raw_pe_images.csv',
'imp': '/kaggle/input/olivera-dataset/Top 1000 PE/top_1000_pe_imports.csv'
}
try:
dfs = {k: pd.read_csv(v) for k, v in [Link]()}
except FileNotFoundError as e:
print(f"Error loading data: {e}. Please ensure file paths are correct.")
exit()
hash_sets = [set(df['hash']) for df in [Link]()]
common_hashes = sorted(reduce(lambda a, b: a & b, hash_sets))
if not common_hashes:
print("Error: No common hashes found across datasets. Check data alignment.")
exit()
for name, df in [Link]():
df_unique = df.drop_duplicates(subset=['hash'], keep='first').set_index('hash')
dfs[name] = df_unique.reindex(common_hashes)
df_imp, df_img, df_hdr, df_api = dfs['imp'], dfs['img'], dfs['hdr'], dfs['api']
print(f"Data loaded for {len(common_hashes)} common samples.")

# --- 2. Feature Extraction & Preprocessing ---


print("\n--- 2. Feature Extraction & Preprocessing ---")
# PE Imports (Assuming already 0/1 or numerical)
X_imp = df_imp.drop(columns=[LABEL_COL]).astype(np.float32)
print(f"X_imp shape: {X_imp.shape}")

# Image Features
img_cols = [c for c in df_img.columns if [Link]('pix_')]
X_img = df_img[img_cols].astype(np.float32)
# Pad if necessary (ensure IMG_FEATURES is correct)
if X_img.shape[1] != IMG_FEATURES:
print(f"Padding/Truncating image features from {X_img.shape[1]} to
{IMG_FEATURES}")
padded_img = [Link]((len(X_img), IMG_FEATURES), dtype=np.float32)
copy_len = min(X_img.shape[1], IMG_FEATURES)
padded_img[:, :copy_len] = X_img.values[:, :copy_len]
X_img = [Link](padded_img, index=X_img.index)
X_img /= 255.0 # Normalize
print(f"X_img shape: {X_img.shape}")
IMG_DIM = int([Link](IMG_FEATURES)) # Recalculate IMG_DIM based on final
IMG_FEATURES
if IMG_DIM * IMG_DIM != IMG_FEATURES:
print(f"Warning: IMG_FEATURES ({IMG_FEATURES}) is not a perfect square. Reshape
layer might behave unexpectedly.")

# PE Header Features
hdr_cols = [c for c in df_hdr.columns if any(p in [Link]() for p in ['entropy',
'section', 'virtual', 'raw'])]
X_hdr = df_hdr[hdr_cols]
# Handle potential non-numeric columns before median/scaling if necessary
X_hdr = X_hdr.apply(pd.to_numeric, errors='coerce') # Convert non-numeric to NaN
X_hdr = X_hdr.fillna(X_hdr.median()) # Fill NaNs with median
scaler = StandardScaler()
X_hdr_scaled = [Link](scaler.fit_transform(X_hdr), index=X_hdr.index,
columns=X_hdr.columns)
print(f"X_hdr_scaled shape: {X_hdr_scaled.shape}")

# API Sequence Features


seq_cols = [f't_{i}' for i in range(MAX_SEQ_LEN)] # Use MAX_SEQ_LEN for column
names
# Ensure columns exist, fill missing with empty string
existing_seq_cols = [c for c in seq_cols if c in df_api.columns]
missing_seq_cols = [c for c in seq_cols if c not in df_api.columns]
for col in missing_seq_cols:
df_api[col] = '' # Add missing columns
api_raw = df_api[seq_cols].fillna('').astype(str).agg(' '.join, axis=1) # Join with
space for Tokenizer

tokenizer = Tokenizer(num_words=MAX_NUM_WORDS, oov_token='<OOV>')


tokenizer.fit_on_texts(api_raw)
api_seq = pad_sequences(tokenizer.texts_to_sequences(api_raw), maxlen=MAX_SEQ_LEN)
VOCAB_SIZE = min(MAX_NUM_WORDS, len(tokenizer.word_index) + 1)
print(f"api_seq shape: {api_seq.shape}, Vocab Size: {VOCAB_SIZE}")

# Target Variable
y = LabelEncoder().fit_transform(df_imp[LABEL_COL])
print(f"Target y shape: {[Link]}, Class distribution: {[Link](y)}")

# --- 3. Initial Data Split (Train 80% / Validation 10% / Test 10%) ---
print("\n--- 3. Initial Data Split ---")
idx = [Link](len(y))

# Split into Train (80%) and Temp (20%)


train_idx, test_idx = train_test_split(idx, test_size=0.1, stratify=y,
random_state=RANDOM_STATE)
y_train = y[train_idx]
y_test = y[test_idx]

print(f"Train indices: {len(train_idx)} ({len(train_idx)/len(idx):.1%}), Class


distribution: {[Link](y_train)}")
print(f"Test indices: {len(test_idx)} ({len(test_idx)/len(idx):.1%}), Class
distribution: {[Link](y_test)}")

# --- Assemble full feature matrix (ensure order matches split_inputs) ---
print("\nAssembling full feature matrix X_full...")
# Use .values to ensure numpy arrays and consistent concatenation
X_full = [Link]([X_imp.values, X_img.values, api_seq])#X_hdr_scaled.values,
api_seq])
print(f"X_full shape: {X_full.shape}")

# --- Define split_inputs function (uses global variables for dimensions) ---
def split_inputs(X):
# Calculate split points based on the *original* dataframes used to build
X_full
n1 = X_imp.shape[1]
n2 = n1 + X_img.shape[1] # Use X_img shape *after* padding/truncating
#n3 = n2 + X_hdr_scaled.shape[1]
# Return list of slices
return [X[:, :n1], X[:, n1:n2], X[:, n2:]]#X[:, n2:n3], X[:, n3:]]

# --- 4. Model builder ---

print("\n--- 4. Defining Model Builder ---")

class CustomBiLSTM([Link]):
def __init__(self, units, return_sequences=False, **kwargs):
super(CustomBiLSTM, self).__init__(**kwargs)

[Link] = units
self.return_sequences = return_sequences

self.return_sequences = return_sequences
self.forward_lstm = [Link](
units, return_sequences=True, return_state=True, name="forward_lstm"
)
self.backward_lstm = [Link](
units, return_sequences=True, return_state=True, go_backwards=True,
name="backward_lstm"
)

self.forward_sequence_outputs = None
self.backward_sequence_outputs = None

def call(self, inputs, training=None, return_sequences_only=False):

f_seq_output, f_h, f_c = self.forward_lstm(inputs, training=training)


b_seq_output, b_h, b_c = self.backward_lstm(inputs, training=training)
b_seq_output = [Link](b_seq_output, axis=[1])

if return_sequences_only:
return f_seq_output, b_seq_output
elif self.return_sequences:
return [Link]([f_seq_output, b_seq_output], axis=-1)
else:
return [Link]([f_h, b_h], axis=-1)

def get_intermediate_outputs(self):
return self.forward_sequence_outputs, self.backward_sequence_outputs

def build_model_alt(dense_units=128, dropout_merged=0.5, lr=1e-4):

imp_shape = X_imp.shape[1]
img_shape = X_img.shape[1]
hdr_shape = X_hdr_scaled.shape[1]
seq_shape = MAX_SEQ_LEN # From padding

in_imp = Input(shape=(imp_shape,), name='imp')


x = Dense(256, activation="relu")(in_imp)
x = BatchNormalization()(x)
x_imp = Dropout(0.4)(x)

in_img = Input(shape=(img_shape,), name='img')


# Only reshape if IMG_FEATURES is a perfect square
if IMG_DIM * IMG_DIM == img_shape:
r2 = Reshape((IMG_DIM, IMG_DIM, 1))(in_img)
x = Conv2D(32, (3,3), activation='relu',padding='same')(r2)
x = BatchNormalization()(x)
x = MaxPooling2D()(x)
x = Dropout(0.4)(x)

x = Flatten()(x)
x = Dropout(0.5)(x)
x = Dense(32, activation='relu')(x)
x_img = Dropout(0.5)(x)
else: # Fallback to Dense if not reshapeable
print("IMG_FEATURES not a perfect square, using Dense layer for images.")
x = Dense(128, activation='relu')(in_img)
x_img = Dropout(0.5)(x)

in_seq = Input(shape=(MAX_SEQ_LEN,), name='seq')


x = Embedding(308, 128)(in_seq)
x = x = CustomBiLSTM(units=128, return_sequences=False, name="custom_bilstm")
(x)
x_seq = Dropout(0.5)(x)
#x = Dropout(0.5)(x)
#x_seq = Dense(1, activation='sigmoid')(x)

merged = concatenate([x_imp, x_img, x_seq])

x = Dense(dense_units, activation='relu')(merged)
x = BatchNormalization()(x)
x = Dropout(dropout_merged)(x)

output = Dense(1, activation='sigmoid')(x)

model = Model(inputs=[in_imp, in_img, in_seq], outputs=output)


[Link](optimizer=[Link](learning_rate=lr),
loss='binary_crossentropy',
metrics=['accuracy', AUC(name='auc')])
return model

# --- 3. Data Split (Train 80% / Validation 10% / Test 10%) ---

print("\n--- 3. Data Split ---")


idx = [Link](len(y))

# First split into Train+Val (90%) and Test (10%)


train_val_idx, test_idx = train_test_split(
idx,
test_size=0.1,
stratify=y,
random_state=RANDOM_STATE
)

train_idx, val_idx = train_test_split(


train_val_idx,
test_size=1/9,
stratify=y[train_val_idx],
random_state=RANDOM_STATE
)

y_train = y[train_idx]
y_val = y[val_idx]
y_test = y[test_idx]

print(f"Train indices: {len(train_idx)} ({len(train_idx)/len(idx):.1%}), Class


distribution: {[Link](y_train)}")
print(f"Validation indices: {len(val_idx)} ({len(val_idx)/len(idx):.1%}), Class
distribution: {[Link](y_val)}")
print(f"Test indices: {len(test_idx)} ({len(test_idx)/len(idx):.1%}), Class
distribution: {[Link](y_test)}")

# --- 4. Feature Preparation ---


print("\n--- 4. Feature Preparation ---")

# Apply ADASYN to the training portion only


print(f"Original training shape: {len(train_idx)}, Class distribution:
{[Link](y_train)}")
adasyn = ADASYN(random_state=RANDOM_STATE, n_neighbors=5)
X_train_resampled, y_train_resampled = adasyn.fit_resample(X_full[train_idx],
y_train)
print(f"Resampled training shape: {X_train_resampled.shape}, Class distribution:
{[Link](y_train_resampled)}")

# Split features for model input


Xs_train = split_inputs(X_train_resampled)
Xs_val = split_inputs(X_full[val_idx])
Xs_test = split_inputs(X_full[test_idx])

# --- 5. Model Training ---

print("\n--- 5. Model Training ---")

# Fixed parameters
params = {
'dense_units': 128,
'dropout_merged': 0.52,
'lr': 9e-5,
'batch_size': 64
}

# Build and train model


print(f"\nTraining model with parameters: {params}")
model = build_model_alt(
dense_units=params['dense_units'],
dropout_merged=params['dropout_merged'],
lr=params['lr']
)

early_stopping = EarlyStopping(
monitor='val_accuracy',
mode='max',
patience=10,
verbose=1,
restore_best_weights=True
)

start_t_train = [Link]()

hist = [Link](
Xs_train, y_train_resampled,
validation_data=(Xs_val, y_val),
epochs=10,
batch_size=params['batch_size'],
callbacks=[early_stopping],
verbose=1
)

train_time = [Link]() - start_t_train

# --- 6. Evaluation ---


print("\n" + "="*30)
print("--- 6. Model Evaluation ---")
print("="*30)

# Validation set evaluation


print("\n--- Validation Set Evaluation ---")
start_t_inf = [Link]()

loss_val, acc_val, auc_val = [Link](Xs_val, y_val, verbose=1)

y_pred_proba_val = [Link](Xs_val).flatten()
inf_time = [Link]() - start_t_inf

y_pred_val = (y_pred_proba_val > 0.5).astype(int)

# Calculate validation metrics


cm_val = confusion_matrix(y_val, y_pred_val)
fpr_val, tpr_val, _ = roc_curve(y_val, y_pred_proba_val)
auc_val = auc(fpr_val, tpr_val)
report_val = classification_report(y_val, y_pred_val, output_dict=True,
zero_division=0)
recall_val = report_val['weighted avg']['recall']
f1_val = report_val['weighted avg']['f1-score']

# Test set evaluation


print("\n--- Test Set Evaluation ---")
start_t_inf_test = [Link]()
loss_test, accuracy_test, auc_test = [Link](Xs_test, y_test, verbose=0)
y_pred_proba_test = [Link](Xs_test).flatten()
inf_time_test = [Link]() - start_t_inf_test

y_pred_test = (y_pred_proba_test > 0.5).astype(int)

# Calculate test metrics


cm_test = confusion_matrix(y_test, y_pred_test)
report_test_dict = classification_report(y_test, y_pred_test, output_dict=True,
zero_division=0)
report_test_str = classification_report(y_test, y_pred_test, zero_division=0)
fpr_test, tpr_test, _ = roc_curve(y_test, y_pred_proba_test)
auc_test = auc(fpr_test, tpr_test)
recall_test = report_test_dict['weighted avg']['recall']
f1_score_test = report_test_dict['weighted avg']['f1-score']

# Print results
print("\n=== Final Results ===")
print("\nValidation Set:")
print(f" Loss: {loss_val:.4f}")
print(f" Accuracy: {acc_val:.4f}")
print(f" AUC: {auc_val:.4f}")
print(f" Recall (Weighted): {recall_val:.4f}")
print(f" F1-score (Weighted): {f1_val:.4f}")
print(f" Confusion Matrix:\n{cm_val}")

print("\nTest Set:")
print(f" Loss: {loss_test:.4f}")
print(f" Accuracy: {accuracy_test:.4f}")
print(f" AUC: {auc_test:.4f}")
print(f" Recall (Weighted): {recall_test:.4f}")
print(f" F1-score (Weighted): {f1_score_test:.4f}")
print(f" Confusion Matrix:\n{cm_test}")

# Plot training history


print("\nPlotting training history...")
history_dict = [Link]
epochs_trained = len(history_dict['loss'])
epochs_range = range(1, epochs_trained + 1)

[Link](figsize=(14, 5))

[Link](1, 2, 1)
[Link](epochs_range, history_dict['loss'], 'bo-', label='Training Loss')
[Link](epochs_range, history_dict['val_loss'], 'ro-', label='Validation Loss')
[Link]('Training and Validation Loss')
[Link]('Epochs')
[Link]('Loss')
[Link]()
[Link](True)

[Link](1, 2, 2)
[Link](epochs_range, history_dict['accuracy'], 'bo-', label='Training Accuracy')
[Link](epochs_range, history_dict['val_accuracy'], 'ro-', label='Validation
Accuracy')
[Link]('Training and Validation Accuracy')
[Link]('Epochs')
[Link]('Accuracy')
[Link]()
[Link](True)

plt.tight_layout()
[Link]()

# Plot confusion matrices


fig, (ax1, ax2) = [Link](1, 2, figsize=(14, 5))

ConfusionMatrixDisplay(confusion_matrix=cm_val, display_labels=[0, 1]).plot(ax=ax1,


cmap=[Link])
ax1.set_title('Validation Set Confusion Matrix')

ConfusionMatrixDisplay(confusion_matrix=cm_test, display_labels=[0,
1]).plot(ax=ax2, cmap=[Link])
ax2.set_title('Test Set Confusion Matrix')

plt.tight_layout()
[Link]()

# Plot ROC curves


fig, (ax1, ax2) = [Link](1, 2, figsize=(14, 5))

[Link](fpr_val, tpr_val, color='darkorange', lw=2, label=f'ROC (AUC =


{auc_val:.4f})')
[Link]([0, 1], [0, 1], color='navy', lw=2, linestyle='--', label='Chance (AUC =
0.50)')
ax1.set_xlim([0.0, 1.0])
ax1.set_ylim([0.0, 1.05])
ax1.set_xlabel('False Positive Rate')
ax1.set_ylabel('True Positive Rate')
ax1.set_title('Validation Set ROC Curve')
[Link](loc="lower right")
[Link](alpha=0.5)

[Link](fpr_test, tpr_test, color='darkorange', lw=2, label=f'ROC (AUC =


{auc_test:.4f})')
[Link]([0, 1], [0, 1], color='navy', lw=2, linestyle='--', label='Chance (AUC =
0.50)')
ax2.set_xlim([0.0, 1.0])
ax2.set_ylim([0.0, 1.05])
ax2.set_xlabel('False Positive Rate')
ax2.set_ylabel('True Positive Rate')
ax2.set_title('Test Set ROC Curve')
[Link](loc="lower right")
[Link](alpha=0.5)

plt.tight_layout()
[Link]()

print("\n--- Complete ---")

You might also like