predict.py

#!/usr/bin/env python3
"""Full pipeline: architecture detection followed by per-byte code/data labelling."""

import os
import sys

_HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, _HERE)
sys.path.insert(0, os.path.join(_HERE, 'arch_classifier'))

from PackedClassifier import PackedClassifier
from models.BidirRNN import BidirRNN

_ARCH_CLF = os.path.join(_HERE, 'arch_classifier', 'packedClassifiers', 'pretrained_classifier.pkl')
_LSTM_DIR = os.path.join(_HERE, 'code_section_identification', 'bilstm-classifiers')

# Arch-classifier label -> (LSTM model filename suffix, arch name for BidirRNN.predict)
_ARCH_MAP = {
    'amd64':  ('lstm_amd64',   'amd64'),
    'arm64':  ('lstm_arm64',   'arm64'),
    'armel':  ('lstm_armel',   'armel'),
    'armhf':  ('lstm_arm32',   'arm32'),
    'i386':   ('lstm_i386',    'i386'),
    'mips':   ('lstm_mips',    'mips'),
    'mipsel': ('lstm_mipsel',  'mipsel'),
    'ppc64el':('lstm_powerpc', 'powerpc'),
}


def predict(binary_path):
    """Detect the architecture of a binary and label each byte as code (1) or data (0).

    Args:
        binary_path: path to the binary file (ELF, PE, or Mach-O).

    Returns:
        arch (str): detected architecture label.
        labels (np.ndarray): integer array of length == file size, 1=code 0=data.

    Raises:
        ValueError: if the detected architecture has no matching LSTM model.
    """
    clf = PackedClassifier.load_classifier(_ARCH_CLF)
    clf.entropy_threshold = 6.667   # not persisted in the pkl; matches original inference script
    with open(binary_path, 'rb') as f:
        blob = f.read()
    arch = clf.predict(clf.gen_feature_vector(blob))[0]

    entry = _ARCH_MAP.get(arch)
    if entry is None:
        raise ValueError(
            f"No LSTM model for detected architecture '{arch}'. "
            f"Supported: {list(_ARCH_MAP)}"
        )
    lstm_file, lstm_arch = entry
    model = BidirRNN.load(os.path.join(_LSTM_DIR, lstm_file))
    labels = model.predict(binary_path, lstm_arch)
    return arch, labels


if __name__ == '__main__':
    import argparse
    import numpy as np

    parser = argparse.ArgumentParser(
        description='Detect architecture and label code/data bytes in a binary.'
    )
    parser.add_argument('binary', help='Path to the binary file')
    args = parser.parse_args()

    arch, labels = predict(args.binary)

    total = len(labels)
    code_bytes = int(np.sum(labels))
    data_bytes = total - code_bytes

    print(f'Architecture : {arch}')
    print(f'Total bytes  : {total}')
    print(f'Code bytes   : {code_bytes}  ({100 * code_bytes / total:.1f}%)')
    print(f'Data bytes   : {data_bytes}  ({100 * data_bytes / total:.1f}%)')