An advanced, offline static malware detection engine that uses Machine Learning (LightGBM) and a comprehensive feature extraction pipeline to classify Windows Portable Executables (PE) as Clean or Malicious without executing them.
Technical Details
This project is built around the extraction of static features from raw binary files, drawing heavy inspiration from the EMBER (Endgame Malware BEnchmark for Research) dataset. The pipeline consists of two major components: Feature Extraction and Classification.
1. Feature Extraction (core/extractor.py)
Instead of analyzing raw bytes, the application relies on the lief parsing library to deconstruct the PE file and extract hundreds of meaningful attributes. The feature topology (currently tuned for EMBER feature version 1) converts the binary into a fixed-length numerical vector (2351 dimensions). The features include:
- Byte Histogram (256-dim): Counts the occurrences of each byte value, helping to identify packed or encrypted files.
- Byte Entropy Histogram (256-dim): Approximates the joint probability of byte value and local entropy, which is highly indicative of obfuscation.
- Section Information: Extracts properties of the PE sections (e.g.,
.text,.rdata,.data), including their count, sizes, virtual sizes, entropy, and access permissions (Read/Write/Execute). This is hashed to a fixed dimension. - Imported Libraries and Functions (1280-dim): Uses the "hashing trick" (via
FeatureHasher) to map the dynamically imported Windows APIs into a fixed-size vector. Malware often relies on specific combinations of APIs (e.g.,CreateRemoteThread,VirtualAllocEx). - Exported Functions (128-dim): Hashes the names of exported functions, which is useful for classifying malicious DLLs.
- General File Information (10-dim): Includes macroscopic characteristics such as file size, virtual size, debug info presence, total imported/exported functions, symbols, and signature presence.
- Header Information (62-dim): Extracts details from the COFF and Optional headers, such as machine architecture, subsystem, linker versions, and timestamps.
- String Information: Analyzes printable strings found in the binary, tracking metrics like average length, entropy, and counts of interesting patterns (URLs, paths, registry keys, and the "MZ" header).
2. Classification (MalwareDetector.py)
Once the binary is vectorized, it is passed to a pre-trained LightGBM gradient-boosting tree model (model.txt).
- The model evaluates the 2351-dimensional feature vector across thousands of decision trees.
- It outputs a decimal probability score between
0.0(Clean) and1.0(Malware). - By default, a threshold of
0.85is used to classify the file as Malicious, heavily reducing False Positives in production environments.
The CLI, powered by rich, presents these technical findings in a visually appealing and easy-to-read format, and even supports exporting Threat Intelligence data in STIX 2.1 format for SIEM integration.
Model Architecture Pipeline
graph TD
Data["EMBER Dataset<br>(1.1M+ PE Files)"]
Data -->|Feature Extraction & Hashing| Vec["Vectorized Database<br>(2351-dimensional arrays)"]
Vec --> |"Histogram algorithm<br>Leaf-wise growth strategy"| DT1["Decision Tree 1"]
Vec -.-> DT2["Decision Tree 2"]
Vec -.-> DTm["Decision Tree m"]
DT1 -->|"Residuals"| DT2
DT2 -->|"Residuals"| DTm
DT1 --> Tr1["Train"]
DT2 --> Tr2["Train"]
DTm --> Trm["Train"]
Tr1 --> Final["PE Malware LightGBM<br>Prediction Model (model.txt)"]
Tr2 --> Final
Trm --> Final
classDef default fill:#ffffff,stroke:#000000,stroke-width:1.5px,color:#000000;
Getting Started: Training and Evaluation
While a pre-trained model.txt may be provided, you can train a state-of-the-art model from scratch using the official EMBER dataset and the scripts included in this repository.
Follow these steps carefully:
Step 1: Download the EMBER Dataset
The EMBER dataset contains pre-extracted JSONL files of features from over 1.1 million Windows PE files.
- Download the EMBER 2017 v1 dataset archive (approx. 1.6 GB compressed): Download EMBER 2017 v1
- Extract the archive using a tool like 7-Zip (Windows) or
tar(Linux).- Note: The extracted folder will be around 5-6 GB and contain multiple
train_features_X.jsonlfiles andtest_features.jsonl. - Keep track of the extraction path (e.g.,
C:\ember_dataset\ember).
- Note: The extracted folder will be around 5-6 GB and contain multiple
Step 2: Install EMBER Utilities & Vectorize Data
To convert the raw JSONL feature files into memory-mapped NumPy arrays, you must install the official EMBER Python library.
- Clone the EMBER repository:
git clone https://github.com/elastic/ember.git cd ember - Install the library:
pip install -r requirements.txt python setup.py install - Open a Python shell and vectorize the JSONL files:
This process will read all JSONL files, extract the 2351 features, and createimport ember ember.create_vectorized_features("C:\\path\\to\\extracted\\ember").datmemory-mapped files (e.g.,X_train.dat,y_train.dat). This requires substantial disk space (around 15-20 GB) and RAM.
Step 3: Train the Model (Train.py)
Now that the data is vectorized, you can train the LightGBM model.
- Open
Train.pyin this repository and update thedata_dirvariable to point to your vectorized EMBER directory:data_dir = r"C:\path\to\extracted\ember" - Run the training script:
python Train.py- The script will load the memory-mapped arrays for both training and testing.
- It uses optimized hyper-parameters (such as tracking
auc, penalizing missed malware withscale_pos_weight, and utilizing all CPU cores withn_jobs=-1). - It employs early stopping to monitor the test set accuracy and prevent overfitting.
- Upon completion, it will save the final weights to
model.txtin your dataset directory.
Step 4: Evaluate the Model (Evaluate.py)
To ensure your trained model is effective against unseen threats, use the evaluation script.
- Open
Evaluate.pyand ensure thedata_dirvariable correctly points to your EMBER directory:data_dir = r"C:\path\to\extracted\ember" - Run the evaluation script:
python Evaluate.py- The script will load your newly trained
model.txt. - It runs predictions on the EMBER test set (200,000 files).
- It outputs a detailed report including:
- Overall Accuracy
- ROC AUC Score (The area under the receiver operating characteristic curve)
- Detailed Classification Metrics (Precision, Recall, F1-Score)
- Confusion Matrix (Showing exact counts of True Positives, False Positives, True Negatives, and False Negatives).
- The script will load your newly trained
Step 5: Deploy
Once you are satisfied with the evaluation metrics, copy the model.txt file from your dataset directory into the root folder of this repository. You can now use MalwareDetector.py to scan real-world executables!
python MalwareDetector.py path\to\suspicious.exe
Comments