Revolutionizing Security Audits with AI: Using Open Source Large Language Models to Analyze Source Code

January 2, 2023

One emerging technology that has the potential to revolutionize security audits is artificial intelligence (AI), specifically machine learning algorithms. In this post, we'll explore how AI can be used to analyze source code and identify security vulnerabilities.

One type of AI model that has gained significant attention in the security field is the large language model.

Large language models are trained on massive datasets of text and are able to understand and generate human-like language. Examples of open source large language models include GPT-3 and BERT.

Let's explore how you can use an open source large language model to conduct security audits on your source code.

This can be especially useful for identifying potential vulnerabilities that might be missed by traditional static analysis tools.

Here is a Python script that uses the Hugging Face transformers library, specifically mrm8488/codebert-base-finetuned-detect-insecure-code to access a large language model and check source code for potential security issues:

from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import numpy as np

# Load the model
model = AutoModelForSequenceClassification.from_pretrained('mrm8488/codebert-base-finetuned-detect-insecure-code')

# Function to classify if code file as vulnerable or not
def classify_vulnerability(file_name, model):
  with open(file_name, "r") as f:
    code = f.read()
  # Tokenize
  tokenizer = AutoTokenizer.from_pretrained('mrm8488/codebert-base-finetuned-detect-insecure-code')
  inputs = tokenizer(code, return_tensors="pt", truncation=True, padding='max_length')
  labels = torch.tensor([1]).unsqueeze(0)  # Batch size 1
  outputs = model(**inputs, labels=labels)
  logits = outputs.logits
  # Check if vulnerable, this model returns 1 for vulnerable and 0 for not vulnerable
  vuln_check = np.argmax(logits.detach().numpy())
  if vuln_check == 1:
    return 'Vulnerable'
  return 'Not Vulnerable'

# Test the model on a file
file_name = 'insecure.java'
vulnerability = classify_vulnerability(file_name, model)
print(f'Vulnerability: {vulnerability}')

This script loads the specified model using the Hugging Face transformers library, reads in the source code file, and uses the model to analyze the source code and make predictions about potential security issues. The results are then printed to the console.

Using a large language model for security audits has several advantages. First, it allows you to scale your security efforts, as the model can analyze a large volume of source code quickly and efficiently. Second, it can potentially identify vulnerabilities that might be missed by traditional static analysis tools, as the model has a deep understanding of language and can identify patterns and anomalies that a human security expert might miss.

However, there are also challenges and limitations to using large language models for security audits.

One challenge is the need for a large and diverse dataset to train the model on. Without a sufficient dataset, the model may not accurately identify vulnerabilities or may miss them altogether. Additionally, there may be limitations to the types of vulnerabilities that the model can detect, as it may not have the same level of understanding and context as a human security expert.

To mitigate the risk of false positives and false negatives, it's important to properly validate the results of an AI-powered security audit