malware_cnn.py

import os
import math
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.model_selection import train_test_split
from sklearn.utils import class_weight
from sklearn.metrics import confusion_matrix
import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Dropout, Flatten, Conv2D, MaxPooling2D, Input
from tensorflow.keras.preprocessing.image import ImageDataGenerator
from tensorflow.keras import backend as K

RAW_DATA_DIR = r"D:\ASSEMBLY\data_raw"
PROCESSED_DATA_DIR = r"D:\ASSEMBLY\multi virusi"
IMG_SIZE = (64, 64)
BATCH_SIZE = 128
EPOCHS = 100
NUM_CLASSES = 22

def convert_bytes_to_image(file_path, save_path):
    try:
        with open(file_path, 'r') as f:
            array = []
            for line in f:
                xx = line.split()
                if len(xx) != 17:
                    continue
                array.append([int(i, 16) if i != '??' else 0 for i in xx[1:]])
        
        if not array:
            return None

        data = np.array(array)
        if data.shape[1] != 16:
            return None

        b = int((data.shape[0] * 16)**(0.5))
        b = 2**(int(math.log(b) / math.log(2)) + 1)
        a = int(data.shape[0] * 16 / b)
        data = data[:a * b // 16, :]
        data = np.reshape(data, (a, b))
        
        im = Image.fromarray(np.uint8(data))
        im.save(save_path, "PNG")
        return im
    except Exception:
        return None

def process_raw_files(source_dir, dest_dir):
    if not os.path.exists(source_dir):
        return

    files = os.listdir(source_dir)
    for name in files:
        if not name.endswith('.bytes'):
            continue
        
        file_path = os.path.join(source_dir, name)
        save_path = os.path.join(dest_dir, name + '.png')
        convert_bytes_to_image(file_path, save_path)

def recall_m(y_true, y_pred):
    true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))
    possible_positives = K.sum(K.round(K.clip(y_true, 0, 1)))
    recall = true_positives / (possible_positives + K.epsilon())
    return recall

def precision_m(y_true, y_pred):
    true_positives = K.sum(K.round(K.clip(y_true * y_pred, 0, 1)))
    predicted_positives = K.sum(K.round(K.clip(y_pred, 0, 1)))
    precision = true_positives / (predicted_positives + K.epsilon())
    return precision

def f1_m(y_true, y_pred):
    precision = precision_m(y_true, y_pred)
    recall = recall_m(y_true, y_pred)
    return 2 * ((precision * recall) / (precision + recall + K.epsilon()))

def create_cnn_model(input_shape, num_classes):
    model = Sequential([
        Input(shape=input_shape),
        Conv2D(30, kernel_size=(3, 3), activation='relu'),
        MaxPooling2D(pool_size=(2, 2)),
        Conv2D(15, (3, 3), activation='relu'),
        MaxPooling2D(pool_size=(2, 2)),
        Dropout(0.25),
        Flatten(),
        Dense(128, activation='relu'),
        Dropout(0.25),
        Dense(128, activation='relu'),
        Dropout(0.5),
        Dense(50, activation='relu'),
        Dense(num_classes, activation='softmax')
    ])
    
    model.compile(
        loss='categorical_crossentropy', 
        optimizer='adam', 
        metrics=['accuracy', f1_m, precision_m, recall_m, tf.keras.metrics.AUC()]
    )
    return model

def load_and_split_data(data_dir, target_size, batch_size):
    datagen = ImageDataGenerator(rescale=1./255)
    batches = datagen.flow_from_directory(
        directory=data_dir, 
        target_size=target_size, 
        batch_size=20000,
        class_mode='categorical'
    )
    
    imgs, labels = next(batches)
    X_train, X_test, y_train, y_test = train_test_split(imgs, labels, test_size=0.3, random_state=42)
    return X_train, X_test, y_train, y_test, batches.class_indices

def plot_confusion_matrix(y_true, y_pred, class_names, output_file="confusion_matrix.png"):
    cm = confusion_matrix(y_true, y_pred)
    plt.figure(figsize=(20, 15))
    sns.heatmap(cm, annot=True, fmt="d", cmap="Blues", xticklabels=class_names, yticklabels=class_names)
    plt.ylabel('True label')
    plt.xlabel('Predicted label')
    plt.title('Confusion Matrix')
    plt.savefig(output_file)
    plt.close()

if __name__ == "__main__":
    if os.path.exists(RAW_DATA_DIR):
        process_raw_files(RAW_DATA_DIR, RAW_DATA_DIR)

    if os.path.exists(PROCESSED_DATA_DIR):
        try:
            X_train, X_test, y_train, y_test, class_indices = load_and_split_data(PROCESSED_DATA_DIR, IMG_SIZE, BATCH_SIZE)
            
            y_train_indices = np.argmax(y_train, axis=1)
            unique_classes = np.unique(y_train_indices)
            class_weights_computed = class_weight.compute_class_weight(
                class_weight='balanced',
                classes=unique_classes,
                y=y_train_indices
            )
            class_weights_dict = {l: c for l, c in zip(unique_classes, class_weights_computed)}

            model = create_cnn_model((IMG_SIZE[0], IMG_SIZE[1], 3), NUM_CLASSES)
            
            model.fit(
                X_train, 
                y_train, 
                validation_data=(X_test, y_test), 
                epochs=EPOCHS, 
                batch_size=BATCH_SIZE,
                class_weight=class_weights_dict,
                verbose=1
            )

            loss, accuracy, f1, precision, recall, auc = model.evaluate(X_test, y_test, verbose=0)
            
            y_pred_probs = model.predict(X_test)
            y_pred_classes = np.argmax(y_pred_probs, axis=-1)
            y_true_classes = np.argmax(y_test, axis=1)
            
            class_names = list(class_indices.keys())
            plot_confusion_matrix(y_true_classes, y_pred_classes, class_names)
            
            model.save("malware_cnn_model.h5")

        except Exception:
            pass