wehub-resource-sync
902 lines
39 KiB
Plaintext
902 lines
39 KiB
Plaintext
{
|
|
"cells": [
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "nUWIDIzvODfw",
|
|
"metadata": {
|
|
"id": "nUWIDIzvODfw"
|
|
},
|
|
"source": [
|
|
"# Code Vulnerability Scanning & Automated Remediation using Gemini API in Vertex AI (Gemini 2.0)\n",
|
|
"\n",
|
|
"---\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "xHjHLyn8m2WF",
|
|
"metadata": {
|
|
"id": "xHjHLyn8m2WF"
|
|
},
|
|
"source": [
|
|
"<table align=\"left\">\n",
|
|
" <td style=\"text-align: center\">\n",
|
|
" <a href=\"https://colab.research.google.com/github/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/code_scanning_and_vulnerability_detection.ipynb\">\n",
|
|
" <img width=\"32px\" src=\"https://www.gstatic.com/pantheon/images/bigquery/welcome_page/colab-logo.svg\" alt=\"Google Colaboratory logo\"><br> Open in Colab\n",
|
|
" </a>\n",
|
|
" </td>\n",
|
|
" <td style=\"text-align: center\">\n",
|
|
" <a href=\"https://console.cloud.google.com/vertex-ai/colab/import/https:%2F%2Fraw.githubusercontent.com%2FGoogleCloudPlatform%2Fgenerative-ai%2Fmain%2Fgemini%2Fuse-cases%2Fcode%2Fcode_scanning_and_vulnerability_detection.ipynb\">\n",
|
|
" <img width=\"32px\" src=\"https://lh3.googleusercontent.com/JmcxdQi-qOpctIvWKgPtrzZdJJK-J3sWE1RsfjZNwshCFgE_9fULcNpuXYTilIR2hjwN\" alt=\"Google Cloud Colab Enterprise logo\"><br> Open in Colab Enterprise\n",
|
|
" </a>\n",
|
|
" </td> \n",
|
|
" <td style=\"text-align: center\">\n",
|
|
" <a href=\"https://console.cloud.google.com/vertex-ai/workbench/deploy-notebook?download_url=https://raw.githubusercontent.com/GoogleCloudPlatform/generative-ai/main/gemini/use-cases/code/code_scanning_and_vulnerability_detection.ipynb\">\n",
|
|
" <img src=\"https://lh3.googleusercontent.com/UiNooY4LUgW_oTvpsNhPpQzsstV5W8F7rYgxgGBD85cWJoLmrOzhVs_ksK_vgx40SHs7jCqkTkCk=e14-rj-sc0xffffff-h130-w32\" alt=\"Vertex AI logo\"><br> Open in Workbench\n",
|
|
" </a>\n",
|
|
" </td>\n",
|
|
" <td style=\"text-align: center\">\n",
|
|
" <a href=\"https://github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/code_scanning_and_vulnerability_detection.ipynb\">\n",
|
|
" <img width=\"32px\" src=\"https://raw.githubusercontent.com/primer/octicons/refs/heads/main/icons/mark-github-24.svg\" alt=\"GitHub logo\"><br> View on GitHub\n",
|
|
" </a>\n",
|
|
" </td>\n",
|
|
" <td style=\"text-align: center\">\n",
|
|
" <a href=\"https://goo.gle/4j91rJr\">\n",
|
|
" <img width=\"32px\" src=\"https://cdn.qwiklabs.com/assets/gcp_cloud-e3a77215f0b8bfa9b3f611c0d2208c7e8708ed31.svg\" alt=\"Google Cloud logo\"><br> Open in Cloud Skills Boost\n",
|
|
" </a>\n",
|
|
" </td>\n",
|
|
"</table>\n",
|
|
"\n",
|
|
"<div style=\"clear: both;\"></div>\n",
|
|
"\n",
|
|
"<b>Share to:</b>\n",
|
|
"\n",
|
|
"<a href=\"https://www.linkedin.com/sharing/share-offsite/?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/code_scanning_and_vulnerability_detection.ipynb\" target=\"_blank\">\n",
|
|
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/8/81/LinkedIn_icon.svg\" alt=\"LinkedIn logo\">\n",
|
|
"</a>\n",
|
|
"\n",
|
|
"<a href=\"https://bsky.app/intent/compose?text=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/code_scanning_and_vulnerability_detection.ipynb\" target=\"_blank\">\n",
|
|
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/7/7a/Bluesky_Logo.svg\" alt=\"Bluesky logo\">\n",
|
|
"</a>\n",
|
|
"\n",
|
|
"<a href=\"https://twitter.com/intent/tweet?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/code_scanning_and_vulnerability_detection.ipynb\" target=\"_blank\">\n",
|
|
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/5a/X_icon_2.svg\" alt=\"X logo\">\n",
|
|
"</a>\n",
|
|
"\n",
|
|
"<a href=\"https://reddit.com/submit?url=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/code_scanning_and_vulnerability_detection.ipynb\" target=\"_blank\">\n",
|
|
" <img width=\"20px\" src=\"https://redditinc.com/hubfs/Reddit%20Inc/Brand/Reddit_Logo.png\" alt=\"Reddit logo\">\n",
|
|
"</a>\n",
|
|
"\n",
|
|
"<a href=\"https://www.facebook.com/sharer/sharer.php?u=https%3A//github.com/GoogleCloudPlatform/generative-ai/blob/main/gemini/use-cases/code/code_scanning_and_vulnerability_detection.ipynb\" target=\"_blank\">\n",
|
|
" <img width=\"20px\" src=\"https://upload.wikimedia.org/wikipedia/commons/5/51/Facebook_f_logo_%282019%29.svg\" alt=\"Facebook logo\">\n",
|
|
"</a> \n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "zBrMoWSrXFI6",
|
|
"metadata": {
|
|
"id": "zBrMoWSrXFI6"
|
|
},
|
|
"source": [
|
|
"| | |\n",
|
|
"|-|-|\n",
|
|
"|Author(s) | [Souvik Mukherjee](https://github.com/talktosauvik)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "R2__rS9iZeVx",
|
|
"metadata": {
|
|
"id": "R2__rS9iZeVx"
|
|
},
|
|
"source": [
|
|
"## Background\n",
|
|
"\n",
|
|
"In today's digital landscape, software security is paramount. With the increasing sophistication of cyber threats, it's more important than ever for developers to proactively identify and address vulnerabilities in their code. Vulnerabilities can lead to data breaches, financial losses, and reputational damage. By harnessing the power of Gemini 2.0, we can help transform code vulnerability detection and remediation, and build a software vulnerability scanning mechanism"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "6j_ed1PcwS3f",
|
|
"metadata": {
|
|
"id": "6j_ed1PcwS3f"
|
|
},
|
|
"source": [
|
|
"## Overview\n",
|
|
"\n",
|
|
"Gemini 2.0, a member of Google Gemini family, is a generative AI model purpose-built for diverse multimodal applications. It's proficiency in understanding and generating content across text, code, and images makes it a powerful asset for intricate codebase analysis. With its expansive 2M token context window, Gemini 2.0 efficiently processes large code volumes in a single call, streamlining large-scale code scanning.Gemini 2.0's deep comprehension of programming languages and security best practices enables it to identify potential vulnerabilities and suggest helpful and contextual modifications. Learn more about [Gemini 2.0](https://deepmind.google/technologies/gemini/pro/).\n",
|
|
"\n",
|
|
"\n",
|
|
"\n",
|
|
"\n",
|
|
"\n",
|
|
"This experimental approach aims to efficiently scan large codebases, analyze multiple files in a single call, and delve deeper into complex code relationships and patternsThe model's deep analysis of code can help ensure comprehensive vulnerability detection, going beyond surface-level flaws. By using this approach, we can accommodate code written in several programming languages. Additionally, we can generate the findings and recommendations as JSON or CSV reports, which we would hypothetically use to make comparisons against established benchmarks and policy checks. With this tutorial, you learn how to use the Gemini API in Vertex AI, Google Cloud Storage API and the Vertex AI SDK to work with the Gemini 2.0 model to build a step by step code vulnerability scanning approach using Gemini 2.0 with the following steps:\n",
|
|
"\n",
|
|
"\n",
|
|
"* Read Python files from a GCS bucket and combining them into a single string\n",
|
|
"* Prompt engineering by crafting a clear and comprehensive prompt for Gemini 2.0, providing instructions for code analysis and output formatting\n",
|
|
"\n",
|
|
"* Submit the consolidated code string to Gemini 2.0 for analysis\n",
|
|
"* Extract vulnerability information, recommendations, and code snippets from the model response\n",
|
|
"\n",
|
|
"* Generate CSV and JSON output reports for further analysis, benchmarking and integration with security tools\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "TpJD9T2_PvPp",
|
|
"metadata": {
|
|
"id": "TpJD9T2_PvPp"
|
|
},
|
|
"source": [
|
|
"### Getting Started"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "k9uwJYsyP7rv",
|
|
"metadata": {
|
|
"id": "k9uwJYsyP7rv"
|
|
},
|
|
"source": [
|
|
"### Install Vertex AI and Google Cloud Storage SDKs for Python\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "N75LKYKXbuHp5cpzkJiq89Sb",
|
|
"metadata": {
|
|
"id": "N75LKYKXbuHp5cpzkJiq89Sb"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"%pip install -q --upgrade --user google-cloud-aiplatform google-cloud-storage"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "Xb3No6mSQFoe",
|
|
"metadata": {
|
|
"id": "Xb3No6mSQFoe"
|
|
},
|
|
"source": [
|
|
"### Restart the Kernel\n",
|
|
"\n",
|
|
"To use the newly installed packages, you must restart the current runtime."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "2l4Q9oy8f2KR",
|
|
"metadata": {
|
|
"id": "2l4Q9oy8f2KR"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# restart the current runtime to be able to access the downloaded packages\n",
|
|
"import IPython\n",
|
|
"\n",
|
|
"app = IPython.Application.instance()\n",
|
|
"app.kernel.do_shutdown(True)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "IIyao4PXQkAm",
|
|
"metadata": {
|
|
"id": "IIyao4PXQkAm"
|
|
},
|
|
"source": [
|
|
"### Authenticate the notebook environment (Colab only)\n",
|
|
"\n",
|
|
"If you are running this notebook on Google Colab, run the following cell to authenticate your environment. This step is not required if you are using [Vertex AI Workbench](https://cloud.google.com/vertex-ai-workbench)."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "pX9rkTlff8l-",
|
|
"metadata": {
|
|
"id": "pX9rkTlff8l-"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"from google.colab import auth\n",
|
|
"\n",
|
|
"auth.authenticate_user()"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "x4A6vhaKQqKT",
|
|
"metadata": {
|
|
"id": "x4A6vhaKQqKT"
|
|
},
|
|
"source": [
|
|
"### Set Google Cloud project information and initialize Vertex AI SDK\n",
|
|
"\n",
|
|
"\n",
|
|
"To get started using Vertex AI, you must have an existing Google Cloud project and [enable the Vertex AI API](https://console.cloud.google.com/flows/enableapi?apiid=aiplatform.googleapis.com).\n",
|
|
"\n",
|
|
"Learn more about [setting up a project and a development environment](https://cloud.google.com/vertex-ai/docs/start/cloud-environment).\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "-c8ZjhqIgA7K",
|
|
"metadata": {
|
|
"id": "-c8ZjhqIgA7K"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# initialize variables\n",
|
|
"\n",
|
|
"PROJECT_ID = \"[your-project-id]\" # @param {type:\"string\"}\n",
|
|
"REGION = \"us-central1\" # @param {type:\"string\"}\n",
|
|
"BUCKET_NAME = \"your-bucket-name\" # @param {type:\"string\"}\n",
|
|
"PREFIX = \"your prefix folder-name/\" # @param {type:\"string\"}"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "TwEE78OKgTRB",
|
|
"metadata": {
|
|
"id": "TwEE78OKgTRB"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# import and initialize Vertex AI\n",
|
|
"import vertexai\n",
|
|
"\n",
|
|
"vertexai.init(project=PROJECT_ID, location=REGION)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "zrgeiLeZRB8m",
|
|
"metadata": {
|
|
"id": "zrgeiLeZRB8m"
|
|
},
|
|
"source": [
|
|
"## Process Python files in batch\n",
|
|
"\n",
|
|
"This block of code reads Python files from the GCS bucket, combines their content, and adds respective `filename` as separator for LLM to better identify each file.\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "lZNJhxiignA5",
|
|
"metadata": {
|
|
"id": "lZNJhxiignA5"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"from google.cloud import storage\n",
|
|
"\n",
|
|
"\n",
|
|
"def process_py_files(bucket_name, prefix):\n",
|
|
" \"\"\"\n",
|
|
" Reads .py files from a GCS bucket and combines them into a single string.\n",
|
|
" Returns:\n",
|
|
" A string containing the combined content of all .py files.\n",
|
|
" \"\"\"\n",
|
|
"\n",
|
|
" storage_client = storage.Client()\n",
|
|
" bucket = storage_client.get_bucket(BUCKET_NAME)\n",
|
|
" blobs = bucket.list_blobs(prefix=PREFIX)\n",
|
|
"\n",
|
|
" combined_text = \"\"\n",
|
|
" for blob in blobs:\n",
|
|
" if blob.name.endswith(\".py\"):\n",
|
|
" file_content = blob.download_as_string().decode(\"utf-8\")\n",
|
|
" combined_text += f\"### File: {blob.name} ###\\n{file_content}\\n\"\n",
|
|
"\n",
|
|
" return combined_text\n",
|
|
"\n",
|
|
"\n",
|
|
"combined_string = process_py_files(BUCKET_NAME, PREFIX)\n",
|
|
"print(combined_string)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "ODdfYDIdRzzX",
|
|
"metadata": {
|
|
"id": "ODdfYDIdRzzX"
|
|
},
|
|
"source": [
|
|
"### Import Generative model library from Vertex AI"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "WCIOjVZYh28Q",
|
|
"metadata": {
|
|
"id": "WCIOjVZYh28Q"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"from IPython.display import display\n",
|
|
"from vertexai.generative_models import GenerationConfig, GenerativeModel"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "zyh4WLTZR42A",
|
|
"metadata": {
|
|
"id": "zyh4WLTZR42A"
|
|
},
|
|
"source": [
|
|
"### Initiate Gemini 2.0"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "sjA145qSh9O8",
|
|
"metadata": {
|
|
"id": "sjA145qSh9O8"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"model = GenerativeModel(\"gemini-2.0-flash\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "QHjEFzCBR_n9",
|
|
"metadata": {
|
|
"id": "QHjEFzCBR_n9"
|
|
},
|
|
"source": [
|
|
"### Setting up model configuration & Prompt template\n",
|
|
"\n",
|
|
"This piece of code sets up the model configurations & prompt template with 1 shot inference. For this specific notebook, the safety filters have not been imported.\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "rX2dng4piDHd",
|
|
"metadata": {
|
|
"id": "rX2dng4piDHd"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# define the prompt template to be passed to gemini.\n",
|
|
"context = combined_string\n",
|
|
"\n",
|
|
"my_prompt = f\"\"\"You are an expert code assistant. Review the following code for vulnerabilities and provide recommendations:\n",
|
|
"{context}\n",
|
|
"\n",
|
|
"Please format your response using markdown and Display with the following structure:\n",
|
|
"\n",
|
|
"file_name: Name of the code file\n",
|
|
"Vulnerability Name: Name of the identified vulnerability\n",
|
|
"Vulnerability: Description of the vulnerability and its potential impact.\n",
|
|
"Recommendations: List of actionable steps to mitigate the vulnerability.\n",
|
|
"Recommended code: Recommended code snippet to remove the suspected vulnerability\n",
|
|
"\n",
|
|
"I am also providing a sample response output that you should follow-\n",
|
|
"-------------------------------------------------\n",
|
|
"**file_name: bulk/example_1.py**\n",
|
|
"\n",
|
|
"**Vulnerability Name:** Information Exposure\n",
|
|
"\n",
|
|
"**Vulnerability:** The `server_bad` function returns the entire traceback when an exception occurs\n",
|
|
".This can leak sensitive information about the application's internals, such as file paths, variable names, and even the type of exception raised. Attackers could\n",
|
|
" exploit this information to gain a deeper understanding of the system and potentially launch further attacks.\n",
|
|
"\n",
|
|
"**Recommendations:**\n",
|
|
"\n",
|
|
"*Catch specific exceptions instead of using a broad `Exception` clause.\n",
|
|
"*Return a generic error message to the user without revealing internal details.\n",
|
|
"*Log the full traceback for debugging purposes, but do not expose it to the user.\n",
|
|
"\n",
|
|
"**Recommended code:**\n",
|
|
"\n",
|
|
"```python\n",
|
|
"sample code\n",
|
|
"```\n",
|
|
"\"\"\" # try your own prompt\n",
|
|
"\n",
|
|
"\n",
|
|
"generation_config = GenerationConfig(\n",
|
|
" temperature=0.5,\n",
|
|
" top_p=0.4,\n",
|
|
" top_k=24,\n",
|
|
" candidate_count=1,\n",
|
|
" max_output_tokens=8192,\n",
|
|
")\n",
|
|
"responses = model.generate_content(\n",
|
|
" contents=my_prompt,\n",
|
|
" generation_config=generation_config,\n",
|
|
" stream=True,\n",
|
|
")\n",
|
|
"\n",
|
|
"for res in responses:\n",
|
|
" print(res.text)"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "cnuNp-heSUMD",
|
|
"metadata": {
|
|
"id": "cnuNp-heSUMD"
|
|
},
|
|
"source": [
|
|
"## Capture the model response into a single variable\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "M_gJqJWAhALS",
|
|
"metadata": {
|
|
"id": "M_gJqJWAhALS"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"response_text = r\"\"\"\n",
|
|
" Security Analysis of Code Snippets\n",
|
|
"\n",
|
|
"### file_name: Bulk/example_1.py\n",
|
|
"\n",
|
|
"**Vulnerability Name:** Information Exposure Through Exception Handling\n",
|
|
"\n",
|
|
"**Vulnerability:** The `server_bad` function returns the entire traceback when an exception occurs. This can leak sensitive information about the application's internals, such\n",
|
|
" as file paths, variable names, and even the type of exception raised. Attackers could exploit this information to gain a deeper understanding of the system and potentially launch\n",
|
|
" further attacks.\n",
|
|
"\n",
|
|
"**Recommendations:**\n",
|
|
"\n",
|
|
"*Catch specific exceptions instead of using a broad `Exception` clause.\n",
|
|
"*Return a generic error message to the user without revealing internal details.\n",
|
|
"*Log the full traceback for debugging purposes, but do not expose it to the user.\n",
|
|
"\n",
|
|
"**Recommended code:**\n",
|
|
"\n",
|
|
"```python\n",
|
|
"@app.route('/bad')\n",
|
|
"def server_bad():\n",
|
|
" try:\n",
|
|
" result = do_computation()\n",
|
|
" return\n",
|
|
" result\n",
|
|
" except SpecificException as e:\n",
|
|
" # Log the error for debugging\n",
|
|
" app.logger.exception(\"An error occurred during computation\")\n",
|
|
" return \"An error occurred. Please try again later.\"\n",
|
|
"```\n",
|
|
"\n",
|
|
"### file_name: Bulk/example_10.py\n",
|
|
"\n",
|
|
"**Vulnerability Name:** Insecure Temporary File Creation\n",
|
|
"\n",
|
|
"**Vulnerability:** The `mktemp()` function creates temporary files with predictable names, which could allow attackers to guess the file names and access sensitive data. Additionally, the file permissions might be insecure, allowing unauthorized access.\n",
|
|
"\n",
|
|
"**Recommendations:**\n",
|
|
"\n",
|
|
"*Use the `tempfile .mkstemp()` function to create temporary files with more secure permissions and unpredictable names.\n",
|
|
"*Ensure that the temporary files are deleted after they are no longer needed.\n",
|
|
"\n",
|
|
"**Recommended code:**\n",
|
|
"\n",
|
|
"```python\n",
|
|
"import tempfile\n",
|
|
"\n",
|
|
"def write_results(results):\n",
|
|
" with tempfile.NamedTemporaryFile(mode=\"\n",
|
|
"w+\", delete=False) as f:\n",
|
|
" filename = f.name\n",
|
|
" f.write(results)\n",
|
|
" print(\"Results written to\", filename)\n",
|
|
" # ... process the file ...\n",
|
|
" os.remove(filename) # Delete the file when done\n",
|
|
"```\n",
|
|
"\n",
|
|
"### file_name: Bulk/example_100.py\n",
|
|
"\n",
|
|
"**Vulnerability Name:** Weak Key Generation Algorithm\n",
|
|
"\n",
|
|
"**Vulnerability:** The `DSA.generate(1024)` function generates a DSA key with a length of 1024 bits, which is considered weak and vulnerable to modern cryptanalysis techniques.\n",
|
|
"\n",
|
|
"\n",
|
|
"**Recommendations:**\n",
|
|
"\n",
|
|
"*Use a stronger key generation algorithm, such as RSA or ECC, with a key length of at least 2048 bits.\n",
|
|
"*Consider using libraries that provide secure defaults for key generation.\n",
|
|
"\n",
|
|
"**Recommended code:**\n",
|
|
"\n",
|
|
"```python\n",
|
|
"from Crypto.PublicKey import RSA\n",
|
|
"\n",
|
|
"\n",
|
|
"def generate_private_key():\n",
|
|
" key = RSA.generate(2048)\n",
|
|
" return key.export_key('PEM') # Export in a standard format\n",
|
|
"```\n",
|
|
"\n",
|
|
"### file_name: Bulk/example_101.py\n",
|
|
"\n",
|
|
"**Vulnerability Name:** Insecure Encryption Mode and Static IV\n",
|
|
"\n",
|
|
"**Vulnerability:** The code uses AES in CBC mode with a static initialization vector (IV). Using a static IV weakens the encryption and makes it vulnerable to certain attacks, especially when encrypting multiple messages with the same key.\n",
|
|
"\n",
|
|
"**Recommendations:**\n",
|
|
"\n",
|
|
"*Use a random IV for each encryption operation.\n",
|
|
"*Consider using more secure modes of operation, such as GCM or CTR, which provide authenticated encryption.\n",
|
|
"\n",
|
|
"**Recommended code:**\n",
|
|
"\n",
|
|
"```python\n",
|
|
"from Crypto.Cipher import AES\n",
|
|
"from Crypto.Random import get_random_bytes\n",
|
|
"\n",
|
|
"def getEncryptionKey(data, key):\n",
|
|
" iv\n",
|
|
" = get_random_bytes(AES.block_size)\n",
|
|
" cipher = AES.new(key, AES.MODE_GCM, iv)\n",
|
|
" ciphertext, tag = cipher.encrypt_and_digest(data)\n",
|
|
" return iv, ciphertext, tag\n",
|
|
"```\n",
|
|
"\n",
|
|
"### file_name: Bulk/example_102.py\n",
|
|
"\n",
|
|
"**Vulnerability Name:** Insecure FTP Usage and Hardcoded Credentials\n",
|
|
"\n",
|
|
"**Vulnerability:** The code uses FTP, which transmits data in plaintext, including usernames and passwords. Additionally, the credentials are hardcoded in the script, making them easily discoverable.\n",
|
|
"\n",
|
|
"**Recommendations:**\n",
|
|
"\n",
|
|
"\n",
|
|
"*Use a secure file transfer protocol, such as SFTP or FTPS, which encrypts the communication.\n",
|
|
"*Avoid hardcoding credentials and instead use environment variables or a secure credential storage mechanism.\n",
|
|
"\n",
|
|
"**Recommended code:**\n",
|
|
"\n",
|
|
"```python\n",
|
|
"import pysftp # Assuming you have the pysftp\n",
|
|
" library installed\n",
|
|
"\n",
|
|
"def downloadFile(hostname, username, password, sensitive_file):\n",
|
|
" with pysftp.Connection(hostname, username=username, password=password) as sftp:\n",
|
|
" sftp.get(sensitive_file, preserve_mtime=True)\n",
|
|
"```\n",
|
|
"\n",
|
|
"### file_name:Bulk/example_103.py\n",
|
|
"\n",
|
|
"**Vulnerability Name:** Potential Issues with Password Hashing\n",
|
|
"\n",
|
|
"**Vulnerability:** While PBKDF2 is a good choice for password hashing, the code uses a static salt, which reduces its effectiveness against rainbow table attacks.\n",
|
|
"\n",
|
|
"**Recommendations:**\n",
|
|
"\n",
|
|
"*Use a unique salt for each password. This can be achieved by generating a random salt for each user and storing it securely alongside the hashed password.\n",
|
|
"\n",
|
|
"**Recommended code:**\n",
|
|
"\n",
|
|
"```python\n",
|
|
"import os\n",
|
|
"\n",
|
|
"def getHash(password):\n",
|
|
" salt = os.urandom(16) # Generate a random\n",
|
|
" salt\n",
|
|
" hash = pbkdf2_hmac('sha256', password, salt, 100000)\n",
|
|
" return salt, hash # Store the salt with the hash\n",
|
|
"```\n",
|
|
"\n",
|
|
"### file_name: Bulk/example_104.py\n",
|
|
"\n",
|
|
"**Vulnerability Name:** Hardcoded and Weak Admin Password\n",
|
|
"\n",
|
|
"**Vulnerability:** The code uses a hardcoded and easily guessable password for the admin user. This makes the system extremely vulnerable to unauthorized access.\n",
|
|
"\n",
|
|
"**Recommendations:**\n",
|
|
"\n",
|
|
"*Never hardcode passwords, especially for privileged accounts.\n",
|
|
"*Enforce strong password policies that require a minimum length, complexity, and regular password changes.\n",
|
|
"*Consider using multi-factor authentication for additional security.\n",
|
|
"\n",
|
|
"**Recommended code:**\n",
|
|
"\n",
|
|
"```python\n",
|
|
"# Implement a secure password storage and verification mechanism, such as bcrypt or Argon2\n",
|
|
"# and avoid hardcoding passwords.\n",
|
|
"```\n",
|
|
"\n",
|
|
"### file_name: Bulk/example_105.py\n",
|
|
"\n",
|
|
"**Vulnerability Name:** Open Redirect Vulnerability\n",
|
|
"\n",
|
|
"**Vulnerability:** The `flask_redirect` function redirects the user to any URL provided in the `next` parameter without validation. This can be exploited by attackers to redirect users to malicious websites.\n",
|
|
"\n",
|
|
"**Recommendations:**\n",
|
|
"\n",
|
|
"*Validate the `next` parameter to ensure it points to a trusted domain or a specific set of allowed URLs.\n",
|
|
"\n",
|
|
"**Recommended code:**\n",
|
|
"\n",
|
|
"```python\n",
|
|
"from urllib.parse import urlparse\n",
|
|
"\n",
|
|
"@app.route('flask_redirect')\n",
|
|
"def flask_redirect():\n",
|
|
" url = request\n",
|
|
".args.get(\"next\")\n",
|
|
" parsed_url = urlparse(url)\n",
|
|
" if parsed_url.netloc in allowed_domains: # Replace with your allowed domains\n",
|
|
" return redirect(url)\n",
|
|
" else:\n",
|
|
" return \"Invalid redirect URL\"\n",
|
|
"```\n",
|
|
"\n",
|
|
"### file_name: Bulk/example_106.py\n",
|
|
"\n",
|
|
"**Vulnerability Name:** Potential XXE (XML External Entity) Vulnerability\n",
|
|
"\n",
|
|
"**Vulnerability:** The code parses XML data from the request without disabling external entity processing. This could allow attackers to inject malicious entities that can lead to information disclosure, denial of service\n",
|
|
", or server-side request forgery (SSRF) attacks.\n",
|
|
"\n",
|
|
"**Recommendations:**\n",
|
|
"\n",
|
|
"* Disable external entity processing when parsing XML data.\n",
|
|
"* Use a library that provides secure XML parsing options.\n",
|
|
"\n",
|
|
"**Recommended code:**\n",
|
|
"\n",
|
|
"```python\n",
|
|
"from defusedxml import lxml\n",
|
|
"\n",
|
|
"@app.post(\"/\n",
|
|
"upload\")\n",
|
|
"def upload():\n",
|
|
" xml_src = request.get_data()\n",
|
|
" parser = lxml.XMLParser(resolve_entities=False)\n",
|
|
" doc = lxml.fromstring(xml_src, parser=parser)\n",
|
|
" return lxml.tostring(doc)\n",
|
|
"\n",
|
|
"```\n",
|
|
"\n",
|
|
"### file_name: Bulk/example_107.py\n",
|
|
"\n",
|
|
"**Vulnerability Name:** Potential XXE and Insecure Password Handling\n",
|
|
"\n",
|
|
"**Vulnerability:** Similar to example_106.py, this code is vulnerable to XXE attacks due to insecure XML parsing. Additionally, the `check\n",
|
|
"Password` function is not implemented, and it's unclear how passwords are stored and verified.\n",
|
|
"\n",
|
|
"**Recommendations:**\n",
|
|
"\n",
|
|
"*Address the XXE vulnerability as described in example_106.py.\n",
|
|
"*Implement a secure password storage and verification mechanism, such as bcrypt or Argon2.\n",
|
|
"*Avoid storing passwords in plaintext or using weak hashing algorithms.\n",
|
|
"\n",
|
|
"**Recommended code:**\n",
|
|
"\n",
|
|
"```python\n",
|
|
"# Implement secure XML parsing as shown in example_106.py\n",
|
|
"# Implement a secure password storage and verification mechanism\n",
|
|
"```\n",
|
|
"\"\"\""
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "1tK5AtEdIeep",
|
|
"metadata": {
|
|
"id": "1tK5AtEdIeep"
|
|
},
|
|
"source": [
|
|
"### Parse the response and export it into a csv for further analysis.\n",
|
|
"\n",
|
|
"To parse the response we have used regular expression pattern matching. This is easy to configure and maintain\n"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "yVethA0vhsiv",
|
|
"metadata": {
|
|
"id": "yVethA0vhsiv"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"import re\n",
|
|
"\n",
|
|
"import pandas as pd\n",
|
|
"\n",
|
|
"\n",
|
|
"def extract_vulnerability_data(text):\n",
|
|
" # Regular expressions for extracting data\n",
|
|
" file_pattern = r\"###\\s*file_name:\\s*(.*\\.py)\" # Adjusted for \"file_name\"\n",
|
|
" vulnerability_name_pattern = r\"\\*\\*Vulnerability Name:\\*\\*\\s*(.*)\"\n",
|
|
" vulnerability_pattern = r\"\\*\\*Vulnerability:\\*\\*\\s*(.*?)(?=\\*\\*Recommendations)\"\n",
|
|
" recommendation_pattern = r\"\\*\\*Recommendations:\\*\\*\\s*((?:\\*.*\\n)+)\"\n",
|
|
" code_pattern = r\"```python(.*?)```\" # Added pattern for recommended code\n",
|
|
"\n",
|
|
" data = []\n",
|
|
"\n",
|
|
" # Iterate through each vulnerability report\n",
|
|
" for match in re.finditer(r\"###.*?(?=###|$)\", text, re.DOTALL):\n",
|
|
" report = match.group(0)\n",
|
|
"\n",
|
|
" file_name = re.search(file_pattern, report).group(1)\n",
|
|
" vulnerability_name = re.search(vulnerability_name_pattern, report).group(1)\n",
|
|
" vulnerability = (\n",
|
|
" re.search(vulnerability_pattern, report, re.DOTALL).group(1).strip()\n",
|
|
" )\n",
|
|
" recommendations = (\n",
|
|
" re.search(recommendation_pattern, report, re.DOTALL).group(1).strip()\n",
|
|
" )\n",
|
|
"\n",
|
|
" # Extract recommended code, handling potential absence\n",
|
|
" code_match = re.search(code_pattern, report, re.DOTALL)\n",
|
|
" recommended_code = code_match.group(1).strip() if code_match else \"N/A\"\n",
|
|
"\n",
|
|
" data.append(\n",
|
|
" {\n",
|
|
" \"File Number\": file_name.split(\"_\")[-1].split(\".\")[0],\n",
|
|
" \"File Name\": file_name,\n",
|
|
" \"Vulnerability Name\": vulnerability_name,\n",
|
|
" \"Description\": vulnerability,\n",
|
|
" \"Recommendations\": recommendations,\n",
|
|
" \"Recommended Code\": recommended_code, # Include recommended code\n",
|
|
" }\n",
|
|
" )\n",
|
|
"\n",
|
|
" return pd.DataFrame(data)\n",
|
|
"\n",
|
|
"\n",
|
|
"# Extract data and create DataFrame\n",
|
|
"df = extract_vulnerability_data(response_text)\n",
|
|
"\n",
|
|
"# Display DataFrame in Colab\n",
|
|
"display(df)\n",
|
|
"\n",
|
|
"# Save DataFrame to CSV file\n",
|
|
"df.to_csv(\"vulnerability_report_BULK.csv\", index=False)\n",
|
|
"\n",
|
|
"print(\"CSV file 'vulnerability_report_BULK.csv' created successfully.\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "ipY3TXNUieVo",
|
|
"metadata": {
|
|
"id": "ipY3TXNUieVo"
|
|
},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"application/javascript": "\n async function download(id, filename, size) {\n if (!google.colab.kernel.accessAllowed) {\n return;\n }\n const div = document.createElement('div');\n const label = document.createElement('label');\n label.textContent = `Downloading \"${filename}\": `;\n div.appendChild(label);\n const progress = document.createElement('progress');\n progress.max = size;\n div.appendChild(progress);\n document.body.appendChild(div);\n\n const buffers = [];\n let downloaded = 0;\n\n const channel = await google.colab.kernel.comms.open(id);\n // Send a message to notify the kernel that we're ready.\n channel.send({})\n\n for await (const message of channel.messages) {\n // Send a message to notify the kernel that we're ready.\n channel.send({})\n if (message.buffers) {\n for (const buffer of message.buffers) {\n buffers.push(buffer);\n downloaded += buffer.byteLength;\n progress.value = downloaded;\n }\n }\n }\n const blob = new Blob(buffers, {type: 'application/binary'});\n const a = document.createElement('a');\n a.href = window.URL.createObjectURL(blob);\n a.download = filename;\n div.appendChild(a);\n a.click();\n div.remove();\n }\n ",
|
|
"text/plain": [
|
|
"<IPython.core.display.Javascript object>"
|
|
]
|
|
},
|
|
"metadata": {},
|
|
"output_type": "display_data"
|
|
},
|
|
{
|
|
"data": {
|
|
"application/javascript": "download(\"download_aab82d58-53a1-40c6-be64-e5bb8ff671f5\", \"vulnerability_report_BULK.csv\", 9660)",
|
|
"text/plain": [
|
|
"<IPython.core.display.Javascript object>"
|
|
]
|
|
},
|
|
"metadata": {},
|
|
"output_type": "display_data"
|
|
}
|
|
],
|
|
"source": [
|
|
"# Display the created CSV file (assuming you have the 'google.colab' library installed)\n",
|
|
"from google.colab import files\n",
|
|
"\n",
|
|
"files.download(\"vulnerability_report_BULK.csv\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "7-kVfA_VIVv8",
|
|
"metadata": {
|
|
"id": "7-kVfA_VIVv8"
|
|
},
|
|
"source": [
|
|
"## Parse the response and export it into JSON output.\n",
|
|
"\n",
|
|
"NOTE: The following process is just a way to capture the response text and parse it further with regular expression. With Gemini 2.0, we can force the model to respond in JSON structure as well."
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "F3hgDQcbn5YO",
|
|
"metadata": {
|
|
"id": "F3hgDQcbn5YO"
|
|
},
|
|
"outputs": [],
|
|
"source": [
|
|
"# Extract data and write to JSON file (assuming you have 'response_text' defined)\n",
|
|
"data = extract_vulnerability_data(response_text)\n",
|
|
"\n",
|
|
"data.to_json(\"vulnerabilities.json\", indent=4)\n",
|
|
"\n",
|
|
"print(\"Vulnerability data extracted and saved to vulnerabilities.json\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "code",
|
|
"execution_count": null,
|
|
"id": "u-jhc3_Qpx9c",
|
|
"metadata": {
|
|
"id": "u-jhc3_Qpx9c"
|
|
},
|
|
"outputs": [
|
|
{
|
|
"data": {
|
|
"application/javascript": "\n async function download(id, filename, size) {\n if (!google.colab.kernel.accessAllowed) {\n return;\n }\n const div = document.createElement('div');\n const label = document.createElement('label');\n label.textContent = `Downloading \"${filename}\": `;\n div.appendChild(label);\n const progress = document.createElement('progress');\n progress.max = size;\n div.appendChild(progress);\n document.body.appendChild(div);\n\n const buffers = [];\n let downloaded = 0;\n\n const channel = await google.colab.kernel.comms.open(id);\n // Send a message to notify the kernel that we're ready.\n channel.send({})\n\n for await (const message of channel.messages) {\n // Send a message to notify the kernel that we're ready.\n channel.send({})\n if (message.buffers) {\n for (const buffer of message.buffers) {\n buffers.push(buffer);\n downloaded += buffer.byteLength;\n progress.value = downloaded;\n }\n }\n }\n const blob = new Blob(buffers, {type: 'application/binary'});\n const a = document.createElement('a');\n a.href = window.URL.createObjectURL(blob);\n a.download = filename;\n div.appendChild(a);\n a.click();\n div.remove();\n }\n ",
|
|
"text/plain": [
|
|
"<IPython.core.display.Javascript object>"
|
|
]
|
|
},
|
|
"metadata": {},
|
|
"output_type": "display_data"
|
|
},
|
|
{
|
|
"data": {
|
|
"application/javascript": "download(\"download_7ee96007-491b-4a8f-b52a-6b7d5e63e4b7\", \"vulnerabilities.json\", 11541)",
|
|
"text/plain": [
|
|
"<IPython.core.display.Javascript object>"
|
|
]
|
|
},
|
|
"metadata": {},
|
|
"output_type": "display_data"
|
|
}
|
|
],
|
|
"source": [
|
|
"from google.colab import files\n",
|
|
"\n",
|
|
"files.download(\"vulnerabilities.json\")"
|
|
]
|
|
},
|
|
{
|
|
"cell_type": "markdown",
|
|
"id": "qbk8xpGJUASo",
|
|
"metadata": {
|
|
"id": "qbk8xpGJUASo"
|
|
},
|
|
"source": [
|
|
"## Conclusions\n",
|
|
"\n",
|
|
"In this notebook we have successfully leveraged Gemini 2.0's code scanning and code generation capability to\n",
|
|
"1. Analyze multiple python files for potential vulnerabilities\n",
|
|
"2. Used Gemini 2.0 to provide recommendations (both in wordings and in the form of code)\n",
|
|
"3. Export the details into csv and json for further analysis\n",
|
|
"\n",
|
|
"The scope of this experiment is limited to identifying issues and providing helpful and contextual modification. Automating remediations or fitting the findings into a review workflow would exist in a more mature tool, and hasn't been considered as part of the experiment. While Gemini 2.0 demonstrates promising capabilities in code analysis, it's important to note that this approach is still experimental. We believe that it is important to explore the potential of this technology for vulnerability detection, and continue development and validation efforts before it can be considered a robust security tool"
|
|
]
|
|
}
|
|
],
|
|
"metadata": {
|
|
"colab": {
|
|
"name": "code_scanning_and_vulnerability_detection.ipynb",
|
|
"toc_visible": true
|
|
},
|
|
"kernelspec": {
|
|
"display_name": "Python 3",
|
|
"name": "python3"
|
|
}
|
|
},
|
|
"nbformat": 4,
|
|
"nbformat_minor": 0
|
|
}
|