{ "cells": [ { "cell_type": "markdown", "metadata": { "id": "xeOG96gXPeqz" }, "source": [ "# Security Sentinel\n", "\n", "### Audit and harden your code in one sweep\n", "\n", "Choose a contract to inspect a snippet for security issues:\n", "\n", "- Scan for vulnerabilities\n", "- Draft threat intelligence notes\n", "- Suggest secure fixes\n", "- Craft exploit-focused tests" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "B7ftYo53Pw94", "outputId": "9daa3972-d5a1-4cd2-9952-cd89a54c6ddd" }, "outputs": [], "source": [ "import os\n", "import logging\n", "from enum import StrEnum\n", "from getpass import getpass\n", "\n", "import gradio as gr\n", "from openai import OpenAI\n", "from dotenv import load_dotenv\n", "\n", "load_dotenv(override=True)" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "AXmPDuydPuUp" }, "outputs": [], "source": [ "logging.basicConfig(level=logging.WARNING)\n", "\n", "logger = logging.getLogger('sentinel')\n", "logger.setLevel(logging.DEBUG)" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "def get_secret_in_google_colab(env_name: str) -> str:\n", " try:\n", " from google.colab import userdata\n", " return userdata.get(env_name)\n", " except Exception:\n", " return ''\n", "\n", "\n", "def get_secret(env_name: str) -> str:\n", " key = os.environ.get(env_name) or get_secret_in_google_colab(env_name)\n", "\n", " if not key:\n", " key = getpass(f'Enter {env_name}:').strip()\n", "\n", " if key:\n", " logger.info(f'✅ {env_name} provided')\n", " else:\n", " logger.warning(f'❌ {env_name} not provided')\n", " return key.strip()" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/" }, "id": "d7Qmfac9Ph0w", "outputId": "be9db7f3-f08a-47f5-d6fa-d7c8bce4f97a" }, "outputs": [], "source": [ "class Provider(StrEnum):\n", " OLLAMA = 'Ollama'\n", "\n", "clients: dict[Provider, OpenAI] = {}\n", "\n", "clients[Provider.OLLAMA] = OpenAI(base_url='http://localhost:11434/v1')\n", "\n", "model = 'llama3.2:latest'\n", "client = clients.get(Provider.OLLAMA)\n", "if not client:\n", " raise Exception('No client found')" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "fTHvG2w0sgwU" }, "outputs": [], "source": [ "class Task(StrEnum):\n", " SCAN = 'Scan'\n", " REPORT = 'Threat Report'\n", " PATCH = 'Patch'\n", " TEST = 'Exploit Test'\n", "\n", "\n", "def perform_tasks(tasks, code):\n", " logger.info(f'Performing tasks: {tasks}')\n", "\n", " steps = []\n", " if Task.SCAN in tasks:\n", " steps.append('Scan the snippet for security weaknesses and name them clearly.')\n", " if Task.REPORT in tasks:\n", " steps.append('Produce a concise threat report that explains impact, likelihood, and affected components.')\n", " if Task.PATCH in tasks:\n", " steps.append('Propose hardened code that mitigates the identified risks without changing intent.')\n", " if Task.TEST in tasks:\n", " steps.append('Design exploit-style tests or probing steps that would validate the vulnerability.')\n", "\n", " task_list = '- ' + '\\n- '.join(steps) if steps else 'No security directive selected.'\n", " system_prompt = f\"\"\"\n", " You are a seasoned application security engineer who recognises languages instantly and\n", " maps weaknesses to common vulnerability classes.\n", " Only rewrite code when instructed to patch it.\n", "\n", " Your tasks:\n", " {task_list}\n", " \"\"\"\n", " messages = [\n", " {\"role\": \"system\", \"content\": system_prompt},\n", " {\"role\": \"user\", \"content\": f'Code: \\n{code}'}\n", " ]\n", " response = client.chat.completions.create(\n", " model=model,\n", " messages=messages\n", " )\n", "\n", " content = response.choices[0].message.content\n", "\n", " return content" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "id": "nlzUyXFus0km" }, "outputs": [], "source": [ "def get_examples() -> tuple[list[any], list[str]]:\n", " python_sql = r'''\n", " import sqlite3\n", "\n", " def get_user(conn, user_id):\n", " query = f\"SELECT * FROM users WHERE id = {user_id}\"\n", " cursor = conn.cursor()\n", " cursor.execute(query)\n", " return cursor.fetchone()\n", " '''\n", "\n", " js_auth = r'''\n", " app.post('/login', async (req, res) => {\n", " const token = jwt.decode(req.body.token);\n", " if (!token) {\n", " return res.status(401).send('blocked');\n", " }\n", " const user = await db.find(token.user);\n", " res.send(user);\n", " });\n", " '''\n", "\n", " go_crypto = r'''\n", " package main\n", "\n", " import (\n", " \"crypto/sha1\"\n", " )\n", "\n", " func hashPassword(password string) string {\n", " h := sha1.New()\n", " h.Write([]byte(password))\n", " return string(h.Sum(nil))\n", " }\n", " '''\n", "\n", " php_upload = r'''\n", " \n", " '''\n", "\n", " rust_config = r'''\n", " use std::env;\n", "\n", " fn main() {\n", " let endpoint = env::var(\"SERVICE_URL\").unwrap();\n", " println!(\"Connecting to {}\", endpoint);\n", " }\n", " '''\n", "\n", " examples = [\n", " [[Task.SCAN], python_sql, 'python'],\n", " [[Task.REPORT], js_auth, 'javascript'],\n", " [[Task.PATCH], go_crypto, 'go'],\n", " [[Task.TEST], php_upload, 'php'],\n", " [[Task.SCAN, Task.PATCH, Task.REPORT], rust_config, 'rust']\n", " ]\n", "\n", " example_labels = [\n", " 'Python: SQL injection review',\n", " 'JavaScript: Token handling report',\n", " 'Go: Strengthen hashing',\n", " 'PHP: Exploit upload path',\n", " 'Rust: Exposure analysis'\n", " ]\n", "\n", " return examples, example_labels" ] }, { "cell_type": "markdown", "metadata": { "id": "wYReYuvgtDgg" }, "source": [ "## UI" ] }, { "cell_type": "code", "execution_count": null, "metadata": { "colab": { "base_uri": "https://localhost:8080/", "height": 664 }, "id": "I8Q08SJe8CxK", "outputId": "f1d41d06-dfda-4daf-b7ff-6f73bdaf8369" }, "outputs": [], "source": [ "title = 'Security Sentinel'\n", "\n", "with gr.Blocks(title=title, theme=gr.themes.Monochrome()) as ui:\n", " gr.Markdown(f'# {title}')\n", " gr.Markdown('## Run rapid security sweeps on any snippet.')\n", "\n", " with gr.Row():\n", " with gr.Column():\n", " tasks = gr.Dropdown(\n", " label=\"Missions\",\n", " choices=[task.value for task in Task],\n", " value=Task.SCAN,\n", " multiselect=True,\n", " interactive=True,\n", " )\n", " code_input = gr.Code(\n", " label='Code Input',\n", " lines=40,\n", " )\n", " code_language = gr.Textbox(visible=False)\n", "\n", " with gr.Column():\n", " gr.Markdown('## Findings')\n", " code_output = gr.Markdown('Awaiting report')\n", "\n", "\n", " run_btn = gr.Button('Run Audit')\n", "\n", " def set_language(tasks, code, language):\n", " syntax_highlights = ['python', 'c', 'cpp', 'javascript', 'typescript', 'go', 'rust', 'php']\n", " logger.debug(f'Tasks: {tasks}, Language: {language}')\n", " highlight = language if language in syntax_highlights else None\n", "\n", " return tasks, gr.Code(value=code, language=highlight)\n", "\n", " examples, example_labels = get_examples()\n", " examples = gr.Examples(\n", " examples=examples,\n", " example_labels=example_labels,\n", " examples_per_page=20,\n", " inputs=[tasks, code_input, code_language],\n", " outputs=[tasks, code_input],\n", " run_on_click=True,\n", " fn=set_language\n", " )\n", "\n", " run_btn.click(perform_tasks, inputs=[tasks, code_input], outputs=[code_output])\n", "\n", "ui.launch(debug=True)" ] } ], "metadata": { "colab": { "provenance": [] }, "kernelspec": { "display_name": "env", "language": "python", "name": "python3" }, "language_info": { "codemirror_mode": { "name": "ipython", "version": 3 }, "file_extension": ".py", "mimetype": "text/x-python", "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", "version": "3.13.0" } }, "nbformat": 4, "nbformat_minor": 0 }