PE Malware Detection Research

Research repository for Windows PE malware detection using ML/DL.

Abstract

As malware becomes increasingly sophisticated, deep learning—particularly Convolutional Neural Networks (CNNs)—has emerged as a powerful tool for automated classification.

Because CNNs are inherently optimized for spatial data, the format in which malware features are presented can significantly impact model performance.

This ongoing research investigates the efficacy of different data representation modalities in CNN-based malware classification. Utilizing the EMBER dataset, with planned expansions into multi-source datasets, we extract standard malware feature vectors and evaluate them across three distinct input formats: raw numerical feature vectors, converted 2D images, and converted WAV audio files.

The primary objective is to determine whether translating 1D vector data into spatial (image) or temporal (audio) formats improves the classification accuracy and feature extraction capabilities of a 2D CNN compared to baseline vector inputs.

Data policy

No malware binaries or large datasets are committed. See data/README.md for dataset notes.

Setup

Install Miniforge or another Conda distribution, then create the environment and install the local package:

conda env create -f environment.yml
conda activate malware_project
python -m ipykernel install --user --name malware_project --display-name "Python (malware_project)"

On Windows, make sure conda is available in a new terminal before running these commands. If conda env create fails while installing EMBER, confirm that Git is installed and available on PATH; this environment installs EMBER from https://github.com/elastic/ember.git because the malware research package is not published as a normal PyPI package.

The notebooks and conversion scripts expect vectorized EMBER 2018 feature files under data/ember2018/:

data/ember2018/
  X_train.dat
  X_test.dat
  y_train.dat
  y_test.dat

Download the EMBER 2018 dataset separately from Elastic and keep it local; this repository intentionally does not track datasets or generated representations.

Tech stack

• Python 3.11 (Conda)
• EMBER 2018 feature vectors (ember from GitHub)
• PyTorch + torchvision with CUDA (cu128 wheels on Windows)
• scikit-learn, NumPy, Pillow, matplotlib, tqdm
• pytest for converter tests

Project layout

src/malware_pck/
  common/                 # EMBER loading, paths, split manifests
  images/                 # PNG generation, inspection, image CNN
  audio/                  # WAV generation, EmberAudioDataset, 1D CNN
scripts/
  images/                 # generate, generate_test, train_cnn, evaluate_cnn
  audio/                  # generate, generate_test, train_cnn, evaluate_cnn
  README.md               # Script index
models/                   # Saved checkpoints (local only)
data/                     # Local datasets and generated artifacts
tests/
  common/  images/  audio/

Primary workflow: image CNN (PyTorch)

Run the scripts under scripts/images/ (see scripts/README.md). For metric tables and plots, use notebooks/evaluate_image_cnn.ipynb after training and test-image generation.

Methodology (leakage-safe)

1. Sample labeled train vectors with generation_seed (default 42).
2. Split train/val with a separate split_seed (default 43) and write data/images/split_manifest.json.
3. Fit MinMaxScaler on train indices only, then generate train+val PNGs.
4. Train CNN using manifest indices; validation is used only for early stopping on AUC (no LR schedule on val).
5. Tune threshold on training predictions (not validation).
6. Generate official test images into data/images_test/ (scaler still fit on train only).
7. Evaluate once on test with scripts/images/evaluate_cnn.py.

Grouped/family-held-out splits are not available from vectorized EMBER features alone; the code uses stratified index splits.

Commands

conda activate malware_project

# Regenerate 500k train/val images (requires --clear-existing if old dataset exists)
python scripts/images/generate.py 500000 --generation-seed 42 --split-seed 43 --clear-existing

# Train (reads split_manifest.json from data/images/)
python scripts/images/train_cnn.py --batch-size 256 --workers 8

# Official EMBER test images + one-shot test metrics
python scripts/images/generate_test.py --clear-existing
python scripts/images/evaluate_cnn.py

Generate representations

python scripts/images/generate.py 1000 --generation-seed 42 --split-seed 43
python scripts/audio/generate.py   # reuses data/images/split_manifest.json
python scripts/images/generate_test.py
python scripts/audio/generate_test.py

generation_seed and split_seed are intentionally independent. Pass --clear-existing to replace an existing image set. Audio generation is manifest-driven by default so the audio and image datasets share the same EMBER indices; use --samples N to override with a standalone seed-based draw. Audio 1D CNN eval plots: notebooks/evaluate_audio_cnn.ipynb.

AI session continuity

Local session notes live in .sessions/ (gitignored). Cursor loads them automatically via .cursor/rules/session-continuity.mdc at the start of each chat. When you close a session, ask the agent to append a new dated breakdown file there.

Tests

pytest