JLB-JLB/android-ransomware-gat

Android ransomware detection GATs (thesis models)

Graph attention networks over class call graphs of Android APKs, trained for the Master's thesis Ransomware Detection using Graph Neural Networks Enhanced by Large Language Models (Joscha Lasse Bisping, TU Berlin). The node features are produced by a large language model that reads each decompiled class (default variants) or by a code embedding model (comparison arm).

Layout

<variant>/<held_out_family>/model.ckpt     Lightning checkpoint (state_dict + hyper_parameters)
<variant>/<held_out_family>/bundle.json    feature layout, model/training config, provenance
<variant>/<held_out_family>/metrics.json   test metrics of that fold

The thesis evaluated with family-leave-one-out: for each of the six ransomware families a model was trained on the other five (plus benign apps) and tested on the held-out family with an equally large benign sample. Every variant therefore comes as six models; each is the best of three seeds for its fold (test macro-F1, ties: AUROC, then loss). To score apps of a family the model must not have seen, take that family's fold; for general use, the pletor fold saw the largest share of the ransomware corpus.

verdict_behavior: enlarged thesis corpus, 2,166 benign and 213 ransomware apps (10.2:1)

Model	Seed	Test macro-F1	Test AUROC	n test
verdict_behavior/wipelocker	42	1.000	1.000	140
verdict_behavior/simplelocker	44	0.824	0.940	128
verdict_behavior/wannalocker	43	0.961	0.988	102
verdict_behavior/blackroselucy	44	0.333	0.059	34
verdict_behavior/pletor	44	1.000	1.000	12
verdict_behavior/filecoder	42	1.000	1.000	10

verdict: enlarged thesis corpus, 2,166 benign and 213 ransomware apps (10.2:1)

Model	Seed	Test macro-F1	Test AUROC	n test
verdict/wipelocker	42	1.000	1.000	140
verdict/simplelocker	44	0.764	0.909	128
verdict/wannalocker	44	0.922	0.963	102
verdict/blackroselucy	44	1.000	1.000	34
verdict/pletor	43	1.000	1.000	12
verdict/filecoder	44	1.000	1.000	10

embedding: base thesis corpus, 502 benign and 213 ransomware apps

Model	Seed	Test macro-F1	Test AUROC	n test
embedding/wipelocker	44	0.483	0.830	140
embedding/simplelocker	44	0.772	0.775	128
embedding/wannalocker	42	0.434	0.755	102
embedding/blackroselucy	43	0.333	0.872	34
embedding/pletor	42	1.000	1.000	12
embedding/filecoder	42	0.333	1.000	10

Test sets are the held-out family plus the same number of benign apps (n in the tables), so folds with few APKs carry little evidence. macro-F1 0.333 means the model flagged nothing (or everything) on that fold.

Intended use and limitations

These checkpoints are research artefacts that reproduce the thesis' family-leave-one-out experiments; they are meant for re-running or extending that evaluation, not for screening apps. Each model was trained without one ransomware family and scored on a balanced hold-out of that family, so none of the eighteen is a calibrated detector for real app traffic, where benign apps outnumber ransomware by orders of magnitude and the 0.5 threshold was never tuned. The verdict variants only work on node features produced by the same upstream step, Gemma 4 E4B with the prompt below; features from another model or prompt, or raw code, give meaningless outputs. The training data covers six ransomware families from 2014 to 2020 and Google Play apps of the same period, so other malware kinds and newer ransomware are outside what the models were trained to recognise.

Architecture (all variants)

GAT-large: 4 GATConv layers (8 heads x 32 = 256 channels, concat=True), each followed by BatchNorm1d(256) and ReLU; global_mean_pool over the nodes; head Linear(256,256) -> ReLU -> Dropout(0.5) -> Linear(256,2). Class-weighted cross-entropy, AdamW (lr 1e-3, weight decay 1e-4), early stopping on validation macro-F1. Output index 1 is ransomware; decision threshold 0.5.

Input

One graph per APK, the class call graph:

• Nodes: one per outer class with bytecode in the APK (inner and anonymous classes Foo$Bar fold into Foo). No library filter: every class is a node.
• Edges: directed, caller class -> callee class, whenever any method of the caller invokes any method of the callee (Androguard method cross-references aggregated to class level; intra-class calls dropped). Shape [2, n_edges], long.
• Node features x, shape [n_nodes, in_dim], float32, in this exact column order:

Variant	in_dim	Columns
verdict	2	potentially_malicious, potentially_ransomware (0/1)
verdict_behavior	7	the two verdict bits, then device_admin, screen_lock_or_overlay, sms_abuse, file_enumeration, anti_analysis (0/1)
embedding	768	nomic-ai/CodeRankEmbed embedding of the class source, L2-normalised

Verdict features come from google/gemma-4-E4B-it served with vLLM (temperature 0, structured JSON output, max_model_len 8192), one call per class with the class source truncated to 16,000 characters. Comments are stripped from the JADX source first. The JSON schema has nine booleans, crypto_use, file_enumeration, screen_lock_or_overlay, c2_network, sms_abuse, device_admin, anti_analysis, potentially_malicious, potentially_ransomware; crypto_use and c2_network are not fed to the models. The prompt (v3_evidence) is:

You are an Android security analyst. Examine the decompiled class below. First record
which concrete behaviours are present, then give an overall verdict. Base every field
strictly on code that is actually shown — do not guess. potentially_ransomware should
be true only when the class encrypts or locks files/the device and/or shows a ransom
demand, and it implies potentially_malicious. Respond with JSON only.

Class:
```
<class source>
```

Embedding features: SentenceTransformer("nomic-ai/CodeRankEmbed", trust_remote_code=True) with max_seq_length = 8192 and normalize_embeddings=True on the comment-stripped class source (JADX Java/Kotlin; Androguard DAD pseudo-Java or smali when JADX has no file for the class).

Minimal code (no pipeline needed)

Requires torch>=2.8, torch-geometric>=2.7, lightning>=2.6. The checkpoint is a PyTorch Lightning checkpoint. The architecture hyperparameters it stores (in_dim, hidden, num_layers, heads, dropout) are handed to the module below by load_from_checkpoint, so nothing has to be typed in by hand.

import lightning as L
import torch, torch.nn as nn
from torch_geometric.data import Batch, Data
from torch_geometric.nn import GATConv, global_mean_pool


class RansomwareGAT(L.LightningModule):
    def __init__(self, in_dim, hidden, num_layers, heads, dropout, **training_hparams):
        super().__init__()
        self.save_hyperparameters()
        convs, norms = [], []
        for i in range(num_layers):
            convs.append(GATConv(in_dim if i == 0 else hidden, hidden // heads, heads=heads))
            norms.append(nn.BatchNorm1d(hidden))
        self.model = nn.Module()
        self.model.convs = nn.ModuleList(convs)
        self.model.norms = nn.ModuleList(norms)
        self.model.head = nn.Sequential(nn.Linear(hidden, hidden), nn.ReLU(), nn.Dropout(dropout),
                                        nn.Linear(hidden, 2))
        # Training-fold class weights of the loss; stored in the checkpoint, unused at inference.
        self.register_buffer("class_weights", torch.ones(2))

    def forward(self, batch):
        x = batch.x
        for conv, norm in zip(self.model.convs, self.model.norms):
            x = torch.relu(norm(conv(x, batch.edge_index)))
        return self.model.head(global_mean_pool(x, batch.batch))

    def predict_step(self, batch, batch_idx=0):
        return torch.softmax(self(batch), dim=1)[:, 1]  # P(ransomware) per graph


model = RansomwareGAT.load_from_checkpoint("verdict_behavior/pletor/model.ckpt", map_location="cpu")
model.eval()

# One APK: n_nodes x in_dim features in the column order above, directed class-call edges.
graph = Data(
    x=torch.tensor(node_features, dtype=torch.float32),        # [n_nodes, 7] for verdict_behavior
    edge_index=torch.tensor(edges, dtype=torch.long).t(),      # [2, n_edges], rows = (caller, callee)
)
with torch.no_grad():
    p_ransomware = model.predict_step(Batch.from_data_list([graph])).item()
print("ransomware" if p_ransomware >= 0.5 else "benign", p_ransomware)

Several APKs can be scored in one call by passing a list of Data objects to Batch.from_data_list, or with L.Trainer().predict(model, loader) where loader is a torch_geometric.loader.DataLoader.

With the thesis' detection_pipeline repository the whole path from an APK file to this prediction (JADX decompilation, graph construction, Gemma verdicts) is python -m detection_pipeline fetch --variant verdict_behavior followed by scripts/run_predict.sh --bundle models/verdict_behavior/pletor app.apk.