Detecting Malware Using Dynamic Memory Heatmaps and Video Classification

Introduction

The cybersecurity landscape is constantly evolving, with malware authors employing sophisticated techniques to evade detection by traditional antivirus (AV) solutions. As attackers leverage advanced methods to obfuscate and disguise malicious code, defenders must adopt innovative approaches to detect and mitigate threats effectively.

One such innovative approach is the use of dynamic memory heatmaps combined with video classification techniques for malware detection. By visualizing the memory usage patterns of executables during runtime and treating these patterns as frames in a video, we can leverage deep learning models designed for video classification to identify malicious behavior.

In this post, we will explore how to build a malware detection system using dynamic memory heatmaps and video classification. We will cover the following steps:

Setting up the environment and tools.
Building a dataset of dynamic memory heatmaps.
Training a video classification model.
Implementing the detection method.
Evaluating the system’s performance.

We will use Python along with libraries such as LIEF for parsing executable files and Qiling for emulating and monitoring program execution. Let’s dive in!

Background

Malware Detection Challenges

Traditional malware detection techniques rely heavily on signature-based methods and heuristic analysis. However, these methods can be circumvented by malware employing code obfuscation, encryption, polymorphism, and other evasion tactics. As highlighted in the paper “Bypassing Antivirus Detection: Old-School Malware, New Tricks,” attackers can modify malware code using publicly available techniques to evade detection by popular AV solutions.

Dynamic Analysis and Memory Heatmaps

Dynamic analysis involves executing a program in a controlled environment (sandbox) to observe its behavior. By monitoring memory usage patterns during execution, we can create visual representations (heatmaps) that capture the program’s runtime characteristics. These heatmaps can reveal anomalies indicative of malicious activity.

Video Classification for Malware Detection

Treating a sequence of memory heatmaps as frames in a video allows us to apply video classification techniques for malware detection. Deep learning models like Convolutional Neural Networks (CNNs) and Long Short-Term Memory (LSTM) networks can learn temporal and spatial features from these sequences, enabling the detection of malicious patterns.

Prerequisites

Before we begin, ensure you have the following installed:

Python 3.7 or higher
LIEF library (pip install lief)
Qiling Framework (pip install qiling)
PyTorch (pip install torch torchvision)
NumPy (pip install numpy)
Matplotlib (pip install matplotlib)
OpenCV (pip install opencv-python)

Setting Up the Environment

First, let’s import the necessary libraries and set up the environment.

import os
import lief
import numpy as np
import matplotlib.pyplot as plt
import cv2
from qiling import Qiling
from qiling.const import QL_VERBOSE
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader

Building the Dataset

Collecting Malware and Benign Samples

To build a robust dataset, we need a collection of malware and benign executable files.

For this example, we’ll assume you have two directories:

malware_samples/: Contains malicious executables.
benign_samples/: Contains benign executables.

Generating Dynamic Memory Heatmaps

We will execute each sample in an emulated environment using Qiling and capture memory snapshots at regular intervals.

Function to Generate Heatmaps

def generate_memory_heatmaps(file_path, label, output_dir, num_frames=30, interval=0.1):
    ql = Qiling([file_path], rootfs="", verbose=QL_VERBOSE.OFF)
    heatmaps = []

    def memory_snapshot(ql):
        mem_map = ql.mem.get_mem_range()
        mem_image = np.zeros((512, 512), dtype=np.uint8)
        for start, end, perm in mem_map:
            addr = start % (512 * 512)
            x = addr % 512
            y = addr // 512
            mem_image[y % 512, x % 512] += 1
        heatmaps.append(mem_image)

    ql.hook_code(memory_snapshot)

    try:
        ql.run(timeout=interval * num_frames)
    except Exception as e:
        print(f"Error executing {file_path}: {e}")

    for i, heatmap in enumerate(heatmaps):
        plt.imsave(f"{output_dir}/{label}_{os.path.basename(file_path)}_frame_{i}.png", heatmap, cmap='hot')

Generating Heatmaps for All Samples

def generate_dataset(sample_dir, label, output_dir):
    os.makedirs(output_dir, exist_ok=True)
    for file_name in os.listdir(sample_dir):
        file_path = os.path.join(sample_dir, file_name)
        if os.path.isfile(file_path):
            print(f"Processing {file_path}")
            generate_memory_heatmaps(file_path, label, output_dir)

Running the Data Generation

# Generate heatmaps for malware samples
generate_dataset('malware_samples', 'malware', 'dataset/malware')

# Generate heatmaps for benign samples
generate_dataset('benign_samples', 'benign', 'dataset/benign')

Preparing the Data for Training

We use Qiling to emulate the execution of each sample.
We capture memory snapshots at each instruction execution (hook_code).
We convert memory addresses into a 2D heatmap.
We save each heatmap as a frame in the output directory.

3. Preparing the Data for Training

Creating a Custom Dataset Class

We need to create a PyTorch Dataset to load our heatmaps and their corresponding labels.

class MalwareDataset(Dataset):
    def __init__(self, dataset_dir, transform=None):
        self.samples = []
        self.labels = []
        self.transform = transform

        for label in ['malware', 'benign']:
            label_dir = os.path.join(dataset_dir, label)
            for file_name in os.listdir(label_dir):
                if file_name.endswith('.png'):
                    self.samples.append(os.path.join(label_dir, file_name))
                    self.labels.append(0 if label == 'benign' else 1)

    def __len__(self):
        return len(self.samples)

    def __getitem__(self, idx):
        video_frames = []
        for i in range(30):  # Assuming 30 frames per sample
            frame_path = self.samples[idx].replace('_frame_0.png', f'_frame_{i}.png')
            if os.path.exists(frame_path):
                frame = cv2.imread(frame_path, cv2.IMREAD_GRAYSCALE)
                if self.transform:
                    frame = self.transform(frame)
                video_frames.append(frame)
        video = np.stack(video_frames, axis=0)
        label = self.labels[idx]
        return torch.tensor(video, dtype=torch.float32), torch.tensor(label, dtype=torch.long)

Explanation:

The MalwareDataset class loads the heatmap frames for each sample.
Each sample consists of a sequence of frames (heatmaps).
Labels are assigned based on the directory (0 for benign, 1 for malware).

Creating DataLoaders

from sklearn.model_selection import train_test_split

dataset = MalwareDataset('dataset')

train_indices, val_indices = train_test_split(range(len(dataset)), test_size=0.2, random_state=42)

train_dataset = torch.utils.data.Subset(dataset, train_indices)
val_dataset = torch.utils.data.Subset(dataset, val_indices)

train_loader = DataLoader(train_dataset, batch_size=8, shuffle=True)
val_loader = DataLoader(val_dataset, batch_size=8, shuffle=False)

Building the Video Classification Model

We will build a simple 3D CNN model for video classification.

Defining the Model

class MalwareClassifier(nn.Module):
    def __init__(self):
        super(MalwareClassifier, self).__init__()
        self.conv1 = nn.Conv3d(1, 16, kernel_size=(3, 3, 3), padding=1)
        self.pool1 = nn.MaxPool3d(kernel_size=(1, 2, 2))
        self.conv2 = nn.Conv3d(16, 32, kernel_size=(3, 3, 3), padding=1)
        self.pool2 = nn.MaxPool3d(kernel_size=(1, 2, 2))
        self.fc1 = nn.Linear(32 * 30 * 128 * 128, 128)
        self.fc2 = nn.Linear(128, 2)  # Two classes: benign and malware

    def forward(self, x):
        x = x.unsqueeze(1)  # Add channel dimension
        x = self.pool1(F.relu(self.conv1(x)))
        x = self.pool2(F.relu(self.conv2(x)))
        x = x.view(x.size(0), -1)
        x = F.relu(self.fc1(x))
        x = self.fc2(x)
        return x

Explanation:

The model consists of 3D convolutional layers to capture spatial and temporal features.
We use max pooling to reduce spatial dimensions.
Fully connected layers map the features to class scores.

Training the Model

import torch.nn.functional as F

model = MalwareClassifier()
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.001)

num_epochs = 10

for epoch in range(num_epochs):
    model.train()
    running_loss = 0.0
    for inputs, labels in train_loader:
        optimizer.zero_grad()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        running_loss += loss.item() * inputs.size(0)
    epoch_loss = running_loss / len(train_loader.dataset)
    print(f"Epoch {epoch+1}/{num_epochs}, Loss: {epoch_loss:.4f}")

    # Validation
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for inputs, labels in val_loader:
            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)
            correct += (preds == labels).sum().item()
            total += labels.size(0)
    val_acc = correct / total
    print(f"Validation Accuracy: {val_acc:.4f}")

Implementing the Detection Method

Once the model is trained, we can use it to detect malware in new samples.

Detection Function

def detect_malware(file_path, model):
    generate_memory_heatmaps(file_path, 'test', 'test_data')
    video_frames = []
    for i in range(30):
        frame_path = f"test_data/test_{os.path.basename(file_path)}_frame_{i}.png"
        if os.path.exists(frame_path):
            frame = cv2.imread(frame_path, cv2.IMREAD_GRAYSCALE)
            video_frames.append(frame)
    if len(video_frames) == 0:
        print("No frames generated.")
        return
    video = np.stack(video_frames, axis=0)
    inputs = torch.tensor(video, dtype=torch.float32).unsqueeze(0)
    outputs = model(inputs)
    _, preds = torch.max(outputs, 1)
    if preds.item() == 1:
        print(f"{file_path} is detected as Malware.")
    else:
        print(f"{file_path} is detected as Benign.")

Testing the Detection

# Load the trained model (if saved)
# model.load_state_dict(torch.load('malware_classifier.pth'))

# Detect malware in a new sample
detect_malware('unknown_sample.exe', model)

Evaluating the System’s Performance

To evaluate the system’s performance, we can compute metrics such as accuracy, precision, recall, and F1-score on a test set.

Computing Metrics

from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score

def evaluate_model(model, dataloader):
    model.eval()
    all_preds = []
    all_labels = []
    with torch.no_grad():
        for inputs, labels in dataloader:
            outputs = model(inputs)
            _, preds = torch.max(outputs, 1)
            all_preds.extend(preds.tolist())
            all_labels.extend(labels.tolist())
    acc = accuracy_score(all_labels, all_preds)
    prec = precision_score(all_labels, all_preds)
    rec = recall_score(all_labels, all_preds)
    f1 = f1_score(all_labels, all_preds)
    print(f"Accuracy: {acc:.4f}, Precision: {prec:.4f}, Recall: {rec:.4f}, F1-Score: {f1:.4f}")

# Evaluate on validation set
evaluate_model(model, val_loader)

Conclusion

In this post, we’ve explored how to use dynamic memory heatmaps and video classification techniques for malware detection. By visualizing memory usage patterns during program execution and leveraging deep learning models, we can detect malicious behavior that might evade traditional detection methods.

This approach can be further enhanced by:

Collecting a larger and more diverse dataset to improve model generalization.
Implementing more advanced models, such as pre-trained CNNs or Transformers.
Incorporating additional features, like API call sequences or system event logs.