[ { "path": ".gitignore", "content": "# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\n\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\nwheels/\npip-wheel-metadata/\nshare/python-wheels/\n*.egg-info/\n.installed.cfg\n*.egg\nMANIFEST\n\n# PyInstaller\n# Usually these files are written by a python script from a template\n# before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.nox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*.cover\n*.py,cover\n.hypothesis/\n.pytest_cache/\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\ndb.sqlite3\ndb.sqlite3-journal\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\ntarget/\n\n# Jupyter Notebook\n.ipynb_checkpoints\n\n# IPython\nprofile_default/\nipython_config.py\n\n# pyenv\n.python-version\n\n# pipenv\n# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.\n# However, in case of collaboration, if having platform-specific dependencies or dependencies\n# having no cross-platform support, pipenv may install dependencies that don't work, or not\n# install all needed dependencies.\n#Pipfile.lock\n\n# PEP 582; used by e.g. github.com/David-OConnor/pyflow\n__pypackages__/\n\n# Celery stuff\ncelerybeat-schedule\ncelerybeat.pid\n\n# SageMath parsed files\n*.sage.py\n\n# Environments\n.env\n.venv\nenv/\nvenv/\nENV/\nenv.bak/\nvenv.bak/\n\n# Spyder project settings\n.spyderproject\n.spyproject\n\n# Rope project settings\n.ropeproject\n\n# mkdocs documentation\n/site\n\n# mypy\n.mypy_cache/\n.dmypy.json\ndmypy.json\n\n# Pyre type checker\n.pyre/\n\n# Pyproject\nbots/triage-slackbot/pyproject.toml\nbots/incident-response-slackbot/pyproject.toml\n\n.trufflehog\n*.pkl\n*.json\nexclude.txt\n" }, { "path": ".pre-commit-config.yaml", "content": "repos:\n- repo: local\n hooks:\n - id: trufflehog\n name: TruffleHog\n description: Detect secrets in your data.\n entry: bash -c 'trufflehog git file://. --since-commit HEAD --fail'\n language: system\n stages: [\"commit\", \"push\"]\n\n- repo: https://github.com/hauntsaninja/black-pre-commit-mirror\n rev: 1836df4ee440d6637f6a4f3d4e0727f1d75ba0eb # 23.10.1\n hooks:\n - id: black\n args: [--line-length=100, --workers=6]\n\n- repo: https://github.com/pycqa/isort\n rev: e44834b7b294701f596c9118d6c370f86671a50d # 5.12.0\n hooks:\n - id: isort\n name: isort (python)\n" }, { "path": "CODEOWNERS", "content": "* @openai/security-team\n" }, { "path": "LICENSE", "content": "MIT License\n\nCopyright (c) 2024 OpenAI\n\nPermission is hereby granted, free of charge, to any person obtaining a copy\nof this software and associated documentation files (the \"Software\"), to deal\nin the Software without restriction, including without limitation the rights\nto use, copy, modify, merge, publish, distribute, sublicense, and/or sell\ncopies of the Software, and to permit persons to whom the Software is\nfurnished to do so, subject to the following conditions:\n\nThe above copyright notice and this permission notice shall be included in all\ncopies or substantial portions of the Software.\n\nTHE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\nIMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\nFITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\nAUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\nLIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\nOUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\nSOFTWARE.\n" }, { "path": "Makefile", "content": "SHELL := /bin/bash\n\nclean-venv: rm-venv\n\tpython3 -m venv venv \n\n\nrm-venv:\n\tif [ -d \"venv\" ]; then rm -rf venv; fi\n\nmaybe-clear-shared:\nifeq ($(SKIP_CLEAR_SHARED), true)\nelse\n\tpip cache remove openai_slackbot\nendif\n\nbuild-shared:\n\tpip install -e ./shared/openai-slackbot\n\n\nbuild-bot: maybe-clear-shared build-shared\n\tcd bots/$(BOT) && $(MAKE) init-pyproject && pip install -e .\n\n\nrun-bot:\n\tpython bots/$(BOT)/$(subst -,_,$(BOT))/bot.py\n\n\nclear:\n\tfind . | grep -E \"(/__pycache__$|\\.pyc$|\\.pyo$)\" | xargs rm -rf\n\n\nbuild-all: \n\t$(MAKE) build-bot BOT=triage-slackbot SKIP_CLEAR_SHARED=true\n\n\ntest-all: \n\tpytest shared/openai-slackbot && \\\n\tpytest bots/triage-slackbot && \\\n\tpytest bots/incident-response-slackbot" }, { "path": "README.md", "content": "# OpenAI Security Bots 🤖\n\nSlack bots integrated with OpenAI APIs to streamline security team's workflows.\n\nAll the bots can be found under `bots/` directory.\n\n```\nshared/\n openai-slackbot/\nbots/\n triage-slackbot/\n incident-response-slackbot/\n sdlc-slackbot/\n```\n\nRefer to each bot's README for more information and setup instruction.\n\n\nIf you wish to contribute, note this repo uses pre-commit to help. In this directory, run:\n```\npip install pre-commit\npre-commit install\n```\n" }, { "path": "SECURITY.md", "content": "# Security Policy\nFor a more in-depth look at our security policy, please check out our [Coordinated Vulnerability Disclosure Policy](https://openai.com/security/disclosure/#:~:text=Disclosure%20Policy,-Security%20is%20essential&text=OpenAI%27s%20coordinated%20vulnerability%20disclosure%20policy,expect%20from%20us%20in%20return.).\n\nOur PGP key can located [at this address.](https://cdn.openai.com/security.txt)\n" }, { "path": "bots/incident-response-slackbot/Makefile", "content": "CWD := $(shell pwd)\nREPO_ROOT := $(shell git rev-parse --show-toplevel)\nESCAPED_REPO_ROOT := $(shell echo $(REPO_ROOT) | sed 's/\\//\\\\\\//g')\n\ninit-env-file:\n\tcp ./incident_response_slackbot/.env.template ./incident_response_slackbot/.env\n\ninit-pyproject:\n\tcat $(CWD)/pyproject.template.toml | \\\n\tsed \"s/\\$$REPO_ROOT/$(ESCAPED_REPO_ROOT)/g\" > $(CWD)/pyproject.toml \n" }, { "path": "bots/incident-response-slackbot/README.md", "content": "

\n \"triage-slackbot-logo\"\n

Incident Response Slackbot

\n

\n\nIncident Response Slackbot automatically chats with users who have been part of an incident alert.\n\n\n\n## Prerequisites\n\nYou will need:\n1. A Slack application (aka your triage bot) with Socket Mode enabled\n2. OpenAI API key\n\nGrab your `SLACK_BOT_TOKEN` by Oauth & Permissions tab in your Slack App page.\n\nGenerate an App-level token for your Slack app, by going to:\n```\nYour Slack App > Basic Information > App-Level Tokens > Generate Token and Scopes\n```\nCreate a new token with `connections:write` scope. This is your `SOCKET_APP_TOKEN` token.\n\nOnce you have them, from the current directory, run:\n```\n$ make init-env-file\n```\nand fill in the right values.\n\nYour Slack App needs the following scopes:\n\n - users:read\n - channels:history\n - chat:write\n - groups:history\n\n## Setup\n\nFrom the current directory, run:\n```\nmake init-pyproject\n```\n\nFrom the repo root, run:\n```\nmake clean-venv\nsource venv/bin/activate\nmake build-bot BOT=incident-response-slackbot\n```\n\n## Run bot with example configuration\n\nThe example configuration is `config.toml`. Replace the configuration values as needed. In particular, the bot will post to channel `feed_channel_id`, and will take an OpenAI Organization ID associated with your OpenAI API key.\n\n⚠️ *Make sure that the bot is added to the channels it needs to read from and post to.* ⚠️\n\nWe will need to add example alerts to `./scripts/alerts.toml` Replace with alert information and user_id. To get the user_id:\n1. Click on the desired user name within Slack. \n2. Click on the ellpises (three dots).\n3. Click on \"Copy Member ID\".\n\n⚠️ *These are mock alerts. In real-world scenarios, this will be integrated with alert feed/database* ⚠️\n\nTo generate an axample alert, in this directory, run:\n```\npython ./scripts/send_alert.py\n```\n\nAn example alert will be sent to the channel.\n\n\nhttps://github.com/openai/openai-security-bots/assets/124844323/b919639c-b691-4b01-aa0c-7be987c9a70b\n\n\nTo have the bot start listening, run the following from the repo root:\n\n```\nmake run-bot BOT=incident-response-slackbot\n```\n\nNow you can start a chat with a user, or do nothing. \nWhen you start a chat, \n\n1. The bot will reach out to the user involved with the alert\n2. Post a message to the original thread in monitoring channel what was sent to the user (message generated with OpenAI API)\n3. Post any messages the user sends to original thread\n4. Checks to see if the user has answered the question using OpenAI's API.\n - If yes, end the chat and provide a summary to the original thread\n - If no, continues sending a message to the user, and repeats this step\n \nLet's start a chat:\n\nhttps://github.com/openai/openai-security-bots/assets/124844323/4b5dd292-b4d3-437a-9809-d6d80e824a9d\n\n\n\n## Alert Details\n\nIn practice, the app will connect with a database or queuing system that monitors alerts. We provide a mock alert system here, and a mock database to hold the state of users and their alerts.\n\nIn the `alerts.toml` file:\n\n```\n[[ alerts ]]\nid = \"pivot\"\n...\nuser_id = ID of person to start chat with (@harold user)\n\n[alerts.properties]\nsource_host = \"source.machine.org\"\ndestination_host = \"destination.machine.org\"\n\n[[ alerts ]]\nid = \"privesc\"\n...\nuser_id = ID of person to start chat with (@harold user)\n```\n" }, { "path": "bots/incident-response-slackbot/incident_response_slackbot/bot.py", "content": "import asyncio\nimport os\n\nfrom incident_response_slackbot.config import load_config, get_config\nfrom incident_response_slackbot.handlers import (\n InboundDirectMessageHandler,\n InboundIncidentDoNothingHandler,\n InboundIncidentEndChatHandler,\n InboundIncidentStartChatHandler,\n)\nfrom openai_slackbot.bot import start_bot\n\nif __name__ == \"__main__\":\n current_dir = os.path.dirname(os.path.abspath(__file__))\n load_config(os.path.join(current_dir, \"config.toml\"))\n\n message_handler = InboundDirectMessageHandler\n action_handlers = [\n InboundIncidentStartChatHandler,\n InboundIncidentDoNothingHandler,\n InboundIncidentEndChatHandler,\n ]\n\n template_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), \"templates\")\n\n config = get_config()\n asyncio.run(\n start_bot(\n openai_organization_id=config.openai_organization_id,\n slack_message_handler=message_handler,\n slack_action_handlers=action_handlers,\n slack_template_path=template_path,\n )\n )\n" }, { "path": "bots/incident-response-slackbot/incident_response_slackbot/config.py", "content": "import os\nimport typing as t\n\nimport toml\nfrom dotenv import load_dotenv\nfrom pydantic import BaseModel\n\n_CONFIG = None\n\n\nclass Config(BaseModel):\n # OpenAI organization ID associated with OpenAI API key.\n openai_organization_id: str\n\n # Slack channel where triage alerts are posted.\n feed_channel_id: str\n\n\ndef load_config(config_path: str = None) -> Config:\n load_dotenv()\n\n if config_path is None:\n config_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), \"config.toml\")\n with open(config_path) as f:\n cfg = toml.loads(f.read())\n config = Config(**cfg)\n\n global _CONFIG\n _CONFIG = config\n return _CONFIG\n\n\ndef get_config() -> Config:\n global _CONFIG\n if _CONFIG is None:\n raise Exception(\"config not initialized, call load_config() first\")\n return _CONFIG\n" }, { "path": "bots/incident-response-slackbot/incident_response_slackbot/config.toml", "content": "# Organization ID associated with OpenAI API key.\nopenai_organization_id = \"\"\n\n# Where the alerts will be posted.\nfeed_channel_id = \"\"\n\n\n" }, { "path": "bots/incident-response-slackbot/incident_response_slackbot/db/database.py", "content": "import os\nimport pickle\n\n\nclass Database:\n \"\"\"\n This class represents a database for storing user messages.\n The data is stored in a pickle file.\n \"\"\"\n\n def __init__(self):\n \"\"\"\n Initialize the database. Load data from the pickle file if it exists,\n otherwise create an empty dictionary.\n \"\"\"\n current_dir = os.path.dirname(os.path.realpath(__file__))\n self.file_path = os.path.join(current_dir, \"data.pkl\")\n self.data = {}\n\n def _load_data(self):\n \"\"\"\n Load data from the pickle file if it exists,\n otherwise return an empty dictionary.\n \"\"\"\n if os.path.exists(self.file_path):\n with open(self.file_path, \"rb\") as f:\n return pickle.load(f)\n else:\n return {}\n\n def _save(self):\n \"\"\"\n Save the current state of the database to the pickle file.\n \"\"\"\n with open(self.file_path, \"wb\") as f:\n pickle.dump(self.data, f)\n\n # Add a new entry to the database\n def add(self, user_id, message_ts):\n \"\"\"\n Add a new entry to the database. If the user_id already exists,\n update the message timestamp. Otherwise, create a new entry.\n \"\"\"\n self.data = self._load_data()\n self.data.setdefault(user_id, {})[\"message_ts\"] = message_ts\n self._save()\n\n # Delete an entry from the database\n def delete(self, user_id):\n \"\"\"\n Delete an entry from the database using the user_id as the key.\n \"\"\"\n self.data = self._load_data()\n self.data.pop(user_id, None)\n self._save()\n\n # Check if user_id exists in the database\n def user_exists(self, user_id):\n \"\"\"\n Check if the user_id exists in the database.\n \"\"\"\n self.data = self._load_data()\n return user_id in self.data\n\n # Return the message timestamp for a given user_id\n def get_ts(self, user_id):\n \"\"\"\n Return the message timestamp for a given user_id.\n \"\"\"\n self.data = self._load_data()\n return self.data[user_id][\"message_ts\"]\n\n # Return the user_id given a message_ts\n def get_user_id(self, message_ts):\n \"\"\"\n Return the user_id given a message_ts.\n \"\"\"\n self.data = self._load_data()\n for user_id, data in self.data.items():\n if data[\"message_ts\"] == message_ts:\n return user_id\n return None\n" }, { "path": "bots/incident-response-slackbot/incident_response_slackbot/handlers.py", "content": "import os\nimport pickle\nimport typing as t\nfrom enum import Enum\nfrom logging import getLogger\n\nfrom incident_response_slackbot.config import load_config, get_config\nfrom incident_response_slackbot.db.database import Database\nfrom incident_response_slackbot.openai_utils import (\n create_greeting,\n generate_awareness_question,\n get_thread_summary,\n get_user_awareness,\n messages_to_string,\n)\nfrom openai_slackbot.handlers import BaseActionHandler, BaseMessageHandler\n\nlogger = getLogger(__name__)\n\nDATABASE = Database()\n\nclass InboundDirectMessageHandler(BaseMessageHandler):\n \"\"\"\n Handles Direct Messages for incident response use cases\n \"\"\"\n\n def __init__(self, slack_client):\n super().__init__(slack_client)\n self.config = get_config()\n\n async def should_handle(self, args):\n return True\n\n async def handle(self, args):\n event = args.event\n user_id = event.get(\"user\")\n\n if not DATABASE.user_exists(user_id):\n # If the user_id does not exist, they're not part of an active chat\n return\n\n message_ts = DATABASE.get_ts(user_id)\n await self.send_message_to_channel(event, message_ts)\n\n user_awareness = await get_user_awareness(event[\"text\"])\n logger.info(f\"User awareness decision: {user_awareness}\")\n\n if user_awareness[\"has_answered\"]:\n await self.handle_user_response(user_id, message_ts)\n else:\n await self.nudge_user(user_id, message_ts)\n\n async def send_message_to_channel(self, event, message_ts):\n # Send the received message to the monitoring channel\n await self._slack_client.post_message(\n channel=self.config.feed_channel_id,\n text=f\"Received message from <@{event['user']}>:\\n> {event['text']}\",\n thread_ts=message_ts,\n )\n\n async def handle_user_response(self, user_id, message_ts):\n # User has answered the question\n messages = await self._slack_client.get_thread_messages(\n channel=self.config.feed_channel_id,\n thread_ts=message_ts,\n )\n\n # Send the end message to the user\n thank_you = \"Thanks for your time!\"\n await self._slack_client.post_message(\n channel=user_id,\n text=thank_you,\n )\n\n # Send message to the channel\n await self._slack_client.post_message(\n channel=self.config.feed_channel_id,\n text=f\"Sent message to <@{user_id}>:\\n> {thank_you}\",\n thread_ts=message_ts,\n )\n\n summary = await get_thread_summary(messages)\n\n # Send message to the channel\n await self._slack_client.post_message(\n channel=self.config.feed_channel_id,\n text=f\"Here is the summary of the chat:\\n> {summary}\",\n thread_ts=message_ts,\n )\n\n DATABASE.delete(user_id)\n\n await self.end_chat(message_ts)\n\n async def end_chat(self, message_ts):\n original_blocks = await self._slack_client.get_original_blocks(\n message_ts, self.config.feed_channel_id\n )\n\n # Remove action buttons and add \"Chat has ended\" text\n new_blocks = [block for block in original_blocks if block.get(\"type\") != \"actions\"]\n\n # Add the \"Chat has ended\" text\n new_blocks.append(\n {\n \"type\": \"section\",\n \"block_id\": \"end_chat_automatically\",\n \"text\": {\n \"type\": \"mrkdwn\",\n \"text\": f\"The chat was automatically ended from SecurityBot review. :done_:\",\n \"verbatim\": True,\n },\n }\n )\n\n await self._slack_client.update_message(\n channel=self.config.feed_channel_id,\n blocks=new_blocks,\n ts=message_ts,\n text=\"Ended chat automatically\",\n )\n\n async def nudge_user(self, user_id, message_ts):\n # User has not answered the question\n\n nudge_message = await generate_awareness_question()\n # Send the greeting message to the user\n await self._slack_client.post_message(\n channel=user_id,\n text=nudge_message,\n )\n\n # Send message to the channel\n await self._slack_client.post_message(\n channel=self.config.feed_channel_id,\n text=f\"Sent message to <@{user_id}>:\\n> {nudge_message}\",\n thread_ts=message_ts,\n )\n\n\nclass InboundIncidentStartChatHandler(BaseActionHandler):\n def __init__(self, slack_client):\n super().__init__(slack_client)\n self.config = get_config()\n\n @property\n def action_id(self):\n return \"start_chat_submit_action\"\n\n async def handle(self, args):\n body = args.body\n original_message = body[\"container\"]\n original_message_ts = original_message[\"message_ts\"]\n alert_user_id = DATABASE.get_user_id(original_message_ts)\n user = body[\"user\"]\n\n name = user[\"name\"]\n first_name = name.split(\".\")[1]\n\n logger.info(f\"Handling inbound incident start chat action from {user['name']}\")\n\n # Update the blocks and elements\n blocks = self.update_blocks(body, alert_user_id)\n\n # Add the \"Started a chat\" text\n blocks.append(self.create_chat_start_section(user[\"id\"]))\n\n messages = await self._slack_client.get_thread_messages(\n channel=self.config.feed_channel_id,\n thread_ts=original_message_ts,\n )\n\n message = await self._slack_client.update_message(\n channel=self.config.feed_channel_id,\n blocks=blocks,\n ts=original_message_ts,\n text=messages[0][\"text\"],\n )\n\n text_messages = messages_to_string(messages)\n logger.info(f\"Alert and detail: {text_messages}\")\n\n username = await self._slack_client.get_user_display_name(alert_user_id)\n\n greeting_message = await create_greeting(first_name, text_messages)\n logger.info(f\"generated greeting message: {greeting_message}\")\n\n # Send the greeting message to the user and to the channel\n await self.send_greeting_message(alert_user_id, greeting_message, original_message_ts)\n\n logger.info(f\"Succesfully started chat with user: {username}\")\n\n return message\n\n def update_blocks(self, body, alert_user_id):\n body_copy = body.copy()\n new_elements = []\n for block in body_copy.get(\"message\", {}).get(\"blocks\", []):\n if block.get(\"type\") == \"actions\":\n for element in block.get(\"elements\", []):\n if element.get(\"action_id\") == \"do_nothing_submit_action\":\n element[\"action_id\"] = \"end_chat_submit_action\"\n element[\"text\"][\"text\"] = \"End Chat\"\n element[\"value\"] = alert_user_id\n new_elements.append(element)\n block[\"elements\"] = new_elements\n return body_copy.get(\"message\", {}).get(\"blocks\", [])\n\n def create_chat_start_section(self, user_id):\n return {\n \"type\": \"section\",\n \"block_id\": \"started_chat\",\n \"text\": {\n \"type\": \"mrkdwn\",\n \"text\": f\"<@{user_id}> started a chat.\",\n \"verbatim\": True,\n },\n }\n\n async def send_greeting_message(self, alert_user_id, greeting_message, original_message_ts):\n # Send the greeting message to the user\n await self._slack_client.post_message(\n channel=alert_user_id,\n text=greeting_message,\n )\n\n # Send message to the channel\n await self._slack_client.post_message(\n channel=self.config.feed_channel_id,\n text=f\"Sent message to <@{alert_user_id}>:\\n> {greeting_message}\",\n thread_ts=original_message_ts,\n )\n\n\nclass InboundIncidentDoNothingHandler(BaseActionHandler):\n \"\"\"\n Handles incoming alerts and decides whether to take no action.\n This will close the alert and mark the status as complete.\n \"\"\"\n\n def __init__(self, slack_client):\n super().__init__(slack_client)\n self.config = get_config()\n\n @property\n def action_id(self):\n return \"do_nothing_submit_action\"\n\n async def handle(self, args):\n body = args.body\n user_id = body[\"user\"][\"id\"]\n original_message_ts = body[\"message\"][\"ts\"]\n\n # Remove action buttons and add \"Chat has ended\" text\n new_blocks = [\n block\n for block in body.get(\"message\", {}).get(\"blocks\", [])\n if block.get(\"type\") != \"actions\"\n ]\n\n # Add the \"Chat has ended\" text\n new_blocks.append(\n {\n \"type\": \"section\",\n \"block_id\": \"do_nothing\",\n \"text\": {\n \"type\": \"mrkdwn\",\n \"text\": f\"<@{user_id}> decided that no action was necessary :done_:\",\n \"verbatim\": True,\n },\n }\n )\n\n await self._slack_client.update_message(\n channel=self.config.feed_channel_id,\n blocks=new_blocks,\n ts=original_message_ts,\n text=\"Do Nothing action selected\",\n )\n\n\nclass InboundIncidentEndChatHandler(BaseActionHandler):\n \"\"\"\n Ends the chat manually\n \"\"\"\n\n def __init__(self, slack_client):\n super().__init__(slack_client)\n self.config = get_config()\n\n @property\n def action_id(self):\n return \"end_chat_submit_action\"\n\n async def handle(self, args):\n body = args.body\n user_id = body[\"user\"][\"id\"]\n message_ts = body[\"message\"][\"ts\"]\n\n alert_user_id = DATABASE.get_user_id(message_ts)\n\n original_blocks = await self._slack_client.get_original_blocks(\n message_ts, self.config.feed_channel_id\n )\n\n # Remove action buttons and add \"Chat has ended\" text\n new_blocks = [block for block in original_blocks if block.get(\"type\") != \"actions\"]\n\n # Add the \"Chat has ended\" text\n new_blocks.append(\n {\n \"type\": \"section\",\n \"block_id\": \"end_chat_manually\",\n \"text\": {\n \"type\": \"mrkdwn\",\n \"text\": f\"<@{user_id}> has ended the chat :done_:\",\n \"verbatim\": True,\n },\n }\n )\n\n await self._slack_client.update_message(\n channel=self.config.feed_channel_id,\n blocks=new_blocks,\n ts=message_ts,\n text=\"Ended chat automatically\",\n )\n\n # User has answered the question\n messages = await self._slack_client.get_thread_messages(\n channel=self.config.feed_channel_id,\n thread_ts=message_ts,\n )\n\n thank_you = \"Thanks for your time!\"\n await self._slack_client.post_message(\n channel=alert_user_id,\n text=thank_you,\n )\n\n # Send message to the channel\n await self._slack_client.post_message(\n channel=self.config.feed_channel_id,\n text=f\"Sent message to <@{alert_user_id}>:\\n> {thank_you}\",\n thread_ts=message_ts,\n )\n\n summary = await get_thread_summary(messages)\n\n # Send message to the channel\n await self._slack_client.post_message(\n channel=self.config.feed_channel_id,\n text=f\"Here is the summary of the chat:\\n> {summary}\",\n thread_ts=message_ts,\n )\n\n DATABASE.delete(user_id)\n" }, { "path": "bots/incident-response-slackbot/incident_response_slackbot/openai_utils.py", "content": "import json\n\nimport openai\nfrom incident_response_slackbot.config import load_config, get_config\n\nload_config()\nconfig = get_config()\n\n# Convert slack threaded messages to string\ndef messages_to_string(messages):\n text_messages = \" \".join([message[\"text\"] for message in messages if \"text\" in message])\n return text_messages\n\n\nasync def get_clean_output(completion: str) -> str:\n return completion.choices[0].message.content\n\n\nasync def create_greeting(username, details):\n if not openai.api_key:\n raise Exception(\"OpenAI API key not found.\")\n\n prompt = f\"\"\"\n You are a helpful cybersecurity AI analyst assistant to the security team that wants to keep\n your company secure. You just received an alert with the following details:\n {details}\n Without being accusatory, gently ask the user, whose name is {username} in a casual tone if they were aware\n about the topic of the alert.\n Keep the message brief, not more than 3 or 4 sentences.\n Do not end with a signature. End with a question.\n \"\"\"\n\n messages = [\n {\"role\": \"system\", \"content\": prompt},\n {\"role\": \"user\", \"content\": \"\"},\n ]\n\n completion = openai.chat.completions.create(\n model=\"gpt-4-32k\",\n messages=messages,\n temperature=0.3,\n stream=False,\n )\n response = await get_clean_output(completion)\n return response\n\n\naware_decision_function = [\n {\n \"name\": \"is_user_aware\",\n \"description\": \"Determines if the user has answered whether they were aware, and what that response is.\",\n \"parameters\": {\n \"type\": \"object\",\n \"properties\": {\n \"has_answered\": {\n \"type\": \"boolean\",\n \"description\": \"Determine whether user answered the quesiton of whether they were aware.\",\n },\n \"is_aware\": {\n \"type\": \"boolean\",\n \"description\": \"Determine whether user was aware of the alert details.\",\n },\n },\n \"required\": [\"has_answered\", \"is_aware\"],\n },\n }\n]\n\n\nasync def get_user_awareness(inbound_direct_message: str) -> str:\n \"\"\"\n This function uses the OpenAI Chat Completion API to determine whether user was aware.\n \"\"\"\n # Define the prompt\n prompt = f\"\"\"\n You are a helpful cybersecurity AI analyst assistant to the security team that wants to keep\n your company secure. You just received an alert and are having a chat with the user whether\n they were aware about the details of an alert. Based on the chat so far, determine whether\n the user has answered the question of whether they were aware of the alert details, and whether\n they were aware or not.\n \"\"\"\n\n messages = [\n {\"role\": \"system\", \"content\": prompt},\n {\"role\": \"user\", \"content\": inbound_direct_message},\n ]\n\n # Call the API\n response = openai.chat.completions.create(\n model=\"gpt-4-32k\",\n messages=messages,\n temperature=0,\n stream=False,\n functions=aware_decision_function,\n function_call={\"name\": \"is_user_aware\"},\n )\n\n function_args = json.loads(response.choices[0].message.function_call.arguments) # type: ignore\n return function_args\n\n\nasync def get_thread_summary(messages):\n if not openai.api_key:\n raise Exception(\"OpenAI API key not found.\")\n\n text_messages = messages_to_string(messages)\n\n prompt = f\"\"\"\n You are a helpful cybersecurity AI analyst assistant to the security team that wants to keep\n your company secure. The following is a conversation that you had with the user.\n Please summarize the following conversation, and note whether the user was aware or not aware\n of the alert, and whether they acted suspiciously when answering:\n {text_messages}\n \"\"\"\n\n messages = [\n {\"role\": \"system\", \"content\": prompt},\n {\"role\": \"user\", \"content\": \"\"},\n ]\n\n completion = openai.chat.completions.create(\n model=\"gpt-4-32k\",\n messages=messages,\n temperature=0.3,\n stream=False,\n )\n response = await get_clean_output(completion)\n return response\n\n\nasync def generate_awareness_question():\n if not openai.api_key:\n raise Exception(\"OpenAI API key not found.\")\n\n prompt = f\"\"\"\n You are a helpful cybersecurity AI analyst assistant to the security team that wants to keep\n your company secure. You have received an alert regarding the user you're chatting with, and\n you have asked whether the user was aware of the alert. The user has not answered the question,\n so now you are asking the user again whether they were aware of the alert. You ask in a gentle,\n kind, and casual tone. You keep it short, to two sentences at most. You end with a question.\n \"\"\"\n\n messages = [\n {\"role\": \"system\", \"content\": prompt},\n {\"role\": \"user\", \"content\": \"\"},\n ]\n\n completion = openai.chat.completions.create(\n model=\"gpt-4-32k\",\n messages=messages,\n temperature=0.5,\n stream=False,\n )\n response = await get_clean_output(completion)\n return response\n" }, { "path": "bots/incident-response-slackbot/incident_response_slackbot/templates/messages/incident_alert.j2", "content": "[\n\t{\n\t\t\"type\": \"section\",\n\t\t\"text\": {\n\t\t\t\"type\": \"mrkdwn\",\n\t\t\t\"text\": \"Incident: {{ alert_name }} with user <@{{ user_id }}>\"\n\t\t}\n\t},\n\t{\n\t\t\"type\": \"context\",\n\t\t\"elements\": [\n\t\t\t{\n\t\t\t\t\"type\": \"plain_text\",\n\t\t\t\t\"text\": \"Details in :thread:\",\n\t\t\t\t\"emoji\": true\n\t\t\t}\n\t\t]\n\t},\n\t{\n\t\t\"type\": \"actions\",\n\t\t\"elements\": [\n\t\t\t{\n\t\t\t\t\"type\": \"button\",\n\t\t\t\t\"text\": {\n\t\t\t\t\t\"type\": \"plain_text\",\n\t\t\t\t\t\"emoji\": true,\n\t\t\t\t\t\"text\": \"Start Chat\"\n\t\t\t\t},\n\t\t\t\t\"style\": \"primary\",\n\t\t\t\t\"value\": \"{{ user_id }}\",\n\t\t\t\t\"action_id\": \"start_chat_submit_action\"\n\t\t\t},\n\t\t\t{\n\t\t\t\t\"type\": \"button\",\n\t\t\t\t\"text\": {\n\t\t\t\t\t\"type\": \"plain_text\",\n\t\t\t\t\t\"emoji\": true,\n\t\t\t\t\t\"text\": \"Do Nothing\"\n\t\t\t\t},\n\t\t\t\t\"style\": \"danger\",\n\t\t\t\t\"value\": \"recategorize\",\n\t\t\t\t\"action_id\": \"do_nothing_submit_action\"\n\t\t\t}\n\t\t]\n\t}\n]\n" }, { "path": "bots/incident-response-slackbot/pyproject.template.toml", "content": "[project]\nname = \"openai-incident-response-slackbot\"\nrequires-python = \">=3.8\"\nversion = \"1.0.0\"\ndependencies = [\n \"toml\",\n \"openai_slackbot @ file://$REPO_ROOT/shared/openai-slackbot\",\n]\n\n[build-system]\nrequires = [\"setuptools>=64.0\"]\nbuild-backend = \"setuptools.build_meta\"\n\n[tool.pytest.ini_options]\nasyncio_mode = \"auto\"\nenv = [\n \"SLACK_BOT_TOKEN=mock-token\",\n \"SOCKET_APP_TOKEN=mock-token\",\n \"OPENAI_API_KEY=mock-key\",\n]" }, { "path": "bots/incident-response-slackbot/scripts/alert_feed.py", "content": "import os\nfrom logging import getLogger\n\nfrom incident_response_slackbot.config import load_config, get_config\nfrom incident_response_slackbot.db.database import Database\nfrom openai_slackbot.clients.slack import CreateSlackMessageResponse, SlackClient\nfrom openai_slackbot.utils.envvars import string\nfrom slack_bolt.app.async_app import AsyncApp\n\nlogger = getLogger(__name__)\n\nDATABASE = Database()\n\nload_config()\nconfig = get_config()\n\nasync def post_alert(alert):\n \"\"\"\n This function posts an alert to the Slack channel.\n It first initializes the Slack client with the bot token and template path.\n Then, it extracts the user_id, alert_name, and properties from the alert.\n Finally, it posts the alert to the Slack channel and sends the initial details.\n\n Args:\n alert (dict): The alert to be posted. It should contain 'user_id', 'name', and 'properties'.\n \"\"\"\n\n slack_bot_token = string(\"SLACK_BOT_TOKEN\")\n app = AsyncApp(token=slack_bot_token)\n slack_template_path = os.path.join(\n os.path.dirname(os.path.abspath(__file__)),\n \"../incident_response_slackbot/templates\",\n )\n slack_client = SlackClient(app.client, slack_template_path)\n\n # Extracting the user_id, alert_name, and properties from the alert\n user_id = alert.get(\"user_id\")\n alert_name = alert.get(\"name\")\n properties = alert.get(\"properties\")\n\n message = await incident_feed_begin(\n slack_client=slack_client, user_id=user_id, alert_name=alert_name\n )\n\n DATABASE.add(user_id, message.ts)\n\n await initial_details(slack_client=slack_client, message=message, properties=properties)\n\n\nasync def incident_feed_begin(\n *, slack_client: SlackClient, user_id: str, alert_name: str\n) -> CreateSlackMessageResponse:\n \"\"\"\n This function begins the incident feed by posting the initial alert message.\n It first renders the blocks from the template with the user_id and alert_name.\n Then, it posts the message to the Slack channel.\n\n Args:\n slack_client (SlackClient): The Slack client.\n user_id (str): The Slack user ID.\n alert_name (str): The name of the alert.\n\n Returns:\n CreateSlackMessageResponse: The response from creating the Slack message.\n\n Raises:\n Exception: If the initial alert message fails to post.\n \"\"\"\n\n try:\n blocks = slack_client.render_blocks_from_template(\n \"messages/incident_alert.j2\",\n {\n \"user_id\": user_id,\n \"alert_name\": alert_name,\n },\n )\n message = await slack_client.post_message(\n channel=config.feed_channel_id,\n blocks=blocks,\n text=f\"{alert_name} via <@{user_id}>\",\n )\n return message\n\n except Exception:\n logger.exception(\"Initial alert feed message failed\")\n\n\ndef get_alert_details(**kwargs) -> str:\n \"\"\"\n This function returns the alert details for each key in the\n property. Each alert could have different properties.\n \"\"\"\n content = \"\"\n for key, value in kwargs.items():\n line = f\"The value of {key} for this alert is {value}. \"\n content += line\n if content:\n return content\n return \"No details available for this alert.\"\n\n\nasync def initial_details(*, slack_client: SlackClient, message, properties):\n \"\"\"\n This function posts the initial details of an alert to a Slack thread.\n\n Args:\n slack_client (SlackClient): The Slack client.\n message: The initial alert message.\n properties: The properties of the alert.\n \"\"\"\n thread_ts = message.ts\n details = get_alert_details(**properties)\n\n await slack_client.post_message(\n channel=config.feed_channel_id, text=f\"{details}\", thread_ts=thread_ts\n )\n" }, { "path": "bots/incident-response-slackbot/scripts/alerts.toml", "content": "# Alert Examples - These are the alerts that will be sent to the feed channel.\n[[alerts]]\nid = \"pivot\"\nname = \"Pivoting\"\ndescription = \"User was found pivoting from one host to another\"\nuser_id = \"\"\n\n[alerts.properties]\nsource_host = \"source.machine.org\"\ndestination_host = \"destination.machine.org\"\n\n[[alerts]]\nid = \"privesc\"\nname = \"Privileged Escalation\"\ndescription = \"Privileged escalation was detected\"\nuser_id = \"\"\n\n[alerts.properties]\nprevious_role = \"reader\"\nnew_role = \"admin\"\n" }, { "path": "bots/incident-response-slackbot/scripts/send_alert.py", "content": "import asyncio\nimport os\nimport random\nimport time\n\nimport toml\nfrom alert_feed import post_alert\n\n\ndef load_alerts():\n alerts_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), \"alerts.toml\")\n with open(alerts_path, \"r\") as file:\n data = toml.load(file)\n return data\n\n\ndef generate_random_alert(alerts):\n random_alert = random.choice([0, 1])\n print(alerts[\"alerts\"][random_alert])\n return alerts[\"alerts\"][random_alert]\n\n\nasync def main():\n alerts = load_alerts()\n\n alert = generate_random_alert(alerts)\n await post_alert(alert)\n\n\nif __name__ == \"__main__\":\n asyncio.run(main())\n" }, { "path": "bots/incident-response-slackbot/tests/__init__.py", "content": "" }, { "path": "bots/incident-response-slackbot/tests/conftest.py", "content": "import os\nfrom unittest.mock import AsyncMock, MagicMock, patch\n\nimport pytest\nimport toml\nfrom incident_response_slackbot.config import load_config\nfrom pydantic import ValidationError\n\n####################\n##### FIXTURES #####\n####################\n\n\n@pytest.fixture(autouse=True)\ndef mock_config():\n # Load the test config\n current_dir = os.path.dirname(os.path.abspath(__file__))\n config_path = os.path.join(current_dir, \"test_config.toml\")\n try:\n config = load_config(config_path)\n except ValidationError as e:\n print(f\"Error validating the config: {e}\")\n raise\n return config\n\n\n@pytest.fixture()\ndef mock_slack_client():\n # Mock the Slack client\n slack_client = MagicMock()\n slack_client.post_message = AsyncMock()\n slack_client.update_message = AsyncMock()\n slack_client.get_original_blocks = AsyncMock()\n slack_client.get_thread_messages = AsyncMock()\n\n return slack_client\n\n\n@pytest.fixture(autouse=True)\n@patch(\"openai.ChatCompletion.create\")\ndef mock_chat_completion(mock_create):\n mock_create.return_value = {\n \"id\": \"chatcmpl-1234567890\",\n \"object\": \"chat.completion\",\n \"created\": 1640995200,\n \"model\": \"gpt-4-32k\",\n \"usage\": {\"prompt_tokens\": 10, \"completion_tokens\": 20, \"total_tokens\": 30},\n \"choices\": [\n {\n \"message\": {\n \"role\": \"assistant\",\n \"content\": \"This is a mock response from the OpenAI API.\",\n },\n \"finish_reason\": \"stop\",\n \"index\": 0,\n }\n ],\n }\n yield\n\n\n@pytest.fixture\ndef mock_generate_awareness_question():\n with patch(\n \"incident_response_slackbot.handlers.generate_awareness_question\",\n new_callable=AsyncMock,\n ) as mock_generate_question:\n mock_generate_question.return_value = \"Mock question\"\n yield mock_generate_question\n\n\n@pytest.fixture\ndef mock_get_thread_summary():\n with patch(\n \"incident_response_slackbot.handlers.get_thread_summary\",\n new_callable=AsyncMock,\n ) as mock_get_summary:\n mock_get_summary.return_value = \"Mock summary\"\n yield mock_get_summary\n" }, { "path": "bots/incident-response-slackbot/tests/test_config.toml", "content": "# Organization ID associated with OpenAI API key.\nopenai_organization_id = \"mock_openai_organization_id\"\n\n# Where the alerts will be posted.\nfeed_channel_id = \"mock_feed_channel_id\"\n\n" }, { "path": "bots/incident-response-slackbot/tests/test_handlers.py", "content": "# in tests/test_handlers.py\nfrom unittest.mock import AsyncMock, MagicMock, patch\nfrom collections import namedtuple\n\nimport pytest\nfrom incident_response_slackbot.handlers import (\n InboundDirectMessageHandler,\n InboundIncidentStartChatHandler,\n InboundIncidentDoNothingHandler,\n InboundIncidentEndChatHandler,\n)\n\n\n@pytest.mark.asyncio\nasync def test_send_message_to_channel(mock_slack_client, mock_config):\n # Arrange\n handler = InboundDirectMessageHandler(slack_client=mock_slack_client)\n mock_event = {\"text\": \"mock_event_text\", \"user_profile\": {\"name\": \"mock_user_name\"}}\n mock_message_ts = \"mock_message_ts\"\n\n # Act\n await handler.send_message_to_channel(mock_event, mock_message_ts)\n\n # Assert\n mock_slack_client.post_message.assert_called_once_with(\n channel=mock_config.feed_channel_id,\n text=\"Received message from <@mock_user_name>:\\n> mock_event_text\",\n thread_ts=mock_message_ts,\n )\n\n\n@pytest.mark.asyncio\nasync def test_end_chat(mock_slack_client, mock_config):\n # Define the return value for get_original_blocks\n mock_slack_client.get_original_blocks.return_value = [\n {\"type\": \"section\", \"block_id\": \"block1\"},\n {\"type\": \"actions\", \"block_id\": \"block2\"},\n {\"type\": \"section\", \"block_id\": \"block3\"},\n ]\n\n # Create an instance of the handler\n handler = InboundDirectMessageHandler(slack_client=mock_slack_client)\n\n # Call the end_chat method\n await handler.end_chat(\"12345\")\n\n # Assert that update_message was called with the correct arguments\n mock_slack_client.update_message.assert_called_once()\n\n # Get the actual call arguments\n args, kwargs = mock_slack_client.update_message.call_args\n\n # Check the blocks argument\n assert kwargs[\"blocks\"] == [\n {\"type\": \"section\", \"block_id\": \"block1\"},\n {\"type\": \"section\", \"block_id\": \"block3\"},\n {\n \"type\": \"section\",\n \"block_id\": \"end_chat_automatically\",\n \"text\": {\n \"type\": \"mrkdwn\",\n \"text\": \"The chat was automatically ended from SecurityBot review. :done_:\",\n \"verbatim\": True,\n },\n },\n ]\n\n\n@pytest.mark.asyncio\nasync def test_nudge_user(mock_slack_client, mock_config, mock_generate_awareness_question):\n # Create an instance of the handler\n handler = InboundDirectMessageHandler(slack_client=mock_slack_client)\n\n # Call the nudge_user method\n await handler.nudge_user(\"user123\", \"12345\")\n\n # Assert that post_message was called twice with the correct arguments\n assert mock_slack_client.post_message.call_count == 2\n mock_slack_client.post_message.assert_any_call(channel=\"user123\", text=\"Mock question\")\n mock_slack_client.post_message.assert_any_call(\n channel=handler.config.feed_channel_id,\n text=\"Sent message to <@user123>:\\n> Mock question\",\n thread_ts=\"12345\",\n )\n\n\n@pytest.mark.asyncio\nasync def test_incident_start_chat_handle(mock_slack_client, mock_config):\n # Create an instance of the handler\n handler = InboundIncidentStartChatHandler(slack_client=mock_slack_client)\n\n # Create a mock args object\n args = MagicMock()\n args.body = {\n \"container\": {\"message_ts\": \"12345\"},\n \"user\": {\"name\": \"test_user\", \"id\": \"user123\"},\n }\n\n # Mock the DATABASE.get_user_id method\n with patch(\n \"incident_response_slackbot.handlers.DATABASE.get_user_id\", return_value=\"alert_user123\"\n ) as mock_get_user_id, patch(\n \"incident_response_slackbot.handlers.create_greeting\",\n new_callable=AsyncMock,\n return_value=\"greeting message\",\n ) as mock_create_greeting, patch.object(\n handler._slack_client, \"get_thread_messages\", new_callable=AsyncMock\n ), patch.object(\n handler._slack_client, \"update_message\", new_callable=AsyncMock\n ), patch.object(\n handler._slack_client,\n \"get_user_display_name\",\n new_callable=AsyncMock,\n return_value=\"username\",\n ), patch.object(\n handler._slack_client, \"post_message\", new_callable=AsyncMock\n ):\n # Call the handle method\n await handler.handle(args)\n\n # Assert that the slack client methods were called with the correct arguments\n handler._slack_client.get_thread_messages.assert_called_once_with(\n channel=mock_config.feed_channel_id, thread_ts=\"12345\"\n )\n handler._slack_client.update_message.assert_called_once()\n handler._slack_client.get_user_display_name.assert_called_once_with(\"alert_user123\")\n\n # Assert that post_message was called twice\n assert handler._slack_client.post_message.call_count == 2\n\n # Assert that post_message was called with the correct arguments\n handler._slack_client.post_message.assert_any_call(\n channel=\"alert_user123\", text=\"greeting message\"\n )\n\n\n@pytest.mark.asyncio\nasync def test_do_nothing_handle(mock_slack_client, mock_config):\n # Create an instance of the handler\n handler = InboundIncidentDoNothingHandler(slack_client=mock_slack_client)\n\n # Create a mock args object\n args = MagicMock()\n args.body = {\n \"user\": {\"id\": \"user123\"},\n \"message\": {\"ts\": \"12345\", \"blocks\": [{\"type\": \"actions\"}, {\"type\": \"section\"}]},\n }\n\n # Call the handle method\n await handler.handle(args)\n\n # Assert that the slack client update_message method was called with the correct arguments\n mock_slack_client.update_message.assert_called_once_with(\n channel=mock_config.feed_channel_id,\n blocks=[\n {\"type\": \"section\"},\n {\n \"type\": \"section\",\n \"block_id\": \"do_nothing\",\n \"text\": {\n \"type\": \"mrkdwn\",\n \"text\": \"<@user123> decided that no action was necessary :done_:\",\n \"verbatim\": True,\n },\n },\n ],\n ts=\"12345\",\n text=\"Do Nothing action selected\",\n )\n\n\n@pytest.mark.asyncio\nasync def test_end_chat_handle(mock_slack_client, mock_config, mock_get_thread_summary):\n # Mock the Slack client and the database\n with patch(\n \"incident_response_slackbot.handlers.Database\", new_callable=AsyncMock\n ) as MockDatabase:\n # Instantiate the handler\n handler = InboundIncidentEndChatHandler(slack_client=mock_slack_client)\n\n # Define a namedtuple for args\n Args = namedtuple(\"Args\", [\"body\"])\n\n # Instantiate the args object\n args = Args(body={\"user\": {\"id\": \"user_id\"}, \"message\": {\"ts\": \"message_ts\"}})\n\n # Mock the get_user_id method of the database to return a user id\n MockDatabase.get_user_id.return_value = \"alert_user_id\"\n\n # Call the handle method\n await handler.handle(args)\n\n # Assert that the correct methods were called with the expected arguments\n mock_slack_client.get_original_blocks.assert_called_once_with(\n \"message_ts\", mock_config.feed_channel_id\n )\n mock_slack_client.update_message.assert_called()\n mock_slack_client.post_message.assert_called()\n" }, { "path": "bots/incident-response-slackbot/tests/test_openai.py", "content": "# in tests/test_openai_utils.py\nfrom unittest.mock import patch\n\nimport pytest\nfrom incident_response_slackbot.openai_utils import get_user_awareness\n\n\n@pytest.mark.asyncio\n@patch(\"openai.ChatCompletion.create\")\nasync def test_get_user_awareness(mock_create):\n # Arrange\n mock_create.return_value = {\n \"choices\": [\n {\n \"message\": {\n \"function_call\": {\"arguments\": '{\"has_answered\": true, \"is_aware\": false}'}\n }\n }\n ]\n }\n inbound_direct_message = \"mock_inbound_direct_message\"\n\n # Act\n result = await get_user_awareness(inbound_direct_message)\n\n # Assert\n assert result == {\"has_answered\": True, \"is_aware\": False}\n" }, { "path": "bots/sdlc-slackbot/Makefile", "content": "CWD := $(shell pwd)\nREPO_ROOT := $(shell git rev-parse --show-toplevel)\nESCAPED_REPO_ROOT := $(shell echo $(REPO_ROOT) | sed 's/\\//\\\\\\//g')\n\ninit-env-file:\n\tcp ./sdlc_slackbot/.env.template ./sdlc_slackbot/.env\n\ninit-pyproject:\n\tcat $(CWD)/pyproject.template.toml | \\\n\tsed \"s/\\$$REPO_ROOT/$(ESCAPED_REPO_ROOT)/g\" > $(CWD)/pyproject.toml \n" }, { "path": "bots/sdlc-slackbot/README.md", "content": "\n

\n \"sdlc-slackbot-logo\"\n

SDLC Slackbot

\n

\n\nSDLC Slackbot decides if a project merits a security review.\n\n## Prerequisites\n\nYou will need:\n1. A Slack application (aka your sdlc bot) with Socket Mode enabled\n2. OpenAI API key\n\nGenerate an App-level token for your Slack app, by going to:\n```\nYour Slack App > Basic Information > App-Level Tokens > Generate Token and Scopes\n```\nCreate a new token with `connections:write` scope. This is your `SOCKET_APP_TOKEN` token.\n\nOnce you have them, from the current directory, run:\n```\n$ make init-env-file\n```\nand fill in the right values.\n\nYour Slack App needs the following scopes:\n\n- app\\_mentions:read\n- channels:join\n- channels:read\n- channels:history\n- chat:write\n- groups:history\n- groups:read\n- groups:write\n- usergroups:read\n- users:read\n- users:read.email\n\n\n## Setup\n\nFrom the current directory, run:\n```\nmake init-pyproject\n```\n\nFrom the repo root, run:\n```\nmake clean-venv\nsource venv/bin/activate\nmake build-bot BOT=sdlc-slackbot\n```\n\n## Run bot with example configuration\n\nThe example configuration is `config.toml`. Replace the configuration values as needed.\nYou need to at least replace the `openai_organization_id` and `notification_channel_id`.\n\nFor optional Google Docs integration you'll need a 'credentials.json' file:\n- Go to the Google Cloud Console.\n- Select your project.\n- Navigate to \"APIs & Services\" > \"Credentials\".\n- Under \"OAuth 2.0 Client IDs\", find your client ID and download the JSON file.\n- Save it in the `sdlc-slackbot/sdlc_slackbot` directory as `credentials.json`.\n\n\n\n⚠️ *Make sure that the bot is added to the channels it needs to read from and post to.* ⚠️\n\nFrom the repo root, run:\n\n```\nmake run-bot BOT=sdlc-slackbot\n```\n\n" }, { "path": "bots/sdlc-slackbot/pyproject.template.toml", "content": "[project]\nname = \"openai-sdlc-slackbot\"\nrequires-python = \">=3.8\"\nversion = \"1.0.0\"\ndependencies = [\n \"toml\",\n \"openai_slackbot @ file://$REPO_ROOT/shared/openai-slackbot\",\n \"validators\",\n \"google-auth\",\n \"google-auth-httplib2\",\n \"google-auth-oauthlib\",\n \"google-api-python-client\",\n \"psycopg\",\n \"psycopg2-binary\",\n \"peewee\",\n]\n\n[build-system]\nrequires = [\"setuptools>=64.0\"]\nbuild-backend = \"setuptools.build_meta\"\n\n[tool.pytest.ini_options]\nasyncio_mode = \"auto\"\nenv = [\n \"SLACK_BOT_TOKEN=mock-token\",\n \"SOCKET_APP_TOKEN=mock-token\",\n \"OPENAI_API_KEY=mock-key\",\n]\n" }, { "path": "bots/sdlc-slackbot/requirements.txt", "content": "openai\npython-dotenv\nslack-bolt\nvalidators\ngoogle-auth\ngoogle-auth-httplib2\ngoogle-auth-oauthlib\ngoogle-api-python-client\npsycopg\npsycopg2-binary\npeewee\n" }, { "path": "bots/sdlc-slackbot/sdlc_slackbot/bot.py", "content": "import asyncio\nimport hashlib\nimport json\nimport os\nimport re\nimport threading\nimport time\nimport traceback\nfrom logging import getLogger\n\nimport validate\nimport validators\nfrom database import *\nfrom gdoc import gdoc_get\nfrom openai_slackbot.bot import init_bot, start_app\nfrom openai_slackbot.utils.envvars import string\nfrom peewee import *\nfrom playhouse.db_url import *\nfrom playhouse.shortcuts import model_to_dict\nfrom sdlc_slackbot.config import get_config, load_config\nfrom slack_bolt import App\nfrom slack_bolt.adapter.socket_mode import SocketModeHandler\nfrom slack_sdk import WebClient\nfrom utils import *\n\n\nlogger = getLogger(__name__)\n\n\nasync def send_update_notification(input, response):\n risk_str, confidence_str = risk_and_confidence_to_string(response)\n risk_num = response[\"risk\"]\n confidence_num = response[\"confidence\"]\n\n msg = f\"\"\"\n Project {input['project_name']} has been updated and has a new decision:\n\n This new decision for the project is that it is: *{risk_str}({risk_num})* with *{confidence_str}({confidence_num})*. {response['justification']}.\"\n \"\"\"\n\n await app.client.chat_postMessage(channel=config.notification_channel_id, text=msg)\n\n\ndef hash_content(content):\n return hashlib.sha256(content.encode(\"utf-8\")).hexdigest()\n\n\nurl_pat = re.compile(\n r\"http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\\\\(\\\\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+\\b(?!>)\"\n)\n\n\ndef extract_urls(text):\n logger.info(f\"extracting urls from {text}\")\n urls = re.findall(url_pat, text)\n return [url for url in urls if validators.url(url)]\n\n\nasync def async_fetch_slack(url):\n parts = url.split(\"/\")\n channel = parts[-2]\n ts = parts[-1]\n ts = ts[1:] # trim p\n seconds = ts[:-6]\n nanoseconds = ts[-6:]\n result = await app.client.conversations_replies(channel=channel, ts=f\"{seconds}.{nanoseconds}\")\n return \" \".join(message.get(\"text\", \"\") for message in result.data.get(\"messages\", []))\n\n\ncontent_fetchers = [\n (\n lambda u: u.startswith((\"https://docs.google.com/document\", \"docs.google.com/document\")),\n gdoc_get,\n ),\n (lambda u: \"slack.com/archives\" in u, async_fetch_slack),\n]\n\n\nasync def fetch_content(url):\n for condition, fetcher in content_fetchers:\n if condition(url):\n if asyncio.iscoroutinefunction(fetcher):\n return await fetcher(url) # Await the result if it's a coroutine function\n else:\n return fetcher(url) # Call it directly if it's not a coroutine function\n\n\nform = [\n input_block(\n \"project_name\",\n \"Project Name\",\n field(\"plain_text_input\", \"Enter the project name\"),\n ),\n input_block(\n \"project_description\",\n \"Project Description\",\n field(\"plain_text_input\", \"Enter the project description\", multiline=True),\n ),\n input_block(\n \"links_to_resources\",\n \"Links to Resources\",\n field(\"plain_text_input\", \"Enter links to resources\", multiline=True),\n ),\n input_block(\"point_of_contact\", \"Point of Contact\", field(\"users_select\", \"Select a user\")),\n input_block(\n \"estimated_go_live_date\",\n \"Estimated Go Live Date\",\n field(\"datepicker\", \"Select a date\"),\n ),\n submit_block(\"submit_form\"),\n]\n\n\ndef risk_and_confidence_to_string(decision):\n # Lookup tables for risk and confidence\n risk_lookup = {\n (1, 2): \"extremely low risk\",\n (3, 3): \"low risk\",\n (4, 5): \"medium risk\",\n (6, 7): \"medium-high risk\",\n (8, 9): \"high risk\",\n (10, 10): \"critical risk\",\n }\n\n confidence_lookup = {\n (1, 2): \"extremely low confidence\",\n (3, 3): \"low confidence\",\n (4, 5): \"medium confidence\",\n (6, 7): \"medium-high confidence\",\n (8, 9): \"high confidence\",\n (10, 10): \"extreme confidence\",\n }\n\n # Function to find the appropriate string from a lookup table\n def find_in_lookup(value, lookup):\n for (min_val, max_val), descriptor in lookup.items():\n if min_val <= value <= max_val:\n return descriptor\n return \"unknown\"\n\n # Convert risk and confidence using their respective lookup tables\n risk_str = find_in_lookup(decision[\"risk\"], risk_lookup)\n confidence_str = find_in_lookup(decision[\"confidence\"], confidence_lookup)\n\n return risk_str, confidence_str\n\n\ndef decision_msg(response):\n risk_str, confidence_str = risk_and_confidence_to_string(response)\n risk_num = response[\"risk\"]\n confidence_num = response[\"confidence\"]\n\n return f\"Thanks for your response! Based on this input, we've decided that this project is *{risk_str}({risk_num})* with *{confidence_str}({confidence_num})*. {response['justification']}.\"\n\n\nskip_params = set(\n [\n \"id\",\n \"project_name\",\n \"links_to_resources\",\n \"point_of_contact\",\n \"estimated_go_live_date\",\n ]\n)\n\nmultiple_whitespace_pat = re.compile(r\"\\s+\")\n\n\ndef model_params_to_str(params):\n ss = (v for k, v in params.items() if k not in skip_params)\n return re.sub(multiple_whitespace_pat, \" \", \"\\n\".join(map(str, ss))).strip()\n\n\ndef summarize_params(params):\n summary = {}\n for k, v in params.items():\n if k not in skip_params:\n summary[k] = ask_ai(\n config.base_prompt + config.summary_prompt, v[: config.context_limit]\n )\n else:\n summary[k] = v\n\n return summary\n\n\nasync def handle_app_mention_events(say, event):\n logger.info(\"App mention event received:\", event)\n await say(blocks=form, thread_ts=event[\"ts\"])\n\n\nasync def handle_message_events(say, message):\n logger.info(\"message: \", message)\n if message[\"channel_type\"] == \"im\":\n await say(blocks=form, thread_ts=message[\"ts\"])\n\n\ndef get_response_with_retry(prompt, context, max_retries=1):\n prompt = prompt.strip().replace(\"\\n\", \" \")\n retries = 0\n while retries <= max_retries:\n try:\n response = ask_ai(prompt, context)\n return response\n except json.JSONDecodeError as e:\n logger.error(f\"JSON error on attempt {retries + 1}: {e}\")\n retries += 1\n if retries > max_retries:\n return {}\n\n\ndef normalize_response(response):\n if isinstance(response, list):\n return [json.loads(block.text) for block in response]\n elif isinstance(response, dict):\n return [response]\n else:\n raise TypeError(\"Unsupported response type\")\n\n\ndef clean_normalized_response(normalized_responses):\n \"\"\"\n Remove the 'decision' key from each dictionary in a list of dictionaries.\n Break it down into 'risk' and 'confidence'\n\n :param normalized_responses: A list of dictionaries.\n :return: The list of dictionaries with 'decision' key broken down.\n \"\"\"\n for response in normalized_responses:\n if \"decision\" in response:\n decision = response[\"decision\"]\n response[\"risk\"] = decision.get(\"risk\")\n response[\"confidence\"] = decision.get(\"confidence\")\n response.pop(\"decision\", None)\n\n return normalized_responses\n\n\nasync def submit_form(ack, body, say):\n await ack()\n\n try:\n ts = body[\"container\"][\"message_ts\"]\n values = body[\"state\"][\"values\"]\n params = get_form_input(\n values,\n \"project_name\",\n \"project_description\",\n \"links_to_resources\",\n \"point_of_contact\",\n \"estimated_go_live_date\",\n )\n\n validate.required(params, \"project_name\", \"project_description\", \"point_of_contact\")\n\n await say(text=config.reviewing_message, thread_ts=ts)\n\n try:\n assessment = Assessment.create(**params, user_id=body[\"user\"][\"id\"])\n except IntegrityError as e:\n raise validate.ValidationError(\"project_name\", \"must be unique\")\n\n resources = []\n for url in extract_urls(params.get(\"links_to_resources\", \"\")):\n content = await fetch_content(url)\n if content:\n params[url] = content\n resources.append(\n dict(\n assessment=assessment,\n url=url,\n content_hash=hash_content(content),\n )\n )\n Resource.insert_many(resources).execute()\n\n context = model_params_to_str(params)\n if len(context) > config.context_limit:\n logger.info(f\"context too long: {len(context)}. Summarizing...\")\n summarized_context = summarize_params(params)\n context = model_params_to_str(summarized_context)\n # FIXME: is there a better way to handle this? currently, if the summary is still too long\n # we just give up and cut it off\n if len(context) > config.context_limit:\n logger.info(f\"Summarized context too long: {len(context)}. Cutting off...\")\n context = context[: config.context_limit]\n\n response = get_response_with_retry(config.base_prompt + config.initial_prompt, context)\n if not response:\n return\n\n normalized_response = normalize_response(response)\n clean_response = clean_normalized_response(normalized_response)\n\n for item in clean_response:\n if item[\"outcome\"] == \"decision\":\n assessment.update(**item).execute()\n await say(text=decision_msg(item), thread_ts=ts)\n elif item[\"outcome\"] == \"followup\":\n db_questions = [dict(assessment=assessment, question=q) for q in item[\"questions\"]]\n Question.insert_many(db_questions).execute()\n\n form = []\n for i, q in enumerate(item[\"questions\"]):\n form.append(\n input_block(\n f\"question_{i}\",\n q,\n field(\"plain_text_input\", \"...\", multiline=True),\n )\n )\n form.append(submit_block(f\"submit_followup_questions_{assessment.id}\"))\n\n await say(blocks=form, thread_ts=ts)\n except validate.ValidationError as e:\n await say(text=f\"{e.field}: {e.issue}\", thread_ts=ts)\n except Exception as e:\n import traceback\n\n traceback.print_exc()\n await say(text=config.irrecoverable_error_message, thread_ts=ts)\n\n\nasync def submit_followup_questions(ack, body, say):\n await ack()\n\n try:\n assessment_id = int(body[\"actions\"][0][\"action_id\"].split(\"_\")[-1])\n ts = body[\"container\"][\"message_ts\"]\n assessment = Assessment.get(Assessment.id == assessment_id)\n params = model_to_dict(assessment)\n followup_questions = [q.question for q in assessment.questions]\n except Exception as e:\n logger.error(f\"Failed to find params for user {body['user']['id']}\", e)\n await say(text=config.recoverable_error_message, thread_ts=ts)\n return\n\n try:\n await say(text=config.reviewing_message, thread_ts=ts)\n\n values = body[\"state\"][\"values\"]\n for i, q in enumerate(followup_questions):\n params[q] = values[f\"question_{i}\"][f\"question_{i}_input\"][\"value\"]\n\n for question in assessment.questions:\n question.answer = params[question.question]\n question.save()\n\n context = model_params_to_str(params)\n\n response = ask_ai(config.base_prompt, context)\n text_to_update = response\n if (\n isinstance(response, dict)\n and \"text\" in response\n and \"type\" in response\n and response[\"type\"] == \"text\"\n ):\n # Extract the text from the content block\n text_to_update = response.text\n\n normalized_response = normalize_response(text_to_update)\n clean_response = clean_normalized_response(normalized_response)\n\n for item in clean_response:\n if item[\"outcome\"] == \"decision\":\n assessment.update(**item).execute()\n await say(text=decision_msg(item), thread_ts=ts)\n\n except Exception as e:\n logger.error(f\"error: {e} processing followup questions: {json.dumps(body, indent=2)}\")\n await say(text=config.irrecoverable_error_message, thread_ts=ts)\n\n\ndef update_resources():\n while True:\n time.sleep(monitor_thread_sleep_seconds)\n try:\n for assessment in Assessment.select():\n logger.info(f\"checking {assessment.project_name} for updates\")\n\n assessment_params = model_to_dict(assessment)\n new_params = assessment_params.copy()\n\n changed = False\n\n previous_content = \"\"\n\n for resource in assessment.resources:\n new_content = asyncio.run(fetch_content(resource.url))\n\n if resource.content_hash != hash_content(new_content):\n # just save previous content in memory temporarily\n previous_content = resource.content\n resource.content = new_content\n new_params[resource.url] = new_content\n changed = True\n\n if not changed:\n continue\n\n old_context = model_params_to_str(assessment_params)\n new_context = model_params_to_str(new_params)\n\n context = {\n \"previous_context\": previous_content,\n \"previous_decision\": {\n \"risk\": assessment.risk,\n \"confidence\": assessment.confidence,\n \"justification\": assessment.justification,\n },\n \"new_context\": new_content,\n }\n\n context_json = json.dumps(context, indent=2)\n\n new_response = ask_ai(config.base_prompt + config.update_prompt, context_json)\n\n resource.content_hash = hash_content(new_content)\n resource.save()\n\n if new_response[\"outcome\"] == \"unchanged\":\n continue\n\n normalized_response = normalize_response(new_response)\n clean_response = clean_normalized_response(normalized_response)\n\n for item in clean_response:\n assessment.update(**item).execute()\n\n asyncio.run(send_update_notification(assessment_params, new_response))\n except Exception as e:\n logger.error(f\"error: {e} updating resources\")\n traceback.print_exc()\n\n\nmonitor_thread_sleep_seconds = 6\n\nif __name__ == \"__main__\":\n current_dir = os.path.dirname(os.path.abspath(__file__))\n load_config(os.path.join(current_dir, \"config.toml\"))\n\n template_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), \"templates\")\n\n config = get_config()\n\n message_handler = []\n action_handlers = []\n view_submission_handlers = []\n\n app = asyncio.run(\n init_bot(\n openai_organization_id=config.openai_organization_id,\n slack_message_handler=message_handler,\n slack_action_handlers=action_handlers,\n slack_template_path=template_path,\n )\n )\n\n # Register your custom event handlers\n app.event(\"app_mention\")(handle_app_mention_events)\n app.message()(handle_message_events)\n\n app.action(\"submit_form\")(submit_form)\n app.action(re.compile(\"submit_followup_questions.*\"))(submit_followup_questions)\n\n t = threading.Thread(target=update_resources)\n t.start()\n\n # Start the app\n asyncio.run(start_app(app))\n" }, { "path": "bots/sdlc-slackbot/sdlc_slackbot/config.py", "content": "import os\nimport typing as t\n\nimport toml\nfrom dotenv import load_dotenv\nfrom pydantic import BaseModel, ValidationError, field_validator, model_validator\nfrom pydantic.functional_validators import AfterValidator, BeforeValidator\n\n_CONFIG = None\n\n\ndef validate_channel(channel_id: str) -> str:\n if not channel_id.startswith(\"C\"):\n raise ValueError(\"channel ID must start with 'C'\")\n return channel_id\n\n\nclass Config(BaseModel):\n # OpenAI organization ID associated with OpenAI API key.\n openai_organization_id: str\n\n context_limit: int\n\n # OpenAI prompts\n base_prompt: str\n initial_prompt: str\n update_prompt: str\n summary_prompt: str\n\n reviewing_message: str\n recoverable_error_message: str\n irrecoverable_error_message: str\n\n # Slack channel for notifications\n notification_channel_id: t.Annotated[str, AfterValidator(validate_channel)]\n\n\ndef load_config(path: str):\n load_dotenv()\n\n with open(path) as f:\n cfg = toml.loads(f.read())\n config = Config(**cfg)\n\n global _CONFIG\n _CONFIG = config\n return _CONFIG\n\n\ndef get_config() -> Config:\n global _CONFIG\n if _CONFIG is None:\n raise Exception(\"config not initialized, call load_config() first\")\n return _CONFIG\n" }, { "path": "bots/sdlc-slackbot/sdlc_slackbot/config.toml", "content": "# Organization ID associated with OpenAI API key.\nopenai_organization_id = \"\"\n\nnotification_channel_id = \"\"\n\ncontext_limit = 31_500\n\nbase_prompt = \"\"\"\nYou're a highly skilled security analyst who is excellent at asking the right questions to determine the true risk of a development project to your organization.\nYou work at a small company with a small security team with limited resources. You ruthlessly prioritize your team's time to ensure that you can reduce\nthe greatest amount of security risk with your limited resources. Your bar for reviewing things before launch is high. They should have the potential to introduce significant security risk to your company.\nYour bar for putting things in your backlog is lower but also still high. Projects in the backlog should have the potential to be high leverage security work.\nYou should base your decision on a variety of factors including but not limited to:\n- if changes would affect any path to model weights or customer data \n- if changes are accessible from the internet\n- if changes affect end users\n- if changes affect security critical parts of the system, like authentication, authorization, encryption\n- if changes deal with historically risky technology like xml parsing\n- if changes will likely involve interpolating user input into a dynamic language like html, sql, or javascript\n\nIf changes affect model weights and customer data, the risk should definitely increase. Model weights should never be exposed. Customer data should be handled extremely safely. \n\nBe conservative about how you rate the risk score in general though. There are tons of projects and there's not enough bandwidth to cover everything in depth.\nYou've been asked to analyze a new project that is being developed by another team at your company \nand determine if and when it should be reviewed by your team. Your decision option should be two numeric scores: \nOne score for the risk: score with values between 1 and 10, where 1 means zero risk, while 10 means extremely risky and needs a security review. \nThe second score is your confidence: how confident are you about your decision, with 1 meaning very low confidence, while 10 meaning super confident.\nPut both number in the \"decision\" as follows: \n\ndecision: { \"risk\": \n \"confidence\": \n\nYou should base your decision on how risky you think the project is to the company.\nYou should also provide a brief justification for your decision. You should only respond with a json object.\nThe decision object should look like this: {\"outcome\": \"decision\", \"decision\": { \"risk\": <1 to 10>, \"confidence\": <1 to 10>}, \"justification\": \"I think this project is risky because...\"}.\n\nDon't send any other responses. Our team has very limited resources and only wants to review the most important projects, so you\nshould enforce a high bar for go live reviews.\n\"\"\"\n\n\ninitial_prompt = \"\"\"\nYou should ask as many questions as you need to make an informed, accurate decision. Don't hesitate at all to ask followup questions.\nAsk for clarification for any critical vague language in the fields below. If the project description doesn't contain information about\nfactors that are critical to your decision, ask about them.\nIf you need to ask a followup question, respond with {\"outcome\": \"followup\", \"questions\": [\"What is the project's budget?\", \"What is the project's timeline?\"]}.\n\"\"\"\n\nupdate_prompt = \"\"\"\nYou've already reviewed this project before, but some information has changed. Below you'll find the previous project context\nyour previous decision, a justification for your previous decision and the new content. If your decision still makes sense\nrespond with a json object with a single property named \"outcome\" set to \"unchanged\". If your decision no longer makes sense\nrespond with a new json object containing the outcome and decision. Carefuly compare the \"previous_context\" part with the \"new_context\" part and detect any changes that might be affecting security components.\n\"\"\"\n\nsummary_prompt = \"\"\"\nYou're a highly skilled security analyst who is excellent at asking the right questions to determine the true risk of a development project to your organization.\nYou work at a small company with a small security team with limited resources. You ruthlessly prioritize your team's time to ensure that you can reduce\nthe greatest amount of security risk with your limited resources.\nPlease provide a summary of the key security design elements, potential vulnerabilities, and recommended mitigation strategies presented in the following project document. Highlight any areas of particular concern and emphasize best practices that have been implemented. Also outline all key technical aspects of the project that you assess would require a security review. Anything that deals with data, end users, authentication, authorization, encryption, untrusted user input, internet exposure, new features or risky technologies like file processing, xml parsing and so on\"\n\"\"\"\n\nreviewing_message = \"Thanks for your submission! We're currently reviewing it and will let you know if we need more information and if / when you'll need a review\"\n\nrecoverable_error_message = \"Something went wrong. We've been notified and will fix it as soon as possible. Start a new conversation to try again\"\n\nirrecoverable_error_message = \"Something went wrong. We've been notified and will fix it as soon as possible. Start a thread in #security if you need help immediately.\"\n" }, { "path": "bots/sdlc-slackbot/sdlc_slackbot/database.py", "content": "import os\n\nfrom peewee import *\nfrom playhouse.db_url import *\n\ndb_url = os.getenv(\"DATABASE_URL\") or \"postgres://postgres:postgres@localhost:5432/postgres\"\ndb = connect(db_url)\n\n\nclass BaseModel(Model):\n class Meta:\n database = db\n\n\nclass Assessment(BaseModel):\n project_name = CharField(unique=True)\n project_description = TextField()\n links_to_resources = TextField(null=True)\n point_of_contact = CharField()\n estimated_go_live_date = CharField(null=True)\n outcome = CharField(null=True)\n risk = IntegerField(null=True) # Storing risk as an integer\n confidence = IntegerField(null=True) # Storing confidence as an integer\n justification = TextField(null=True)\n\n\nclass Question(Model):\n question = TextField()\n answer = TextField(null=True)\n assessment = ForeignKeyField(Assessment, backref=\"questions\")\n\n class Meta:\n database = db\n indexes = (((\"question\", \"assessment\"), True),)\n\n\nclass Resource(BaseModel):\n url = TextField()\n content_hash = CharField()\n content = TextField(null=True)\n assessment = ForeignKeyField(Assessment, backref=\"resources\")\n\n\ndb.connect()\ndb.create_tables([Assessment, Question, Resource])\n" }, { "path": "bots/sdlc-slackbot/sdlc_slackbot/gdoc.py", "content": "from __future__ import print_function\n\nimport os.path\nimport re\nfrom logging import getLogger\n\nfrom google.auth.transport.requests import Request\nfrom google.oauth2.credentials import Credentials\nfrom google_auth_oauthlib.flow import InstalledAppFlow\nfrom googleapiclient.discovery import build\nfrom googleapiclient.errors import HttpError\n\n# If modifying these scopes, delete the file token.json.\nSCOPES = [\"https://www.googleapis.com/auth/documents.readonly\"]\n\nlogger = getLogger(__name__)\n\n\ndef read_paragraph_element(element):\n \"\"\"Returns the text in the given ParagraphElement.\n\n Args:\n element: a ParagraphElement from a Google Doc.\n \"\"\"\n text_run = element.get(\"textRun\")\n if not text_run:\n return \"\"\n return text_run.get(\"content\")\n\n\ndef read_structural_elements(elements):\n \"\"\"Recurses through a list of Structural Elements to read a document's text where text may be\n in nested elements.\n\n Args:\n elements: a list of Structural Elements.\n \"\"\"\n text = \"\"\n for value in elements:\n if \"paragraph\" in value:\n elements = value.get(\"paragraph\").get(\"elements\")\n for elem in elements:\n text += read_paragraph_element(elem)\n elif \"table\" in value:\n # The text in table cells are in nested Structural Elements and tables may be\n # nested.\n table = value.get(\"table\")\n for row in table.get(\"tableRows\"):\n cells = row.get(\"tableCells\")\n for cell in cells:\n text += read_structural_elements(cell.get(\"content\"))\n elif \"tableOfContents\" in value:\n # The text in the TOC is also in a Structural Element.\n toc = value.get(\"tableOfContents\")\n text += read_structural_elements(toc.get(\"content\"))\n return text\n\n\ndef gdoc_creds():\n creds = None\n # The file token.json stores the user's access and refresh tokens, and is\n # created automatically when the authorization flow completes for the first\n # time.\n creds_path = \"./bots/sdlc-slackbot/sdlc_slackbot/\"\n\n if os.path.exists(creds_path + \"token.json\"):\n creds = Credentials.from_authorized_user_file(creds_path + \"token.json\", SCOPES)\n\n # If there are no (valid) credentials available, let the user log in.\n if not creds or not creds.valid:\n if creds and creds.expired and creds.refresh_token:\n creds.refresh(Request())\n else:\n flow = InstalledAppFlow.from_client_secrets_file(\n creds_path + \"credentials.json\", SCOPES\n )\n creds = flow.run_local_server(port=0)\n # Save the credentials for the next run\n with open(creds_path + \"token.json\", \"w\") as token:\n token.write(creds.to_json())\n\n return creds\n\n\ndef gdoc_get(gdoc_url):\n # https://docs.google.com/document/d//edit\n\n result = None\n logger.info(gdoc_url)\n if not gdoc_url.startswith(\"https://docs.google.com/document\") and not gdoc_url.startswith(\n \"docs.google.com/document\"\n ):\n logger.error(\"Invalid google doc url\")\n return result\n\n # This regex captures the ID after \"/d/\" and before an optional \"/edit\", \"/\" or the end of the string.\n pattern = r\"/d/([^/]+)\"\n match = re.search(pattern, gdoc_url)\n\n if match:\n document_id = match.group(1)\n logger.info(document_id)\n else:\n logger.error(\"No ID found in the URL\")\n return result\n\n creds = gdoc_creds()\n try:\n service = build(\"docs\", \"v1\", credentials=creds)\n\n # Retrieve the documents contents from the Docs service.\n document = service.documents().get(documentId=document_id).execute()\n\n logger.info(\"The title of the document is: {}\".format(document.get(\"title\")))\n\n doc_content = document.get(\"body\").get(\"content\")\n result = read_structural_elements(doc_content)\n\n except HttpError as err:\n logger.error(err)\n\n return result\n" }, { "path": "bots/sdlc-slackbot/sdlc_slackbot/utils.py", "content": "import json\nimport os\nfrom logging import getLogger\n\n# import anthropic\nimport openai\n\nlogger = getLogger(__name__)\n\n\ndef get_form_input(values, *fields):\n ret = {}\n for f in fields:\n container = values[f][f + \"_input\"]\n value = container.get(\"value\")\n if value:\n ret[f] = container[\"value\"]\n else:\n for key, item in container.items():\n if key.startswith(\"selected_\") and item:\n ret[f] = item\n break\n return ret\n\n\ndef plain_text(text):\n return dict(type=\"plain_text\", text=text)\n\n\ndef field(type, placeholder, **kwargs):\n return dict(type=type, placeholder=plain_text(placeholder), **kwargs)\n\n\ndef input_block(block_id, label, element):\n if \"action_id\" not in element:\n element[\"action_id\"] = block_id + \"_input\"\n\n return dict(\n type=\"input\",\n block_id=block_id,\n label=plain_text(label),\n element=element,\n )\n\n\ndef submit_block(action_id):\n return dict(\n type=\"actions\",\n elements=[\n dict(\n type=\"button\",\n text=plain_text(\"Submit\"),\n action_id=action_id,\n style=\"primary\",\n )\n ],\n )\n\n\ndef ask_ai(prompt, context):\n # return ask_claude(prompt, context) # YOU CAN USE CLAUDE HERE\n response = ask_gpt(prompt, context)\n\n # Removing leading and trailing backticks and whitespace\n clean_response = response.strip(\"`\\n \")\n\n # Check if 'json' is at the beginning and remove it\n if clean_response.lower().startswith(\"json\"):\n clean_response = clean_response[4:].strip()\n\n # Remove a trailing } if it exists\n if clean_response.endswith(\"}}\"):\n clean_response = clean_response[:-1] # Remove the last character\n\n logger.info(clean_response)\n\n try:\n parsed_response = json.loads(clean_response)\n return parsed_response\n except json.JSONDecodeError as e:\n logger.error(f\"Failed to parse JSON response from ask_gpt: {response}\\nError: {e}\")\n return None\n\n\ndef ask_gpt(prompt, context):\n response = openai.chat.completions.create(\n model=\"gpt-4-32k\",\n messages=[\n {\"role\": \"system\", \"content\": prompt},\n {\"role\": \"user\", \"content\": context},\n ],\n )\n return response.choices[0].message.content\n\n\ndef ask_claude(prompt, context):\n client = anthropic.Anthropic(api_key=os.environ[\"CLAUDE_API_KEY\"])\n message = client.messages.create(\n model=\"claude-3-opus-20240229\",\n max_tokens=4096,\n messages=[{\"role\": \"user\", \"content\": prompt}],\n )\n return message.content\n" }, { "path": "bots/sdlc-slackbot/sdlc_slackbot/validate.py", "content": "class ValidationError(Exception):\n def __init__(self, field, issue):\n self.field = field\n self.issue = issue\n super().__init__(f\"{field} {issue}\")\n\n\ndef required(values, *fields):\n for f in fields:\n if f not in values:\n raise ValidationError(f, \"required\")\n if values[f] == \"\":\n raise ValidationError(f, \"required\")\n" }, { "path": "bots/sdlc-slackbot/setup.py", "content": "from setuptools import setup, find_packages\n\nwith open(\"requirements.txt\", \"r\") as f:\n requirements = f.read().splitlines()\n\nsetup(\n name=\"sdlc_bot\",\n version=\"0.1\",\n packages=find_packages(),\n install_requires=requirements,\n)\n" }, { "path": "bots/triage-slackbot/Makefile", "content": "CWD := $(shell pwd)\nREPO_ROOT := $(shell git rev-parse --show-toplevel)\nESCAPED_REPO_ROOT := $(shell echo $(REPO_ROOT) | sed 's/\\//\\\\\\//g')\n\ninit-env-file:\n\tcp ./triage_slackbot/.env.template ./triage_slackbot/.env\n\ninit-pyproject:\n\tcat $(CWD)/pyproject.template.toml | \\\n\tsed \"s/\\$$REPO_ROOT/$(ESCAPED_REPO_ROOT)/g\" > $(CWD)/pyproject.toml \n" }, { "path": "bots/triage-slackbot/README.md", "content": "

\n \"triage-slackbot-logo\"\n

Triage Slackbot

\n

\n\nTriage Slackbot triages inbound requests in a Slack channel to different sub-teams within your organization.\n\n## Prerequisites\n\nYou will need:\n1. A Slack application (aka your triage bot) with Socket Mode enabled\n2. OpenAI API key\n\nGenerate an App-level token for your Slack app, by going to:\n```\nYour Slack App > Basic Information > App-Level Tokens > Generate Token and Scopes\n```\nCreate a new token with `connections:write` scope. This is your `SOCKET_APP_TOKEN` token.\n\nOnce you have them, from the current directory, run:\n```\n$ make init-env-file\n```\nand fill in the right values.\n\nYour Slack App needs the following scopes:\n\n- channels:history\n- chat:write\n- groups:history\n- reactions:read\n- reactions:write\n\n## Setup\n\nFrom the current directory, run:\n```\nmake init-pyproject\n```\n\nFrom the repo root, run:\n```\nmake clean-venv\nsource venv/bin/activate\nmake build-bot BOT=triage-slackbot\n```\n\n## Run bot with example configuration\n\nThe example configuration is `config.toml`. Replace the configuration values as needed.\n\n⚠️ *Make sure that the bot is added to the channels it needs to read from and post to.* ⚠️\n\nFrom the repo root, run:\n\n```\nmake run-bot BOT=triage-slackbot\n```\n\n## Demo\n\nThis demo is run with the provided `config.toml`. In this demo:\n\n```\ninbound_request_channel_id = ID of #inbound-security-requests channel\nfeed_channel_id = ID of #inbound-security-requests-feed channel\n\n[[ categories ]]\nkey = \"appsec\"\n...\noncall_slack_id = ID of #appsec-requests channel\n\n[[ categories ]]\nkey = \"privacy\"\n...\noncall_slack_id = ID of @tiffany user\n```\n\nThe following triage scenarios are supported: \n\nFirst, the bot categorizes the inbound requests accurately, and on-call acknowledges this prediction.\n\nhttps://github.com/openai/openai-security-bots/assets/10287796/2bb8b301-41b6-450f-a578-482e89a75050\n\nSecondly, the bot categorizes the request into a category that it can autorespond to, e.g. Physical Security, \nand there is no manual action from on-call required.\n\nhttps://github.com/openai/openai-security-bots/assets/10287796/e77bacf0-e16d-4ed3-9567-6f3caaab02ad\n\nFinally, on-call can re-route an inbound request to another category's on-call if the initial predicted \ncategory is not accurate. Additionally, if `other_category_enabled` is set to true, on-call can select any\nchannels it can route the user to:\n\nhttps://github.com/openai/openai-security-bots/assets/10287796/04247a29-f904-42bc-82d8-12b7f2b7e170\n\nThe bot will reply to the thread with this:\n\n\"autorespond\"\n" }, { "path": "bots/triage-slackbot/pyproject.template.toml", "content": "[project]\nname = \"openai-triage-slackbot\"\nrequires-python = \">=3.8\"\nversion = \"1.0.0\"\ndependencies = [\n \"toml\",\n \"openai_slackbot @ file://$REPO_ROOT/shared/openai-slackbot\",\n]\n\n[build-system]\nrequires = [\"setuptools>=64.0\"]\nbuild-backend = \"setuptools.build_meta\"\n\n[tool.pytest.ini_options]\nasyncio_mode = \"auto\"\nenv = [\n \"SLACK_BOT_TOKEN=mock-token\",\n \"SOCKET_APP_TOKEN=mock-token\",\n \"OPENAI_API_KEY=mock-key\",\n]" }, { "path": "bots/triage-slackbot/tests/__init__.py", "content": "" }, { "path": "bots/triage-slackbot/tests/conftest.py", "content": "import os\nfrom unittest.mock import AsyncMock, MagicMock, patch\n\nimport pytest\nfrom openai_slackbot.clients.slack import SlackClient\nfrom triage_slackbot.category import RequestCategory\nfrom triage_slackbot.config import load_config\nfrom triage_slackbot.handlers import MessageTemplatePath\n\n##########################\n##### HELPER METHODS #####\n##########################\n\n\ndef bot_message_extra_data():\n return {\n \"bot_id\": \"bot_id\",\n \"bot_profile\": {\"id\": \"bot_profile_id\"},\n \"team\": \"team\",\n \"type\": \"type\",\n \"user\": \"user\",\n }\n\n\ndef recategorize_message_data(\n *,\n ts,\n channel_id,\n user,\n message,\n category,\n conversation=None,\n):\n return MagicMock(\n ack=AsyncMock(),\n body={\n \"actions\": [\"recategorize_submit_action\"],\n \"container\": {\n \"message_ts\": ts,\n \"channel_id\": channel_id,\n },\n \"user\": user,\n \"message\": message,\n \"state\": {\n \"values\": {\n \"recategorize_select_category_block\": {\n \"recategorize_select_category_action\": {\n \"selected_option\": {\n \"value\": category,\n }\n }\n },\n \"recategorize_select_conversation_block\": {\n \"recategorize_select_conversation_action\": {\n \"selected_conversation\": conversation,\n }\n },\n }\n },\n },\n )\n\n\n####################\n##### FIXTURES #####\n####################\n\n\n@pytest.fixture(autouse=True)\ndef mock_config():\n current_dir = os.path.dirname(os.path.abspath(__file__))\n config_path = os.path.join(current_dir, \"test_config.toml\")\n return load_config(config_path)\n\n\n@pytest.fixture()\ndef mock_post_message_response():\n return AsyncMock(\n return_value=MagicMock(\n ok=True,\n data={\n \"ok\": True,\n \"channel\": \"\",\n \"ts\": \"\",\n \"message\": {\n \"blocks\": [],\n \"text\": \"\",\n \"ts\": \"\",\n **bot_message_extra_data(),\n },\n },\n )\n )\n\n\n@pytest.fixture()\ndef mock_generic_slack_response():\n return AsyncMock(return_value=MagicMock(ok=True, data={\"ok\": True}))\n\n\n@pytest.fixture()\ndef mock_conversations_history_response():\n return AsyncMock(\n return_value=MagicMock(\n ok=True,\n data={\n \"ok\": True,\n \"messages\": [\n {\n \"blocks\": [],\n \"text\": \"\",\n \"ts\": \"\",\n **bot_message_extra_data(),\n }\n ],\n },\n )\n )\n\n\n@pytest.fixture()\ndef mock_get_permalink_response():\n return AsyncMock(\n return_value={\n \"ok\": True,\n \"permalink\": \"mockpermalink\",\n },\n )\n\n\n@pytest.fixture\ndef mock_slack_asyncwebclient(\n mock_conversations_history_response,\n mock_generic_slack_response,\n mock_post_message_response,\n mock_get_permalink_response,\n):\n with patch(\"slack_sdk.web.async_client.AsyncWebClient\", autospec=True) as mock_client:\n wc = mock_client.return_value\n wc.reactions_add = mock_generic_slack_response\n wc.chat_update = mock_generic_slack_response\n wc.conversations_history = mock_conversations_history_response\n wc.chat_postMessage = mock_post_message_response\n wc.chat_getPermalink = mock_get_permalink_response\n yield wc\n\n\n@pytest.fixture\ndef mock_slack_client(mock_slack_asyncwebclient):\n template_path = os.path.join(\n os.path.dirname(os.path.abspath(__file__)), \"../triage_slackbot/templates\"\n )\n return SlackClient(mock_slack_asyncwebclient, template_path)\n\n\n@pytest.fixture\ndef mock_inbound_request_channel_id(mock_config):\n return mock_config.inbound_request_channel_id\n\n\n@pytest.fixture\ndef mock_feed_channel_id(mock_config):\n return mock_config.feed_channel_id\n\n\n@pytest.fixture\ndef mock_appsec_oncall_slack_channel_id(mock_config):\n return mock_config.categories[\"appsec\"].oncall_slack_id\n\n\n@pytest.fixture\ndef mock_privacy_oncall_slack_user_id(mock_config):\n return mock_config.categories[\"privacy\"].oncall_slack_id\n\n\n@pytest.fixture\ndef mock_appsec_oncall_slack_user():\n return {\"id\": \"U1234567890\"}\n\n\n@pytest.fixture\ndef mock_appsec_oncall_slack_user_id(mock_appsec_oncall_slack_user):\n return mock_appsec_oncall_slack_user[\"id\"]\n\n\n@pytest.fixture\ndef mock_inbound_request_ts():\n return \"t0\"\n\n\n@pytest.fixture\ndef mock_feed_message_ts():\n return \"t1\"\n\n\n@pytest.fixture\ndef mock_notify_appsec_oncall_message_ts():\n return \"t2\"\n\n\n@pytest.fixture\ndef mock_appsec_oncall_recategorize_ts():\n return \"t3\"\n\n\n@pytest.fixture\ndef mock_inbound_request(mock_inbound_request_channel_id, mock_inbound_request_ts):\n return MagicMock(\n ack=AsyncMock(),\n event={\n \"channel\": mock_inbound_request_channel_id,\n \"text\": \"sample inbound request\",\n \"thread_ts\": None,\n \"ts\": mock_inbound_request_ts,\n },\n )\n\n\n@pytest.fixture\ndef mock_inbound_request_permalink(mock_inbound_request_channel_id):\n return f\"https://myorg.slack.com/archives/{mock_inbound_request_channel_id}/p1234567890\"\n\n\n@pytest.fixture\nasync def mock_notify_appsec_oncall_message_data(\n mock_slack_client,\n mock_config,\n mock_inbound_request_channel_id,\n mock_inbound_request_permalink,\n mock_inbound_request_ts,\n mock_feed_channel_id,\n mock_feed_message_ts,\n mock_appsec_oncall_slack_channel_id,\n mock_notify_appsec_oncall_message_ts,\n):\n appsec_key = \"appsec\"\n remaining_categories = [c for c in mock_config.categories.values() if c.key != appsec_key]\n blocks = mock_slack_client.render_blocks_from_template(\n MessageTemplatePath.notify_oncall_channel.value,\n {\n \"inbound_message_url\": mock_inbound_request_permalink,\n \"inbound_message_channel\": mock_inbound_request_channel_id,\n \"predicted_category\": mock_config.categories[appsec_key].display_name,\n \"options\": RequestCategory.to_block_options(remaining_categories),\n },\n )\n\n return {\n \"ok\": True,\n \"channel\": mock_appsec_oncall_slack_channel_id,\n \"ts\": mock_notify_appsec_oncall_message_ts,\n \"message\": {\n \"blocks\": blocks,\n \"text\": \"Notify on-call for new inbound request\",\n \"metadata\": {\n \"event_type\": \"notify_oncall\",\n \"event_payload\": {\n \"inbound_message_channel\": mock_inbound_request_channel_id,\n \"inbound_message_ts\": mock_inbound_request_ts,\n \"feed_message_channel\": mock_feed_channel_id,\n \"feed_message_ts\": mock_feed_message_ts,\n \"inbound_message_url\": mock_inbound_request_permalink,\n \"predicted_category\": appsec_key,\n },\n },\n \"ts\": mock_notify_appsec_oncall_message_ts,\n **bot_message_extra_data(),\n },\n }\n\n\n@pytest.fixture\ndef mock_notify_appsec_oncall_message(\n mock_notify_appsec_oncall_message_data,\n mock_appsec_oncall_slack_channel_id,\n mock_appsec_oncall_slack_user,\n):\n return MagicMock(\n ack=AsyncMock(),\n body={\n \"actions\": [\"acknowledge_submit_action\"],\n \"container\": {\n \"message_ts\": mock_notify_appsec_oncall_message_data[\"ts\"],\n \"channel_id\": mock_appsec_oncall_slack_channel_id,\n },\n \"user\": mock_appsec_oncall_slack_user,\n \"message\": mock_notify_appsec_oncall_message_data[\"message\"],\n },\n )\n\n\n@pytest.fixture\ndef mock_appsec_oncall_recategorize_to_privacy_message(\n mock_appsec_oncall_recategorize_ts,\n mock_appsec_oncall_slack_channel_id,\n mock_appsec_oncall_slack_user,\n mock_notify_appsec_oncall_message_data,\n):\n return recategorize_message_data(\n ts=mock_appsec_oncall_recategorize_ts,\n channel_id=mock_appsec_oncall_slack_channel_id,\n user=mock_appsec_oncall_slack_user,\n message=mock_notify_appsec_oncall_message_data[\"message\"],\n category=\"privacy\",\n )\n\n\n@pytest.fixture\ndef mock_appsec_oncall_recategorize_to_other_message(\n mock_appsec_oncall_recategorize_ts,\n mock_appsec_oncall_slack_channel_id,\n mock_appsec_oncall_slack_user,\n mock_notify_appsec_oncall_message_data,\n):\n return recategorize_message_data(\n ts=mock_appsec_oncall_recategorize_ts,\n channel_id=mock_appsec_oncall_slack_channel_id,\n user=mock_appsec_oncall_slack_user,\n message=mock_notify_appsec_oncall_message_data[\"message\"],\n category=\"other\",\n conversation=\"C11111\",\n )\n" }, { "path": "bots/triage-slackbot/tests/test_config.toml", "content": "# Organization ID associated with OpenAI API key.\nopenai_organization_id = \"org-1234\"\n\n# Prompt to use for categorizing inbound requests.\nopenai_prompt = \"\"\"\nYou are currently an on-call engineer for a security team at a tech company. \nYour goal is to triage the following incoming Slack message into three categories:\n1. Privacy, return \"privacy\"\n2. Application security, return \"appsec\"\n3. Physical security, return \"physical_security\"\n\"\"\"\ninbound_request_channel_id = \"C12345\"\nfeed_channel_id = \"C23456\"\nother_category_enabled = true\n\n[[ categories ]] \nkey = \"appsec\"\ndisplay_name = \"Application Security\"\noncall_slack_id = \"C34567\"\nautorespond = false\n\n[[ categories ]] \nkey = \"privacy\"\ndisplay_name = \"Privacy\"\noncall_slack_id = \"U12345\"\nautorespond = false\n\n[[ categories ]] \nkey = \"physical_security\"\ndisplay_name = \"Physical Security\"\nautorespond = true\nautorespond_message = \"Looking for Physical or Office Security? You can reach out to physical-security@company.com.\"" }, { "path": "bots/triage-slackbot/tests/test_handlers.py", "content": "import json\nfrom unittest.mock import call, patch\n\nimport pytest\nfrom triage_slackbot.handlers import (\n InboundRequestAcknowledgeHandler,\n InboundRequestHandler,\n InboundRequestRecategorizeHandler,\n)\nfrom triage_slackbot.openai_utils import openai\n\n\ndef get_mock_chat_completion_response(category: str):\n category_args = json.dumps({\"category\": category})\n return {\n \"choices\": [\n {\n \"message\": {\n \"function_call\": {\n \"arguments\": category_args,\n }\n }\n }\n ]\n }\n\n\ndef assert_chat_completion_called(mock_chat_completion, mock_config):\n mock_chat_completion.create.assert_called_once_with(\n model=\"gpt-4-32k-0613\",\n messages=[\n {\n \"role\": \"system\",\n \"content\": mock_config.openai_prompt,\n },\n {\"role\": \"user\", \"content\": \"sample inbound request\"},\n ],\n temperature=0,\n stream=False,\n functions=[\n {\n \"name\": \"get_predicted_category\",\n \"description\": \"Predicts the category of an inbound request.\",\n \"parameters\": {\n \"type\": \"object\",\n \"properties\": {\n \"category\": {\n \"type\": \"string\",\n \"enum\": [\"appsec\", \"privacy\", \"physical_security\"],\n \"description\": \"Predicted category of the inbound request\",\n }\n },\n \"required\": [\"category\"],\n },\n }\n ],\n function_call={\"name\": \"get_predicted_category\"},\n )\n\n\n@patch.object(openai, \"ChatCompletion\")\nasync def test_inbound_request_handler_handle(\n mock_chat_completion,\n mock_config,\n mock_slack_client,\n mock_inbound_request,\n):\n # Setup mocks\n mock_chat_completion.create.return_value = get_mock_chat_completion_response(\"appsec\")\n\n # Call handler\n handler = InboundRequestHandler(mock_slack_client)\n await handler.maybe_handle(mock_inbound_request)\n\n # Assert that handler calls OpenAI API\n assert_chat_completion_called(mock_chat_completion, mock_config)\n\n mock_slack_client._client.assert_has_calls(\n [\n call.chat_getPermalink(channel=\"C12345\", message_ts=\"t0\"),\n call.chat_postMessage(\n channel=\"C23456\",\n blocks=[\n {\n \"type\": \"section\",\n \"text\": {\n \"type\": \"mrkdwn\",\n \"text\": \"Received an in <#C12345>:\",\n },\n },\n {\n \"type\": \"context\",\n \"elements\": [\n {\n \"type\": \"plain_text\",\n \"text\": \"Predicted category: Application Security\",\n \"emoji\": True,\n },\n {\"type\": \"mrkdwn\", \"text\": \"Triaged to: <#C34567>\"},\n {\n \"type\": \"plain_text\",\n \"text\": \"Triage updates in the :thread:\",\n \"emoji\": True,\n },\n ],\n },\n ],\n text=\"New inbound request received\",\n ),\n call.chat_postMessage(\n channel=\"C34567\",\n thread_ts=None,\n blocks=[\n {\n \"type\": \"section\",\n \"text\": {\n \"type\": \"mrkdwn\",\n \"text\": \":wave: Hi, we received an in <#C12345>, which was categorized as Application Security. Is this accurate?\\n\\n\",\n },\n },\n {\n \"type\": \"context\",\n \"elements\": [\n {\n \"type\": \"plain_text\",\n \"text\": \":thumbsup: Acknowledge this message and response directly to the inbound request.\",\n \"emoji\": True,\n },\n {\n \"type\": \"plain_text\",\n \"text\": \":thumbsdown: Recategorize this message, and if defined, I will route it to the appropriate on-call. If none applies, select Other and pick a channel that I will route the user to.\",\n \"emoji\": True,\n },\n ],\n },\n {\n \"type\": \"actions\",\n \"elements\": [\n {\n \"type\": \"button\",\n \"text\": {\n \"type\": \"plain_text\",\n \"emoji\": True,\n \"text\": \"Acknowledge\",\n },\n \"style\": \"primary\",\n \"value\": \"Application Security\",\n \"action_id\": \"acknowledge_submit_action\",\n },\n {\n \"type\": \"button\",\n \"text\": {\n \"type\": \"plain_text\",\n \"emoji\": True,\n \"text\": \"Inaccurate, recategorize\",\n },\n \"style\": \"danger\",\n \"value\": \"recategorize\",\n \"action_id\": \"recategorize_submit_action\",\n },\n ],\n },\n {\n \"type\": \"section\",\n \"block_id\": \"recategorize_select_category_block\",\n \"text\": {\n \"type\": \"mrkdwn\",\n \"text\": \"*Select a category from the dropdown list, or*\",\n },\n \"accessory\": {\n \"type\": \"static_select\",\n \"placeholder\": {\n \"type\": \"plain_text\",\n \"text\": \"Select an item\",\n \"emoji\": True,\n },\n \"options\": [\n {\n \"text\": {\n \"type\": \"plain_text\",\n \"text\": \"Privacy\",\n \"emoji\": True,\n },\n \"value\": \"privacy\",\n },\n {\n \"text\": {\n \"type\": \"plain_text\",\n \"text\": \"Physical Security\",\n \"emoji\": True,\n },\n \"value\": \"physical_security\",\n },\n {\n \"text\": {\n \"type\": \"plain_text\",\n \"text\": \"Other\",\n \"emoji\": True,\n },\n \"value\": \"other\",\n },\n ],\n \"action_id\": \"recategorize_select_category_action\",\n },\n },\n ],\n metadata={\n \"event_type\": \"notify_oncall\",\n \"event_payload\": {\n \"inbound_message_channel\": \"C12345\",\n \"inbound_message_ts\": \"t0\",\n \"feed_message_channel\": \"\",\n \"feed_message_ts\": \"\",\n \"inbound_message_url\": \"mockpermalink\",\n \"predicted_category\": \"appsec\",\n },\n },\n text=\"Notify on-call for new inbound request\",\n ),\n ]\n )\n\n\n@patch.object(openai, \"ChatCompletion\")\nasync def test_inbound_request_handler_handle_autorespond(\n mock_chat_completion,\n mock_config,\n mock_slack_client,\n mock_inbound_request,\n):\n # Setup mocks\n mock_chat_completion.create.return_value = get_mock_chat_completion_response(\n \"physical_security\"\n )\n\n # Call handler\n handler = InboundRequestHandler(mock_slack_client)\n await handler.maybe_handle(mock_inbound_request)\n\n # Assert that handler calls OpenAI API\n assert_chat_completion_called(mock_chat_completion, mock_config)\n\n mock_slack_client._client.assert_has_calls(\n [\n call.chat_getPermalink(channel=\"C12345\", message_ts=\"t0\"),\n call.chat_postMessage(\n channel=\"C23456\",\n blocks=[\n {\n \"type\": \"section\",\n \"text\": {\n \"type\": \"mrkdwn\",\n \"text\": \"Received an in <#C12345>:\",\n },\n },\n {\n \"type\": \"context\",\n \"elements\": [\n {\n \"type\": \"plain_text\",\n \"text\": \"Predicted category: Physical Security\",\n \"emoji\": True,\n },\n {\n \"type\": \"mrkdwn\",\n \"text\": \"Triaged to: No on-call assigned\",\n },\n {\n \"type\": \"plain_text\",\n \"text\": \"Triage updates in the :thread:\",\n \"emoji\": True,\n },\n ],\n },\n ],\n text=\"New inbound request received\",\n ),\n call.chat_postMessage(\n channel=\"C12345\",\n thread_ts=\"t0\",\n text=\"Hi, thanks for reaching out! Looking for Physical or Office Security? You can reach out to physical-security@company.com.\",\n blocks=[\n {\n \"type\": \"section\",\n \"text\": {\n \"type\": \"mrkdwn\",\n \"text\": \"Hi, thanks for reaching out! Looking for Physical or Office Security? You can reach out to physical-security@company.com.\",\n },\n },\n {\n \"type\": \"context\",\n \"elements\": [\n {\n \"type\": \"plain_text\",\n \"text\": \"If you feel strongly this is a Security issue, respond to this thread and someone from our team will get back to you.\",\n \"emoji\": True,\n }\n ],\n },\n ],\n ),\n call.chat_getPermalink(channel=\"\", message_ts=\"\"),\n call.chat_postMessage(\n channel=\"\",\n thread_ts=\"\",\n text=\" to inbound request.\",\n ),\n ]\n )\n\n\nasync def test_inbound_request_acknowledge_handler(\n mock_slack_client,\n mock_notify_appsec_oncall_message,\n):\n handler = InboundRequestAcknowledgeHandler(mock_slack_client)\n await handler.maybe_handle(mock_notify_appsec_oncall_message)\n mock_slack_client._client.assert_has_calls(\n [\n call.reactions_add(\n blocks=[],\n channel=\"C34567\",\n ts=\"t2\",\n text=\":thumbsup: <@U1234567890> acknowledged the triaged to Application Security.\",\n ),\n call.chat_postMessage(\n blocks=[],\n channel=\"C23456\",\n thread_ts=\"t1\",\n text=\":thumbsup: <@U1234567890> acknowledged the inbound message triaged to Application Security.\",\n ),\n call.conversations_history(channel=\"C23456\", inclusive=True, latest=\"t1\", limit=1),\n call.reactions_add(channel=\"C23456\", name=\"thumbsup\", timestamp=\"t1\"),\n ]\n )\n\n\nasync def test_inbound_request_recategorize_to_listed_category_handler(\n mock_slack_client,\n mock_appsec_oncall_recategorize_to_privacy_message,\n):\n handler = InboundRequestRecategorizeHandler(mock_slack_client)\n await handler.maybe_handle(mock_appsec_oncall_recategorize_to_privacy_message)\n mock_slack_client._client.assert_has_calls(\n [\n call.reactions_add(\n blocks=[],\n channel=\"C34567\",\n ts=\"t3\",\n text=\":thumbsdown: <@U1234567890> reassigned the from Application Security to: Privacy.\",\n ),\n call.reactions_add(channel=\"C23456\", name=\"thumbsdown\", timestamp=\"t1\"),\n call.chat_postMessage(\n blocks=[],\n channel=\"C23456\",\n thread_ts=\"t1\",\n text=\":thumbsdown: <@U1234567890> reassigned the inbound message from Application Security to: Privacy.\",\n ),\n call.chat_postMessage(\n channel=\"C23456\",\n thread_ts=\"t1\",\n blocks=[\n {\n \"type\": \"section\",\n \"text\": {\n \"type\": \"mrkdwn\",\n \"text\": \":wave: Hi <@U12345>, is this assignment accurate?\\n\\n\",\n },\n },\n {\n \"type\": \"context\",\n \"elements\": [\n {\n \"type\": \"plain_text\",\n \"text\": \":thumbsup: Acknowledge this message and response directly to the inbound request.\",\n \"emoji\": True,\n },\n {\n \"type\": \"plain_text\",\n \"text\": \":thumbsdown: Recategorize this message, and if defined, I will route it to the appropriate on-call. If none applies, select Other and pick a channel that I will route the user to.\",\n \"emoji\": True,\n },\n ],\n },\n {\n \"type\": \"actions\",\n \"elements\": [\n {\n \"type\": \"button\",\n \"text\": {\n \"type\": \"plain_text\",\n \"emoji\": True,\n \"text\": \"Acknowledge\",\n },\n \"style\": \"primary\",\n \"value\": \"Privacy\",\n \"action_id\": \"acknowledge_submit_action\",\n },\n {\n \"type\": \"button\",\n \"text\": {\n \"type\": \"plain_text\",\n \"emoji\": True,\n \"text\": \"Inaccurate, recategorize\",\n },\n \"style\": \"danger\",\n \"value\": \"recategorize\",\n \"action_id\": \"recategorize_submit_action\",\n },\n ],\n },\n {\n \"type\": \"section\",\n \"block_id\": \"recategorize_select_category_block\",\n \"text\": {\n \"type\": \"mrkdwn\",\n \"text\": \"*Select a category from the dropdown list, or*\",\n },\n \"accessory\": {\n \"type\": \"static_select\",\n \"placeholder\": {\n \"type\": \"plain_text\",\n \"text\": \"Select an item\",\n \"emoji\": True,\n },\n \"options\": [\n {\n \"text\": {\n \"type\": \"plain_text\",\n \"text\": \"Physical Security\",\n \"emoji\": True,\n },\n \"value\": \"physical_security\",\n },\n {\n \"text\": {\n \"type\": \"plain_text\",\n \"text\": \"Other\",\n \"emoji\": True,\n },\n \"value\": \"other\",\n },\n ],\n \"action_id\": \"recategorize_select_category_action\",\n },\n },\n ],\n metadata={\n \"event_type\": \"notify_oncall\",\n \"event_payload\": {\n \"inbound_message_channel\": \"C12345\",\n \"inbound_message_ts\": \"t0\",\n \"feed_message_channel\": \"C23456\",\n \"feed_message_ts\": \"t1\",\n \"inbound_message_url\": \"https://myorg.slack.com/archives/C12345/p1234567890\",\n \"predicted_category\": \"privacy\",\n },\n },\n text=\"Notify on-call for new inbound request\",\n ),\n ]\n )\n\n\nasync def test_inbound_request_recategorize_to_other_category_handler(\n mock_slack_client,\n mock_appsec_oncall_recategorize_to_other_message,\n):\n handler = InboundRequestRecategorizeHandler(mock_slack_client)\n await handler.maybe_handle(mock_appsec_oncall_recategorize_to_other_message)\n mock_slack_client._client.assert_has_calls(\n [\n call.reactions_add(\n blocks=[],\n channel=\"C34567\",\n ts=\"t3\",\n text=\":thumbsdown: <@U1234567890> reassigned the from Application Security to: Other.\",\n ),\n call.reactions_add(channel=\"C23456\", name=\"thumbsdown\", timestamp=\"t1\"),\n call.chat_postMessage(\n blocks=[],\n channel=\"C23456\",\n thread_ts=\"t1\",\n text=\":thumbsdown: <@U1234567890> reassigned the inbound message from Application Security to: Other.\",\n ),\n call.chat_postMessage(\n channel=\"C12345\",\n thread_ts=\"t0\",\n text=\"Hi, thanks for reaching out! Our team looked at your request, and this is actually something that we don't own. We recommend reaching out to <#C11111> instead.\",\n blocks=[\n {\n \"type\": \"section\",\n \"text\": {\n \"type\": \"mrkdwn\",\n \"text\": \"Hi, thanks for reaching out! Our team looked at your request, and this is actually something that we don't own. We recommend reaching out to <#C11111> instead.\",\n },\n },\n {\n \"type\": \"context\",\n \"elements\": [\n {\n \"type\": \"plain_text\",\n \"text\": \"If you feel strongly this is a Security issue, respond to this thread and someone from our team will get back to you.\",\n \"emoji\": True,\n }\n ],\n },\n ],\n ),\n call.chat_getPermalink(channel=\"\", message_ts=\"\"),\n call.chat_postMessage(\n channel=\"C23456\",\n thread_ts=\"t1\",\n text=\" to inbound request.\",\n ),\n ]\n )\n\n\n@pytest.mark.parametrize(\n \"event_args_override\",\n [\n # Channel is not inbound request channel\n {\"channel\": \"c0\"},\n # No text\n {\"text\": \"\"},\n # Bot message\n {\"subtype\": \"bot_message\"},\n # Thread response, not broadcasted\n {\"thread_ts\": \"t0\"},\n ],\n)\n@patch.object(openai, \"ChatCompletion\")\nasync def test_inbound_request_handler_skip_handle(\n mock_chat_completion, event_args_override, mock_slack_client, mock_inbound_request\n):\n mock_inbound_request.event = {**mock_inbound_request.event, **event_args_override}\n handler = InboundRequestHandler(mock_slack_client)\n\n await handler.maybe_handle(mock_inbound_request)\n mock_chat_completion.create.assert_not_called()\n" }, { "path": "bots/triage-slackbot/triage_slackbot/bot.py", "content": "import asyncio\nimport os\n\nfrom openai_slackbot.bot import start_bot\nfrom triage_slackbot.config import get_config, load_config\nfrom triage_slackbot.handlers import (\n InboundRequestAcknowledgeHandler,\n InboundRequestHandler,\n InboundRequestRecategorizeHandler,\n InboundRequestRecategorizeSelectConversationHandler,\n InboundRequestRecategorizeSelectHandler,\n)\n\nif __name__ == \"__main__\":\n current_dir = os.path.dirname(os.path.abspath(__file__))\n load_config(os.path.join(current_dir, \"config.toml\"))\n\n message_handler = InboundRequestHandler\n action_handlers = [\n InboundRequestAcknowledgeHandler,\n InboundRequestRecategorizeHandler,\n InboundRequestRecategorizeSelectHandler,\n InboundRequestRecategorizeSelectConversationHandler,\n ]\n\n template_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), \"templates\")\n\n config = get_config()\n asyncio.run(\n start_bot(\n openai_organization_id=config.openai_organization_id,\n slack_message_handler=message_handler,\n slack_action_handlers=action_handlers,\n slack_template_path=template_path,\n )\n )\n" }, { "path": "bots/triage-slackbot/triage_slackbot/category.py", "content": "import typing as t\n\nfrom pydantic import BaseModel, ValidationError, model_validator\n\nOTHER_KEY = \"other\"\n\n\nclass RequestCategory(BaseModel):\n # Key used to identify the category in the config.\n key: str\n\n # Display name of the category.\n display_name: str\n\n # Slack ID of the user or channel to route the request to.\n # If user is specified, user will be tagged on the message\n # in the feed channel.\n oncall_slack_id: t.Optional[str] = None\n\n # If true, no manual triage is required for this category\n # and that the bot will autorespond to the inbound request.\n autorespond: bool = False\n\n # Message to send when autoresponding to the inbound request.\n autorespond_message: t.Optional[str] = None\n\n @model_validator(mode=\"after\")\n def check_autorespond(self) -> \"RequestCategory\":\n if self.autorespond and not self.autorespond_message:\n raise ValidationError(\"autorespond_message must be set if autorespond is True\")\n return self\n\n @property\n def route_to_channel(self) -> bool:\n return (self.oncall_slack_id or \"\").startswith(\"C\")\n\n @classmethod\n def to_block_options(cls, categories: t.List[\"RequestCategory\"]) -> t.Dict[str, str]:\n return dict((c.key, c.display_name) for c in categories)\n\n def is_other(self) -> bool:\n return self.key == OTHER_KEY\n" }, { "path": "bots/triage-slackbot/triage_slackbot/config.py", "content": "import os\nimport typing as t\n\nimport toml\nfrom dotenv import load_dotenv\nfrom pydantic import BaseModel, ValidationError, field_validator, model_validator\nfrom pydantic.functional_validators import AfterValidator, BeforeValidator\nfrom triage_slackbot.category import OTHER_KEY, RequestCategory\n\n_CONFIG = None\n\n\ndef convert_categories(v: t.List[t.Dict]):\n categories = {}\n\n for category in v:\n categories[category[\"key\"]] = category\n return categories\n\n\ndef validate_channel(channel_id: str) -> str:\n if not channel_id.startswith(\"C\"):\n raise ValueError(\"channel ID must start with 'C'\")\n return channel_id\n\n\nclass Config(BaseModel):\n # OpenAI organization ID associated with OpenAI API key.\n openai_organization_id: str\n\n # OpenAI prompt to categorize the request.\n openai_prompt: str\n\n # Slack channel where inbound requests are received.\n inbound_request_channel_id: t.Annotated[str, AfterValidator(validate_channel)]\n\n # Slack channel where triage updates are posted.\n feed_channel_id: t.Annotated[str, AfterValidator(validate_channel)]\n\n # Valid categories for inbound requests to be triaged into.\n categories: t.Annotated[t.Dict[str, RequestCategory], BeforeValidator(convert_categories)]\n\n # Enables \"Other\" category, which will allow triager to\n # route the request to a specific conversation.\n other_category_enabled: bool\n\n @model_validator(mode=\"after\")\n def check_category_keys(config: \"Config\") -> \"Config\":\n if config.other_category_enabled:\n if OTHER_KEY in config.categories:\n raise ValidationError(\"other category is reserved and cannot be used\")\n\n category_keys = set(config.categories.keys())\n if len(category_keys) != len(config.categories):\n raise ValidationError(\"category keys must be unique\")\n\n return config\n\n\ndef load_config(path: str):\n load_dotenv()\n\n with open(path) as f:\n cfg = toml.loads(f.read())\n config = Config(**cfg)\n\n if config.other_category_enabled:\n other_category = RequestCategory(\n key=OTHER_KEY,\n display_name=OTHER_KEY.capitalize(),\n oncall_slack_id=None,\n autorespond=True,\n autorespond_message=\"Our team looked at your request, and this is actually something that we don't own. We recommend reaching out to {} instead.\",\n )\n config.categories[other_category.key] = other_category\n\n global _CONFIG\n _CONFIG = config\n return _CONFIG\n\n\ndef get_config() -> Config:\n global _CONFIG\n if _CONFIG is None:\n raise Exception(\"config not initialized, call load_config() first\")\n return _CONFIG\n" }, { "path": "bots/triage-slackbot/triage_slackbot/config.toml", "content": "# Organization ID associated with OpenAI API key.\nopenai_organization_id = \"\"\n\n# Prompt to use for categorizing inbound requests.\nopenai_prompt = \"\"\"\nYou are currently an on-call engineer for a security team at a tech company. \nYour goal is to triage the following incoming Slack message into three categories: \n1. Privacy, return \"privacy\"\n2. Application security, return \"appsec\"\n3. Physical security, return \"physical_security\"\n\"\"\"\ninbound_request_channel_id = \"\"\nfeed_channel_id = \"\"\nother_category_enabled = true\n\n[[ categories ]] \nkey = \"appsec\"\ndisplay_name = \"Application Security\"\noncall_slack_id = \"\"\nautorespond = false\n\n[[ categories ]] \nkey = \"privacy\"\ndisplay_name = \"Privacy\"\noncall_slack_id = \"\"\nautorespond = false\n\n[[ categories ]] \nkey = \"physical_security\"\ndisplay_name = \"Physical Security\"\nautorespond = true\nautorespond_message = \"Looking for Physical or Office Security? You can reach out to physical-security@company.com.\"\n" }, { "path": "bots/triage-slackbot/triage_slackbot/handlers.py", "content": "import typing as t\nfrom enum import Enum\nfrom logging import getLogger\n\nfrom openai_slackbot.clients.slack import CreateSlackMessageResponse, SlackClient\nfrom openai_slackbot.handlers import BaseActionHandler, BaseHandler, BaseMessageHandler\nfrom openai_slackbot.utils.slack import (\n RenderedSlackBlock,\n block_id_exists,\n extract_text_from_event,\n get_block_by_id,\n remove_block_id_if_exists,\n render_slack_id_to_mention,\n render_slack_url,\n)\nfrom triage_slackbot.category import RequestCategory\nfrom triage_slackbot.config import get_config\nfrom triage_slackbot.openai_utils import get_predicted_category\n\nlogger = getLogger(__name__)\n\n\nclass BlockId(str, Enum):\n # Block that will be rendered if on-call recagorizes inbound request but doesn't select a new category.\n empty_category_warning = \"empty_category_warning_block\"\n\n # Block that will be rendreed if on-call recategorizes inbound request to \"Other\" but doesn't select a conversation.\n empty_conversation_warning = \"empty_conversation_warning_block\"\n\n # Block that will be rendered to show on-call all the remaining categories to route to.\n recategorize_select_category = \"recategorize_select_category_block\"\n\n # Block that will be rendered to show on-call all the conversations they can reroute the user\n # to, if they select \"Other\" as the category.\n recategorize_select_conversation = \"recategorize_select_conversation_block\"\n\n\nclass MessageTemplatePath(str, Enum):\n # Template for feed channel message that summarizes triage updates.\n feed = \"messages/feed.j2\"\n\n # Template for message that notifies oncall about inbound request in the same channel as the feed channel.\n notify_oncall_in_feed = \"messages/notify_oncall_in_feed.j2\"\n\n # Template for message that notifies oncall about inbound request in a different channel from the feed channel.\n notify_oncall_channel = \"messages/notify_oncall_channel.j2\"\n\n # Template for message that will autorespond to inbound requests.\n autorespond = \"messages/autorespond.j2\"\n\n\nBlockIdToTemplatePath: t.Dict[BlockId, str] = {\n BlockId.empty_category_warning: \"blocks/empty_category_warning.j2\",\n BlockId.empty_conversation_warning: \"blocks/empty_conversation_warning.j2\",\n BlockId.recategorize_select_conversation: \"blocks/select_conversation.j2\",\n}\n\n\nclass InboundRequestHandlerMixin(BaseHandler):\n def __init__(self, slack_client: SlackClient) -> None:\n super().__init__(slack_client)\n self.config = get_config()\n\n def render_block_if_not_exists(\n self, *, block_id: BlockId, blocks: t.List[RenderedSlackBlock]\n ) -> t.List[RenderedSlackBlock]:\n if not block_id_exists(blocks, block_id):\n template_path = BlockIdToTemplatePath[block_id]\n block = self._slack_client.render_blocks_from_template(template_path)\n blocks.append(block)\n return blocks\n\n def get_selected_category(self, body: t.Dict[str, t.Any]) -> t.Optional[RequestCategory]:\n category = (\n body[\"state\"]\n .get(\"values\", {})\n .get(BlockId.recategorize_select_category, {})\n .get(\"recategorize_select_category_action\", {})\n .get(\"selected_option\", {})\n or {}\n ).get(\"value\")\n\n if not category:\n return None\n\n return self.config.categories[category]\n\n def get_selected_conversation(self, body: t.Dict[str, t.Any]) -> t.Optional[str]:\n return (\n body[\"state\"]\n .get(\"values\", {})\n .get(BlockId.recategorize_select_conversation, {})\n .get(\"recategorize_select_conversation_action\", {})\n .get(\"selected_conversation\")\n )\n\n async def notify_oncall(\n self,\n *,\n predicted_category: RequestCategory,\n selected_conversation: t.Optional[str],\n remaining_categories: t.List[RequestCategory],\n inbound_message_channel: str,\n inbound_message_ts: str,\n feed_message_channel: str,\n feed_message_ts: str,\n inbound_message_url: str,\n ) -> None:\n autoresponded = await self._maybe_autorespond(\n predicted_category,\n selected_conversation,\n inbound_message_channel,\n inbound_message_ts,\n feed_message_channel,\n feed_message_ts,\n )\n\n if autoresponded:\n logger.info(f\"Autoresponded to inbound request: {inbound_message_url}\")\n return\n\n # This metadata will continue to be passed along to the subsequent\n # notify on-call messages.\n metadata = {\n \"event_type\": \"notify_oncall\",\n \"event_payload\": {\n \"inbound_message_channel\": inbound_message_channel,\n \"inbound_message_ts\": inbound_message_ts,\n \"feed_message_channel\": feed_message_channel,\n \"feed_message_ts\": feed_message_ts,\n \"inbound_message_url\": inbound_message_url,\n \"predicted_category\": predicted_category.key,\n },\n }\n\n block_args = {\n \"predicted_category\": predicted_category,\n \"remaining_categories\": remaining_categories,\n \"inbound_message_channel\": inbound_message_channel,\n }\n\n if predicted_category.route_to_channel:\n channel = predicted_category.oncall_slack_id\n thread_ts = None # This will be a new message, not a thread.\n blocks = await self._get_notify_oncall_channel_blocks(\n **block_args,\n inbound_message_url=inbound_message_url,\n )\n else:\n channel = feed_message_channel\n thread_ts = feed_message_ts # Post this as a thread reply to the original feed message.\n blocks = await self._get_notify_oncall_in_feed_blocks(**block_args)\n\n await self._slack_client.post_message(\n channel=channel,\n thread_ts=thread_ts,\n blocks=blocks,\n metadata=metadata,\n text=\"Notify on-call for new inbound request\",\n )\n\n async def _get_notify_oncall_in_feed_blocks(\n self,\n *,\n predicted_category: RequestCategory,\n remaining_categories: t.List[RequestCategory],\n inbound_message_channel: str,\n ):\n oncall_mention = self._get_oncall_mention(predicted_category)\n predicted_category_display_name = predicted_category.display_name\n oncall_greeting = (\n f\":wave: Hi {oncall_mention}\"\n if oncall_mention\n else f\"No on-call defined for {predicted_category_display_name}\"\n )\n\n return self._slack_client.render_blocks_from_template(\n MessageTemplatePath.notify_oncall_in_feed.value,\n {\n \"predicted_category\": predicted_category_display_name,\n \"oncall_greeting\": oncall_greeting,\n \"options\": RequestCategory.to_block_options(remaining_categories),\n \"inbound_message_channel\": inbound_message_channel,\n },\n )\n\n async def _get_notify_oncall_channel_blocks(\n self,\n *,\n predicted_category: RequestCategory,\n remaining_categories: t.List[RequestCategory],\n inbound_message_channel: str,\n inbound_message_url: str,\n ):\n return self._slack_client.render_blocks_from_template(\n MessageTemplatePath.notify_oncall_channel.value,\n {\n \"inbound_message_url\": inbound_message_url,\n \"inbound_message_channel\": inbound_message_channel,\n \"predicted_category\": predicted_category.display_name,\n \"options\": RequestCategory.to_block_options(remaining_categories),\n },\n )\n\n def _get_oncall_mention(self, predicted_category: RequestCategory) -> t.Optional[str]:\n oncall_slack_id = predicted_category.oncall_slack_id\n return render_slack_id_to_mention(oncall_slack_id) if oncall_slack_id else None\n\n async def _maybe_autorespond(\n self,\n predicted_category: RequestCategory,\n selected_conversation: t.Optional[str],\n inbound_message_channel: str,\n inbound_message_ts: str,\n feed_message_channel: str,\n feed_message_ts: str,\n ) -> bool:\n if not predicted_category.autorespond:\n return False\n\n text = \"Hi, thanks for reaching out!\"\n if predicted_category.autorespond_message:\n rendered_selected_conversation = (\n render_slack_id_to_mention(selected_conversation) if selected_conversation else None\n )\n text += (\n f\" {predicted_category.autorespond_message.format(rendered_selected_conversation)}\"\n )\n\n blocks = self._slack_client.render_blocks_from_template(\n MessageTemplatePath.autorespond.value, {\"text\": text}\n )\n message = await self._slack_client.post_message(\n channel=inbound_message_channel,\n thread_ts=inbound_message_ts,\n text=text,\n blocks=blocks,\n )\n message_link = await self._slack_client.get_message_link(\n channel=message.channel, message_ts=message.ts\n )\n\n # Post an update to the feed channel.\n feed_message = (\n f\"{render_slack_url(url=message_link, text='Autoresponded')} to inbound request.\"\n )\n await self._slack_client.post_message(\n channel=feed_message_channel, thread_ts=feed_message_ts, text=feed_message\n )\n\n return True\n\n\nclass InboundRequestHandler(BaseMessageHandler, InboundRequestHandlerMixin):\n \"\"\"\n Handles inbound requests in inbound request channel.\n \"\"\"\n\n async def handle(self, args):\n event = args.event\n\n channel = event.get(\"channel\")\n ts = event.get(\"ts\")\n\n logging_extra = self.logging_extra(args)\n\n text = extract_text_from_event(event)\n if not text:\n logger.info(\"No text in event, done processing\", extra=logging_extra)\n return\n\n predicted_category = await self._predict_category(text)\n logger.info(f\"Predicted category: {predicted_category}\", extra=logging_extra)\n\n message_link = await self._slack_client.get_message_link(channel=channel, message_ts=ts)\n feed_message = await self._update_feed(\n predicted_category=predicted_category,\n message_channel=channel,\n message_link=message_link,\n )\n logger.info(\n f\"Updated feed channel for inbound message link: {message_link}\",\n extra=logging_extra,\n )\n\n remaining_categories = [\n r for r in self.config.categories.values() if r != predicted_category\n ]\n await self.notify_oncall(\n predicted_category=predicted_category,\n selected_conversation=None,\n remaining_categories=remaining_categories,\n inbound_message_channel=channel,\n inbound_message_ts=ts,\n feed_message_channel=feed_message.channel,\n feed_message_ts=feed_message.ts,\n inbound_message_url=message_link,\n )\n logger.info(\"Notified on-call\", extra=logging_extra)\n\n async def should_handle(self, args):\n event = args.event\n\n return (\n event[\"channel\"] == self.config.inbound_request_channel_id\n and\n # Don't respond to messages in threads (with the exception of thread replies\n # that are also sent to the channel)\n (\n (\n event.get(\"thread_ts\") is None\n and (not event.get(\"subtype\") or event.get(\"subtype\") == \"file_share\")\n )\n or event.get(\"subtype\") == \"thread_broadcast\"\n )\n )\n\n async def _predict_category(self, body) -> RequestCategory:\n predicted_category = await get_predicted_category(body)\n return self.config.categories[predicted_category]\n\n async def _update_feed(\n self,\n *,\n predicted_category: RequestCategory,\n message_channel: str,\n message_link: str,\n ) -> CreateSlackMessageResponse:\n oncall_mention = self._get_oncall_mention(predicted_category) or \"No on-call assigned\"\n blocks = self._slack_client.render_blocks_from_template(\n MessageTemplatePath.feed.value,\n {\n \"predicted_category\": predicted_category.display_name,\n \"inbound_message_channel\": message_channel,\n \"inbound_message_url\": message_link,\n \"oncall_mention\": oncall_mention,\n },\n )\n\n message = await self._slack_client.post_message(\n channel=self.config.feed_channel_id,\n blocks=blocks,\n text=\"New inbound request received\",\n )\n return message\n\n\nclass InboundRequestAcknowledgeHandler(BaseActionHandler, InboundRequestHandlerMixin):\n \"\"\"\n Once InboundRequestHandler has predicted the category of an inbound request\n and notifies the corresponding on-call, this handler will be called if on-call\n acknowledges the prediction, i.e. they think the prediction is accurate.\n \"\"\"\n\n @property\n def action_id(self):\n return \"acknowledge_submit_action\"\n\n async def handle(self, args):\n body = args.body\n\n notify_oncall_msg = body[\"container\"]\n notify_oncall_msg_ts = notify_oncall_msg[\"message_ts\"]\n notify_oncall_msg_channel = notify_oncall_msg[\"channel_id\"]\n\n feed_message_metadata = body[\"message\"].get(\"metadata\", {}).get(\"event_payload\", {})\n feed_message_ts = feed_message_metadata[\"feed_message_ts\"]\n feed_message_channel = feed_message_metadata[\"feed_message_channel\"]\n inbound_message_url = feed_message_metadata[\"inbound_message_url\"]\n predicted_category = feed_message_metadata[\"predicted_category\"]\n\n # Oncall that was notified.\n user = body[\"user\"]\n\n await self._slack_client.update_message(\n blocks=[],\n channel=notify_oncall_msg_channel,\n ts=notify_oncall_msg_ts,\n # If oncall is notified in the feed channel, don't need to include\n # the inbound message URL since oncall will be notified in the feed\n # message thread, and the URL is already in the original message.\n text=self._get_message(\n user=user,\n category=predicted_category,\n inbound_message_url=inbound_message_url,\n with_url=notify_oncall_msg_channel != feed_message_channel,\n ),\n )\n\n # If oncall gets notified in a separate channel and not the feed channel,\n # update the feed thread with the acknowledgment.\n if notify_oncall_msg_channel != feed_message_channel:\n await self._slack_client.post_message(\n blocks=[],\n channel=feed_message_channel,\n thread_ts=feed_message_ts,\n text=self._get_message(\n user=user,\n category=predicted_category,\n inbound_message_url=inbound_message_url,\n with_url=False,\n ),\n )\n\n feed_message = await self._slack_client.get_message(\n channel=feed_message_channel, ts=feed_message_ts\n )\n if feed_message:\n # If the original message has been thumbs-downed, this means\n # that the bot's original prediction is wrong, so don't thumbs\n # up the feed message.\n wrong_original_prediction = any(\n [r[\"name\"] == \"-1\" for r in feed_message.get(\"reactions\", [])]\n )\n\n if not wrong_original_prediction:\n await self._slack_client.add_reaction(\n channel=feed_message_channel,\n name=\"thumbsup\",\n timestamp=feed_message_ts,\n )\n\n def _get_message(\n self, user: t.Dict, category: str, inbound_message_url: str, with_url: bool\n ) -> str:\n message = f\":thumbsup: {render_slack_id_to_mention(user['id'])} acknowledged the \"\n if with_url:\n message += render_slack_url(url=inbound_message_url, text=\"inbound message\")\n else:\n message += \"inbound message\"\n\n return f\"{message} triaged to {self.config.categories[category].display_name}.\"\n\n\nclass InboundRequestRecategorizeHandler(BaseActionHandler, InboundRequestHandlerMixin):\n \"\"\"\n This handler will be called if on-call wants to recategorize the request\n that they get notified about.\n \"\"\"\n\n @property\n def action_id(self):\n return \"recategorize_submit_action\"\n\n async def handle(self, args):\n body = args.body\n\n notify_oncall_msg = body[\"container\"]\n notify_oncall_msg_ts = notify_oncall_msg[\"message_ts\"]\n notify_oncall_msg_channel = notify_oncall_msg[\"channel_id\"]\n\n msg_metadata = body[\"message\"].get(\"metadata\", {}).get(\"event_payload\", {})\n feed_message_ts = msg_metadata[\"feed_message_ts\"]\n feed_message_channel = msg_metadata[\"feed_message_channel\"]\n inbound_message_url = msg_metadata[\"inbound_message_url\"]\n\n # Predicted category that turned out to be incorrect\n # and wanted to be recategorized.\n predicted_category = self.config.categories[msg_metadata.pop(\"predicted_category\")]\n assert predicted_category\n\n user: t.Dict = body[\"user\"]\n\n notify_oncall_msg_blocks = body[\"message\"][\"blocks\"]\n selection_block = get_block_by_id(\n notify_oncall_msg_blocks, BlockId.recategorize_select_category\n )\n remaining_category_keys: t.List[str] = [\n o[\"value\"] for o in selection_block[\"accessory\"][\"options\"]\n ]\n\n selected_category: t.Optional[RequestCategory] = self.get_selected_category(body)\n selected_conversation: t.Optional[str] = self.get_selected_conversation(body)\n valid, notify_oncall_msg_blocks = await self._validate_selection(\n selected_category, selected_conversation, notify_oncall_msg_blocks\n )\n if valid:\n assert selected_category, \"selected_category should be set if valid\"\n message_kwargs = {\n \"user\": user,\n \"predicted_category\": predicted_category,\n \"selected_category\": selected_category,\n \"selected_conversation\": selected_conversation,\n \"inbound_message_url\": inbound_message_url,\n }\n\n await self._slack_client.update_message(\n blocks=[],\n channel=notify_oncall_msg_channel,\n ts=notify_oncall_msg_ts,\n # If the feed message is in the same channel as the notify on-call message, don't need to include\n # the URL since it's already in the original feed message.\n text=self._get_message(\n **message_kwargs,\n with_url=notify_oncall_msg_channel != feed_message_channel,\n ),\n )\n\n # Indicate that the previous predicted category is not accurate.\n await self._slack_client.add_reaction(\n channel=feed_message_channel,\n name=\"thumbsdown\",\n timestamp=feed_message_ts,\n )\n\n # If the feed message is in a different channel than the notify on-call message,\n # post recategorization update to the feed channel.\n if notify_oncall_msg_channel != feed_message_channel:\n await self._slack_client.post_message(\n blocks=[],\n channel=feed_message_channel,\n thread_ts=feed_message_ts,\n text=self._get_message(**message_kwargs, with_url=False),\n )\n\n remaining_categories = [\n self.config.categories[category_key]\n for category_key in remaining_category_keys\n if category_key != selected_category.key\n ]\n\n # Route this to the next oncall.\n await self.notify_oncall(\n predicted_category=selected_category,\n selected_conversation=selected_conversation,\n remaining_categories=remaining_categories,\n **msg_metadata,\n )\n else:\n # Display warning.\n await self._slack_client.update_message(\n blocks=notify_oncall_msg_blocks,\n channel=notify_oncall_msg_channel,\n ts=notify_oncall_msg_ts,\n text=\"\",\n )\n\n def _get_message(\n self,\n *,\n user: t.Dict,\n predicted_category: RequestCategory,\n selected_category: RequestCategory,\n selected_conversation: t.Optional[str],\n inbound_message_url: str,\n with_url: bool,\n ) -> str:\n rendered_selected_conversation = (\n render_slack_id_to_mention(selected_conversation) if selected_conversation else None\n )\n selected_category_display_name = selected_category.display_name.format(\n rendered_selected_conversation\n )\n\n message_text = f\"<{inbound_message_url}|inbound message>\" if with_url else \"inbound message\"\n return f\":thumbsdown: {render_slack_id_to_mention(user['id'])} reassigned the {message_text} from {predicted_category.display_name} to: {selected_category_display_name}.\"\n\n async def _validate_selection(\n self,\n selected_category: t.Optional[RequestCategory],\n selected_conversation: t.Optional[str],\n blocks: t.List[RenderedSlackBlock],\n ) -> t.Tuple[bool, t.List[RenderedSlackBlock]]:\n if not selected_category:\n return False, self.render_block_if_not_exists(\n block_id=BlockId.empty_category_warning, blocks=blocks\n )\n elif selected_category.is_other() and not selected_conversation:\n return False, self.render_block_if_not_exists(\n block_id=BlockId.empty_conversation_warning, blocks=blocks\n )\n\n return True, blocks\n\n\nclass InboundRequestRecategorizeSelectHandler(BaseActionHandler, InboundRequestHandlerMixin):\n \"\"\"\n This handler will be called if on-call selects a new category for a request they\n get notififed about.\n \"\"\"\n\n @property\n def action_id(self):\n return \"recategorize_select_category_action\"\n\n async def handle(self, args):\n body = args.body\n\n notify_oncall_msg = body[\"container\"]\n notify_oncall_msg_ts = notify_oncall_msg[\"message_ts\"]\n notify_oncall_msg_channel = notify_oncall_msg[\"channel_id\"]\n\n notify_oncall_msg_blocks = body[\"message\"][\"blocks\"]\n notify_oncall_msg_blocks = remove_block_id_if_exists(\n notify_oncall_msg_blocks, BlockId.empty_category_warning\n )\n\n selected_category = self.get_selected_category(body)\n if selected_category.is_other():\n # Prompt on-call to select a conversation if Other category is selected.\n notify_oncall_msg_blocks = self.render_block_if_not_exists(\n block_id=BlockId.recategorize_select_conversation,\n blocks=notify_oncall_msg_blocks,\n )\n else:\n # Remove warning if on-call updates their selection from Other to non-Other.\n notify_oncall_msg_blocks = remove_block_id_if_exists(\n notify_oncall_msg_blocks, BlockId.recategorize_select_conversation\n )\n\n # Update message with warnings, if any.\n await self._slack_client.update_message(\n blocks=notify_oncall_msg_blocks,\n channel=notify_oncall_msg_channel,\n ts=notify_oncall_msg_ts,\n )\n\n\nclass InboundRequestRecategorizeSelectConversationHandler(BaseActionHandler):\n \"\"\"\n This handler will be called if on-call selects a conversation to route the request to.\n \"\"\"\n\n @property\n def action_id(self):\n return \"recategorize_select_conversation_action\"\n\n async def handle(self, args):\n pass\n" }, { "path": "bots/triage-slackbot/triage_slackbot/openai_utils.py", "content": "import json\nfrom functools import cache\n\nimport openai\nfrom triage_slackbot.category import OTHER_KEY, RequestCategory\nfrom triage_slackbot.config import get_config\n\n\n@cache\ndef predict_category_functions(categories: list[RequestCategory]) -> list[dict]:\n return [\n {\n \"name\": \"get_predicted_category\",\n \"description\": \"Predicts the category of an inbound request.\",\n \"parameters\": {\n \"type\": \"object\",\n \"properties\": {\n \"category\": {\n \"type\": \"string\",\n \"enum\": [\n category.key for category in categories if category.key != OTHER_KEY\n ],\n \"description\": \"Predicted category of the inbound request\",\n },\n },\n \"required\": [\"category\"],\n },\n }\n ]\n\n\nasync def get_predicted_category(inbound_request_content: str) -> str:\n \"\"\"\n This function uses the OpenAI Chat Completion API to predict the category of an inbound request.\n \"\"\"\n config = get_config()\n\n # Define the prompt\n messages = [\n {\"role\": \"system\", \"content\": config.openai_prompt},\n {\"role\": \"user\", \"content\": inbound_request_content},\n ]\n\n # Call the API\n response = openai.chat.completions.create(\n model=\"gpt-4-32k\",\n messages=messages,\n temperature=0,\n stream=False,\n functions=predict_category_functions(config.categories.values()),\n function_call={\"name\": \"get_predicted_category\"},\n )\n\n function_args = json.loads(response.choices[0].message.function_call.arguments) # type: ignore\n return function_args[\"category\"]\n" }, { "path": "bots/triage-slackbot/triage_slackbot/templates/blocks/empty_category_warning.j2", "content": "{\n \"type\": \"context\",\n \"block_id\": \"empty_category_warning_block\",\n \"elements\": [{\"type\": \"plain_text\", \"text\": \"Category is required.\"}]\n}" }, { "path": "bots/triage-slackbot/triage_slackbot/templates/blocks/empty_conversation_warning.j2", "content": "{\n \"type\": \"context\",\n \"block_id\": \"empty_conversation_warning_block\",\n \"elements\": [{\"type\": \"plain_text\", \"text\": \"Conversation is required.\"}]\n}\n" }, { "path": "bots/triage-slackbot/triage_slackbot/templates/blocks/select_conversation.j2", "content": "{\n \"type\": \"section\",\n \"text\": {\"type\": \"mrkdwn\", \"text\": \"*Select a channel*\"},\n \"accessory\": {\n \"type\": \"conversations_select\",\n \"placeholder\": {\n \"type\": \"plain_text\",\n \"text\": \"Select conversations\",\n \"emoji\": true\n },\n \"action_id\": \"recategorize_select_conversation_action\"\n },\n \"block_id\": \"recategorize_select_conversation_block\"\n}" }, { "path": "bots/triage-slackbot/triage_slackbot/templates/messages/_notify_oncall_body.j2", "content": "{\n\t\"type\": \"context\",\n\t\"elements\": [\n\t\t{\n\t\t\t\"type\": \"plain_text\",\n\t\t\t\"text\": \":thumbsup: Acknowledge this message and response directly to the inbound request.\",\n\t\t\t\"emoji\": true\n\t\t},\n\t\t{\n\t\t\t\"type\": \"plain_text\",\n\t\t\t\"text\": \":thumbsdown: Recategorize this message, and if defined, I will route it to the appropriate on-call. If none applies, select Other and pick a channel that I will route the user to.\",\n\t\t\t\"emoji\": true\n\t\t}\n\t]\n},\n{\n\t\"type\": \"actions\",\n\t\"elements\": [\n\t\t{\n\t\t\t\"type\": \"button\",\n\t\t\t\"text\": {\n\t\t\t\t\"type\": \"plain_text\",\n\t\t\t\t\"emoji\": true,\n\t\t\t\t\"text\": \"Acknowledge\"\n\t\t\t},\n\t\t\t\"style\": \"primary\",\n\t\t\t\"value\": \"{{ predicted_category }}\",\n\t\t\t\"action_id\": \"acknowledge_submit_action\"\n\t\t},\n\t\t{\n\t\t\t\"type\": \"button\",\n\t\t\t\"text\": {\n\t\t\t\t\"type\": \"plain_text\",\n\t\t\t\t\"emoji\": true,\n\t\t\t\t\"text\": \"Inaccurate, recategorize\"\n\t\t\t},\n\t\t\t\"style\": \"danger\",\n\t\t\t\"value\": \"recategorize\",\n\t\t\t\"action_id\": \"recategorize_submit_action\"\n\t\t}\n\t]\n},\n{\n\t\"type\": \"section\",\n\t\"block_id\": \"recategorize_select_category_block\",\n\t\"text\": {\n\t\t\"type\": \"mrkdwn\",\n\t\t\"text\": \"*Select a category from the dropdown list, or*\"\n\t},\n\t\"accessory\": {\n\t\t\"type\": \"static_select\",\n\t\t\"placeholder\": {\n\t\t\t\"type\": \"plain_text\",\n\t\t\t\"text\": \"Select an item\",\n\t\t\t\"emoji\": true\n\t\t},\n\t\t\"options\": [\n\t\t\t{% for value, text in options.items() %}\n\t\t\t{\n\t\t\t\t\"text\": {\n\t\t\t\t\t\"type\": \"plain_text\",\n\t\t\t\t\t\"text\": \"{{ text }}\",\n\t\t\t\t\t\"emoji\": true\n\t\t\t\t},\n\t\t\t\t\"value\": \"{{ value }}\"\n\t\t\t}{% if not loop.last %},{% endif %}\n\t\t\t{% endfor %}\n\t\t],\n\t\t\"action_id\": \"recategorize_select_category_action\"\n\t}\n}\n" }, { "path": "bots/triage-slackbot/triage_slackbot/templates/messages/autorespond.j2", "content": "[\n {\n \"type\": \"section\",\n \"text\": {\n \"type\": \"mrkdwn\",\n \"text\": \"{{ text }}\"\n }\n },\n {\n \"type\": \"context\",\n \"elements\": [\n {\n \"type\": \"plain_text\",\n \"text\": \"If you feel strongly this is a Security issue, respond to this thread and someone from our team will get back to you.\",\n \"emoji\": true\n }\n ]\n }\n]\n" }, { "path": "bots/triage-slackbot/triage_slackbot/templates/messages/feed.j2", "content": "[\n\t{\n\t\t\"type\": \"section\",\n\t\t\"text\": {\n\t\t\t\"type\": \"mrkdwn\",\n\t\t\t\"text\": \"Received an <{{ inbound_message_url }}|inbound message> in <#{{ inbound_message_channel }}>:\"\n\t\t}\n\t},\n\t{\n\t\t\"type\": \"context\",\n\t\t\"elements\": [\n\t\t\t{\n\t\t\t\t\"type\": \"plain_text\",\n\t\t\t\t\"text\": \"Predicted category: {{ predicted_category }}\",\n\t\t\t\t\"emoji\": true\n\t\t\t},\n\t\t\t{\n\t\t\t\t\"type\": \"mrkdwn\",\n\t\t\t\t\"text\": \"Triaged to: {{ oncall_mention }}\"\n\t\t\t},\n\t\t\t{\n\t\t\t\t\"type\": \"plain_text\",\n\t\t\t\t\"text\": \"Triage updates in the :thread:\",\n\t\t\t\t\"emoji\": true\n\t\t\t}\n\t\t]\n\t}\n]\n" }, { "path": "bots/triage-slackbot/triage_slackbot/templates/messages/notify_oncall_channel.j2", "content": "[\n\t{\n\t\t\"type\": \"section\",\n\t\t\"text\": {\n\t\t\t\"type\": \"mrkdwn\",\n\t\t\t\"text\": \":wave: Hi, we received an <{{ inbound_message_url }}|inbound message> in <#{{ inbound_message_channel }}>, which was categorized as {{ predicted_category }}. Is this accurate?\\n\\n\"\n\t\t}\n\t},\n\t{% include 'messages/_notify_oncall_body.j2' %}\n]\n" }, { "path": "bots/triage-slackbot/triage_slackbot/templates/messages/notify_oncall_in_feed.j2", "content": "[\n\t{\n\t\t\"type\": \"section\",\n\t\t\"text\": {\n\t\t\t\"type\": \"mrkdwn\",\n\t\t\t\"text\": \"{{ oncall_greeting }}, is this assignment accurate?\\n\\n\"\n\t\t}\n\t},\n\t{% include 'messages/_notify_oncall_body.j2' %}\n]\n" }, { "path": "shared/openai-slackbot/openai_slackbot/__init__.py", "content": "" }, { "path": "shared/openai-slackbot/openai_slackbot/bot.py", "content": "import typing as t\nfrom logging import getLogger\n\nimport openai\nfrom openai_slackbot.clients.slack import SlackClient\nfrom openai_slackbot.handlers import BaseActionHandler, BaseMessageHandler\nfrom openai_slackbot.utils.envvars import string\nfrom slack_bolt.adapter.socket_mode.async_handler import AsyncSocketModeHandler\nfrom slack_bolt.app.async_app import AsyncApp\n\nlogger = getLogger(__name__)\n\n\nasync def register_app_handlers(\n *,\n app: AsyncApp,\n message_handler: t.Type[BaseMessageHandler],\n action_handlers: t.List[t.Type[BaseActionHandler]],\n slack_client: SlackClient,\n):\n if message_handler:\n app.event(\"message\")(message_handler(slack_client).maybe_handle)\n\n if action_handlers:\n for action_handler in action_handlers:\n handler = action_handler(slack_client)\n app.action(handler.action_id)(handler.maybe_handle)\n\n\nasync def init_bot(\n *,\n openai_organization_id: str,\n slack_message_handler: t.Type[BaseMessageHandler],\n slack_action_handlers: t.List[t.Type[BaseActionHandler]],\n slack_template_path: str,\n):\n slack_bot_token = string(\"SLACK_BOT_TOKEN\")\n openai_api_key = string(\"OPENAI_API_KEY\")\n\n # Init OpenAI API\n openai.organization = openai_organization_id\n openai.api_key = openai_api_key\n\n # Init slack bot\n app = AsyncApp(token=slack_bot_token)\n slack_client = SlackClient(app.client, slack_template_path)\n await register_app_handlers(\n app=app,\n message_handler=slack_message_handler,\n action_handlers=slack_action_handlers,\n slack_client=slack_client,\n )\n\n return app\n\n\nasync def start_app(app):\n socket_app_token = string(\"SOCKET_APP_TOKEN\")\n handler = AsyncSocketModeHandler(app, socket_app_token)\n await handler.start_async()\n\n\nasync def start_bot(\n *,\n openai_organization_id: str,\n slack_message_handler: t.Type[BaseMessageHandler],\n slack_action_handlers: t.List[t.Type[BaseActionHandler]],\n slack_template_path: str,\n):\n app = await init_bot(\n openai_organization_id=openai_organization_id,\n slack_message_handler=slack_message_handler,\n slack_action_handlers=slack_action_handlers,\n slack_template_path=slack_template_path,\n )\n\n await start_app(app)\n" }, { "path": "shared/openai-slackbot/openai_slackbot/clients/__init__.py", "content": "" }, { "path": "shared/openai-slackbot/openai_slackbot/clients/slack.py", "content": "import json\nimport os\nimport typing as t\nfrom logging import getLogger\n\nfrom jinja2 import Environment, FileSystemLoader\nfrom pydantic import BaseModel\nfrom slack_sdk.errors import SlackApiError\nfrom slack_sdk.web.async_client import AsyncWebClient\n\nlogger = getLogger(__name__)\n\n\nclass SlackMessage(BaseModel):\n app_id: t.Optional[str] = None\n blocks: t.Optional[t.List[t.Any]] = None\n bot_id: t.Optional[str] = None\n bot_profile: t.Optional[t.Dict[str, t.Any]] = None\n team: str\n text: str\n ts: str\n type: str\n user: t.Optional[str] = None\n\n\nclass CreateSlackMessageResponse(BaseModel):\n ok: bool\n channel: str\n ts: str\n message: SlackMessage\n\n\nclass SlackClient:\n \"\"\"\n SlackClient wraps the Slack AsyncWebClient implementation and\n provides some additional functionality specific to the Slackbot\n implementation.\n \"\"\"\n\n def __init__(self, client: AsyncWebClient, template_path: str) -> None:\n self._client = client\n self._jinja = self._init_jinja(template_path)\n\n async def get_message_link(self, **kwargs) -> str:\n response = await self._client.chat_getPermalink(**kwargs)\n if not response[\"ok\"]:\n raise Exception(f\"Failed to get Slack message link: {response['error']}\")\n return response[\"permalink\"]\n\n async def get_message(self, channel: str, ts: str) -> t.Optional[t.Dict[str, t.Any]]:\n \"\"\"Follows: https://api.slack.com/messaging/retrieving.\"\"\"\n result = await self._client.conversations_history(\n channel=channel,\n inclusive=True,\n latest=ts,\n limit=1,\n )\n return result[\"messages\"][0] if result[\"messages\"] else None\n\n async def post_message(self, **kwargs) -> CreateSlackMessageResponse:\n response = await self._client.chat_postMessage(**kwargs)\n if not response[\"ok\"]:\n raise Exception(f\"Failed to post Slack message: {response['error']}\")\n\n assert isinstance(response.data, dict)\n return CreateSlackMessageResponse(**response.data)\n\n async def update_message(self, **kwargs) -> t.Dict[str, t.Any]:\n response = await self._client.chat_update(**kwargs)\n if not response[\"ok\"]:\n raise Exception(f\"Failed to update Slack message: {response['error']}\")\n\n assert isinstance(response.data, dict)\n return response.data\n\n async def add_reaction(self, **kwargs) -> t.Dict[str, t.Any]:\n try:\n response = await self._client.reactions_add(**kwargs)\n except SlackApiError as e:\n if e.response[\"error\"] == \"already_reacted\":\n return {}\n raise e\n\n assert isinstance(response.data, dict)\n return response.data\n\n async def get_thread_messages(self, channel: str, thread_ts: str) -> t.List[t.Dict[str, t.Any]]:\n response = await self._client.conversations_replies(channel=channel, ts=thread_ts)\n if not response[\"ok\"]:\n raise Exception(f\"Failed to get thread messages: {response['error']}\")\n\n assert isinstance(response.data, dict)\n return response.data[\"messages\"]\n\n async def get_user_display_name(self, user_id: str) -> str:\n response = await self._client.users_info(user=user_id)\n if not response[\"ok\"]:\n raise Exception(f\"Failed to get user info: {response['error']}\")\n return response[\"user\"][\"profile\"][\"display_name\"]\n\n async def get_original_blocks(self, thread_ts: str, channel: str) -> None:\n \"\"\"Given a thread_ts, get original message block\"\"\"\n response = await self._client.conversations_replies(\n channel=channel,\n ts=thread_ts,\n )\n try:\n messages = response.get(\"messages\", [])\n if not messages:\n raise ValueError(f\"Error fetching original message for thread_ts {thread_ts}\")\n blocks = messages[0].get(\"blocks\")\n if not blocks:\n raise ValueError(f\"Error fetching original message for thread_ts {thread_ts}\")\n return blocks\n except Exception as e:\n logger.exception(f\"Error fetching original message for thread_ts {thread_ts}: {e}\")\n\n def render_blocks_from_template(self, template_filename: str, context: t.Dict = {}) -> t.Any:\n rendered_template = self._jinja.get_template(template_filename).render(context)\n return json.loads(rendered_template)\n\n def _init_jinja(self, template_path: str):\n templates_dir = os.path.join(template_path)\n return Environment(loader=FileSystemLoader(templates_dir))\n" }, { "path": "shared/openai-slackbot/openai_slackbot/handlers.py", "content": "import abc\nimport typing as t\nfrom logging import getLogger\n\nfrom openai_slackbot.clients.slack import SlackClient\n\nlogger = getLogger(__name__)\n\n\nclass BaseHandler(abc.ABC):\n def __init__(self, slack_client: SlackClient) -> None:\n self._slack_client = slack_client\n\n async def maybe_handle(self, args):\n await args.ack()\n\n logging_extra = self.logging_extra(args)\n try:\n should_handle = await self.should_handle(args)\n logger.info(\n f\"Handler: {self.__class__.__name__}, should handle: {should_handle}\",\n extra=logging_extra,\n )\n if should_handle:\n await self.handle(args)\n except Exception:\n logger.exception(\"Failed to handle event\", extra=logging_extra)\n\n @abc.abstractmethod\n async def should_handle(self, args) -> bool:\n ...\n\n @abc.abstractmethod\n async def handle(self, args):\n ...\n\n @abc.abstractmethod\n def logging_extra(self, args) -> t.Dict[str, t.Any]:\n ...\n\n\nclass BaseMessageHandler(BaseHandler):\n def logging_extra(self, args) -> t.Dict[str, t.Any]:\n fields = {}\n for field in [\"type\", \"subtype\", \"channel\", \"ts\"]:\n fields[field] = args.event.get(field)\n return fields\n\n\nclass BaseActionHandler(BaseHandler):\n @abc.abstractproperty\n def action_id(self) -> str:\n ...\n\n async def should_handle(self, args) -> bool:\n return True\n\n def logging_extra(self, args) -> t.Dict[str, t.Any]:\n return {\n \"action_type\": args.body.get(\"type\"),\n \"action\": args.body.get(\"actions\", [])[0],\n }\n" }, { "path": "shared/openai-slackbot/openai_slackbot/utils/__init__.py", "content": "" }, { "path": "shared/openai-slackbot/openai_slackbot/utils/envvars.py", "content": "import os\nimport typing as t\n\n\ndef string(key: str, default: t.Optional[str] = None) -> str:\n val = os.environ.get(key)\n if not val:\n if default is None:\n raise ValueError(f\"Missing required environment variable: {key}\")\n return default\n return val\n" }, { "path": "shared/openai-slackbot/openai_slackbot/utils/slack.py", "content": "import typing as t\n\nRenderedSlackBlock = t.NewType(\"RenderedSlackBlock\", t.Dict[str, t.Any])\n\n\ndef block_id_exists(blocks: t.List[RenderedSlackBlock], block_id: str) -> bool:\n return any([block.get(\"block_id\") == block_id for block in blocks])\n\n\ndef remove_block_id_if_exists(blocks: t.List[RenderedSlackBlock], block_id: str) -> t.List:\n return [block for block in blocks if block.get(\"block_id\") != block_id]\n\n\ndef get_block_by_id(blocks: t.Dict, block_id: str) -> t.Dict:\n for block in blocks:\n if block.get(\"block_id\") == block_id:\n return block\n return {}\n\n\ndef extract_text_from_event(event) -> str:\n \"\"\"Extracts text from either plaintext and block message.\"\"\"\n\n # Extract text from plaintext message.\n text = event.get(\"text\")\n if text:\n return text\n\n # Extract text from message blocks.\n texts = []\n attachments = event.get(\"attachments\", [])\n for attachment in attachments:\n attachment_message_blocks = attachment.get(\"message_blocks\", [])\n for amb in attachment_message_blocks:\n message_blocks = amb.get(\"message\", {}).get(\"blocks\", [])\n for mb in message_blocks:\n mb_elements = mb.get(\"elements\", [])\n for mbe in mb_elements:\n mbe_elements = mbe.get(\"elements\", [])\n for mbee in mbe_elements:\n if mbee.get(\"type\") == \"text\":\n texts.append(mbee[\"text\"])\n\n return \" \".join(texts).strip()\n\n\ndef render_slack_id_to_mention(id: str):\n \"\"\"Render a usergroup or user ID to a mention.\"\"\"\n\n if not id:\n return \"\"\n elif id.startswith(\"U\"):\n return f\"<@{id}>\"\n elif id.startswith(\"S\"):\n return f\"\"\n elif id.startswith(\"C\"):\n return f\"<#{id}>\"\n else:\n raise ValueError(f\"Unsupported/invalid ID type: {id}\")\n\n\ndef render_slack_url(*, url: str, text: str) -> str:\n \"\"\"Render a URL to a clickable link.\"\"\"\n return f\"<{url}|{text}>\"\n" }, { "path": "shared/openai-slackbot/pyproject.toml", "content": "[project]\nname = \"openai-slackbot\"\nrequires-python = \">=3.8\"\nversion = \"1.0.0\"\ndependencies = [\n \"aiohttp\",\n \"Jinja2\",\n \"openai\",\n \"pydantic\",\n \"python-dotenv\",\n \"slack-bolt\",\n \"slack-sdk\",\n \"pytest\",\n \"pytest-env\",\n \"pytest-asyncio\",\n \"aiohttp\",\n]\n\n[build-system]\nrequires = [\"setuptools>=64.0\"]\nbuild-backend = \"setuptools.build_meta\"\n\n[tool.pytest.ini_options]\nasyncio_mode = \"auto\"\nenv = [\n \"SLACK_BOT_TOKEN=mock-token\",\n \"SOCKET_APP_TOKEN=mock-token\",\n \"OPENAI_API_KEY=mock-key\",\n]\n" }, { "path": "shared/openai-slackbot/setup.cfg", "content": "" }, { "path": "shared/openai-slackbot/tests/__init__.py", "content": "" }, { "path": "shared/openai-slackbot/tests/clients/__init__.py", "content": "" }, { "path": "shared/openai-slackbot/tests/clients/test_slack.py", "content": "from unittest.mock import AsyncMock, MagicMock\n\nimport pytest\nfrom openai_slackbot.clients.slack import CreateSlackMessageResponse\nfrom slack_sdk.errors import SlackApiError\n\n\nasync def test_get_message_link_success(mock_slack_client):\n mock_slack_client._client.chat_getPermalink = AsyncMock(\n return_value={\n \"ok\": True,\n \"channel\": \"C123456\",\n \"permalink\": \"https://myorg.slack.com/archives/C123456/p1234567890\",\n }\n )\n link = await mock_slack_client.get_message_link(channel=\"channel\", message_ts=\"message_ts\")\n mock_slack_client._client.chat_getPermalink.assert_called_once_with(\n channel=\"channel\", message_ts=\"message_ts\"\n )\n assert link == \"https://myorg.slack.com/archives/C123456/p1234567890\"\n\n\nasync def test_get_message_link_failed(mock_slack_client):\n mock_slack_client._client.chat_getPermalink = AsyncMock(\n return_value={\"ok\": False, \"error\": \"failed\"}\n )\n with pytest.raises(Exception):\n await mock_slack_client.get_message_link(channel=\"channel\", message_ts=\"message_ts\")\n mock_slack_client._client.chat_getPermalink.assert_called_once_with(\n channel=\"channel\", message_ts=\"message_ts\"\n )\n\n\nasync def test_post_message_success(mock_slack_client):\n mock_message_data = {\n \"ok\": True,\n \"channel\": \"C234567\",\n \"ts\": \"ts\",\n \"message\": {\n \"bot_id\": \"bot_id\",\n \"bot_profile\": {\"id\": \"bot_profile_id\"},\n \"team\": \"team\",\n \"text\": \"text\",\n \"ts\": \"ts\",\n \"type\": \"type\",\n \"user\": \"user\",\n },\n }\n mock_response = MagicMock(data=mock_message_data)\n mock_response.__getitem__.side_effect = mock_message_data.__getitem__\n mock_slack_client._client.chat_postMessage = AsyncMock(return_value=mock_response)\n\n response = await mock_slack_client.post_message(channel=\"C234567\", text=\"text\")\n assert response == CreateSlackMessageResponse(**mock_message_data)\n\n\nasync def test_post_message_failed(mock_slack_client):\n mock_slack_client._client.chat_postMessage = AsyncMock(\n return_value={\"ok\": False, \"error\": \"failed\"}\n )\n with pytest.raises(Exception):\n await mock_slack_client.post_message(channel=\"channel\", text=\"text\")\n mock_slack_client._client.chat_postMessage.assert_called_once_with(\n channel=\"channel\", text=\"text\"\n )\n\n\nasync def test_update_message_success(mock_slack_client):\n mock_message_data = {\n \"ok\": True,\n \"channel\": \"C234567\",\n \"ts\": \"ts\",\n \"message\": {\n \"bot_id\": \"bot_id\",\n \"bot_profile\": {\"id\": \"bot_profile_id\"},\n \"team\": \"team\",\n \"text\": \"text\",\n \"ts\": \"ts\",\n \"type\": \"type\",\n \"user\": \"user\",\n },\n }\n mock_response = MagicMock(data=mock_message_data)\n mock_response.__getitem__.side_effect = mock_message_data.__getitem__\n mock_slack_client._client.chat_update = AsyncMock(return_value=mock_response)\n\n response = await mock_slack_client.update_message(channel=\"C234567\", ts=\"ts\", text=\"text\")\n assert response == mock_message_data\n\n\nasync def test_update_message_failed(mock_slack_client):\n mock_slack_client._client.chat_update = AsyncMock(return_value={\"ok\": False, \"error\": \"failed\"})\n with pytest.raises(Exception):\n await mock_slack_client.update_message(channel=\"channel\", ts=\"ts\", text=\"text\")\n mock_slack_client._client.chat_update.assert_called_once_with(\n channel=\"channel\", ts=\"ts\", text=\"text\"\n )\n\n\nasync def test_add_reaction_success(mock_slack_client):\n mock_response_data = {\"ok\": True}\n mock_response = MagicMock(data=mock_response_data)\n mock_response.__getitem__.side_effect = mock_response_data.__getitem__\n mock_slack_client._client.reactions_add = AsyncMock(return_value=mock_response)\n await mock_slack_client.add_reaction(channel=\"channel\", name=\"thumbsup\", timestamp=\"timestamp\")\n\n\nasync def test_add_reaction_already_reacted(mock_slack_client):\n mock_slack_client._client.reactions_add = AsyncMock(\n side_effect=SlackApiError(\"already_reacted\", {\"error\": \"already_reacted\"})\n )\n response = await mock_slack_client.add_reaction(\n channel=\"channel\", name=\"thumbsup\", timestamp=\"timestamp\"\n )\n assert response == {}\n\n\nasync def test_add_reaction_failed(mock_slack_client):\n mock_slack_client._client.reactions_add = AsyncMock(\n side_effect=SlackApiError(\"failed\", {\"error\": \"invalid_reaction\"})\n )\n with pytest.raises(Exception):\n await mock_slack_client.add_reaction(\n channel=\"channel\", name=\"thumbsup\", timestamp=\"timestamp\"\n )\n" }, { "path": "shared/openai-slackbot/tests/conftest.py", "content": "from unittest.mock import AsyncMock, MagicMock, patch\n\nimport pytest\nfrom openai_slackbot.clients.slack import SlackClient\nfrom openai_slackbot.handlers import BaseActionHandler, BaseMessageHandler\n\n\n@pytest.fixture\ndef mock_slack_app():\n with patch(\"slack_bolt.app.async_app.AsyncApp\") as mock_app:\n yield mock_app.return_value\n\n\n@pytest.fixture\ndef mock_socket_mode_handler():\n with patch(\n \"slack_bolt.adapter.socket_mode.async_handler.AsyncSocketModeHandler\"\n ) as mock_handler:\n mock_handler_object = mock_handler.return_value\n mock_handler_object.start_async = AsyncMock()\n yield mock_handler_object\n\n\n@pytest.fixture\ndef mock_openai():\n mock_openai = MagicMock()\n with patch.dict(\"sys.modules\", openai=mock_openai):\n yield mock_openai\n\n\n@pytest.fixture\ndef mock_slack_asyncwebclient():\n with patch(\"slack_sdk.web.async_client.AsyncWebClient\") as mock_client:\n yield mock_client.return_value\n\n\n@pytest.fixture\ndef mock_slack_client(mock_slack_asyncwebclient):\n return SlackClient(mock_slack_asyncwebclient, \"template_path\")\n\n\n@pytest.fixture\ndef mock_message_handler(mock_slack_client):\n return MockMessageHandler(mock_slack_client)\n\n\n@pytest.fixture\ndef mock_action_handler(mock_slack_client):\n return MockActionHandler(mock_slack_client)\n\n\nclass MockMessageHandler(BaseMessageHandler):\n def __init__(self, slack_client):\n super().__init__(slack_client)\n self.mock_handler = AsyncMock()\n\n async def should_handle(self, args):\n return args.event[\"subtype\"] != \"bot_message\"\n\n async def handle(self, args):\n await self.mock_handler(args)\n\n\nclass MockActionHandler(BaseActionHandler):\n def __init__(self, slack_client):\n super().__init__(slack_client)\n self.mock_handler = AsyncMock()\n\n async def handle(self, args):\n await self.mock_handler(args)\n\n @property\n def action_id(self):\n return \"mock_action\"\n" }, { "path": "shared/openai-slackbot/tests/test_bot.py", "content": "import pytest\n\n\nasync def test_start_bot(\n mock_slack_app, mock_socket_mode_handler, mock_message_handler, mock_action_handler\n):\n from openai_slackbot.bot import start_bot\n\n await start_bot(\n openai_organization_id=\"org-id\",\n slack_message_handler=mock_message_handler.__class__,\n slack_action_handlers=[mock_action_handler.__class__],\n slack_template_path=\"/path/to/templates\",\n )\n\n mock_slack_app.event.assert_called_once_with(\"message\")\n mock_slack_app.action.assert_called_once_with(\"mock_action\")\n mock_socket_mode_handler.start_async.assert_called_once()\n" }, { "path": "shared/openai-slackbot/tests/test_handlers.py", "content": "from unittest.mock import AsyncMock, MagicMock\n\nimport pytest\n\n\n@pytest.mark.parametrize(\"subtype, should_handle\", [(\"message\", True), (\"bot_message\", False)])\nasync def test_message_handler(mock_message_handler, subtype, should_handle):\n args = MagicMock(\n ack=AsyncMock(),\n event={\"type\": \"message\", \"subtype\": subtype, \"channel\": \"channel\", \"ts\": \"ts\"},\n )\n\n await mock_message_handler.maybe_handle(args)\n args.ack.assert_awaited_once()\n if should_handle:\n mock_message_handler.mock_handler.assert_awaited_once_with(args)\n else:\n mock_message_handler.mock_handler.assert_not_awaited()\n\n assert mock_message_handler.logging_extra(args) == {\n \"type\": \"message\",\n \"subtype\": subtype,\n \"channel\": \"channel\",\n \"ts\": \"ts\",\n }\n\n\nasync def test_action_handler(mock_action_handler):\n args = MagicMock(\n ack=AsyncMock(),\n body={\n \"type\": \"type\",\n \"actions\": [\"action\"],\n },\n )\n\n await mock_action_handler.maybe_handle(args)\n args.ack.assert_awaited_once()\n mock_action_handler.mock_handler.assert_awaited_once_with(args)\n\n assert mock_action_handler.logging_extra(args) == {\n \"action_type\": \"type\",\n \"action\": \"action\",\n }\n" } ]