Original implementation for paper Augmenting Decompiler Output with Learned Variable Names and Types.
DIRTY is a Transformer-based model which improves the quality of decompiler outputs by automatically generating meaningful variable names and types, assigning variable names that agree with those written by developers 66.4% of the time and types 75.8% of the time. We also release a large real-world dataset DIRT for this task, which consists of 75K+ programs and 1M+ human-written C functions mined from GitHub paired with their decompiler outputs.
- Linux with Python 3.6/3.7/3.8
- PyTorch โฅ 1.5.1
pip install -r requirements.txt
The first step to train DIRTY is to download the preprocessed DIRT dataset. If you wish to obtain the original unpreprocessed dataset, please open an issue.
cd dirty/
wget cmu-itl.s3.amazonaws.com/dirty/dirt.tar.gz -O dirt.tar.gz
tar -xzf dirt.tar.gz
The command would automatically download and decompress the dataset from Amazon S3.
If your machine does not have access to AWS, please manually download from the above link and untar it to data1/
.
We have setup configuration files for different models reported in our paper:
file | model | time (estimated hours) |
---|---|---|
multitask.xfmr.jsonnet | DIRTY-Multitask | 120 |
rename.xfmr.jsonnet | DIRTY-Rename | 80 |
retype.xfmr.jsonnet | DIRTY-Retype | 80 |
retype_nomem.xfmr.jsonnet | DIRTY_NDL | 80 |
retype_base.xfmr.jsonnet | DIRTY_S | 40 |
Training a models is as easy as specifying the name of the experimental run and the config file. Suppose we want to reproduce the Multi-task model in Table~7 in the paper:
cd dirty/
python exp.py train --cuda --expname=dirty_mt multitask.xfmr.jsonnet
Then, please watch for the line wandb: Run data is saved locally in ...
in the output.
This is where the logs and models are to be saved.
You can also monitor the automatically uploaded training and validation status (e.g., losses, accuracy) in your browser in real-time with the link printed after wandb: ๐ View run at ...
.
Feel free to adjust the hyperparameters in *.jsonnet
config files to train your own model.
As an alternative to train the model by yourself, you can download our trained DIRTY model.
cd dirty/
mkdir exp_runs/
wget cmu-itl.s3.amazonaws.com/dirty/dirty_mt.ckpt -O exp_runs/dirty_mt.ckpt
First, run your trained/downloaded model to produce predictions on the DIRE test set.
python exp.py train --cuda --expname=eval_dirty_mt multitask.xfmr.jsonnet --eval-ckpt <ckpt_path>
<ckpt_path>
is either exp_runs/dirty_mt.ckpt
if you download our trained model,
or saved during training at wandb/run-YYYYMMDD_HHMMSS-XXXXXXXX/files/dire/XXXXXXXX/checkpoints/epoch=N.ckpt
.
We sugguest changing beam_size
in config files to 0
to switch to greedy decoding, which is significantly faster.
The default configuration of beam_size = 5
can take hours.
The predictions will be saved to pred_XXX.json
.
This filename depends on models and can be modified in config files.
You can inspect the prediction results, which is in the following format.
{
binary: {
func_name: {
var1: [var1_retype, var1_rename], ...
}, ...
}, ...
}
Finally, use our standalone benchmark script:
python -m utils.evaluate --pred-file pred_mt.json --config-file multitask.xfmr.jsonnet
Where do I find the DIRTY paper?
We apologize for the inconvenience. It is currently under peer review.
Here is a walk-through of the code files of this repo.
The dirty/
folder contains the main code for the DIRTY model and DIRT dataset.
The entry point for running DIRTY experiments. It loads a configuration file, constructs a dataset instance, a model instance, and launches into a Trainer which runs training or inference according to configuration, and save logs and results in wandb.
Configuration files for running DIRTY experiments.
This folder contains neural models consisting of the DIRTY model.
โโโ dirty
โ โโโ model
โ โ โโโ beam.py # Beam search
โ โ โโโ decoder.py # factory class for building Decoders from configs
โ โ โโโ encoder.py # factory class for building Encoders from configs
โ โ โโโ model.py # training and evaluation step and metric logging
โ โ โโโ simple_decoder.py # A `decoder' consists of a linear layer,
# used for producing a soft mask from Data Layout Encoder
โ โ โโโ xfmr_decoder.py # Type/Multitask Decoder
โ โ โโโ xfmr_mem_encoder.py # Data Layout Encoder
โ โ โโโ xfmr_sequential_encoder.py # Code Encoder
โ โ โโโ xfmr_subtype_decoder.py # Not used in the current version
This folder contains code for the DIRT dataset, data preprocessing, evaluation, helper functions, and demos in the paper.
โโโ dirty
โ โโโ utils
โ โโโ case_study.py # Generate results for Table 3 and Table 6 in the paper
โ โโโ code_processing.py # Code canonicalization such as converting literals
โ โโโ compute_mi.py # Compute the mutual information between variables and types as a proof-of-concept for MT
โ โโโ dataset.py # A parallelized data loading class for preparing batched samples from DIRT for DIRTY
โ โโโ dataset_statistics.py # Compute dataset statistics
โ โโโ dire_types.py -> ../../binary/dire_types.py
โ โโโ evaluate.py # Evaluate final scores from json files saved from different methods for fair comparison
โ โโโ function.py -> ../../binary/function.py
โ โโโ ida_ast.py -> ../../binary/ida_ast.py
โ โโโ lexer.py
โ โโโ preprocess.py # Preprocess data produced from `dataset-gen/` into the DIRT dataset
โ โโโ util.py
โ โโโ variable.py -> ../../binary/variable.py
โ โโโ vocab.py
Empirical baselines included in the paper.
Use python -m baselines.<xxxxxx>
to run.
Results are saved to corresponding json files and can be evaluated with python -m utils.evaluate
.
โโโ dirty
โ โโโ baselines
โ โ โโโ copy_decompiler.py
โ โ โโโ most_common.py
โ โ โโโ most_common_decomp.py
The binary/
folder contains definitions for classes, including types, variables, and functions, constructed from decompiler outputs from binaries.
โโโ binary
โ โโโ __init__.py
โ โโโ dire_types.py # constructing types and a type library
โ โโโ function.py # definition and serialization for function instances
โ โโโ ida_ast.py # constructing ASTs from IDA-Pro outputs
โ โโโ variable.py # definition and serialization for variable instances
The idastubs/
folder contains helper functions used by the ida_ast.py
file.
The dataset-gen/
folder contains producing unpreprocessed data from binaries using IDA-Pro (required).
Legacy code for the DIRE paper.
If you use DIRTY/DIRT in your research or wish to refer to the baseline results, please use the following BibTeX.
@inproceedings {chen2021augmenting,
title = {Augmenting Decompiler Output with Learned Variable Names and Types},
booktitle = {31st USENIX Security Symposium (USENIX Security 22)},
year = {2022},
address = {Boston, MA},
url = {https://www.usenix.org/conference/usenixsecurity22/presentation/chen-qibin},
publisher = {USENIX Association},
month = aug,
}
@inproceedings {lacomis2019dire,
author={Lacomis, Jeremy and Yin, Pengcheng and Schwartz, Edward and Allamanis, Miltiadis and Le Goues, Claire and Neubig, Graham and Vasilescu, Bogdan},
booktitle={2019 34th IEEE/ACM International Conference on Automated Software Engineering (ASE)},
title={DIRE: A Neural Approach to Decompiled Identifier Naming},
year={2019},
volume={},
number={},
pages={628-639},
doi={10.1109/ASE.2019.00064}
}