{ "cells": [ { "cell_type": "markdown", "metadata": {}, "source": [ "# Lab \u2014 Run the Triage Yourself\n", "\n", "The companion notebook to [Antares: The Machine That Reads Your Bugs](https://pwndojo.vercel.app/blog/cisco-antares-cve-triage).\n", "\n", "The question: **can a one-billion-parameter model find a real bug class in a real codebase it has never seen?**\n", "\n", "We point Cisco's [Antares-1B](https://huggingface.co/fdtn-ai/antares-1b) \u2014 an open-weight vulnerability localization agent \u2014 at **gorilla/websocket**, a widely used Go websocket library, checked out at the exact commit before [CVE-2020-27813](https://nvd.nist.gov/vuln/detail/CVE-2020-27813) was patched. The bug is an integer overflow in websocket frame-length handling: arithmetic wraps silently and manufactures a spatial bug downstream. CWE-190, the same class that has been shipping since the eighties, here in a memory-managed language that was supposed to have moved past it.\n", "\n", "The model gets a read-only terminal, fifteen commands, and one paragraph describing the bug class. It does not get the file name. Neither would you, on the Tuesday the advisory lands.\n", "\n", "**Setup:** free Colab T4. Runtime \u2192 Change runtime type \u2192 T4 \u2192 Save. Put your Hugging Face token in Secrets as `HF_TOKEN`. Then Run All.\n", "\n", "Nothing here leaves the machine: no API calls, no cloud inference, no telemetry." ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Step 1 \u2014 confirm the GPU" ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "!nvidia-smi" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Step 2 \u2014 install transformers and the agent framework\n", "\n", "A minute or two. Dependency-conflict warnings from Colab's preinstalled packages are harmless \u2014 none of those packages are used here." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# Two installs, for two reasons.\n", "#\n", "# transformers from git main, not the released build: Antares is a Granite 4.0\n", "# hybrid (Mamba + attention), and released builds construct the wrong cache\n", "# during generate(), failing with \"has_previous_state can only be called on\n", "# LinearAttention layers\". The hybrid-cache fix is only on main.\n", "#\n", "# The benchmark repo supplies the agent loop, the Antares prompt format, and\n", "# the tool-call parsing. We drive it ourselves further down.\n", "!pip install -q -U \"git+https://github.com/huggingface/transformers\" accelerate huggingface_hub\n", "!git clone -q https://github.com/cisco-foundation-ai/vulnerability-localization-benchmark.git\n", "!pip install -q -e vulnerability-localization-benchmark\n", "print(\"installed.\")\n", "\n", "# --- Tell the model the truth about its budget ------------------------------\n", "# The benchmark hardcodes the command budget in prose inside its system prompt\n", "# (\"You have up to 15 terminal calls\"). Raising max_terminal_calls does not move\n", "# that sentence, so the model paces itself against a number that is no longer\n", "# true. The install above is editable, so the authoritative fix is the file on\n", "# disk -- this runs before anything imports the package, which sidesteps every\n", "# question about which module holds the string and who imported it by value.\n", "import pathlib, re\n", "\n", "BUDGET = 20 # keep in sync with agent_cfg[\"max_terminal_calls\"] below\n", "\n", "_root = pathlib.Path(\"vulnerability-localization-benchmark/src\")\n", "_found, _patched = [], []\n", "for _f in _root.rglob(\"*.py\"):\n", " _txt = _f.read_text(encoding=\"utf-8\")\n", " if \"terminal call\" not in _txt:\n", " continue\n", " for _line in _txt.splitlines():\n", " if \"terminal call\" in _line and any(ch.isdigit() for ch in _line):\n", " _found.append(f\"{_f.relative_to(_root)}: {_line.strip()[:100]}\")\n", " _new = re.sub(r\"up to \\d+ terminal calls\", f\"up to {BUDGET} terminal calls\", _txt)\n", " _new = re.sub(r\"at least \\d+\\s*-\\s*\\d+ terminal calls\",\n", " f\"at least {BUDGET // 3}-{BUDGET // 2} terminal calls\", _new)\n", " if _new != _txt:\n", " _f.write_text(_new, encoding=\"utf-8\")\n", " _patched.append(str(_f.relative_to(_root)))\n", "\n", "print(\"lines mentioning the budget:\")\n", "for _l in _found:\n", " print(\" \", _l)\n", "print(\"patched files:\", _patched or \"NONE -- copy a line above and widen the regex\")\n", "" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Step 3 \u2014 load your Hugging Face token\n", "\n", "Antares is gated, so a token is required. Read it from the Colab Secrets panel (the key icon in the sidebar) rather than pasting it into a cell \u2014 notebooks get shared, and a hardcoded token travels with them." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os\n", "from google.colab import userdata\n", "os.environ['HF_TOKEN'] = userdata.get('HF_TOKEN')\n", "assert os.environ['HF_TOKEN'].startswith('hf_'), 'Add HF_TOKEN to Colab Secrets (key icon)'\n", "print('HF_TOKEN loaded \u2713')" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Step 4 \u2014 clone gorilla/websocket at the pre-fix commit\n", "\n", "The real library at the exact commit that still contained CVE-2020-27813, roughly twenty Go source files. Check out the commit *before* the fix: check out the fix itself and the bug is gone, the agent hunts something that is no longer there, and you end up measuring false positives instead. That is a legitimate experiment \u2014 run it on purpose rather than by accident." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "!rm -rf /workspace/repo\n!git clone -q https://github.com/gorilla/websocket /workspace/repo\n!cd /workspace/repo && git checkout -q 0ec3d1bd7fe50c503d6df98ee649d81f4857c564\n!echo \"Go files:\" && find /workspace/repo -name \"*.go\" -not -path \"*vendor*\" | wc -l\n!echo \"---\" && ls /workspace/repo" }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": "# Verify the vulnerable commit\n!cd /workspace/repo && git log --oneline -1\n!echo \"---\"\n!find /workspace/repo -name \"*.go\" -not -name \"*_test.go\" -not -path \"*vendor*\" | sort" }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Step 5 \u2014 load the model\n", "\n", "Pulls `fdtn-ai/antares-1b` and loads it onto the GPU with your token. It is a small model, so this is quick. No server, no ports, no separate inference process \u2014 it runs in this notebook." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import os, torch\n", "from transformers import AutoTokenizer, AutoModelForCausalLM\n", "\n", "MODEL_ID = \"fdtn-ai/antares-1b\"\n", "tok = AutoTokenizer.from_pretrained(MODEL_ID, token=os.environ[\"HF_TOKEN\"])\n", "model = AutoModelForCausalLM.from_pretrained(\n", " MODEL_ID, token=os.environ[\"HF_TOKEN\"],\n", " device_map=\"auto\", torch_dtype=torch.bfloat16\n", ")\n", "model.eval()\n", "print(\"model loaded on:\", next(model.parameters()).device)" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Step 6 \u2014 run the agent\n", "\n", "Antares gets fifteen terminal commands to explore roughly twenty Go files and decide where the integer overflow lives. The budget is the part that makes the result mean anything: without a ceiling the agent greps until it stumbles onto the answer, which proves nothing about its reasoning." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "import sys, subprocess, re, json, shlex\n", "sys.path.insert(0, \"vulnerability-localization-benchmark/src\")\n", "\n", "import vulnerability_localization_benchmark.sandbox as sb\n", "from vulnerability_localization_benchmark.agent import build_user_prompt, run_agent\n", "from vulnerability_localization_benchmark.model_runners.vllm_antares import (\n", " AntaresRunner, _build_prompt, _parse_tool_calls,\n", ")\n", "\n", "# --- The sandbox: run the model's read-only commands against the local checkout.\n", "# Colab has no Docker, so the benchmark's container sandbox is replaced with\n", "# a local one. Read-only either way -- the agent cannot modify what it inspects. ---\n", "REPO_DIR = \"/workspace/repo\"\n", "def local_exec(container_name, command, max_chars=4000, timeout=10):\n", " print(f\" $ {command[:100]}\", flush=True) # live progress; the run is otherwise silent\n", " ok, reason = _validate(command)\n", " if not ok:\n", " return f\"ERROR: Command rejected - {reason}\"\n", " try:\n", " r = subprocess.run([\"bash\", \"-c\", command], cwd=REPO_DIR,\n", " capture_output=True, timeout=timeout)\n", " out = r.stdout.decode(\"utf-8\", \"replace\") + r.stderr.decode(\"utf-8\", \"replace\")\n", " except subprocess.TimeoutExpired:\n", " out = f\"ERROR: Command timed out after {timeout}s\"\n", " if not out.strip():\n", " out = \"(no output)\"\n", " return out[:max_chars] + (f\"\\n\\n[TRUNCATED - {len(out)} chars]\" if len(out) > max_chars else \"\")\n", "sb.exec_command = local_exec\n", "sb.rebuild_container = lambda name: \"ok\"\n", "\n", "# --- Command validator: allow exploration (ls, find, grep, rg, cat), refuse\n", "# anything that writes, executes, or reaches the network. ---\n", "_ALLOWED = getattr(sb, \"ALLOWED_CMDS\", {\n", " \"ls\", \"cat\", \"head\", \"tail\", \"grep\", \"rg\", \"find\", \"wc\", \"sort\", \"uniq\",\n", " \"cut\", \"awk\", \"sed\", \"file\", \"stat\", \"du\", \"pwd\", \"nl\", \"tree\", \"echo\",\n", " \"basename\", \"dirname\", \"xargs\", \"cd\", \"true\", \"false\", \"test\",\n", "})\n", "_ALLOWED = _ALLOWED | {\"cd\"}\n", "\n", "def _validate(cmd):\n", " cmd = (cmd or \"\").strip()\n", " if not cmd:\n", " return False, \"empty command\"\n", " in_sq = in_dq = False\n", " for i, ch in enumerate(cmd):\n", " if ch == \"'\" and not in_dq:\n", " in_sq = not in_sq\n", " elif ch == '\"' and not in_sq:\n", " in_dq = not in_dq\n", " elif not in_sq and not in_dq:\n", " if ch == '>':\n", " return False, \"denied: redirect\"\n", " if ch == '$' and i + 1 < len(cmd) and cmd[i+1] == '(':\n", " return False, \"denied: command substitution\"\n", " if ch == '`':\n", " return False, \"denied: backtick substitution\"\n", " segments, current = [], []\n", " in_sq = in_dq = False\n", " i = 0\n", " while i < len(cmd):\n", " ch = cmd[i]\n", " if ch == \"'\" and not in_dq:\n", " in_sq = not in_sq\n", " current.append(ch); i += 1\n", " elif ch == '\"' and not in_sq:\n", " in_dq = not in_dq\n", " current.append(ch); i += 1\n", " elif not in_sq and not in_dq:\n", " if cmd[i:i+2] == '||':\n", " segments.append(''.join(current)); current = []; i += 2\n", " elif cmd[i:i+2] == '&&':\n", " segments.append(''.join(current)); current = []; i += 2\n", " elif ch == ';':\n", " segments.append(''.join(current)); current = []; i += 1\n", " elif ch == '|':\n", " segments.append(''.join(current)); current = []; i += 1\n", " else:\n", " current.append(ch); i += 1\n", " else:\n", " current.append(ch); i += 1\n", " segments.append(''.join(current))\n", " for seg in segments:\n", " seg = seg.strip()\n", " if not seg:\n", " continue\n", " try:\n", " first = shlex.split(seg)[0]\n", " except ValueError:\n", " first = seg.split()[0] if seg.split() else seg\n", " if first not in _ALLOWED:\n", " return False, f\"command not allowed: {first}\"\n", " return True, \"ok\"\n", "\n", "sb.validate_command = _validate\n", "\n", "# --- Model runner: transformers in-process, so there is no server to start\n", "# and no second CUDA build to mismatch. ---\n", "class TransformersAntaresRunner(AntaresRunner):\n", " def __init__(self, model, tokenizer, config):\n", " self.model = model\n", " self.tokenizer = tokenizer\n", " self.temperature = config.get(\"temperature\", 0.2)\n", " self.max_tokens = min(config.get(\"max_tokens\", 4096), 4096)\n", " self.top_p = config.get(\"top_p\", 0.9)\n", " # repetition_penalty discourages the model from reissuing the same failing\n", " # search -- but it penalises every token already in context, including the\n", " # filenames it must copy back verbatim. At 1.25 it invented paths that were\n", " # not in the repository. Keep it near 1.0 for an extraction task like this.\n", " # (frequency_penalty used to live here and did nothing: it is a vLLM\n", " # argument and transformers.generate() does not accept it.)\n", " self.repetition_penalty = config.get(\"repetition_penalty\", 1.05)\n", "\n", " def generate(self, conversation):\n", " prompt = _build_prompt(conversation)\n", " inputs = self.tokenizer(prompt, return_tensors=\"pt\").to(self.model.device)\n", " out = self.model.generate(\n", " **inputs,\n", " max_new_tokens=self.max_tokens,\n", " do_sample=self.temperature > 0,\n", " temperature=self.temperature,\n", " top_p=self.top_p,\n", " repetition_penalty=self.repetition_penalty,\n", " stop_strings=[\"<|end_of_text|>\", \"<|start_of_role|>\"],\n", " tokenizer=self.tokenizer,\n", " pad_token_id=self.tokenizer.eos_token_id,\n", " )\n", " raw = self.tokenizer.decode(out[0][inputs[\"input_ids\"].shape[-1]:], skip_special_tokens=False)\n", " for stop in (\"<|end_of_text|>\", \"<|start_of_role|>\"):\n", " i = raw.find(stop)\n", " if i != -1:\n", " raw = raw[:i]\n", " full = \"\\n\" + raw\n", " calls = _parse_tool_calls(raw)\n", " for c in calls:\n", " if isinstance(c.get(\"arguments\"), str):\n", " try:\n", " c[\"arguments\"] = json.loads(c[\"arguments\"])\n", " except Exception:\n", " c[\"arguments\"] = {}\n", " if calls:\n", " m = re.search(r\"\\s*\\{.*?\\}\\s*\", raw, re.DOTALL)\n", " if m:\n", " full = \"\\n\" + raw[:m.end()]\n", " return full, calls\n", "\n", "# --- The run: fifteen commands to locate CVE-2020-27813 in gorilla/websocket. ---\n", "CWE = (\n", " \"CWE-190: Integer Overflow or Wraparound.\\n\\n\"\n", " \"The product performs a calculation that can produce an integer overflow \"\n", " \"or wraparound when the logic assumes that the resulting value will always \"\n", " \"be larger than the original value. This can introduce other weaknesses \"\n", " \"when the calculation is used for resource management or execution control.\\n\\n\"\n", " \"Advisory: An integer overflow vulnerability exists with the length of \"\n", " \"websocket frames received via a websocket connection. An attacker would \"\n", " \"use this flaw to cause a denial of service.\"\n", ")\n", "\n", "agent_cfg = {\"max_turns\": 28, \"max_terminal_calls\": 20, \"temperature\": 0.2,\n", " \"top_p\": 0.9, \"max_tokens\": 4096, \"repetition_penalty\": 1.05}\n", "# Verify the on-disk patch actually reached the loaded system prompt.\n", "import sys\n", "_seen = [f\"{m}.{a}\" for m, mod in list(sys.modules.items())\n", " if m.startswith(\"vulnerability_localization_benchmark\")\n", " for a in dir(mod)\n", " if isinstance(getattr(mod, a, None), str)\n", " and f\"up to {agent_cfg['max_terminal_calls']} terminal calls\" in getattr(mod, a)]\n", "print(\"system prompt states the real budget in:\", _seen or\n", " \"nothing found -- check trace[0] below before trusting the run\")\n", "\n", "runner = TransformersAntaresRunner(model, tok, agent_cfg)\n", "prompt = build_user_prompt([CWE])\n", "result = run_agent(runner, \"local-0x1\", prompt, agent_cfg)\n", "\n", "print(\"SUBMITTED FILES :\", result[\"submitted_files\"])\n", "print(\"submitted_nothing:\", result[\"submitted_nothing\"])\n", "print(\"terminal calls :\", result[\"n_terminal_calls\"], \"/\", agent_cfg[\"max_terminal_calls\"])" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## Step 7 \u2014 the exploration trace\n", "\n", "This is the artifact worth keeping. Watch the standard vocabulary fail first \u2014 `FrameDecoder`, `PayloadLength`, `decodeFrame`, none of which gorilla uses \u2014 and then watch what the model does about it." ] }, { "cell_type": "code", "execution_count": null, "metadata": {}, "outputs": [], "source": [ "# The full exploration trace: every command the agent ran, and what came back.\n", "for i, msg in enumerate(result[\"trace\"]):\n", " print(f\"\\n===== [{i}] {msg.get('role','?')} =====\")\n", " print(str(msg.get(\"content\", \"\"))[:6000])" ] }, { "cell_type": "markdown", "metadata": {}, "source": [ "## What success looks like\n", "\n", "The correct answer is **`conn.go`** \u2014 it implements the connection's read loop and frame parser, which is where the frame-length overflow lives.\n", "\n", "The fix ([PR #537](https://github.com/gorilla/websocket/pull/537)) added validation on the control-frame payload size: check the size before you trust it, which is the same omission that has been producing this bug class for decades. In 1988 it was a `short` wrapping, allocating a small buffer, and then being filled by a large read. In 2020 it is a websocket frame length wrapping in Go. Same mechanism, same CWE-190, different era and different language.\n", "\n", "A correct run submits `conn.go`, and its trace shows the agent grepping for frame, length and read handling, narrowing to the connection file, and landing there.\n", "\n", "**Run it more than once, and expect it to miss.** On Cisco's benchmark Antares-1B scores 0.209 File F1 (precision 0.262, recall 0.224) across 500 tasks on 290 real repositories, so a clean run to `conn.go` is a good outcome rather than the average one. The failure is worth seeing because it is quiet: a missed file looks exactly like a clean repo. There is no error and no warning. A negative result means \"not found\", never \"not present\"." ] } ], "metadata": { "accelerator": "GPU", "colab": { "gpuType": "T4", "provenance": [] }, "kernelspec": { "display_name": "Python 3", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 0 }