Tetteh-Apotey/malware-classifier

Malware Classifier using Machine Learning


Project Overview

This project is submitted in partial fulfillment of the requirements for the Introduction to Machine Learning course at Quantic School of Business and Technology.

The project builds and evaluates machine learning models to classify software as malware or goodware using static analysis features from PE executables.

Academic Context

• —Institution: Quantic School of Business and Technology
• —Program: Master of Science in Software Engineering (MSSE)
• —Course: Introduction to Machine Learning
• —Student: Isaac Tetteh-Apotey
• —GitHub Repository: https://github.com/life2allsofts/malware-classifier

Dataset

• —Source: Brazilian Malware Dataset (Ceschin et al., 2018)
• —Citation: F. Ceschin, F. Pinege, M. Castilho, D. Menotti, L.S. Oliveira and A. Gregio, "The Need for Speed: An Analysis of Brazilian Malware Classifiers," IEEE Security & Privacy, vol. 16, no. 6, pp. 31–41, Nov.–Dec. 2018
• —Samples: 50,181 PE files (40,144 after cleaning)
• —Features: 17 PE header features selected after leakage removal
• —Target: Binary (0 = Goodware, 1 = Malware)
• —Class Distribution: 58% Malware, 42% Goodware

Project Components

Model Development

• —Data Cleaning: Removal of identifier columns, string conversion, date feature extraction
• —Feature Engineering: Ratio features (CodeToImageRatio, SectionsToSizeRatio)
• —Leakage Detection: 14 leakage sources identified and removed
• —Models Evaluated: Logistic Regression, Decision Tree, Random Forest, XGBoost, LightGBM, CatBoost
• —Best Model: XGBoost (98.03% accuracy, 99.59% AUC-ROC)
• —Bias Correction: Threshold 0.6 with triage categories (BENIGN/SUSPICIOUS/MALWARE)

Web Application

• —Framework: Flask
• —Features: File upload, batch processing, manual input
• —API Endpoints: /health, /api/predict, /model_info
• —Local Access: http://localhost:5000

Documentation

• —Evaluation and Design
• —AI Tooling Strategy
• —Deployment Information
• —Model Results & Metrics

Project Structure

malware-classifier/
├── .github/
│   └── workflows/          # CI/CD pipeline configurations
├── app/                    # Flask web application
│   ├── app.py
│   └── templates/          # HTML templates
├── data/                   # Dataset files (gitignored)
├── docs/                   # Project documentation
├── models/                 # Saved models and pipelines
│   └── clean/              # Cleaned models (no leakage)
├── notebooks/              # EDA notebooks
├── results/                # Evaluation results and figures
│   └── clean/              # Results from cleaned models
├── src/                    # Source code
│   ├── fix_dtypes.py
│   ├── preprocess.py
│   ├── train_clean.py
│   ├── evaluate_clean.py
│   ├── check_data_leakage.py
│   └── fix_data_leakage.py
├── tests/                  # Unit tests
├── requirements.txt
└── README.md

Installation and Setup

1. Clone the repository

git clone https://github.com/life2allsofts/malware-classifier.git
cd malware-classifier

2. Create and activate virtual environment

python -m venv venv
# Linux/Mac
source venv/bin/activate
# Windows
.\\venv\\Scripts\\activate

3. Install dependencies

pip install -r requirements.txt

4. Download and prepare data

python download_data.py
python src/fix_dtypes.py
python src/preprocess.py

5. Train model

python src/train_clean.py

6. Run web application locally

python app/app.py

Access the application at: http://localhost:5000

Model Performance

Confusion Matrix (Test Set)

API Usage

Health Check

curl http://localhost:5000/health

Prediction

curl -X POST http://localhost:5000/api/predict \
  -H "Content-Type: application/json" \
  -d '{
    "BaseOfData": 0.5,
    "Characteristics": 256,
    "DllCharacteristics": 0,
    "Entropy": 6.2,
    "FileAlignment": 512,
    "ImageBase": 4194304,
    "Machine": 332,
    "NumberOfRvaAndSizes": 16,
    "NumberOfSections": 5,
    "NumberOfSymbols": 0,
    "PointerToSymbolTable": 0,
    "Size": 102400,
    "SizeOfCode": 51200,
    "SizeOfHeaders": 1024,
    "SizeOfImage": 204800,
    "SizeOfInitializedData": 25600,
    "SizeOfUninitializedData": 0
  }'

Expected Response

{
  "success": true,
  "prediction": "MALWARE",
  "confidence": "96.75%",
  "probabilities": {
    "benign": 0.0325,
    "malware": 0.9675
  }
}

Testing

Run the test suite:

pytest tests/ -v

Tests include:

• —Application functionality
• —Preprocessing pipeline
• —Data leakage detection
• —Model performance validation
• —API endpoint testing

Key Findings

Data Leakage Detection: Identified and removed 15 leakage sources including:

• —5 constant features
• —8 time-based features
• —2 highly correlated features

Model Performance: After cleaning, XGBoost achieved 98.03% accuracy with excellent precision-recall balance.

Bias Analysis: The model shows slight bias toward malware (75.7% probability on all-zero features), addressed through confidence thresholding.

Academic Integrity

This project represents my own work in accordance with Quantic's Academic Integrity Policy. The Brazilian Malware Dataset is used as permitted by the authors for research purposes.

Citation

If referencing this work:

Tetteh-Apotey, I. (2026). Malware Classifier using Machine Learning.
Quantic School of Business and Technology, MSSE Program.

Acknowledgments

• —Quantic School of Business and Technology for the project framework and guidance
• —Ceschin et al. for the Brazilian Malware Dataset
• —Open-source libraries: scikit-learn, XGBoost, LightGBM, CatBoost, Flask, PyTorch

License

This project is for educational purposes only. The dataset remains the property of its original authors.

Contact

📞 CONTACT INFORMATION Isaac Tetteh-Apotey

GitHub: https://github.com/life2allsofts

Hugging Face: https://huggingface.co/Tetteh-Apotey

LinkedIn: https://www.linkedin.com/in/isaac-tetteh-apotey-67408b89/