AI Bug Hunter VSCode Extension
AI-based software vulnerability prediction IDE plug-in that was trained using data from National Vulnerability Database (NVD), which can generate vulnerability predictions for a source code function.
This is a very early preview of the extension with a lot of improvements to be made in the future. As of version 1.0.0, the extension does not include one of the major features: repair, although this feature is scheduled to be released soon.
The extension will highlight the vulnerable line in the source code and the ML models will predict the following:
- The potentially vulnerable line in a function
- The types of the predicted vulnerability (i.e., CWE-ID, CWE Type)
- The severity of the predicted vulnerability
- The potential fix to the vulnerability
AIBugHunter will provide accurate vulnerability predictions that enable software developers to detect vulnerabilities in their IDE.
Our tool will have an interface that looks like the image below.
All you have to do is open any C/C++ file, and the tool with automatically start analysing in real time - the analysis should only take couple of seconds.
The file is analysed every time you modify the file, and it will start analysis after defined amount of time after stopping typing. (Refer to Extension Settings)
This extension will use the local inference mode by default, which runs a generic Python script containing the inference code. This requires some packages to be installed.
- Navigate to the Local Inference Repository
- Run the following command:
pip3 install -r requirements.txt
# or manually
pip3 install numpy onnxruntime torch transformers
# If you intend to use the GPU and you are not on MacOS or ARM processor
# You may replace onnxruntime with onnxruntime-gpu
Refer to the ONNX Runtime Website to determine whether to install
onnxruntime on your machine
There are two ways you can deploy an on-premise inference server:
- Running generic Python script with Uvicorn (aibughunter/remote-inference-py)
- Deploy using Docker Image provided <-- Recommended method (aibughunter/remote-inference-docker)
Refer to the respective repository's README on the deployment guide.
This extension provides the following settings:
AiBugHunter.diagnostics.delayBeforeAnalysis: Delay after paused typing before analysis (in milliseconds)
AiBugHunter.diagnostics.diagnosticMessageInformation: Select specific information to be displayed on the diagnostics UI
AiBugHunter.diagnostics.highlightSeverityType: Select the highlight severity type (i.e. the colour of the underline and it's classification)
AiBugHunter.diagnostics.informationLevel: Fluent or Verbose way of displaying the message. Fluent would be easier to read, verbose provides more information.
AiBugHunter.diagnostics.maxNumberOfLines: Top-n number of vulnerable lines to be highlighted
AiBugHunter.diagnostics.showDescription: Boolean to choose to show CWE description or not
AiBugHunter.inference.inferenceMode: Select Inference Mode. Local or On-Premise inference. On-Premise allows you to self-host inference server in a separate machine/server
AiBugHunter.inference.inferenceServerURL: Define the URL/IP of the on-premise server
AiBugHunter.inference.useCUDA: Option to use CUDA as an inference accelerator (NVIDIA GPUs only)
- Overly sensitive predictions
- CUDA mode on local inference does not provide any speed improvements
- Repair feature not implemented yet
- Data sent to the model can be reduced by selectively sending functions (Only functions modified)
Initial release of AIBugHunter
This AIBugHunter VSCode Extension Project is using Machine Learning Models published in the following publications:
Fu, Michael & Tantithamthavorn, Chakkrit. (2022). LineVul: A Transformer-based Line-Level Vulnerability Prediction. 10.1145/3524842.3528452.
Fu, Michael & Tantithamthavorn, Chakkrit & Le, Trung & Le, Trung & Nguyen, Van & Phung, Dinh. (2022). VulRepair: A T5-Based Automated Software Vulnerability Repair. 10.1145/3540250.3549098.
The replication packages of which are available here: