DEV Community

Gaganpreet Singh
Gaganpreet Singh

Posted on Originally published at willeysingh.wordpress.com

Classifying Malware Variants with Graph Neural Networks

#cybersecurity #machinelearning #python #kubernetes

Malware authors rarely write a new family from scratch. They recompile, repack and patch, so the bytes change while the program’s structure stays much the same. Hash-based detection misses those variants. Structure-based classification can catch them. This project is my build of that idea: a Graph Neural Network that classifies malware variants from their control flow graphs and call graphs.

Code: github.com/Willey2003/malware-variant-classifier

How it works

  1. Feature extraction. Each binary is disassembled with Capstone, and its control flow graph and call graph are built with NetworkX. Each node gets a small feature vector describing its basic block.
  2. Graph dataset. The NetworkX graphs are converted to PyTorch Geometric Data objects, labelled by family.
  3. Model. A configurable GNN encoder (GCN, GAT or GraphSAGE, 3 layers, 128 hidden units, 64-dimensional embedding) with pooling and a classification head over 10 families.
  4. Training. A training loop with validation split, checkpoints and MLflow experiment tracking.
  5. Inference API. A FastAPI service with /health, /predict and /predict/batch endpoints that accept uploaded binaries and return the predicted family with a confidence score.

Deployment

The service ships with a Dockerfile, a docker-compose file and Kubernetes manifests, so the same image runs on a laptop or a cluster.

pip install -r requirements.txt
python src/training/train.py --config configs/train_config.yaml

docker-compose up --build     # local stack
kubectl apply -f k8s/         # Kubernetes
Enter fullscreen mode Exit fullscreen mode

Why graphs

  • A repacked or recompiled variant changes its hash and many of its bytes, but its control flow tends to survive.
  • GNNs learn from the shape of the graph, so two samples with similar structure land close together in embedding space even when their bytes differ.
  • The embedding is also useful on its own: nearest-neighbour search over embeddings groups unknown samples with their closest known family.

Status and next steps

This is an early build. The pipeline, model, API and deployment files are in place. Next is training on a public labelled dataset, publishing accuracy per family, and comparing the three GNN architectures. Only handle live malware samples in an isolated analysis environment.

Top comments (0)