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) and 1.0 (Malware).
  • By default, a threshold of 0.85 is 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.

  1. Download the EMBER 2017 v1 dataset archive (approx. 1.6 GB compressed): Download EMBER 2017 v1
  2. 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.jsonl files and test_features.jsonl.
    • Keep track of the extraction path (e.g., C:\ember_dataset\ember).

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.

  1. Clone the EMBER repository:
    git clone https://github.com/elastic/ember.git
    cd ember
    
  2. Install the library:
    pip install -r requirements.txt
    python setup.py install
    
  3. Open a Python shell and vectorize the JSONL files:
    import ember
    ember.create_vectorized_features("C:\\path\\to\\extracted\\ember")
    
    This process will read all JSONL files, extract the 2351 features, and create .dat memory-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.

  1. Open Train.py in this repository and update the data_dir variable to point to your vectorized EMBER directory:
    data_dir = r"C:\path\to\extracted\ember"
    
  2. 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 with scale_pos_weight, and utilizing all CPU cores with n_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.txt in your dataset directory.

Step 4: Evaluate the Model (Evaluate.py)

To ensure your trained model is effective against unseen threats, use the evaluation script.

  1. Open Evaluate.py and ensure the data_dir variable correctly points to your EMBER directory:
    data_dir = r"C:\path\to\extracted\ember"
    
  2. 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).

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