Auto-Remediation Security Agent

Python / CodeBERT / PyGithub / Gemini API

A pipeline that hunts for vulnerable C code in the wild. It scrapes functions from popular GitHub repositories, embeds them with CodeBERT, and flags any that are semantically similar to known vulnerabilities from the Devign dataset. Flagged functions are then handed to an LLM, which names the likely vulnerability, rates its severity, and proposes a patched version.

Data Ingestion & Parsing

The pipeline begins by authenticating with the GitHub API to download C repositories with more than 150 stars, sorted ascending so the corpus reflects commonly used (but not enormous) projects. A run against 5 repositories yielded 2,929 C files. To analyze this code at function granularity, we needed to isolate individual function bodies from files that often lacked headers or dependencies, rendering standard AST compilers (like Clang) ineffective.

An initial regex-based extractor proved too slow per file, so we replaced it with a state-machine parser. Instead of full syntax tree generation, this lightweight algorithm iterates through file lines, tracking brace depth to extract syntactically complete C functions from "dirty" or partial source files in O(N). Against the scraped corpus it extracted 4,799 functions.

def parse_functions_from_file(file_path):
    inside_function = False
    brace_count = 0
    
    # Iterate line by line, tracking scope depth
    for line in lines:
        stripped = line.strip()
        
        # Detect function start
        if not inside_function and "(" in stripped and stripped.endswith("{"):
             inside_function = True
             brace_count = 1
        
        elif inside_function:
             function_body.append(line)
             # Track nested blocks to find true end of function
             brace_count += line.count("{") - line.count("}")
             
             if brace_count == 0:
                 functions.append("".join(function_body))

Vector Detection & LLM Remediation

We map each extracted function to a 768-dimensional vector using the Microsoft CodeBERT base model, and do the same for the 12,460 vulnerable functions in the Devign dataset (the CVElist database originally proposed turned out to contain only prose descriptions, no code). Both embedding sets are cached as .npy arrays so reruns skip regeneration. The detection engine computes cosine similarity in PyTorch; any GitHub function exceeding a 0.95 similarity threshold against a known vulnerable function is flagged. Flagged matches included integer overflows, format string risks, and out-of-bounds writes, at similarity scores between 0.96 and 0.99.

For remediation, we do not simply ask the LLM to fix the bug. Each flagged pair is injected into a prompt containing both the suspect GitHub function and the Devign function it matched, and the model (gemini-exp-1206) is asked to reason about their similarity before naming the specific vulnerability, assessing severity (Low to Critical), and generating a patched version of the code. Grounding the prompt in the retrieved exploit cuts down on hallucination and yields specific patches rather than generic advice. The prompt also states its educational purpose twice, since earlier runs got refused under Gemini's content policies for containing insecure code.

"""
Analyze the following two code snippets:

GitHub Code (Potentially Vulnerable):
```c
{github_code} ```
Known Vulnerable Code (For Comparison):
```c
{vulnerable_code} ```
Identify the most probable vulnerability in the GitHub code based on similarity.
Assess severity (Low/Medium/High/Critical).
Provide an improved version of the GitHub code that mitigates the vulnerability.
"""