Building an RL System to Generate Realistic Malware Samples for Defensive Training

Introduction

In our previous posts, we explored:

Using Reinforcement Learning (RL) for Malware Evasion: How attackers might use RL to craft malware that adapts to evade detection.
Detecting Malware Using Dynamic Memory Heatmaps and Video Classification: Leveraging dynamic analysis and deep learning to detect malicious behavior.

Now, we’ll bring these concepts together to advance cybersecurity research. Our goal is to:

Embed the dynamic classifier into a malware gym framework.
Train an RL agent to bypass this dynamic classifier using custom actions.
Generate realistic malware samples to train defensive systems like malware classifiers.

Overview

Malware Gym Framework: An environment where malware samples can be modified and evaluated against detection mechanisms.
Dynamic Classifier Integration: Incorporate our dynamic memory heatmap-based classifier into the gym environment.
Reinforcement Learning Agent: Train an agent to apply transformations to malware samples to evade detection.
Custom Actions: Define actions that represent realistic modifications malware authors might use, inspired by the paper “Bypassing Antivirus Detection: Old-School Malware, New Tricks.”

Setting Up the Environment

Prerequisites

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

Import Libraries

import os
import gymnasium as gym
from gymnasium import spaces
import numpy as np
import lief
from qiling import Qiling
from qiling.const import QL_VERBOSE
import torch
import torch.nn as nn
import torch.optim as optim
import cv2
import random

Malware Gym Environment

Defining the Environment

We’ll create a custom gym environment where the agent can apply transformations to a malware sample and receive feedback based on whether the modified sample is detected by the dynamic classifier.

class MalwareEnv(gym.Env):
    metadata = {'render.modes': ['human']}

    def __init__(self, malware_sample_path):
        super(MalwareEnv, self).__init__()
        self.original_sample = lief.parse(malware_sample_path)
        self.sample_path = malware_sample_path
        self.action_space = spaces.Discrete(7)  # Seven custom actions
        self.observation_space = spaces.Box(low=0, high=255, shape=(30, 512, 512), dtype=np.uint8)
        self.current_sample = None
        self.steps = 0
        self.max_steps = 10  # Limit the number of modifications

    def reset(self):
        self.current_sample = self.original_sample.copy()
        self.steps = 0
        observation = self.get_observation()
        return observation, {}

    def step(self, action):
        self.steps += 1
        self.apply_action(action)
        observation = self.get_observation()
        done = self.steps >= self.max_steps
        reward = self.get_reward()
        info = {}
        return observation, reward, done, False, info

    def apply_action(self, action):
        # Define the action transformations
        if action == 0:
            self.add_junk_section()
        elif action == 1:
            self.modify_header()
        elif action == 2:
            self.encrypt_sections()
        elif action == 3:
            self.pack_executable()
        elif action == 4:
            self.inject_code()
        elif action == 5:
            self.manipulate_imports()
        elif action == 6:
            self.change_entry_point()
        else:
            pass  # No action

    def get_observation(self):
        # Generate dynamic memory heatmaps
        heatmaps = generate_memory_heatmaps_from_binary(self.current_sample)
        return heatmaps

    def get_reward(self):
        # Use the dynamic classifier to determine detection
        detected = dynamic_classifier(self.current_sample)
        reward = -1 if detected else 1
        return reward

    def render(self, mode='human'):
        pass

    def close(self):
        pass

Explanation:

Action Space: Seven discrete actions representing possible modifications.
Observation Space: A sequence of memory heatmaps (treated as frames).
Reward Function: Positive reward if the sample is not detected, negative if detected.
Termination Condition: Episode ends after a maximum number of steps or if certain conditions are met.

Custom Actions

Based on the paper “Bypassing Antivirus Detection: Old-School Malware, New Tricks,” we define the following actions:

Add Junk Section: Append a section with random data to the binary.
Modify Header: Change non-essential header fields.
Encrypt Sections: Encrypt code or data sections and add a decryption stub.
Pack Executable: Use a packer to compress and obfuscate the binary.
Inject Code: Insert benign-looking code that doesn’t alter functionality.
Manipulate Imports: Alter the import table to obfuscate API calls.
Change Entry Point: Redirect the entry point to a custom loader.

Dynamic Classifier Integration

We will integrate the dynamic memory heatmap-based classifier into the environment.

Dynamic Classifier Function

def dynamic_classifier(binary):
    # Simulate execution and generate heatmaps
    heatmaps = generate_memory_heatmaps_from_binary(binary)
    # Preprocess and pass through the trained model
    inputs = preprocess_heatmaps(heatmaps)
    outputs = model(inputs)
    _, preds = torch.max(outputs, 1)
    detected = preds.item() == 1  # 1 indicates malware
    return detected

Generating Heatmaps from Binary

def generate_memory_heatmaps_from_binary(binary):
    # Save the binary to a temporary file
    temp_path = 'temp_sample.exe'
    binary.write(temp_path)
    # Use Qiling to emulate execution
    ql = Qiling([temp_path], rootfs="", verbose=QL_VERBOSE.OFF)
    heatmaps = []
    def memory_snapshot(ql):
        # Similar to previous implementation
        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=0.1 * 30)
    except Exception as e:
        pass
    os.remove(temp_path)
    return np.stack(heatmaps, axis=0)

Preprocessing Heatmaps

def preprocess_heatmaps(heatmaps):
    # Convert to tensor and add batch dimension
    inputs = torch.tensor(heatmaps, dtype=torch.float32).unsqueeze(0)
    return inputs

Loading the Trained Model

Assuming you have a trained model from the previous blog post:

model = MalwareClassifier()
model.load_state_dict(torch.load('malware_classifier.pth'))
model.eval()

Training the RL Agent

Defining the RL Algorithm

We’ll use a simple Deep Q-Network (DQN) for demonstration purposes.

class DQNAgent:
    def __init__(self, state_size, action_size):
        self.model = nn.Sequential(
            nn.Flatten(),
            nn.Linear(state_size, 128),
            nn.ReLU(),
            nn.Linear(128, action_size)
        )
        self.optimizer = optim.Adam(self.model.parameters(), lr=0.001)
        self.criterion = nn.MSELoss()
        self.memory = []
        self.gamma = 0.9

    def act(self, state):
        with torch.no_grad():
            q_values = self.model(torch.tensor(state, dtype=torch.float32).unsqueeze(0))
        action = torch.argmax(q_values).item()
        return action

    def remember(self, state, action, reward, next_state, done):
        self.memory.append((state, action, reward, next_state, done))
        if len(self.memory) > 1000:
            self.memory.pop(0)

    def replay(self, batch_size):
        if len(self.memory) < batch_size:
            return
        batch = random.sample(self.memory, batch_size)
        for state, action, reward, next_state, done in batch:
            state_t = torch.tensor(state, dtype=torch.float32).unsqueeze(0)
            next_state_t = torch.tensor(next_state, dtype=torch.float32).unsqueeze(0)
            target = reward
            if not done:
                target += self.gamma * torch.max(self.model(next_state_t)).item()
            output = self.model(state_t)[0][action]
            loss = self.criterion(output, torch.tensor(target))
            self.optimizer.zero_grad()
            loss.backward()
            self.optimizer.step()

Training Loop

env = MalwareEnv('path_to_malware_sample.exe')
state_size = env.observation_space.shape[0] * env.observation_space.shape[1] * env.observation_space.shape[2]
action_size = env.action_space.n

agent = DQNAgent(state_size, action_size)
episodes = 100
batch_size = 32

for e in range(episodes):
    state, _ = env.reset()
    state = state.flatten()
    total_reward = 0
    for time in range(env.max_steps):
        action = agent.act(state)
        next_state, reward, done, _, _ = env.step(action)
        next_state = next_state.flatten()
        agent.remember(state, action, reward, next_state, done)
        state = next_state
        total_reward += reward
        if done:
            break
    agent.replay(batch_size)
    print(f"Episode {e+1}/{episodes}, Total Reward: {total_reward}")

Explanation:

Experience Replay: The agent stores experiences and samples batches for training.
Target Update: The target is calculated based on the reward and future estimated rewards.
Training: The agent minimizes the difference between the predicted Q-values and the target.

Generating Realistic Malware Samples

After training, the agent should have learned a sequence of actions that modify the malware sample to evade the dynamic classifier.

Applying the Learned Policy

state, _ = env.reset()
state = state.flatten()
done = False
while not done:
    action = agent.act(state)
    next_state, reward, done, _, _ = env.step(action)
    state = next_state.flatten()

Extracting the Modified Sample

The modified sample is stored in env.current_sample. You can save it for analysis.

env.current_sample.write('modified_sample.exe')

Using Generated Samples for Defensive Training

Augmenting the Dataset

Add the modified samples to your training dataset to improve the classifier’s robustness.

# Assume you have a function to add samples to your dataset
add_sample_to_dataset('modified_sample.exe', label='malware')

Retraining the Classifier

Retrain your dynamic classifier with the augmented dataset to enhance its ability to detect new evasion techniques.

Conclusion

By integrating the dynamic memory heatmap classifier into a malware gym environment, we can simulate an adversarial setting where an RL agent learns to modify malware samples to evade detection. This process helps us:

Understand Potential Evasion Techniques: Anticipate how attackers might adapt.
Improve Defensive Measures: Use generated samples to train more robust classifiers.
Advance Cybersecurity Research: Explore the interplay between attackers and defenders using RL.

References

Bypassing Antivirus Detection: Old-School Malware, New Tricks: Link to Paper
LIEF Library: https://lief.quarkslab.com/
Qiling Framework: https://qiling.io/
OpenAI Gymnasium: https://gymnasium.farama.org/
PyTorch: https://pytorch.org/