Replication package developed during the conduction of our publication on which we examined different input representations and model architectures for the automated categorization of software vulnerabilities into Common Weakness Enumeration (CWE) categories directly from source code.
The high-level overview of the overall methodology that we followed in our study is illustrated in the figure below:
Replication Package of our research work entitled "Context Matters: Vulnerability Categorization with Dual-Input Code Representations"
To replicate the analysis and reproduce the results:
git clone https://github.com/annavou/Vulnerability-Categorization.git
and navigate to the cloned repository.
Inside the Vulnerability Categorization-main folder, there is a folder:
• ml: The folder with the ML pipelines that perform the experiments of the analysis. It contains five scripts, one per approach examined in the paper and a sub-folder:
-
data: The folder that contains the dataset files used in the experiments. It includes:
- train.csv, test.csv, val.csv: The training, test, and validation splits of the Big-Vul dataset. Download the Big-Vul dataset from the official repository and place the CSV files in this folder, following the data splitting strategy of LineVul.
- cwe_type_mapping.json: A JSON file containing the mapping of each CWE-ID to its corresponding CWE-Type, used in the dual-task learning experiment. This file is included in the repository.
-
functionsOnly.ipynb: performs the fine-tuning of a single CodeBERT encoder on the full function body only
-
flawLinewOnly.ipynb: performs the fine-tuning of a single CodeBERT encoder on the vulnerable lines only
-
dualInput.ipynb: performs the fine-tuning of a single CodeBERT encoder on the combined dual-input of function body and vulnerable lines, using the text pair tokenization strategy ([CLS] A [SEP] B [SEP])
-
dualEncoder.ipynb: performs the fine-tuning of two independent CodeBERT encoders, one dedicated to the function body and one to the vulnerable lines, whose embeddings are concatenated and passed to a shared classifier
-
dualTaskTraining.ipynb: extends the dual-encoder approach with a dual-task learning strategy, training the model simultaneously on CWE-ID and CWE-Type classification
Special thanks to HuggingFace for providing the Transformers library.
Special thanks to the paper entitled "A C/C++ Code Vulnerability Dataset with Code Changes and CVE Summaries" for providing the Big-Vul dataset, which was utilized as the basis for our analysis. If you use Big-Vul, please cite:
J. Fan, Y. Li, S. Wang and T. N. Nguyen, "A C/C++ Code Vulnerability Dataset with Code Changes and CVE Summaries," 2020 IEEE/ACM 17th International Conference on Mining Software Repositories (MSR), Seoul, Korea, Republic of, 2020, pp. 508-512, doi: 10.1145/3379597.3387501
Special thanks to the paper entitled "LineVul: A Transformer-Based Line-Level Vulnerability Prediction" for the data splitting strategy adopted in our experiments. If you use this splitting strategy, please cite:
M. Fu and C. Tantithamthavorn, "LineVul: A Transformer-Based Line-Level Vulnerability Prediction," 2022 IEEE/ACM 19th International Conference on Mining Software Repositories (MSR), 2022.
Special thanks to the paper entitled "AIBugHunter: A Practical Tool for Predicting, Classifying and Repairing Software Vulnerabilities" for the dual-task learning strategy adopted in our experiments. If you use this strategy, please cite:
M. Fu, C. Tantithamthavorn, T. Le, Y. Kume, V. Nguyen, D. Phung, and J. Grundy, "AIBugHunter: A Practical Tool for Predicting, Classifying and Repairing Software Vulnerabilities," Empirical Software Engineering, vol. 29, no. 1, p. 4, 2024.