Generative AI Google Brain

gemini-use-cases-code-code-scanning-and-vulnerability-detection

3 min read

Code Vulnerability Scanning & Automated Remediation using Gemini API in Vertex AI (Gemini 2.0)

Source notebook

Repo path: gemini/use-cases/code/code_scanning_and_vulnerability_detection.ipynb · Open on GitHub · intermediate

Scans Python files from GCS with Gemini 2.0 Flash and exports vulnerability reports to CSV and JSON.

Summary

This notebook demonstrates a code vulnerability scanning workflow using Gemini API in Vertex AI. It reads Python files from a Cloud Storage bucket, combines them into one prompt, asks Gemini 2.0 Flash to identify vulnerabilities and recommendations, then parses the markdown response into CSV and JSON reports. It also notes that the approach is experimental and not a mature automated remediation workflow.

Key code patterns

Initialize Vertex AI

import vertexai
 
vertexai.init(project=PROJECT_ID, location=REGION)

Sets the Google Cloud project and region before using Gemini through Vertex AI.

Batch Python Files From GCS

storage_client = storage.Client()
bucket = storage_client.get_bucket(BUCKET_NAME)
blobs = bucket.list_blobs(prefix=PREFIX)
 
combined_text = ""
for blob in blobs:
    if blob.name.endswith(".py"):
        file_content = blob.download_as_string().decode("utf-8")
        combined_text += f"### File: {blob.name} ###\n{file_content}\n"

Combines multiple .py files with filename separators so Gemini can analyze a codebase in one request.

Generate Vulnerability Review

model = GenerativeModel("gemini-2.0-flash")
responses = model.generate_content(
    contents=my_prompt,
    generation_config=generation_config,
    stream=True,
)
 
for res in responses:
    print(res.text)

Uses Gemini 2.0 Flash with a structured prompt and streaming output for vulnerability findings.

Parse Report To DataFrame

for match in re.finditer(r"###.*?(?=###|$)", text, re.DOTALL):
    report = match.group(0)
    file_name = re.search(file_pattern, report).group(1)
    vulnerability_name = re.search(vulnerability_name_pattern, report).group(1)
    recommended_code = code_match.group(1).strip() if code_match else "N/A"
    data.append({...})
 
return pd.DataFrame(data)

Extracts file names, vulnerability descriptions, recommendations, and code snippets from markdown output.

Export Reports

df.to_csv("vulnerability_report_BULK.csv", index=False)
 
data = extract_vulnerability_data(response_text)
data.to_json("vulnerabilities.json", indent=4)

Creates CSV and JSON artifacts for further analysis, benchmarking, or security tool integration.

Models & APIs used

  • Models: gemini-2.0-flash
  • APIs / services: Vertex AI, Cloud Storage
  • SDKs / libraries: google-cloud-aiplatform, google-cloud-storage, vertexai, pandas

When to use this

Use this pattern to prototype Gemini-assisted vulnerability analysis across multiple Python files stored in Cloud Storage.

Gotchas & caveats

  • Requires an existing Google Cloud project with the Vertex AI API enabled.
  • Colab requires auth.authenticate_user(); Vertex AI Workbench does not require that Colab-only step.
  • The notebook installs updated packages and requires a runtime restart before imports work.
  • The prompt says safety filters have not been imported for this notebook.
  • Regex parsing depends on the model following the requested markdown structure.
  • The notebook says Gemini 2.0 can be forced to respond in JSON, but this example parses markdown instead.
  • The approach is experimental and not considered a robust security tool without further validation.

Best practices

  • Add each filename as a separator before its code so the model can identify findings per file.
  • Use a clear prompt that specifies vulnerability name, description, recommendations, and recommended code.
  • Set generation parameters including temperature, top_p, top_k, candidate_count, and max_output_tokens.
  • Export findings to CSV and JSON for further analysis, benchmarking, and integration with security tools.
  • Treat the workflow as experimental and continue validation before using it as a robust security tool.

Graph View

Backlinks