Ed Donner
GitHub
341
community-contributions/abdoul/week_four_exercise.ipynb
Normal file
341
community-contributions/abdoul/week_four_exercise.ipynb
Normal file
@@ -0,0 +1,341 @@
|
||||
{
|
||||
"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",
|
||||
" <?php\n",
|
||||
" if ($_FILES[\"file\"][\"error\"] == 0) {\n",
|
||||
" move_uploaded_file($_FILES[\"file\"][\"tmp_name\"], \"/uploads/\" . $_FILES[\"file\"][\"name\"]);\n",
|
||||
" echo \"done\";\n",
|
||||
" }\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
|
||||
}
|
||||
Reference in New Issue
Block a user