{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "# COAE Practical 03 — AI Agent / MCP Tool Boundary\n",
        "**Track:** Prompt Injection + AI Application & System Security | **Level:** Medium\n",
        "\n",
        "## Mission\n",
        "You are testing a **synthetic, deterministic agent** that reads a third-party tool response and sometimes executes commands. This is **not a real LLM**, and exploiting it does not predict how any particular AI model behaves. It models the key trust-boundary failure: untrusted content controls privileged tool invocation.\n",
        "\n",
        "The lab contains no real network requests, credentials, or enterprise systems.\n",
        "\n",
        "**Objective:** Craft untrusted document text that makes the vulnerable parser change the mock rule `SIG-042` from enabled to disabled. Then demonstrate that an independent authorisation gateway blocks the exact same attempt for an analyst user."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "RULE_ID = \"SIG-042\"\n",
        "def vulnerable_agent(document, rules):\n",
        "    # Deliberately insecure: tool instructions in a document are treated as commands.\n",
        "    for line in document.splitlines():\n",
        "        parts = line.strip().split()\n",
        "        if len(parts) == 3 and parts[0] == \"AGENT_ACTION\" and parts[1] == \"disable_rule\":\n",
        "            if parts[2] in rules:\n",
        "                rules[parts[2]] = False\n",
        "\n",
        "def trusted_tool_gateway(role, action, resource, rules):\n",
        "    policy = {\"analyst\": {\"read_alert\"}, \"admin\": {\"read_alert\", \"disable_rule\"}}\n",
        "    if action not in policy.get(role, set()):\n",
        "        return {\"approved\": False, \"reason\": \"Role is not authorised.\"}\n",
        "    if action == \"disable_rule\" and resource in rules:\n",
        "        rules[resource] = False\n",
        "    return {\"approved\": True}\n",
        "\n",
        "def defended_agent(document, role, rules):\n",
        "    # The untrusted document cannot grant an additional capability.\n",
        "    for line in document.splitlines():\n",
        "        parts = line.strip().split()\n",
        "        if len(parts) == 3 and parts[0] == \"AGENT_ACTION\":\n",
        "            trusted_tool_gateway(role, parts[1], parts[2], rules)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Task A — Indirect instruction injection\n",
        "You control **only** the contents of a third-party tool response. Place text in the response so that `vulnerable_agent` incorrectly disables `SIG-042`. Do not change the parser or the rule ID."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "def make_untrusted_document():\n",
        "    # TODO: Return a string with third-party document contents.\n",
        "    return \"SOC analyst note: review the alert.\"\n",
        "\n",
        "payload = make_untrusted_document()\n",
        "vulnerable_rules = {RULE_ID: True}\n",
        "vulnerable_agent(payload, vulnerable_rules)\n",
        "print(\"Vulnerable agent — rule enabled:\", vulnerable_rules[RULE_ID])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Task B — Show the trust-boundary violation and prevention\n",
        "Both assertions must pass. This validates the *mock implementation*, not real-world LLM prompt reliability."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {},
      "outputs": [],
      "source": [
        "assert vulnerable_rules[RULE_ID] is False, \"Injection did not trigger the vulnerable parser.\"\n",
        "secured_rules = {RULE_ID: True}\n",
        "defended_agent(payload, \"analyst\", secured_rules)\n",
        "assert secured_rules[RULE_ID] is True, \"The policy gateway failed.\"\n",
        "print(\"PASS — vulnerable agent changed the rule; independent gateway prevented the change.\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Task C — Investigation and remediation\n",
        "Document the untrusted entry point, the parsed instruction, the privileged tool action, and the trust boundary violated.\n",
        "\n",
        "**Bonus 1:** Remove the document-to-command parser entirely and require authenticated user requests for sensitive actions.\n",
        "\n",
        "**Bonus 2 (your authorised SOC Agent):** Repeat the trust-boundary test with a harmless *read-only* synthetic tool result. Verify that your real tool gateway rejects any unauthorised write request. Never place secrets in this public notebook.\n",
        "\n",
        "**Deliverable:** The payload, both verification outputs, a root-cause explanation and an authorisation control recommendation."
      ]
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "Python (Pyodide)",
      "language": "python",
      "name": "python"
    },
    "language_info": {
      "name": "python"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 5
}