{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Different models review a set of requirements and architecture in a mermaid file and then do all the steps of security review. Then we use LLM to  rank them and then merge them into a more complete and accurate threat model\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Start with imports \n",
    "\n",
    "import os\n",
    "import json\n",
    "from dotenv import load_dotenv\n",
    "from openai import OpenAI\n",
    "from anthropic import Anthropic\n",
    "from IPython.display import Markdown, display"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Always remember to do this!\n",
    "load_dotenv(override=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print the key prefixes to help with any debugging\n",
    "\n",
    "openai_api_key = os.getenv('OPENAI_API_KEY')\n",
    "anthropic_api_key = os.getenv('ANTHROPIC_API_KEY')\n",
    "google_api_key = os.getenv('GOOGLE_API_KEY')\n",
    "deepseek_api_key = os.getenv('DEEPSEEK_API_KEY')\n",
    "groq_api_key = os.getenv('GROQ_API_KEY')\n",
    "\n",
    "if openai_api_key:\n",
    "    print(f\"OpenAI API Key exists and begins {openai_api_key[:8]}\")\n",
    "else:\n",
    "    print(\"OpenAI API Key not set\")\n",
    "    \n",
    "if anthropic_api_key:\n",
    "    print(f\"Anthropic API Key exists and begins {anthropic_api_key[:7]}\")\n",
    "else:\n",
    "    print(\"Anthropic API Key not set (and this is optional)\")\n",
    "\n",
    "if google_api_key:\n",
    "    print(f\"Google API Key exists and begins {google_api_key[:2]}\")\n",
    "else:\n",
    "    print(\"Google API Key not set (and this is optional)\")\n",
    "\n",
    "if deepseek_api_key:\n",
    "    print(f\"DeepSeek API Key exists and begins {deepseek_api_key[:3]}\")\n",
    "else:\n",
    "    print(\"DeepSeek API Key not set (and this is optional)\")\n",
    "\n",
    "if groq_api_key:\n",
    "    print(f\"Groq API Key exists and begins {groq_api_key[:4]}\")\n",
    "else:\n",
    "    print(\"Groq API Key not set (and this is optional)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "#This is the prompt which asks the LLM to do a security design review and provides a set of requirements and an architectural diagram in mermaid format\n",
    "designreviewrequest = \"\"\"For the following requirements and architectural diagram, please perform a full security design review which includes the following 7 steps\n",
    "1. Define scope and system boundaries.\n",
    "2. Create detailed data flow diagrams.\n",
    "3. Apply threat frameworks (like STRIDE) to identify threats.\n",
    "4. Rate and prioritize identified threats.\n",
    "5. Document-specific security controls and mitigations.\n",
    "6. Rank the threats based on their severity and likelihood of occurrence.\n",
    "7. Provide a summary of the security review and recommendations.\n",
    "\n",
    "Here are the requirements and mermaid architectural diagram:\n",
    "Software Requirements Specification (SRS) - Juice Shop: Secure E-Commerce Platform\n",
    "This document outlines the functional and non-functional requirements for the Juice Shop, a secure online retail platform.\n",
    "\n",
    "1. Introduction\n",
    "\n",
    "1.1 Purpose: To define the requirements for a robust and secure e-commerce platform that allows customers to purchase products online safely and efficiently.\n",
    "1.2 Scope: The system will be a web-based application providing a full range of e-commerce functionalities, from user registration and product browsing to secure payment processing and order management.\n",
    "1.3 Intended Audience: This document is intended for project managers, developers, quality assurance engineers, and stakeholders involved in the development and maintenance of the Juice Shop platform.\n",
    "2. Overall Description\n",
    "\n",
    "2.1 Product Perspective: A customer-facing, scalable, and secure e-commerce website with a comprehensive administrative backend.\n",
    "2.2 Product Features:\n",
    "Secure user registration and authentication with multi-factor authentication (MFA).\n",
    "A product catalog with detailed descriptions, images, pricing, and stock levels.\n",
    "Advanced search and filtering capabilities for products.\n",
    "A secure shopping cart and checkout process integrating with a trusted payment gateway.\n",
    "User profile management, including order history, shipping addresses, and payment information.\n",
    "An administrative dashboard for managing products, inventory, orders, and customer data.\n",
    "2.3 User Classes and Characteristics:\n",
    "Customer: A registered or guest user who can browse products, make purchases, and manage their account.\n",
    "Administrator: An authorized employee who can manage the platform's content and operations.\n",
    "Customer Service Representative: An authorized employee who can assist customers with orders and account issues.\n",
    "3. System Features\n",
    "\n",
    "3.1 Functional Requirements:\n",
    "User Management:\n",
    "Users shall be able to register for a new account with a unique email address and a strong password.\n",
    "The system shall enforce strong password policies (e.g., length, complexity, and expiration).\n",
    "Users shall be able to log in securely and enable/disable MFA.\n",
    "Users shall be able to reset their password through a secure, token-based process.\n",
    "Product Management:\n",
    "The system shall display products with accurate information, including price, description, and availability.\n",
    "Administrators shall be able to add, update, and remove products from the catalog.\n",
    "Order Processing:\n",
    "The system shall process orders through a secure, PCI-compliant payment gateway.\n",
    "The system shall encrypt all sensitive customer and payment data.\n",
    "Customers shall receive email confirmations for orders and shipping updates.\n",
    "3.2 Non-Functional Requirements:\n",
    "Security:\n",
    "All data transmission shall be encrypted using TLS 1.2 or higher.\n",
    "The system shall be protected against common web vulnerabilities, including the OWASP Top 10 (e.g., SQL Injection, XSS, CSRF).\n",
    "Regular security audits and penetration testing shall be conducted.\n",
    "Performance:\n",
    "The website shall load in under 3 seconds on a standard broadband connection.\n",
    "The system shall handle at least 1,000 concurrent users without significant performance degradation.\n",
    "Reliability: The system shall have an uptime of 99.9% or higher.\n",
    "Usability: The user interface shall be intuitive and easy to navigate for all user types.\n",
    "\n",
    "and here is the mermaid architectural diagram:\n",
    "\n",
    "graph TB\n",
    "    subgraph \"Client Layer\"\n",
    "        Browser[Web Browser]\n",
    "        Mobile[Mobile App]\n",
    "    end\n",
    "    \n",
    "    subgraph \"Frontend Layer\"\n",
    "        Angular[Angular SPA Frontend]\n",
    "        Static[Static Assets<br/>CSS, JS, Images]\n",
    "    end\n",
    "    \n",
    "    subgraph \"Application Layer\"\n",
    "        Express[Express.js Server]\n",
    "        Routes[REST API Routes]\n",
    "        Auth[Authentication Module]\n",
    "        Middleware[Security Middleware]\n",
    "        Challenges[Challenge Engine]\n",
    "    end\n",
    "    \n",
    "    subgraph \"Business Logic\"\n",
    "        UserMgmt[User Management]\n",
    "        ProductCatalog[Product Catalog]\n",
    "        OrderSystem[Order System]\n",
    "        Feedback[Feedback System]\n",
    "        FileUpload[File Upload Handler]\n",
    "        Payment[Payment Processing]\n",
    "    end\n",
    "    \n",
    "    subgraph \"Data Layer\"\n",
    "        SQLite[(SQLite Database)]\n",
    "        FileSystem[File System<br/>Uploaded Files]\n",
    "        Memory[In-Memory Storage<br/>Sessions, Cache]\n",
    "    end\n",
    "    \n",
    "    subgraph \"Security Features (Intentionally Vulnerable)\"\n",
    "        XSS[DOM Manipulation]\n",
    "        SQLi[Database Queries]\n",
    "        AuthBypass[Login System]\n",
    "        CSRF[State Changes]\n",
    "        Crypto[Password Hashing]\n",
    "        IDOR[Resource Access]\n",
    "    end\n",
    "    \n",
    "    subgraph \"External Dependencies\"\n",
    "        NPM[NPM Packages]\n",
    "        JWT[JWT Libraries]\n",
    "        Crypto[Crypto Libraries]\n",
    "        Sequelize[Sequelize ORM]\n",
    "    end\n",
    "    \n",
    "    %% Client connections\n",
    "    Browser --> Angular\n",
    "    Mobile --> Routes\n",
    "    \n",
    "    %% Frontend connections\n",
    "    Angular --> Static\n",
    "    Angular --> Routes\n",
    "    \n",
    "    %% Application layer connections\n",
    "    Express --> Routes\n",
    "    Routes --> Auth\n",
    "    Routes --> Middleware\n",
    "    Routes --> Challenges\n",
    "    \n",
    "    %% Business logic connections\n",
    "    Routes --> UserMgmt\n",
    "    Routes --> ProductCatalog\n",
    "    Routes --> OrderSystem\n",
    "    Routes --> Feedback\n",
    "    Routes --> FileUpload\n",
    "    Routes --> Payment\n",
    "    \n",
    "    %% Data layer connections\n",
    "    UserMgmt --> SQLite\n",
    "    ProductCatalog --> SQLite\n",
    "    OrderSystem --> SQLite\n",
    "    Feedback --> SQLite\n",
    "    FileUpload --> FileSystem\n",
    "    Auth --> Memory\n",
    "    \n",
    "    %% Security vulnerabilities (dotted lines indicate vulnerable paths)\n",
    "    Angular -.-> XSS\n",
    "    Routes -.-> SQLi\n",
    "    Auth -.-> AuthBypass\n",
    "    Angular -.-> CSRF\n",
    "    UserMgmt -.-> Crypto\n",
    "    Routes -.-> IDOR\n",
    "    \n",
    "    %% External dependencies\n",
    "    Express --> NPM\n",
    "    Auth --> JWT\n",
    "    UserMgmt --> Crypto\n",
    "    SQLite --> Sequelize\n",
    "    \n",
    "    %% Styling\n",
    "    classDef clientLayer fill:#e1f5fe\n",
    "    classDef frontendLayer fill:#f3e5f5\n",
    "    classDef appLayer fill:#e8f5e8\n",
    "    classDef businessLayer fill:#fff3e0\n",
    "    classDef dataLayer fill:#fce4ec\n",
    "    classDef securityLayer fill:#ffebee\n",
    "    classDef externalLayer fill:#f1f8e9\n",
    "    \n",
    "    class Browser,Mobile clientLayer\n",
    "    class Angular,Static frontendLayer\n",
    "    class Express,Routes,Auth,Middleware,Challenges appLayer\n",
    "    class UserMgmt,ProductCatalog,OrderSystem,Feedback,FileUpload,Payment businessLayer\n",
    "    class SQLite,FileSystem,Memory dataLayer\n",
    "    class XSS,SQLi,AuthBypass,CSRF,Crypto,IDOR securityLayer\n",
    "    class NPM,JWT,Crypto,Sequelize externalLayer\"\"\"\n",
    "\n",
    "\n",
    "messages = [{\"role\": \"user\", \"content\": designreviewrequest}]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "messages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {},
   "outputs": [],
   "source": [
    "openai = OpenAI()\n",
    "competitors = []\n",
    "answers = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We make the first call to the first model\n",
    "model_name = \"gpt-4o-mini\"\n",
    "\n",
    "response = openai.chat.completions.create(model=model_name, messages=messages)\n",
    "answer = response.choices[0].message.content\n",
    "\n",
    "display(Markdown(answer))\n",
    "competitors.append(model_name)\n",
    "answers.append(answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Anthropic has a slightly different API, and Max Tokens is required\n",
    "\n",
    "model_name = \"claude-3-7-sonnet-latest\"\n",
    "\n",
    "claude = Anthropic()\n",
    "response = claude.messages.create(model=model_name, messages=messages, max_tokens=1000)\n",
    "answer = response.content[0].text\n",
    "\n",
    "display(Markdown(answer))\n",
    "competitors.append(model_name)\n",
    "answers.append(answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gemini = OpenAI(api_key=google_api_key, base_url=\"https://generativelanguage.googleapis.com/v1beta/openai/\")\n",
    "model_name = \"gemini-2.0-flash\"\n",
    "\n",
    "response = gemini.chat.completions.create(model=model_name, messages=messages)\n",
    "answer = response.choices[0].message.content\n",
    "\n",
    "display(Markdown(answer))\n",
    "competitors.append(model_name)\n",
    "answers.append(answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "deepseek = OpenAI(api_key=deepseek_api_key, base_url=\"https://api.deepseek.com/v1\")\n",
    "model_name = \"deepseek-chat\"\n",
    "\n",
    "response = deepseek.chat.completions.create(model=model_name, messages=messages)\n",
    "answer = response.choices[0].message.content\n",
    "\n",
    "display(Markdown(answer))\n",
    "competitors.append(model_name)\n",
    "answers.append(answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "groq = OpenAI(api_key=groq_api_key, base_url=\"https://api.groq.com/openai/v1\")\n",
    "model_name = \"llama-3.3-70b-versatile\"\n",
    "\n",
    "response = groq.chat.completions.create(model=model_name, messages=messages)\n",
    "answer = response.choices[0].message.content\n",
    "\n",
    "display(Markdown(answer))\n",
    "competitors.append(model_name)\n",
    "answers.append(answer)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!ollama pull llama3.2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ollama = OpenAI(base_url='http://localhost:11434/v1', api_key='ollama')\n",
    "model_name = \"llama3.2\"\n",
    "\n",
    "response = ollama.chat.completions.create(model=model_name, messages=messages)\n",
    "answer = response.choices[0].message.content\n",
    "\n",
    "display(Markdown(answer))\n",
    "competitors.append(model_name)\n",
    "answers.append(answer)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# So where are we?\n",
    "\n",
    "print(competitors)\n",
    "print(answers)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# It's nice to know how to use \"zip\"\n",
    "for competitor, answer in zip(competitors, answers):\n",
    "    print(f\"Competitor: {competitor}\\n\\n{answer}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Let's bring this together - note the use of \"enumerate\"\n",
    "\n",
    "together = \"\"\n",
    "for index, answer in enumerate(answers):\n",
    "    together += f\"# Response from competitor {index+1}\\n\\n\"\n",
    "    together += answer + \"\\n\\n\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(together)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Now we are going to ask the model to rank the design reviews\n",
    "judge = f\"\"\"You are judging a competition between {len(competitors)} competitors.\n",
    "Each model has been given this question:\n",
    "\n",
    "{designreviewrequest}\n",
    "\n",
    "Your job is to evaluate each response for completeness and accuracy, and rank them in order of best to worst.\n",
    "Respond with JSON, and only JSON, with the following format:\n",
    "{{\"results\": [\"best competitor number\", \"second best competitor number\", \"third best competitor number\", ...]}}\n",
    "\n",
    "Here are the responses from each competitor:\n",
    "\n",
    "{together}\n",
    "\n",
    "Now respond with the JSON with the ranked order of the competitors, nothing else. Do not include markdown formatting or code blocks.\"\"\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(judge)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {},
   "outputs": [],
   "source": [
    "judge_messages = [{\"role\": \"user\", \"content\": judge}]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Judgement time!\n",
    "\n",
    "openai = OpenAI()\n",
    "response = openai.chat.completions.create(\n",
    "    model=\"o3-mini\",\n",
    "    messages=judge_messages,\n",
    ")\n",
    "results = response.choices[0].message.content\n",
    "print(results)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# OK let's turn this into results!\n",
    "\n",
    "results_dict = json.loads(results)\n",
    "ranks = results_dict[\"results\"]\n",
    "for index, result in enumerate(ranks):\n",
    "    competitor = competitors[int(result)-1]\n",
    "    print(f\"Rank {index+1}: {competitor}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Now we have all the design reviews, let's see if LLMs can merge them into a single design review that is more complete and accurate than the individual reviews.\n",
    "mergePrompt = f\"\"\"Here are design reviews from {len(competitors)} LLms. Here are the responses from each one:\n",
    "\n",
    "{together} Your task is to synthesize these reviews into a single, comprehensive design review and threat model that:\n",
    "\n",
    "1. **Includes all identified threats**, consolidating any duplicates with unified wording.\n",
    "2. **Preserves the strongest insights** from each review, especially nuanced or unique observations.\n",
    "3. **Highlights conflicting or divergent findings**, if any, and explains which interpretation seems more likely and why.\n",
    "4. **Organizes the final output** in a clear format, with these sections:\n",
    "   - Scope and System Boundaries\n",
    "   - Data Flow Overview\n",
    "   - Identified Threats (categorized using STRIDE or equivalent)\n",
    "   - Risk Ratings and Prioritization\n",
    "   - Suggested Mitigations\n",
    "   - Final Comments and Open Questions\n",
    "\n",
    "Be concise but thorough. Treat this as a final report for a real-world security audit.\n",
    "\"\"\"\n",
    "\n",
    "\n",
    "openai = OpenAI()\n",
    "response = openai.chat.completions.create(\n",
    "    model=\"gpt-4o-mini\",\n",
    "    messages=[{\"role\": \"user\", \"content\": mergePrompt}],\n",
    ")\n",
    "results = response.choices[0].message.content\n",
    "print(results)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "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.12.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}