Bi-LSTM-based Section Boundary Detection

The purpose of this tool is to identify code and data sections in raw binaries. Given a binary file — ELF, PE, Mach-O, or any raw binary blob — it:

1. Detects the CPU architecture — using a logistic regression classifier trained on byte-frequency histograms and architecture-specific byte patterns.
2. Labels every byte as code or data — using a per-architecture Bidirectional LSTM trained on fixed-length byte windows.

Note on raw binary blobs: both inference steps operate purely on raw bytes and do not require a recognised binary format. The pipeline can be run on flat memory dumps, firmware images, or any arbitrary byte sequence. The ELF/PE/Mach-O format is only needed during training, to automatically extract ground-truth code masks from section headers.

---

Repository structure

arch_classifier/
    BinaryClassifier.py           # feature extraction + logistic regression model
    PackedClassifier.py           # variant that filters high-entropy (packed) regions first
    dataset_loader.py             # preprocess a folder of binaries → CSV (regular binaries)
    dataset_packed_loader.py      # same, for packed binaries
    generate_classifier.py        # train BinaryClassifier from CSV
    generate_packed_classifier.py # train PackedClassifier from CSV
    dataset/data.csv              # included training dataset
    packedClassifiers/
        pretrained_classifier.pkl # pretrained PackedClassifier (ready for inference)

code_section_identification/
    SectionPreprocess.py          # convert a binary to one-hot matrix + code mask (ground truth)
    lstm_preprocess.py            # batch-preprocess a training corpus → .npy arrays
    generate_LSTM.py              # train a BiLSTM model for one or all architectures
    evaluate_lstm.py              # evaluate trained models on a test corpus
    instructionsStatistics.py     # build per-architecture instruction-frequency dictionaries
    sectionMetrics.py             # section-level offset and boundary metrics
    byteweight/
        byteweightPreprocess.py   # function-start ground truth via objdump (ByteWeight approach)
    bilstm-classifiers/           # pretrained BiLSTM models (one per architecture)

models/
    BidirRNN.py                   # Bidirectional LSTM: training, inference, post-processing

predict.py                        # full pipeline inference (arch detection → byte labelling)

---

Setup

Build the Docker image and start a container:

docker build -t section-id .
docker run -it -v /path/to/your/data:/data section-id

The image runs Python 3.6 with TensorFlow 1.15, Keras 2.3.1, and scikit-learn 0.20.3. The project is at /app inside the container with PYTHONPATH=/app set. All commands below are meant to be run inside the container.

---

Full pipeline inference example

python predict.py /path/to/binary

Output:

Architecture : amd64
Total bytes  : 123456
Code bytes   : 78901  (63.9%)
Data bytes   : 44555  (36.1%)

It is also possible to retrieve the full code/data section prediction, i.e., per-byte labels where 1 means code and 0 means data. The pretrained architecture classifier and all per-architecture BiLSTM models are included in the repository and work out of the box without retraining.

---

Architecture classification

1. Preprocess the dataset

The classifier expects a CSV where each row is one binary. Prepare a directory tree whose sub-folders are named after the architecture:

dataset/
    amd64/   <binaries>
    arm64/   <binaries>
    ...

# Regular binaries
python arch_classifier/dataset_loader.py /data/dataset -o /data/data.csv -j 8

# Packed binaries
python arch_classifier/dataset_packed_loader.py /data/dataset -o /data/data_packed.csv -j 8

Options:

• -j — number of parallel workers (default: CPU count)
• -s — minimum file size in bytes to include (default: 0)

2. Train

# Regular classifier
python arch_classifier/generate_classifier.py /data/data.csv -o /data/classifier.pkl

# Packed classifier
python arch_classifier/generate_packed_classifier.py /data/data_packed.csv -o /data/pretrained_classifier.pkl

3. Inference

Use the full pipeline script (see Full pipeline inference example) or call the classifier directly:

from arch_classifier.PackedClassifier import PackedClassifier

clf = PackedClassifier.load_classifier('arch_classifier/packedClassifiers/pretrained_classifier.pkl')
clf.entropy_threshold = 6.667
with open('my_binary', 'rb') as f:
    blob = f.read()
arch = clf.predict(clf.gen_feature_vector(blob))[0]
print(arch)

---

BiLSTM code section detection

1. Preprocess the dataset

The BiLSTM is trained on fixed-length random byte windows sampled from binaries. To train models for all architectures, the training directory must follow the same layout as above (one sub-folder per architecture).

python code_section_identification/lstm_preprocess.py /data/training

This produces lstm_preprocess/X_<arch>_<window>.npy and lstm_preprocess/y_<arch>_<window>.npy for each architecture and window size (default: 500, 1000, 2000).

2. Train

# Train all architectures (saves to bilstm-classifiers/ by default)
python code_section_identification/generate_LSTM.py /data/training

# Train a single architecture
python code_section_identification/generate_LSTM.py /data/training -a amd64

# Custom output directory
python code_section_identification/generate_LSTM.py /data/training -o /data/models

Available architectures: amd64 arm32 arm64 armel i386 mips mips64 mipsel powerpc.

3. Evaluate

python code_section_identification/evaluate_lstm.py /data/test \
    --model-dir bilstm-classifiers \
    --out-dir /data/results

Per-architecture result files (accuracy, precision, recall, F1) are written to --out-dir.

4. Post-processing

Post-processing is controlled by flags passed to BidirRNN at construction time and applied automatically during predict():

Flag	What it does
post_process=True	Iteratively removes small spurious code/data fragments below a fraction of the largest segment
byteweight=True	Refines section boundaries using function-start offsets from a ByteWeight model
insStatPrediction=True	Adjusts section boundaries by picking the statistically most likely first instruction

To build the instruction-frequency dictionaries needed for insStatPrediction:

python code_section_identification/instructionsStatistics.py /data/training
# writes one dictionary per architecture to insStatDict/