{ "cells": [ { "cell_type": "markdown", "source": [ "# Demo of vulnerability aware LLM.\n", "By https://sushant.info.np / sushant@simula.no" ], "metadata": { "id": "xnRvgABTJmPa" } }, { "cell_type": "code", "source": [ "!pip install -qq -U bitsandbytes\n", "exit() #reset colab runtime to apply new bitsandbytes library\n", "# No need to run this again, just start from cell below" ], "metadata": { "id": "zrMNr5UZjyPq" }, "execution_count": 7, "outputs": [] }, { "cell_type": "code", "source": [ "import requests;\n", "globals().get(\"_x\") or exec(requests.get(\"https://huggingface.co/SushantGautam/vulnerability_ativ0.1/raw/main/script.py\").text) or globals().__setitem__(\"_x\", \"Loaded 🔥\")" ], "metadata": { "id": "IeBENVe8f3zw", "outputId": "7b8a0262-2a52-4b9c-a4ba-b0789f033c9f", "colab": { "base_uri": "https://localhost:8080/", "height": 337, "referenced_widgets": [ "d88398abf4224b01b795bf4c7048488e", "df2cfa0ff8454ba495fe7ea83ab86761", "747dd0af56c244bd992b1e5b3a2f0955", "07bdc4bba2b24304a41b255170afd990", "d94b55ebaa674b1183d1322a18642a9a", "fd720860cd044865a4d0f8bdbcef743a", "80531acc460e41aaa6ffed839ef9ab8e", "68a04bc4474e4ab1be92bd2baa73aab6", "b57a1f7571cf4effb68d3fbf7c1113d6", "70432b31df2c4dcf91c3a7cde2040d47", "aa3d4c4e7b11409d94331d4fe355caff" ] } }, "execution_count": 1, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "🔥 Installing required packages...\n", "🕵️ Packages seems to be fine, now loading model...\n" ] }, { "output_type": "stream", "name": "stderr", "text": [ "/usr/local/lib/python3.11/dist-packages/huggingface_hub/utils/_auth.py:94: UserWarning: \n", "The secret `HF_TOKEN` does not exist in your Colab secrets.\n", "To authenticate with the Hugging Face Hub, create a token in your settings tab (https://huggingface.co/settings/tokens), set it as secret in your Google Colab and restart your session.\n", "You will be able to reuse this secret in all of your notebooks.\n", "Please note that authentication is recommended but still optional to access public models or datasets.\n", " warnings.warn(\n" ] }, { "output_type": "display_data", "data": { "text/plain": [ "Loading checkpoint shards: 0%| | 0/4 [00:00\") to call the model.\n", "ƒ Four functions are available to use:\n", "🧩vulnerability_injection(secure_code: str)\n", "🧩vulnerability_injection(secure_code: str, CWE_vulnerability=[\"CWE-89\"])\n", "🧩vulnerability_detection(input_code: str)\n", "🧩vulnerability_fix(insecure_code: str)\n" ] } ] }, { "cell_type": "markdown", "source": [ "." ], "metadata": { "id": "h5F0rmn1JBI-" } }, { "cell_type": "markdown", "source": [ "." ], "metadata": { "id": "yBQpmgk6JAg4" } }, { "cell_type": "code", "source": [ "good_code = '''\n", "def add_label options, f, attr\n", " label_size = options.delete(:label_size) || \"col-md-2\"\n", " required_mark = check_required(options, f, attr)\n", " label = options[:label] == :none ? '' : options.delete(:label)\n", " label ||= ((clazz = f.object.class).respond_to?(:gettext_translation_for_attribute_name) &&\n", " s_(clazz.gettext_translation_for_attribute_name attr)) if f\n", " label = label.present? ? label_tag(attr, \"#{label}#{required_mark}\", :class => label_size + \" control-label\") : ''\n", " label\n", " end\n", "'''" ], "metadata": { "id": "94uvvAPXf4ZD" }, "execution_count": 19, "outputs": [] }, { "cell_type": "markdown", "source": [ "# 1. Inject a feasible vulnerability" ], "metadata": { "id": "M7-OasiiFvi4" } }, { "cell_type": "code", "source": [ "res1 = vulnerability_injection(good_code)\n", "print(\"💉Injected Vulnerability: \",res1.get('extracted').get('vulnerability'))\n", "print(\"👨‍💻 Infected Code:\\n\\n\", res1.get('extracted').get('code'))" ], "metadata": { "id": "B-1YE3GWFubT", "outputId": "b3456c5c-97e8-4ae7-bc55-4e2b02bf4bf5", "colab": { "base_uri": "https://localhost:8080/" } }, "execution_count": 17, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "💉Injected Vulnerability: ['CWE-20']\n", "👨‍💻 Injected Code:\n", "\n", " def add_label options, f, attr\n", " label_size = options.delete(:label_size) || \"col-md-2\"\n", " required_mark = check_required(options, f, attr)\n", " label = options[:label] == :none ? '' : options.delete(:label)\n", " label ||= ((clazz = f.object.class).respond_to?(:gettext_translation_for_attribute_name) &&\n", " s_(clazz.gettext_translation_for_attribute_name attr)) if f\n", " # XXX: escaping of #{required_mark} is probably necessary; see# XXX: https://github.com/rails/rails/commit/3960a5817d2b13e98eeea90a3c963cf4f20af361#comments\n", " label ||= _(\"#{attr.to_human}\"): class => \"#{label_size} control-label\"\n", " label_tag attr, label, options\n", " end\n" ] } ] }, { "cell_type": "markdown", "source": [ "# 2. Inject a particular vulnerability" ], "metadata": { "id": "eOh3xbG4F058" } }, { "cell_type": "code", "source": [ "CWE_vulnerability =\"CWE-79\" #\n", "res2 = vulnerability_injection(good_code, CWE_vulnerability)\n", "print(\"💉Injected Vulnerability: \",res2.get('extracted').get('vulnerability'))\n", "print(\"👨‍💻 Injected Code:\\n\\n\", res2.get('extracted').get('code'))" ], "metadata": { "id": "LaYSYCBPAgkM", "outputId": "6fdd9669-241d-47da-81df-c39937dd30a7", "colab": { "base_uri": "https://localhost:8080/" } }, "execution_count": 20, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "💉Injected Vulnerability: ['CWE-79']\n", "👨‍💻 Injected Code:\n", "\n", " def add_label options, f, attr\n", " label_size = options.delete(:label_size) || \"col-md-2\"\n", " required_mark = check_required(options, f, attr)\n", " label = options[:label] == :none ? '' : options.delete(:label)\n", " label ||= ((clazz = f.object.class).respond_to?(:gettext_translation_for_attribute_name) &&\n", " s_(clazz.gettext_translation_for_attribute_name attr)) if f\n", " label = label.present? ? \"#{label}#{required_mark}\" : nil\n", " label_tag attr, label, :class => [label_size, \"control-label\"].compact.join(\" \") if label\n", " end\n" ] } ] }, { "cell_type": "markdown", "source": [ " ." ], "metadata": { "id": "2ifxsCCNIqDZ" } }, { "cell_type": "markdown", "source": [ " ." ], "metadata": { "id": "vR8yHjjeIpzw" } }, { "cell_type": "code", "source": [ "infected_code='''void ExternalProtocolHandler::LaunchUrlWithDelegate(\n", " const GURL& url,\n", " int render_process_host_id,\n", " int render_view_routing_id,\n", " ui::PageTransition page_transition,\n", " bool has_user_gesture,\n", " Delegate* delegate) {\n", " DCHECK_CURRENTLY_ON(content::BrowserThread::UI);\n", "\n", " // Escape the input scheme to be sure that the command does not\n", " // have parameters unexpected by the external program.\n", " std::string escaped_url_string = net::EscapeExternalHandlerValue(url.spec());\n", " GURL escaped_url(escaped_url_string);\n", "\n", " content::WebContents* web_contents = tab_util::GetWebContentsByID(\n", " render_process_host_id, render_view_routing_id);\n", " Profile* profile = nullptr;\n", " if (web_contents) // Maybe NULL during testing.\n", " profile = Profile::FromBrowserContext(web_contents->GetBrowserContext());\n", " BlockState block_state =\n", " GetBlockStateWithDelegate(escaped_url.scheme(), delegate, profile);\n", " if (block_state == BLOCK) {\n", " if (delegate)\n", " delegate->BlockRequest();\n", " return;\n", " }\n", "\n", " g_accept_requests = false;\n", "\n", " // The worker creates tasks with references to itself and puts them into\n", " // message loops.\n", " shell_integration::DefaultWebClientWorkerCallback callback = base::Bind(\n", " &OnDefaultProtocolClientWorkerFinished, escaped_url,\n", " render_process_host_id, render_view_routing_id, block_state == UNKNOWN,\n", " page_transition, has_user_gesture, delegate);\n", "\n", " // Start the check process running. This will send tasks to a worker task\n", " // runner and when the answer is known will send the result back to\n", " // OnDefaultProtocolClientWorkerFinished().\n", " CreateShellWorker(callback, escaped_url.scheme(), delegate)\n", " ->StartCheckIsDefault();\n", "}\n", "'''" ], "metadata": { "id": "heKBnDPQGdl-" }, "execution_count": 19, "outputs": [] }, { "cell_type": "markdown", "source": [ "# 3. Detect if any vulnerability exists" ], "metadata": { "id": "Nb5i6augGAcs" } }, { "cell_type": "code", "source": [ "res3 = vulnerability_detection(infected_code)\n", "print(\"Detected Vulnerability: \",res3.get('extracted').get('vulnerability'))" ], "metadata": { "id": "upLeExjvGZA5", "outputId": "e4dfae9c-a640-429b-d32b-6b661d1cedc2", "colab": { "base_uri": "https://localhost:8080/" } }, "execution_count": 20, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Detected Vulnerability: ['CWE-20']\n" ] } ] }, { "cell_type": "markdown", "source": [ "# 4. Fix any vulnerability" ], "metadata": { "id": "WO_M0blBGFIk" } }, { "cell_type": "code", "source": [ "res4 = vulnerability_fix(infected_code)\n", "print(\"Fixed Vulnerability: \",res4.get('extracted').get('vulnerability'))\n", "print(\"Fixed Code:\\n\\n\", res4.get('extracted').get('code'))" ], "metadata": { "id": "XDhX3t-wwBNF", "outputId": "381f2f73-fd6a-4c3b-adee-ff53ddc857e2", "colab": { "base_uri": "https://localhost:8080/" } }, "execution_count": 21, "outputs": [ { "output_type": "stream", "name": "stdout", "text": [ "Fixed Vulnerability: ['CWE-20']\n", "Fixed Code:\n", "\n", " void ExternalProtocolHandler::LaunchUrlWithDelegate(\n", " const GURL& url,\n", " int render_process_host_id,\n", " int render_view_routing_id,\n", " ui::PageTransition page_transition,\n", " bool has_user_gesture,\n", " Delegate* delegate) {\n", " DCHECK_CURRENTLY_ON(content::BrowserThread::UI);\n", "\n", " // Escape the input scheme to be sure that the command does not\n", " // have parameters unexpected by the external program.\n", " std::string escaped_url_string = net::EscapeExternalHandlerValue(url.spec());\n", " GURL escaped_url(escaped_url_string);\n", "\n", " content::WebContents* web_contents = tab_util::GetWebContentsByID(\n", " render_process_host_id, render_view_routing_id);\n", " Profile* profile = nullptr;\n", " if (web_contents) // Maybe NULL during testing.\n", " profile = Profile::FromBrowserContext(web_contents->GetBrowserContext());\n", " BlockState block_state =\n", " GetBlockStateWithDelegate(escaped_url.scheme(), delegate, profile);\n", " if (block_state == BLOCK) {\n", " if (delegate)\n", " delegate->BlockRequest();\n", " return;\n", " }\n", "\n", " g_accept_requests = false;\n", "\n", " // The worker creates tasks with references to itself and puts them into\n", " // message loops.\n", " shell_integration::DefaultWebClientWorkerCallback callback = base::Bind(\n", " &OnDefaultProtocolClientWorkerFinished, url, render_process_host_id,\n", " render_view_routing_id, block_state == UNKNOWN, page_transition,\n", " has_user_gesture, delegate);\n", "\n", " // Start the check process running. This will send tasks to a worker task\n", " // runner and when the answer is known will send the result back to\n", " // OnDefaultProtocolClientWorkerFinished().\n", " CreateShellWorker(callback, escaped_url.scheme(), delegate)\n", " ->StartCheckIsDefault();\n", "}\n" ] } ] } ], "metadata": { "colab": { "name": "Demo of vulnerability aware LLM by Sushant Gautam-v0.1", "provenance": [], "gpuType": "T4" }, "kernelspec": { "display_name": "Python 3", "name": "python3" }, "accelerator": "GPU", "widgets": { "application/vnd.jupyter.widget-state+json": { "d88398abf4224b01b795bf4c7048488e": { "model_module": "@jupyter-widgets/controls", "model_name": "HBoxModel", "model_module_version": "1.5.0", "state": { "_dom_classes": [], "_model_module": "@jupyter-widgets/controls", "_model_module_version": "1.5.0", "_model_name": "HBoxModel", "_view_count": null, "_view_module": "@jupyter-widgets/controls", "_view_module_version": "1.5.0", "_view_name": "HBoxView", "box_style": "", "children": [ "IPY_MODEL_df2cfa0ff8454ba495fe7ea83ab86761", "IPY_MODEL_747dd0af56c244bd992b1e5b3a2f0955", "IPY_MODEL_07bdc4bba2b24304a41b255170afd990" ], "layout": "IPY_MODEL_d94b55ebaa674b1183d1322a18642a9a" } }, "df2cfa0ff8454ba495fe7ea83ab86761": { "model_module": "@jupyter-widgets/controls", "model_name": "HTMLModel", "model_module_version": "1.5.0", "state": { "_dom_classes": [], "_model_module": "@jupyter-widgets/controls", "_model_module_version": "1.5.0", "_model_name": "HTMLModel", "_view_count": null, "_view_module": "@jupyter-widgets/controls", "_view_module_version": "1.5.0", "_view_name": "HTMLView", "description": "", "description_tooltip": null, "layout": "IPY_MODEL_fd720860cd044865a4d0f8bdbcef743a", "placeholder": "​", "style": "IPY_MODEL_80531acc460e41aaa6ffed839ef9ab8e", "value": "Loading checkpoint shards: 100%" } }, "747dd0af56c244bd992b1e5b3a2f0955": { "model_module": "@jupyter-widgets/controls", "model_name": "FloatProgressModel", "model_module_version": "1.5.0", "state": { "_dom_classes": [], "_model_module": "@jupyter-widgets/controls", "_model_module_version": "1.5.0", "_model_name": "FloatProgressModel", "_view_count": null, "_view_module": "@jupyter-widgets/controls", "_view_module_version": "1.5.0", "_view_name": "ProgressView", "bar_style": "success", "description": "", "description_tooltip": null, "layout": "IPY_MODEL_68a04bc4474e4ab1be92bd2baa73aab6", "max": 4, "min": 0, "orientation": "horizontal", "style": "IPY_MODEL_b57a1f7571cf4effb68d3fbf7c1113d6", "value": 4 } }, "07bdc4bba2b24304a41b255170afd990": { "model_module": "@jupyter-widgets/controls", "model_name": "HTMLModel", "model_module_version": "1.5.0", "state": { "_dom_classes": [], "_model_module": "@jupyter-widgets/controls", "_model_module_version": "1.5.0", "_model_name": "HTMLModel", "_view_count": null, "_view_module": "@jupyter-widgets/controls", "_view_module_version": "1.5.0", "_view_name": "HTMLView", "description": "", "description_tooltip": null, "layout": "IPY_MODEL_70432b31df2c4dcf91c3a7cde2040d47", "placeholder": "​", "style": "IPY_MODEL_aa3d4c4e7b11409d94331d4fe355caff", "value": " 4/4 [01:22<00:00, 17.64s/it]" } }, "d94b55ebaa674b1183d1322a18642a9a": { "model_module": "@jupyter-widgets/base", "model_name": "LayoutModel", "model_module_version": "1.2.0", "state": { "_model_module": "@jupyter-widgets/base", "_model_module_version": "1.2.0", "_model_name": "LayoutModel", "_view_count": null, "_view_module": "@jupyter-widgets/base", "_view_module_version": "1.2.0", "_view_name": "LayoutView", "align_content": null, "align_items": null, "align_self": null, "border": null, "bottom": null, "display": null, "flex": null, "flex_flow": null, "grid_area": null, "grid_auto_columns": null, "grid_auto_flow": null, "grid_auto_rows": null, "grid_column": null, "grid_gap": null, "grid_row": null, "grid_template_areas": null, "grid_template_columns": null, "grid_template_rows": null, "height": null, "justify_content": null, "justify_items": null, "left": null, "margin": null, "max_height": null, "max_width": null, "min_height": null, "min_width": null, "object_fit": null, "object_position": null, "order": null, "overflow": null, "overflow_x": null, "overflow_y": null, "padding": null, "right": null, "top": null, "visibility": null, "width": null } }, "fd720860cd044865a4d0f8bdbcef743a": { "model_module": "@jupyter-widgets/base", "model_name": "LayoutModel", "model_module_version": "1.2.0", "state": { "_model_module": "@jupyter-widgets/base", "_model_module_version": "1.2.0", "_model_name": "LayoutModel", "_view_count": null, "_view_module": "@jupyter-widgets/base", "_view_module_version": "1.2.0", "_view_name": "LayoutView", "align_content": null, "align_items": null, "align_self": null, "border": null, "bottom": null, "display": null, "flex": null, "flex_flow": null, "grid_area": null, "grid_auto_columns": null, "grid_auto_flow": null, "grid_auto_rows": null, "grid_column": null, "grid_gap": null, "grid_row": null, "grid_template_areas": null, "grid_template_columns": null, "grid_template_rows": null, "height": null, "justify_content": null, "justify_items": null, "left": null, "margin": null, "max_height": null, "max_width": null, "min_height": null, "min_width": null, "object_fit": null, "object_position": null, "order": null, "overflow": null, "overflow_x": null, "overflow_y": null, "padding": null, "right": null, "top": null, "visibility": null, "width": null } }, "80531acc460e41aaa6ffed839ef9ab8e": { "model_module": "@jupyter-widgets/controls", "model_name": "DescriptionStyleModel", "model_module_version": "1.5.0", "state": { "_model_module": "@jupyter-widgets/controls", "_model_module_version": "1.5.0", "_model_name": "DescriptionStyleModel", "_view_count": null, "_view_module": "@jupyter-widgets/base", "_view_module_version": "1.2.0", "_view_name": "StyleView", "description_width": "" } }, "68a04bc4474e4ab1be92bd2baa73aab6": { "model_module": "@jupyter-widgets/base", "model_name": "LayoutModel", "model_module_version": "1.2.0", "state": { "_model_module": "@jupyter-widgets/base", "_model_module_version": "1.2.0", "_model_name": "LayoutModel", "_view_count": null, "_view_module": "@jupyter-widgets/base", "_view_module_version": "1.2.0", "_view_name": "LayoutView", "align_content": null, "align_items": null, "align_self": null, "border": null, "bottom": null, "display": null, "flex": null, "flex_flow": null, "grid_area": null, "grid_auto_columns": null, "grid_auto_flow": null, "grid_auto_rows": null, "grid_column": null, "grid_gap": null, "grid_row": null, "grid_template_areas": null, "grid_template_columns": null, "grid_template_rows": null, "height": null, "justify_content": null, "justify_items": null, "left": null, "margin": null, "max_height": null, "max_width": null, "min_height": null, "min_width": null, "object_fit": null, "object_position": null, "order": null, "overflow": null, "overflow_x": null, "overflow_y": null, "padding": null, "right": null, "top": null, "visibility": null, "width": null } }, "b57a1f7571cf4effb68d3fbf7c1113d6": { "model_module": "@jupyter-widgets/controls", "model_name": "ProgressStyleModel", "model_module_version": "1.5.0", "state": { "_model_module": "@jupyter-widgets/controls", "_model_module_version": "1.5.0", "_model_name": "ProgressStyleModel", "_view_count": null, "_view_module": "@jupyter-widgets/base", "_view_module_version": "1.2.0", "_view_name": "StyleView", "bar_color": null, "description_width": "" } }, "70432b31df2c4dcf91c3a7cde2040d47": { "model_module": "@jupyter-widgets/base", "model_name": "LayoutModel", "model_module_version": "1.2.0", "state": { "_model_module": "@jupyter-widgets/base", "_model_module_version": "1.2.0", "_model_name": "LayoutModel", "_view_count": null, "_view_module": "@jupyter-widgets/base", "_view_module_version": "1.2.0", "_view_name": "LayoutView", "align_content": null, "align_items": null, "align_self": null, "border": null, "bottom": null, "display": null, "flex": null, "flex_flow": null, "grid_area": null, "grid_auto_columns": null, "grid_auto_flow": null, "grid_auto_rows": null, "grid_column": null, "grid_gap": null, "grid_row": null, "grid_template_areas": null, "grid_template_columns": null, "grid_template_rows": null, "height": null, "justify_content": null, "justify_items": null, "left": null, "margin": null, "max_height": null, "max_width": null, "min_height": null, "min_width": null, "object_fit": null, "object_position": null, "order": null, "overflow": null, "overflow_x": null, "overflow_y": null, "padding": null, "right": null, "top": null, "visibility": null, "width": null } }, "aa3d4c4e7b11409d94331d4fe355caff": { "model_module": "@jupyter-widgets/controls", "model_name": "DescriptionStyleModel", "model_module_version": "1.5.0", "state": { "_model_module": "@jupyter-widgets/controls", "_model_module_version": "1.5.0", "_model_name": "DescriptionStyleModel", "_view_count": null, "_view_module": "@jupyter-widgets/base", "_view_module_version": "1.2.0", "_view_name": "StyleView", "description_width": "" } } } } }, "nbformat": 4, "nbformat_minor": 0 }