├── .gitignore ├── .pre-commit-config.yaml ├── CODEOWNERS ├── LICENSE ├── Makefile ├── README.md ├── SECURITY.md ├── bots ├── incident-response-slackbot │ ├── Makefile │ ├── README.md │ ├── incident_response_slackbot │ │ ├── .env.template │ │ ├── bot.py │ │ ├── config.py │ │ ├── config.toml │ │ ├── db │ │ │ └── database.py │ │ ├── handlers.py │ │ ├── openai_utils.py │ │ └── templates │ │ │ └── messages │ │ │ └── incident_alert.j2 │ ├── pyproject.template.toml │ ├── scripts │ │ ├── alert_feed.py │ │ ├── alerts.toml │ │ └── send_alert.py │ └── tests │ │ ├── __init__.py │ │ ├── conftest.py │ │ ├── test_config.toml │ │ ├── test_handlers.py │ │ └── test_openai.py ├── sdlc-slackbot │ ├── Makefile │ ├── README.md │ ├── pyproject.template.toml │ ├── requirements.txt │ ├── sdlc_slackbot │ │ ├── .env.template │ │ ├── bot.py │ │ ├── config.py │ │ ├── config.toml │ │ ├── database.py │ │ ├── gdoc.py │ │ ├── utils.py │ │ └── validate.py │ └── setup.py └── triage-slackbot │ ├── Makefile │ ├── README.md │ ├── pyproject.template.toml │ ├── tests │ ├── __init__.py │ ├── conftest.py │ ├── test_config.toml │ └── test_handlers.py │ └── triage_slackbot │ ├── .env.template │ ├── bot.py │ ├── category.py │ ├── config.py │ ├── config.toml │ ├── handlers.py │ ├── openai_utils.py │ └── templates │ ├── blocks │ ├── empty_category_warning.j2 │ ├── empty_conversation_warning.j2 │ └── select_conversation.j2 │ └── messages │ ├── _notify_oncall_body.j2 │ ├── autorespond.j2 │ ├── feed.j2 │ ├── notify_oncall_channel.j2 │ └── notify_oncall_in_feed.j2 └── shared └── openai-slackbot ├── openai_slackbot ├── __init__.py ├── bot.py ├── clients │ ├── __init__.py │ └── slack.py ├── handlers.py └── utils │ ├── __init__.py │ ├── envvars.py │ └── slack.py ├── pyproject.toml ├── setup.cfg └── tests ├── __init__.py ├── clients ├── __init__.py └── test_slack.py ├── conftest.py ├── test_bot.py └── test_handlers.py /.gitignore: -------------------------------------------------------------------------------- 1 | # Byte-compiled / optimized / DLL files 2 | __pycache__/ 3 | *.py[cod] 4 | *$py.class 5 | 6 | # C extensions 7 | *.so 8 | 9 | # Distribution / packaging 10 | .Python 11 | build/ 12 | develop-eggs/ 13 | dist/ 14 | downloads/ 15 | eggs/ 16 | .eggs/ 17 | lib/ 18 | lib64/ 19 | parts/ 20 | sdist/ 21 | var/ 22 | wheels/ 23 | pip-wheel-metadata/ 24 | share/python-wheels/ 25 | *.egg-info/ 26 | .installed.cfg 27 | *.egg 28 | MANIFEST 29 | 30 | # PyInstaller 31 | # Usually these files are written by a python script from a template 32 | # before PyInstaller builds the exe, so as to inject date/other infos into it. 33 | *.manifest 34 | *.spec 35 | 36 | # Installer logs 37 | pip-log.txt 38 | pip-delete-this-directory.txt 39 | 40 | # Unit test / coverage reports 41 | htmlcov/ 42 | .tox/ 43 | .nox/ 44 | .coverage 45 | .coverage.* 46 | .cache 47 | nosetests.xml 48 | coverage.xml 49 | *.cover 50 | *.py,cover 51 | .hypothesis/ 52 | .pytest_cache/ 53 | 54 | # Translations 55 | *.mo 56 | *.pot 57 | 58 | # Django stuff: 59 | *.log 60 | local_settings.py 61 | db.sqlite3 62 | db.sqlite3-journal 63 | 64 | # Flask stuff: 65 | instance/ 66 | .webassets-cache 67 | 68 | # Scrapy stuff: 69 | .scrapy 70 | 71 | # Sphinx documentation 72 | docs/_build/ 73 | 74 | # PyBuilder 75 | target/ 76 | 77 | # Jupyter Notebook 78 | .ipynb_checkpoints 79 | 80 | # IPython 81 | profile_default/ 82 | ipython_config.py 83 | 84 | # pyenv 85 | .python-version 86 | 87 | # pipenv 88 | # According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. 89 | # However, in case of collaboration, if having platform-specific dependencies or dependencies 90 | # having no cross-platform support, pipenv may install dependencies that don't work, or not 91 | # install all needed dependencies. 92 | #Pipfile.lock 93 | 94 | # PEP 582; used by e.g. github.com/David-OConnor/pyflow 95 | __pypackages__/ 96 | 97 | # Celery stuff 98 | celerybeat-schedule 99 | celerybeat.pid 100 | 101 | # SageMath parsed files 102 | *.sage.py 103 | 104 | # Environments 105 | .env 106 | .venv 107 | env/ 108 | venv/ 109 | ENV/ 110 | env.bak/ 111 | venv.bak/ 112 | 113 | # Spyder project settings 114 | .spyderproject 115 | .spyproject 116 | 117 | # Rope project settings 118 | .ropeproject 119 | 120 | # mkdocs documentation 121 | /site 122 | 123 | # mypy 124 | .mypy_cache/ 125 | .dmypy.json 126 | dmypy.json 127 | 128 | # Pyre type checker 129 | .pyre/ 130 | 131 | # Pyproject 132 | bots/triage-slackbot/pyproject.toml 133 | bots/incident-response-slackbot/pyproject.toml 134 | 135 | .trufflehog 136 | *.pkl 137 | *.json 138 | exclude.txt 139 | -------------------------------------------------------------------------------- /.pre-commit-config.yaml: -------------------------------------------------------------------------------- 1 | repos: 2 | - repo: local 3 | hooks: 4 | - id: trufflehog 5 | name: TruffleHog 6 | description: Detect secrets in your data. 7 | entry: bash -c 'trufflehog git file://. --since-commit HEAD --fail' 8 | language: system 9 | stages: ["commit", "push"] 10 | 11 | - repo: https://github.com/hauntsaninja/black-pre-commit-mirror 12 | rev: 23.10.1 13 | hooks: 14 | - id: black 15 | args: [--line-length=100, --workers=6] 16 | 17 | - repo: https://github.com/pycqa/isort 18 | rev: 5.12.0 19 | hooks: 20 | - id: isort 21 | name: isort (python) 22 | -------------------------------------------------------------------------------- /CODEOWNERS: -------------------------------------------------------------------------------- 1 | * @openai/security-team 2 | -------------------------------------------------------------------------------- /LICENSE: -------------------------------------------------------------------------------- 1 | MIT License 2 | 3 | Copyright (c) 2024 OpenAI 4 | 5 | Permission is hereby granted, free of charge, to any person obtaining a copy 6 | of this software and associated documentation files (the "Software"), to deal 7 | in the Software without restriction, including without limitation the rights 8 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell 9 | copies of the Software, and to permit persons to whom the Software is 10 | furnished to do so, subject to the following conditions: 11 | 12 | The above copyright notice and this permission notice shall be included in all 13 | copies or substantial portions of the Software. 14 | 15 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR 16 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, 17 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE 18 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER 19 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, 20 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE 21 | SOFTWARE. 22 | -------------------------------------------------------------------------------- /Makefile: -------------------------------------------------------------------------------- 1 | SHELL := /bin/bash 2 | 3 | clean-venv: rm-venv 4 | python3 -m venv venv 5 | 6 | 7 | rm-venv: 8 | if [ -d "venv" ]; then rm -rf venv; fi 9 | 10 | maybe-clear-shared: 11 | ifeq ($(SKIP_CLEAR_SHARED), true) 12 | else 13 | pip cache remove openai_slackbot 14 | endif 15 | 16 | build-shared: 17 | pip install -e ./shared/openai-slackbot 18 | 19 | 20 | build-bot: maybe-clear-shared build-shared 21 | cd bots/$(BOT) && $(MAKE) init-pyproject && pip install -e . 22 | 23 | 24 | run-bot: 25 | python bots/$(BOT)/$(subst -,_,$(BOT))/bot.py 26 | 27 | 28 | clear: 29 | find . | grep -E "(/__pycache__$|\.pyc$|\.pyo$)" | xargs rm -rf 30 | 31 | 32 | build-all: 33 | $(MAKE) build-bot BOT=triage-slackbot SKIP_CLEAR_SHARED=true 34 | 35 | 36 | test-all: 37 | pytest shared/openai-slackbot && \ 38 | pytest bots/triage-slackbot && \ 39 | pytest bots/incident-response-slackbot -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | # OpenAI Security Bots 🤖 2 | 3 | Slack bots integrated with OpenAI APIs to streamline security team's workflows. 4 | 5 | All the bots can be found under `bots/` directory. 6 | 7 | ``` 8 | shared/ 9 | openai-slackbot/ 10 | bots/ 11 | triage-slackbot/ 12 | incident-response-slackbot/ 13 | sdlc-slackbot/ 14 | ``` 15 | 16 | Refer to each bot's README for more information and setup instruction. 17 | 18 | 19 | If you wish to contribute, note this repo uses pre-commit to help. In this directory, run: 20 | ``` 21 | pip install pre-commit 22 | pre-commit install 23 | ``` 24 | -------------------------------------------------------------------------------- /SECURITY.md: -------------------------------------------------------------------------------- 1 | # Security Policy 2 | For 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.). 3 | 4 | Our PGP key can located [at this address.](https://cdn.openai.com/security.txt) 5 | -------------------------------------------------------------------------------- /bots/incident-response-slackbot/Makefile: -------------------------------------------------------------------------------- 1 | CWD := $(shell pwd) 2 | REPO_ROOT := $(shell git rev-parse --show-toplevel) 3 | ESCAPED_REPO_ROOT := $(shell echo $(REPO_ROOT) | sed 's/\//\\\//g') 4 | 5 | init-env-file: 6 | cp ./incident_response_slackbot/.env.template ./incident_response_slackbot/.env 7 | 8 | init-pyproject: 9 | cat $(CWD)/pyproject.template.toml | \ 10 | sed "s/\$$REPO_ROOT/$(ESCAPED_REPO_ROOT)/g" > $(CWD)/pyproject.toml 11 | -------------------------------------------------------------------------------- /bots/incident-response-slackbot/README.md: -------------------------------------------------------------------------------- 1 |

2 | triage-slackbot-logo 3 |

Incident Response Slackbot

4 |

5 | 6 | Incident Response Slackbot automatically chats with users who have been part of an incident alert. 7 | 8 | 9 | 10 | ## Prerequisites 11 | 12 | You will need: 13 | 1. A Slack application (aka your triage bot) with Socket Mode enabled 14 | 2. OpenAI API key 15 | 16 | Grab your `SLACK_BOT_TOKEN` by Oauth & Permissions tab in your Slack App page. 17 | 18 | Generate an App-level token for your Slack app, by going to: 19 | ``` 20 | Your Slack App > Basic Information > App-Level Tokens > Generate Token and Scopes 21 | ``` 22 | Create a new token with `connections:write` scope. This is your `SOCKET_APP_TOKEN` token. 23 | 24 | Once you have them, from the current directory, run: 25 | ``` 26 | $ make init-env-file 27 | ``` 28 | and fill in the right values. 29 | 30 | Your Slack App needs the following scopes: 31 | 32 | - users:read 33 | - channels:history 34 | - chat:write 35 | - groups:history 36 | 37 | ## Setup 38 | 39 | From the current directory, run: 40 | ``` 41 | make init-pyproject 42 | ``` 43 | 44 | From the repo root, run: 45 | ``` 46 | make clean-venv 47 | source venv/bin/activate 48 | make build-bot BOT=incident-response-slackbot 49 | ``` 50 | 51 | ## Run bot with example configuration 52 | 53 | The 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. 54 | 55 | ⚠️ *Make sure that the bot is added to the channels it needs to read from and post to.* ⚠️ 56 | 57 | We will need to add example alerts to `./scripts/alerts.toml` Replace with alert information and user_id. To get the user_id: 58 | 1. Click on the desired user name within Slack. 59 | 2. Click on the ellpises (three dots). 60 | 3. Click on "Copy Member ID". 61 | 62 | ⚠️ *These are mock alerts. In real-world scenarios, this will be integrated with alert feed/database* ⚠️ 63 | 64 | To generate an axample alert, in this directory, run: 65 | ``` 66 | python ./scripts/send_alert.py 67 | ``` 68 | 69 | An example alert will be sent to the channel. 70 | 71 | 72 | https://github.com/openai/openai-security-bots/assets/124844323/b919639c-b691-4b01-aa0c-7be987c9a70b 73 | 74 | 75 | To have the bot start listening, run the following from the repo root: 76 | 77 | ``` 78 | make run-bot BOT=incident-response-slackbot 79 | ``` 80 | 81 | Now you can start a chat with a user, or do nothing. 82 | When you start a chat, 83 | 84 | 1. The bot will reach out to the user involved with the alert 85 | 2. Post a message to the original thread in monitoring channel what was sent to the user (message generated with OpenAI API) 86 | 3. Post any messages the user sends to original thread 87 | 4. Checks to see if the user has answered the question using OpenAI's API. 88 | - If yes, end the chat and provide a summary to the original thread 89 | - If no, continues sending a message to the user, and repeats this step 90 | 91 | Let's start a chat: 92 | 93 | https://github.com/openai/openai-security-bots/assets/124844323/4b5dd292-b4d3-437a-9809-d6d80e824a9d 94 | 95 | 96 | 97 | ## Alert Details 98 | 99 | In 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. 100 | 101 | In the `alerts.toml` file: 102 | 103 | ``` 104 | [[ alerts ]] 105 | id = "pivot" 106 | ... 107 | user_id = ID of person to start chat with (@harold user) 108 | 109 | [alerts.properties] 110 | source_host = "source.machine.org" 111 | destination_host = "destination.machine.org" 112 | 113 | [[ alerts ]] 114 | id = "privesc" 115 | ... 116 | user_id = ID of person to start chat with (@harold user) 117 | ``` 118 | -------------------------------------------------------------------------------- /bots/incident-response-slackbot/incident_response_slackbot/.env.template: -------------------------------------------------------------------------------- 1 | SLACK_BOT_TOKEN=xoxb- # Go to your Slack app, Settings > Install App > Bot User OAuth Token 2 | SOCKET_APP_TOKEN=xapp- # Go to your Slack app, Settings > Basic Information > App-Level Tokens 3 | OPENAI_API_KEY= -------------------------------------------------------------------------------- /bots/incident-response-slackbot/incident_response_slackbot/bot.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | import os 3 | 4 | from incident_response_slackbot.config import load_config, get_config 5 | from incident_response_slackbot.handlers import ( 6 | InboundDirectMessageHandler, 7 | InboundIncidentDoNothingHandler, 8 | InboundIncidentEndChatHandler, 9 | InboundIncidentStartChatHandler, 10 | ) 11 | from openai_slackbot.bot import start_bot 12 | 13 | if __name__ == "__main__": 14 | current_dir = os.path.dirname(os.path.abspath(__file__)) 15 | load_config(os.path.join(current_dir, "config.toml")) 16 | 17 | message_handler = InboundDirectMessageHandler 18 | action_handlers = [ 19 | InboundIncidentStartChatHandler, 20 | InboundIncidentDoNothingHandler, 21 | InboundIncidentEndChatHandler, 22 | ] 23 | 24 | template_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "templates") 25 | 26 | config = get_config() 27 | asyncio.run( 28 | start_bot( 29 | openai_organization_id=config.openai_organization_id, 30 | slack_message_handler=message_handler, 31 | slack_action_handlers=action_handlers, 32 | slack_template_path=template_path, 33 | ) 34 | ) 35 | -------------------------------------------------------------------------------- /bots/incident-response-slackbot/incident_response_slackbot/config.py: -------------------------------------------------------------------------------- 1 | import os 2 | import typing as t 3 | 4 | import toml 5 | from dotenv import load_dotenv 6 | from pydantic import BaseModel 7 | 8 | _CONFIG = None 9 | 10 | 11 | class Config(BaseModel): 12 | # OpenAI organization ID associated with OpenAI API key. 13 | openai_organization_id: str 14 | 15 | # Slack channel where triage alerts are posted. 16 | feed_channel_id: str 17 | 18 | 19 | def load_config(config_path: str = None) -> Config: 20 | load_dotenv() 21 | 22 | if config_path is None: 23 | config_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "config.toml") 24 | with open(config_path) as f: 25 | cfg = toml.loads(f.read()) 26 | config = Config(**cfg) 27 | 28 | global _CONFIG 29 | _CONFIG = config 30 | return _CONFIG 31 | 32 | 33 | def get_config() -> Config: 34 | global _CONFIG 35 | if _CONFIG is None: 36 | raise Exception("config not initialized, call load_config() first") 37 | return _CONFIG 38 | -------------------------------------------------------------------------------- /bots/incident-response-slackbot/incident_response_slackbot/config.toml: -------------------------------------------------------------------------------- 1 | # Organization ID associated with OpenAI API key. 2 | openai_organization_id = "" 3 | 4 | # Where the alerts will be posted. 5 | feed_channel_id = "" 6 | 7 | 8 | -------------------------------------------------------------------------------- /bots/incident-response-slackbot/incident_response_slackbot/db/database.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pickle 3 | 4 | 5 | class Database: 6 | """ 7 | This class represents a database for storing user messages. 8 | The data is stored in a pickle file. 9 | """ 10 | 11 | def __init__(self): 12 | """ 13 | Initialize the database. Load data from the pickle file if it exists, 14 | otherwise create an empty dictionary. 15 | """ 16 | current_dir = os.path.dirname(os.path.realpath(__file__)) 17 | self.file_path = os.path.join(current_dir, "data.pkl") 18 | self.data = {} 19 | 20 | def _load_data(self): 21 | """ 22 | Load data from the pickle file if it exists, 23 | otherwise return an empty dictionary. 24 | """ 25 | if os.path.exists(self.file_path): 26 | with open(self.file_path, "rb") as f: 27 | return pickle.load(f) 28 | else: 29 | return {} 30 | 31 | def _save(self): 32 | """ 33 | Save the current state of the database to the pickle file. 34 | """ 35 | with open(self.file_path, "wb") as f: 36 | pickle.dump(self.data, f) 37 | 38 | # Add a new entry to the database 39 | def add(self, user_id, message_ts): 40 | """ 41 | Add a new entry to the database. If the user_id already exists, 42 | update the message timestamp. Otherwise, create a new entry. 43 | """ 44 | self.data = self._load_data() 45 | self.data.setdefault(user_id, {})["message_ts"] = message_ts 46 | self._save() 47 | 48 | # Delete an entry from the database 49 | def delete(self, user_id): 50 | """ 51 | Delete an entry from the database using the user_id as the key. 52 | """ 53 | self.data = self._load_data() 54 | self.data.pop(user_id, None) 55 | self._save() 56 | 57 | # Check if user_id exists in the database 58 | def user_exists(self, user_id): 59 | """ 60 | Check if the user_id exists in the database. 61 | """ 62 | self.data = self._load_data() 63 | return user_id in self.data 64 | 65 | # Return the message timestamp for a given user_id 66 | def get_ts(self, user_id): 67 | """ 68 | Return the message timestamp for a given user_id. 69 | """ 70 | self.data = self._load_data() 71 | return self.data[user_id]["message_ts"] 72 | 73 | # Return the user_id given a message_ts 74 | def get_user_id(self, message_ts): 75 | """ 76 | Return the user_id given a message_ts. 77 | """ 78 | self.data = self._load_data() 79 | for user_id, data in self.data.items(): 80 | if data["message_ts"] == message_ts: 81 | return user_id 82 | return None 83 | -------------------------------------------------------------------------------- /bots/incident-response-slackbot/incident_response_slackbot/handlers.py: -------------------------------------------------------------------------------- 1 | import os 2 | import pickle 3 | import typing as t 4 | from enum import Enum 5 | from logging import getLogger 6 | 7 | from incident_response_slackbot.config import load_config, get_config 8 | from incident_response_slackbot.db.database import Database 9 | from incident_response_slackbot.openai_utils import ( 10 | create_greeting, 11 | generate_awareness_question, 12 | get_thread_summary, 13 | get_user_awareness, 14 | messages_to_string, 15 | ) 16 | from openai_slackbot.handlers import BaseActionHandler, BaseMessageHandler 17 | 18 | logger = getLogger(__name__) 19 | 20 | DATABASE = Database() 21 | 22 | class InboundDirectMessageHandler(BaseMessageHandler): 23 | """ 24 | Handles Direct Messages for incident response use cases 25 | """ 26 | 27 | def __init__(self, slack_client): 28 | super().__init__(slack_client) 29 | self.config = get_config() 30 | 31 | async def should_handle(self, args): 32 | return True 33 | 34 | async def handle(self, args): 35 | event = args.event 36 | user_id = event.get("user") 37 | 38 | if not DATABASE.user_exists(user_id): 39 | # If the user_id does not exist, they're not part of an active chat 40 | return 41 | 42 | message_ts = DATABASE.get_ts(user_id) 43 | await self.send_message_to_channel(event, message_ts) 44 | 45 | user_awareness = await get_user_awareness(event["text"]) 46 | logger.info(f"User awareness decision: {user_awareness}") 47 | 48 | if user_awareness["has_answered"]: 49 | await self.handle_user_response(user_id, message_ts) 50 | else: 51 | await self.nudge_user(user_id, message_ts) 52 | 53 | async def send_message_to_channel(self, event, message_ts): 54 | # Send the received message to the monitoring channel 55 | await self._slack_client.post_message( 56 | channel=self.config.feed_channel_id, 57 | text=f"Received message from <@{event['user']}>:\n> {event['text']}", 58 | thread_ts=message_ts, 59 | ) 60 | 61 | async def handle_user_response(self, user_id, message_ts): 62 | # User has answered the question 63 | messages = await self._slack_client.get_thread_messages( 64 | channel=self.config.feed_channel_id, 65 | thread_ts=message_ts, 66 | ) 67 | 68 | # Send the end message to the user 69 | thank_you = "Thanks for your time!" 70 | await self._slack_client.post_message( 71 | channel=user_id, 72 | text=thank_you, 73 | ) 74 | 75 | # Send message to the channel 76 | await self._slack_client.post_message( 77 | channel=self.config.feed_channel_id, 78 | text=f"Sent message to <@{user_id}>:\n> {thank_you}", 79 | thread_ts=message_ts, 80 | ) 81 | 82 | summary = await get_thread_summary(messages) 83 | 84 | # Send message to the channel 85 | await self._slack_client.post_message( 86 | channel=self.config.feed_channel_id, 87 | text=f"Here is the summary of the chat:\n> {summary}", 88 | thread_ts=message_ts, 89 | ) 90 | 91 | DATABASE.delete(user_id) 92 | 93 | await self.end_chat(message_ts) 94 | 95 | async def end_chat(self, message_ts): 96 | original_blocks = await self._slack_client.get_original_blocks( 97 | message_ts, self.config.feed_channel_id 98 | ) 99 | 100 | # Remove action buttons and add "Chat has ended" text 101 | new_blocks = [block for block in original_blocks if block.get("type") != "actions"] 102 | 103 | # Add the "Chat has ended" text 104 | new_blocks.append( 105 | { 106 | "type": "section", 107 | "block_id": "end_chat_automatically", 108 | "text": { 109 | "type": "mrkdwn", 110 | "text": f"The chat was automatically ended from SecurityBot review. :done_:", 111 | "verbatim": True, 112 | }, 113 | } 114 | ) 115 | 116 | await self._slack_client.update_message( 117 | channel=self.config.feed_channel_id, 118 | blocks=new_blocks, 119 | ts=message_ts, 120 | text="Ended chat automatically", 121 | ) 122 | 123 | async def nudge_user(self, user_id, message_ts): 124 | # User has not answered the question 125 | 126 | nudge_message = await generate_awareness_question() 127 | # Send the greeting message to the user 128 | await self._slack_client.post_message( 129 | channel=user_id, 130 | text=nudge_message, 131 | ) 132 | 133 | # Send message to the channel 134 | await self._slack_client.post_message( 135 | channel=self.config.feed_channel_id, 136 | text=f"Sent message to <@{user_id}>:\n> {nudge_message}", 137 | thread_ts=message_ts, 138 | ) 139 | 140 | 141 | class InboundIncidentStartChatHandler(BaseActionHandler): 142 | def __init__(self, slack_client): 143 | super().__init__(slack_client) 144 | self.config = get_config() 145 | 146 | @property 147 | def action_id(self): 148 | return "start_chat_submit_action" 149 | 150 | async def handle(self, args): 151 | body = args.body 152 | original_message = body["container"] 153 | original_message_ts = original_message["message_ts"] 154 | alert_user_id = DATABASE.get_user_id(original_message_ts) 155 | user = body["user"] 156 | 157 | name = user["name"] 158 | first_name = name.split(".")[1] 159 | 160 | logger.info(f"Handling inbound incident start chat action from {user['name']}") 161 | 162 | # Update the blocks and elements 163 | blocks = self.update_blocks(body, alert_user_id) 164 | 165 | # Add the "Started a chat" text 166 | blocks.append(self.create_chat_start_section(user["id"])) 167 | 168 | messages = await self._slack_client.get_thread_messages( 169 | channel=self.config.feed_channel_id, 170 | thread_ts=original_message_ts, 171 | ) 172 | 173 | message = await self._slack_client.update_message( 174 | channel=self.config.feed_channel_id, 175 | blocks=blocks, 176 | ts=original_message_ts, 177 | text=messages[0]["text"], 178 | ) 179 | 180 | text_messages = messages_to_string(messages) 181 | logger.info(f"Alert and detail: {text_messages}") 182 | 183 | username = await self._slack_client.get_user_display_name(alert_user_id) 184 | 185 | greeting_message = await create_greeting(first_name, text_messages) 186 | logger.info(f"generated greeting message: {greeting_message}") 187 | 188 | # Send the greeting message to the user and to the channel 189 | await self.send_greeting_message(alert_user_id, greeting_message, original_message_ts) 190 | 191 | logger.info(f"Succesfully started chat with user: {username}") 192 | 193 | return message 194 | 195 | def update_blocks(self, body, alert_user_id): 196 | body_copy = body.copy() 197 | new_elements = [] 198 | for block in body_copy.get("message", {}).get("blocks", []): 199 | if block.get("type") == "actions": 200 | for element in block.get("elements", []): 201 | if element.get("action_id") == "do_nothing_submit_action": 202 | element["action_id"] = "end_chat_submit_action" 203 | element["text"]["text"] = "End Chat" 204 | element["value"] = alert_user_id 205 | new_elements.append(element) 206 | block["elements"] = new_elements 207 | return body_copy.get("message", {}).get("blocks", []) 208 | 209 | def create_chat_start_section(self, user_id): 210 | return { 211 | "type": "section", 212 | "block_id": "started_chat", 213 | "text": { 214 | "type": "mrkdwn", 215 | "text": f"<@{user_id}> started a chat.", 216 | "verbatim": True, 217 | }, 218 | } 219 | 220 | async def send_greeting_message(self, alert_user_id, greeting_message, original_message_ts): 221 | # Send the greeting message to the user 222 | await self._slack_client.post_message( 223 | channel=alert_user_id, 224 | text=greeting_message, 225 | ) 226 | 227 | # Send message to the channel 228 | await self._slack_client.post_message( 229 | channel=self.config.feed_channel_id, 230 | text=f"Sent message to <@{alert_user_id}>:\n> {greeting_message}", 231 | thread_ts=original_message_ts, 232 | ) 233 | 234 | 235 | class InboundIncidentDoNothingHandler(BaseActionHandler): 236 | """ 237 | Handles incoming alerts and decides whether to take no action. 238 | This will close the alert and mark the status as complete. 239 | """ 240 | 241 | def __init__(self, slack_client): 242 | super().__init__(slack_client) 243 | self.config = get_config() 244 | 245 | @property 246 | def action_id(self): 247 | return "do_nothing_submit_action" 248 | 249 | async def handle(self, args): 250 | body = args.body 251 | user_id = body["user"]["id"] 252 | original_message_ts = body["message"]["ts"] 253 | 254 | # Remove action buttons and add "Chat has ended" text 255 | new_blocks = [ 256 | block 257 | for block in body.get("message", {}).get("blocks", []) 258 | if block.get("type") != "actions" 259 | ] 260 | 261 | # Add the "Chat has ended" text 262 | new_blocks.append( 263 | { 264 | "type": "section", 265 | "block_id": "do_nothing", 266 | "text": { 267 | "type": "mrkdwn", 268 | "text": f"<@{user_id}> decided that no action was necessary :done_:", 269 | "verbatim": True, 270 | }, 271 | } 272 | ) 273 | 274 | await self._slack_client.update_message( 275 | channel=self.config.feed_channel_id, 276 | blocks=new_blocks, 277 | ts=original_message_ts, 278 | text="Do Nothing action selected", 279 | ) 280 | 281 | 282 | class InboundIncidentEndChatHandler(BaseActionHandler): 283 | """ 284 | Ends the chat manually 285 | """ 286 | 287 | def __init__(self, slack_client): 288 | super().__init__(slack_client) 289 | self.config = get_config() 290 | 291 | @property 292 | def action_id(self): 293 | return "end_chat_submit_action" 294 | 295 | async def handle(self, args): 296 | body = args.body 297 | user_id = body["user"]["id"] 298 | message_ts = body["message"]["ts"] 299 | 300 | alert_user_id = DATABASE.get_user_id(message_ts) 301 | 302 | original_blocks = await self._slack_client.get_original_blocks( 303 | message_ts, self.config.feed_channel_id 304 | ) 305 | 306 | # Remove action buttons and add "Chat has ended" text 307 | new_blocks = [block for block in original_blocks if block.get("type") != "actions"] 308 | 309 | # Add the "Chat has ended" text 310 | new_blocks.append( 311 | { 312 | "type": "section", 313 | "block_id": "end_chat_manually", 314 | "text": { 315 | "type": "mrkdwn", 316 | "text": f"<@{user_id}> has ended the chat :done_:", 317 | "verbatim": True, 318 | }, 319 | } 320 | ) 321 | 322 | await self._slack_client.update_message( 323 | channel=self.config.feed_channel_id, 324 | blocks=new_blocks, 325 | ts=message_ts, 326 | text="Ended chat automatically", 327 | ) 328 | 329 | # User has answered the question 330 | messages = await self._slack_client.get_thread_messages( 331 | channel=self.config.feed_channel_id, 332 | thread_ts=message_ts, 333 | ) 334 | 335 | thank_you = "Thanks for your time!" 336 | await self._slack_client.post_message( 337 | channel=alert_user_id, 338 | text=thank_you, 339 | ) 340 | 341 | # Send message to the channel 342 | await self._slack_client.post_message( 343 | channel=self.config.feed_channel_id, 344 | text=f"Sent message to <@{alert_user_id}>:\n> {thank_you}", 345 | thread_ts=message_ts, 346 | ) 347 | 348 | summary = await get_thread_summary(messages) 349 | 350 | # Send message to the channel 351 | await self._slack_client.post_message( 352 | channel=self.config.feed_channel_id, 353 | text=f"Here is the summary of the chat:\n> {summary}", 354 | thread_ts=message_ts, 355 | ) 356 | 357 | DATABASE.delete(user_id) 358 | -------------------------------------------------------------------------------- /bots/incident-response-slackbot/incident_response_slackbot/openai_utils.py: -------------------------------------------------------------------------------- 1 | import json 2 | 3 | import openai 4 | from incident_response_slackbot.config import load_config, get_config 5 | 6 | load_config() 7 | config = get_config() 8 | 9 | # Convert slack threaded messages to string 10 | def messages_to_string(messages): 11 | text_messages = " ".join([message["text"] for message in messages if "text" in message]) 12 | return text_messages 13 | 14 | 15 | async def get_clean_output(completion: str) -> str: 16 | return completion.choices[0].message.content 17 | 18 | 19 | async def create_greeting(username, details): 20 | if not openai.api_key: 21 | raise Exception("OpenAI API key not found.") 22 | 23 | prompt = f""" 24 | You are a helpful cybersecurity AI analyst assistant to the security team that wants to keep 25 | your company secure. You just received an alert with the following details: 26 | {details} 27 | Without being accusatory, gently ask the user, whose name is {username} in a casual tone if they were aware 28 | about the topic of the alert. 29 | Keep the message brief, not more than 3 or 4 sentences. 30 | Do not end with a signature. End with a question. 31 | """ 32 | 33 | messages = [ 34 | {"role": "system", "content": prompt}, 35 | {"role": "user", "content": ""}, 36 | ] 37 | 38 | completion = openai.chat.completions.create( 39 | model="gpt-4-32k", 40 | messages=messages, 41 | temperature=0.3, 42 | stream=False, 43 | ) 44 | response = await get_clean_output(completion) 45 | return response 46 | 47 | 48 | aware_decision_function = [ 49 | { 50 | "name": "is_user_aware", 51 | "description": "Determines if the user has answered whether they were aware, and what that response is.", 52 | "parameters": { 53 | "type": "object", 54 | "properties": { 55 | "has_answered": { 56 | "type": "boolean", 57 | "description": "Determine whether user answered the quesiton of whether they were aware.", 58 | }, 59 | "is_aware": { 60 | "type": "boolean", 61 | "description": "Determine whether user was aware of the alert details.", 62 | }, 63 | }, 64 | "required": ["has_answered", "is_aware"], 65 | }, 66 | } 67 | ] 68 | 69 | 70 | async def get_user_awareness(inbound_direct_message: str) -> str: 71 | """ 72 | This function uses the OpenAI Chat Completion API to determine whether user was aware. 73 | """ 74 | # Define the prompt 75 | prompt = f""" 76 | You are a helpful cybersecurity AI analyst assistant to the security team that wants to keep 77 | your company secure. You just received an alert and are having a chat with the user whether 78 | they were aware about the details of an alert. Based on the chat so far, determine whether 79 | the user has answered the question of whether they were aware of the alert details, and whether 80 | they were aware or not. 81 | """ 82 | 83 | messages = [ 84 | {"role": "system", "content": prompt}, 85 | {"role": "user", "content": inbound_direct_message}, 86 | ] 87 | 88 | # Call the API 89 | response = openai.chat.completions.create( 90 | model="gpt-4-32k", 91 | messages=messages, 92 | temperature=0, 93 | stream=False, 94 | functions=aware_decision_function, 95 | function_call={"name": "is_user_aware"}, 96 | ) 97 | 98 | function_args = json.loads(response.choices[0].message.function_call.arguments) # type: ignore 99 | return function_args 100 | 101 | 102 | async def get_thread_summary(messages): 103 | if not openai.api_key: 104 | raise Exception("OpenAI API key not found.") 105 | 106 | text_messages = messages_to_string(messages) 107 | 108 | prompt = f""" 109 | You are a helpful cybersecurity AI analyst assistant to the security team that wants to keep 110 | your company secure. The following is a conversation that you had with the user. 111 | Please summarize the following conversation, and note whether the user was aware or not aware 112 | of the alert, and whether they acted suspiciously when answering: 113 | {text_messages} 114 | """ 115 | 116 | messages = [ 117 | {"role": "system", "content": prompt}, 118 | {"role": "user", "content": ""}, 119 | ] 120 | 121 | completion = openai.chat.completions.create( 122 | model="gpt-4-32k", 123 | messages=messages, 124 | temperature=0.3, 125 | stream=False, 126 | ) 127 | response = await get_clean_output(completion) 128 | return response 129 | 130 | 131 | async def generate_awareness_question(): 132 | if not openai.api_key: 133 | raise Exception("OpenAI API key not found.") 134 | 135 | prompt = f""" 136 | You are a helpful cybersecurity AI analyst assistant to the security team that wants to keep 137 | your company secure. You have received an alert regarding the user you're chatting with, and 138 | you have asked whether the user was aware of the alert. The user has not answered the question, 139 | so now you are asking the user again whether they were aware of the alert. You ask in a gentle, 140 | kind, and casual tone. You keep it short, to two sentences at most. You end with a question. 141 | """ 142 | 143 | messages = [ 144 | {"role": "system", "content": prompt}, 145 | {"role": "user", "content": ""}, 146 | ] 147 | 148 | completion = openai.chat.completions.create( 149 | model="gpt-4-32k", 150 | messages=messages, 151 | temperature=0.5, 152 | stream=False, 153 | ) 154 | response = await get_clean_output(completion) 155 | return response 156 | -------------------------------------------------------------------------------- /bots/incident-response-slackbot/incident_response_slackbot/templates/messages/incident_alert.j2: -------------------------------------------------------------------------------- 1 | [ 2 | { 3 | "type": "section", 4 | "text": { 5 | "type": "mrkdwn", 6 | "text": "Incident: {{ alert_name }} with user <@{{ user_id }}>" 7 | } 8 | }, 9 | { 10 | "type": "context", 11 | "elements": [ 12 | { 13 | "type": "plain_text", 14 | "text": "Details in :thread:", 15 | "emoji": true 16 | } 17 | ] 18 | }, 19 | { 20 | "type": "actions", 21 | "elements": [ 22 | { 23 | "type": "button", 24 | "text": { 25 | "type": "plain_text", 26 | "emoji": true, 27 | "text": "Start Chat" 28 | }, 29 | "style": "primary", 30 | "value": "{{ user_id }}", 31 | "action_id": "start_chat_submit_action" 32 | }, 33 | { 34 | "type": "button", 35 | "text": { 36 | "type": "plain_text", 37 | "emoji": true, 38 | "text": "Do Nothing" 39 | }, 40 | "style": "danger", 41 | "value": "recategorize", 42 | "action_id": "do_nothing_submit_action" 43 | } 44 | ] 45 | } 46 | ] 47 | -------------------------------------------------------------------------------- /bots/incident-response-slackbot/pyproject.template.toml: -------------------------------------------------------------------------------- 1 | [project] 2 | name = "openai-incident-response-slackbot" 3 | requires-python = ">=3.8" 4 | version = "1.0.0" 5 | dependencies = [ 6 | "toml", 7 | "openai_slackbot @ file://$REPO_ROOT/shared/openai-slackbot", 8 | ] 9 | 10 | [build-system] 11 | requires = ["setuptools>=64.0"] 12 | build-backend = "setuptools.build_meta" 13 | 14 | [tool.pytest.ini_options] 15 | asyncio_mode = "auto" 16 | env = [ 17 | "SLACK_BOT_TOKEN=mock-token", 18 | "SOCKET_APP_TOKEN=mock-token", 19 | "OPENAI_API_KEY=mock-key", 20 | ] -------------------------------------------------------------------------------- /bots/incident-response-slackbot/scripts/alert_feed.py: -------------------------------------------------------------------------------- 1 | import os 2 | from logging import getLogger 3 | 4 | from incident_response_slackbot.config import load_config, get_config 5 | from incident_response_slackbot.db.database import Database 6 | from openai_slackbot.clients.slack import CreateSlackMessageResponse, SlackClient 7 | from openai_slackbot.utils.envvars import string 8 | from slack_bolt.app.async_app import AsyncApp 9 | 10 | logger = getLogger(__name__) 11 | 12 | DATABASE = Database() 13 | 14 | load_config() 15 | config = get_config() 16 | 17 | async def post_alert(alert): 18 | """ 19 | This function posts an alert to the Slack channel. 20 | It first initializes the Slack client with the bot token and template path. 21 | Then, it extracts the user_id, alert_name, and properties from the alert. 22 | Finally, it posts the alert to the Slack channel and sends the initial details. 23 | 24 | Args: 25 | alert (dict): The alert to be posted. It should contain 'user_id', 'name', and 'properties'. 26 | """ 27 | 28 | slack_bot_token = string("SLACK_BOT_TOKEN") 29 | app = AsyncApp(token=slack_bot_token) 30 | slack_template_path = os.path.join( 31 | os.path.dirname(os.path.abspath(__file__)), 32 | "../incident_response_slackbot/templates", 33 | ) 34 | slack_client = SlackClient(app.client, slack_template_path) 35 | 36 | # Extracting the user_id, alert_name, and properties from the alert 37 | user_id = alert.get("user_id") 38 | alert_name = alert.get("name") 39 | properties = alert.get("properties") 40 | 41 | message = await incident_feed_begin( 42 | slack_client=slack_client, user_id=user_id, alert_name=alert_name 43 | ) 44 | 45 | DATABASE.add(user_id, message.ts) 46 | 47 | await initial_details(slack_client=slack_client, message=message, properties=properties) 48 | 49 | 50 | async def incident_feed_begin( 51 | *, slack_client: SlackClient, user_id: str, alert_name: str 52 | ) -> CreateSlackMessageResponse: 53 | """ 54 | This function begins the incident feed by posting the initial alert message. 55 | It first renders the blocks from the template with the user_id and alert_name. 56 | Then, it posts the message to the Slack channel. 57 | 58 | Args: 59 | slack_client (SlackClient): The Slack client. 60 | user_id (str): The Slack user ID. 61 | alert_name (str): The name of the alert. 62 | 63 | Returns: 64 | CreateSlackMessageResponse: The response from creating the Slack message. 65 | 66 | Raises: 67 | Exception: If the initial alert message fails to post. 68 | """ 69 | 70 | try: 71 | blocks = slack_client.render_blocks_from_template( 72 | "messages/incident_alert.j2", 73 | { 74 | "user_id": user_id, 75 | "alert_name": alert_name, 76 | }, 77 | ) 78 | message = await slack_client.post_message( 79 | channel=config.feed_channel_id, 80 | blocks=blocks, 81 | text=f"{alert_name} via <@{user_id}>", 82 | ) 83 | return message 84 | 85 | except Exception: 86 | logger.exception("Initial alert feed message failed") 87 | 88 | 89 | def get_alert_details(**kwargs) -> str: 90 | """ 91 | This function returns the alert details for each key in the 92 | property. Each alert could have different properties. 93 | """ 94 | content = "" 95 | for key, value in kwargs.items(): 96 | line = f"The value of {key} for this alert is {value}. " 97 | content += line 98 | if content: 99 | return content 100 | return "No details available for this alert." 101 | 102 | 103 | async def initial_details(*, slack_client: SlackClient, message, properties): 104 | """ 105 | This function posts the initial details of an alert to a Slack thread. 106 | 107 | Args: 108 | slack_client (SlackClient): The Slack client. 109 | message: The initial alert message. 110 | properties: The properties of the alert. 111 | """ 112 | thread_ts = message.ts 113 | details = get_alert_details(**properties) 114 | 115 | await slack_client.post_message( 116 | channel=config.feed_channel_id, text=f"{details}", thread_ts=thread_ts 117 | ) 118 | -------------------------------------------------------------------------------- /bots/incident-response-slackbot/scripts/alerts.toml: -------------------------------------------------------------------------------- 1 | # Alert Examples - These are the alerts that will be sent to the feed channel. 2 | [[alerts]] 3 | id = "pivot" 4 | name = "Pivoting" 5 | description = "User was found pivoting from one host to another" 6 | user_id = "" 7 | 8 | [alerts.properties] 9 | source_host = "source.machine.org" 10 | destination_host = "destination.machine.org" 11 | 12 | [[alerts]] 13 | id = "privesc" 14 | name = "Privileged Escalation" 15 | description = "Privileged escalation was detected" 16 | user_id = "" 17 | 18 | [alerts.properties] 19 | previous_role = "reader" 20 | new_role = "admin" 21 | -------------------------------------------------------------------------------- /bots/incident-response-slackbot/scripts/send_alert.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | import os 3 | import random 4 | import time 5 | 6 | import toml 7 | from alert_feed import post_alert 8 | 9 | 10 | def load_alerts(): 11 | alerts_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "alerts.toml") 12 | with open(alerts_path, "r") as file: 13 | data = toml.load(file) 14 | return data 15 | 16 | 17 | def generate_random_alert(alerts): 18 | random_alert = random.choice([0, 1]) 19 | print(alerts["alerts"][random_alert]) 20 | return alerts["alerts"][random_alert] 21 | 22 | 23 | async def main(): 24 | alerts = load_alerts() 25 | 26 | alert = generate_random_alert(alerts) 27 | await post_alert(alert) 28 | 29 | 30 | if __name__ == "__main__": 31 | asyncio.run(main()) 32 | -------------------------------------------------------------------------------- /bots/incident-response-slackbot/tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/openai/openai-security-bots/20fb09044cae19a6b25a192ccb605d24d3743d88/bots/incident-response-slackbot/tests/__init__.py -------------------------------------------------------------------------------- /bots/incident-response-slackbot/tests/conftest.py: -------------------------------------------------------------------------------- 1 | import os 2 | from unittest.mock import AsyncMock, MagicMock, patch 3 | 4 | import pytest 5 | import toml 6 | from incident_response_slackbot.config import load_config 7 | from pydantic import ValidationError 8 | 9 | #################### 10 | ##### FIXTURES ##### 11 | #################### 12 | 13 | 14 | @pytest.fixture(autouse=True) 15 | def mock_config(): 16 | # Load the test config 17 | current_dir = os.path.dirname(os.path.abspath(__file__)) 18 | config_path = os.path.join(current_dir, "test_config.toml") 19 | try: 20 | config = load_config(config_path) 21 | except ValidationError as e: 22 | print(f"Error validating the config: {e}") 23 | raise 24 | return config 25 | 26 | 27 | @pytest.fixture() 28 | def mock_slack_client(): 29 | # Mock the Slack client 30 | slack_client = MagicMock() 31 | slack_client.post_message = AsyncMock() 32 | slack_client.update_message = AsyncMock() 33 | slack_client.get_original_blocks = AsyncMock() 34 | slack_client.get_thread_messages = AsyncMock() 35 | 36 | return slack_client 37 | 38 | 39 | @pytest.fixture(autouse=True) 40 | @patch("openai.ChatCompletion.create") 41 | def mock_chat_completion(mock_create): 42 | mock_create.return_value = { 43 | "id": "chatcmpl-1234567890", 44 | "object": "chat.completion", 45 | "created": 1640995200, 46 | "model": "gpt-4-32k", 47 | "usage": {"prompt_tokens": 10, "completion_tokens": 20, "total_tokens": 30}, 48 | "choices": [ 49 | { 50 | "message": { 51 | "role": "assistant", 52 | "content": "This is a mock response from the OpenAI API.", 53 | }, 54 | "finish_reason": "stop", 55 | "index": 0, 56 | } 57 | ], 58 | } 59 | yield 60 | 61 | 62 | @pytest.fixture 63 | def mock_generate_awareness_question(): 64 | with patch( 65 | "incident_response_slackbot.handlers.generate_awareness_question", 66 | new_callable=AsyncMock, 67 | ) as mock_generate_question: 68 | mock_generate_question.return_value = "Mock question" 69 | yield mock_generate_question 70 | 71 | 72 | @pytest.fixture 73 | def mock_get_thread_summary(): 74 | with patch( 75 | "incident_response_slackbot.handlers.get_thread_summary", 76 | new_callable=AsyncMock, 77 | ) as mock_get_summary: 78 | mock_get_summary.return_value = "Mock summary" 79 | yield mock_get_summary 80 | -------------------------------------------------------------------------------- /bots/incident-response-slackbot/tests/test_config.toml: -------------------------------------------------------------------------------- 1 | # Organization ID associated with OpenAI API key. 2 | openai_organization_id = "mock_openai_organization_id" 3 | 4 | # Where the alerts will be posted. 5 | feed_channel_id = "mock_feed_channel_id" 6 | 7 | -------------------------------------------------------------------------------- /bots/incident-response-slackbot/tests/test_handlers.py: -------------------------------------------------------------------------------- 1 | # in tests/test_handlers.py 2 | from unittest.mock import AsyncMock, MagicMock, patch 3 | from collections import namedtuple 4 | 5 | import pytest 6 | from incident_response_slackbot.handlers import ( 7 | InboundDirectMessageHandler, 8 | InboundIncidentStartChatHandler, 9 | InboundIncidentDoNothingHandler, 10 | InboundIncidentEndChatHandler, 11 | ) 12 | 13 | 14 | @pytest.mark.asyncio 15 | async def test_send_message_to_channel(mock_slack_client, mock_config): 16 | # Arrange 17 | handler = InboundDirectMessageHandler(slack_client=mock_slack_client) 18 | mock_event = {"text": "mock_event_text", "user_profile": {"name": "mock_user_name"}} 19 | mock_message_ts = "mock_message_ts" 20 | 21 | # Act 22 | await handler.send_message_to_channel(mock_event, mock_message_ts) 23 | 24 | # Assert 25 | mock_slack_client.post_message.assert_called_once_with( 26 | channel=mock_config.feed_channel_id, 27 | text="Received message from <@mock_user_name>:\n> mock_event_text", 28 | thread_ts=mock_message_ts, 29 | ) 30 | 31 | 32 | @pytest.mark.asyncio 33 | async def test_end_chat(mock_slack_client, mock_config): 34 | # Define the return value for get_original_blocks 35 | mock_slack_client.get_original_blocks.return_value = [ 36 | {"type": "section", "block_id": "block1"}, 37 | {"type": "actions", "block_id": "block2"}, 38 | {"type": "section", "block_id": "block3"}, 39 | ] 40 | 41 | # Create an instance of the handler 42 | handler = InboundDirectMessageHandler(slack_client=mock_slack_client) 43 | 44 | # Call the end_chat method 45 | await handler.end_chat("12345") 46 | 47 | # Assert that update_message was called with the correct arguments 48 | mock_slack_client.update_message.assert_called_once() 49 | 50 | # Get the actual call arguments 51 | args, kwargs = mock_slack_client.update_message.call_args 52 | 53 | # Check the blocks argument 54 | assert kwargs["blocks"] == [ 55 | {"type": "section", "block_id": "block1"}, 56 | {"type": "section", "block_id": "block3"}, 57 | { 58 | "type": "section", 59 | "block_id": "end_chat_automatically", 60 | "text": { 61 | "type": "mrkdwn", 62 | "text": "The chat was automatically ended from SecurityBot review. :done_:", 63 | "verbatim": True, 64 | }, 65 | }, 66 | ] 67 | 68 | 69 | @pytest.mark.asyncio 70 | async def test_nudge_user(mock_slack_client, mock_config, mock_generate_awareness_question): 71 | # Create an instance of the handler 72 | handler = InboundDirectMessageHandler(slack_client=mock_slack_client) 73 | 74 | # Call the nudge_user method 75 | await handler.nudge_user("user123", "12345") 76 | 77 | # Assert that post_message was called twice with the correct arguments 78 | assert mock_slack_client.post_message.call_count == 2 79 | mock_slack_client.post_message.assert_any_call(channel="user123", text="Mock question") 80 | mock_slack_client.post_message.assert_any_call( 81 | channel=handler.config.feed_channel_id, 82 | text="Sent message to <@user123>:\n> Mock question", 83 | thread_ts="12345", 84 | ) 85 | 86 | 87 | @pytest.mark.asyncio 88 | async def test_incident_start_chat_handle(mock_slack_client, mock_config): 89 | # Create an instance of the handler 90 | handler = InboundIncidentStartChatHandler(slack_client=mock_slack_client) 91 | 92 | # Create a mock args object 93 | args = MagicMock() 94 | args.body = { 95 | "container": {"message_ts": "12345"}, 96 | "user": {"name": "test_user", "id": "user123"}, 97 | } 98 | 99 | # Mock the DATABASE.get_user_id method 100 | with patch( 101 | "incident_response_slackbot.handlers.DATABASE.get_user_id", return_value="alert_user123" 102 | ) as mock_get_user_id, patch( 103 | "incident_response_slackbot.handlers.create_greeting", 104 | new_callable=AsyncMock, 105 | return_value="greeting message", 106 | ) as mock_create_greeting, patch.object( 107 | handler._slack_client, "get_thread_messages", new_callable=AsyncMock 108 | ), patch.object( 109 | handler._slack_client, "update_message", new_callable=AsyncMock 110 | ), patch.object( 111 | handler._slack_client, 112 | "get_user_display_name", 113 | new_callable=AsyncMock, 114 | return_value="username", 115 | ), patch.object( 116 | handler._slack_client, "post_message", new_callable=AsyncMock 117 | ): 118 | # Call the handle method 119 | await handler.handle(args) 120 | 121 | # Assert that the slack client methods were called with the correct arguments 122 | handler._slack_client.get_thread_messages.assert_called_once_with( 123 | channel=mock_config.feed_channel_id, thread_ts="12345" 124 | ) 125 | handler._slack_client.update_message.assert_called_once() 126 | handler._slack_client.get_user_display_name.assert_called_once_with("alert_user123") 127 | 128 | # Assert that post_message was called twice 129 | assert handler._slack_client.post_message.call_count == 2 130 | 131 | # Assert that post_message was called with the correct arguments 132 | handler._slack_client.post_message.assert_any_call( 133 | channel="alert_user123", text="greeting message" 134 | ) 135 | 136 | 137 | @pytest.mark.asyncio 138 | async def test_do_nothing_handle(mock_slack_client, mock_config): 139 | # Create an instance of the handler 140 | handler = InboundIncidentDoNothingHandler(slack_client=mock_slack_client) 141 | 142 | # Create a mock args object 143 | args = MagicMock() 144 | args.body = { 145 | "user": {"id": "user123"}, 146 | "message": {"ts": "12345", "blocks": [{"type": "actions"}, {"type": "section"}]}, 147 | } 148 | 149 | # Call the handle method 150 | await handler.handle(args) 151 | 152 | # Assert that the slack client update_message method was called with the correct arguments 153 | mock_slack_client.update_message.assert_called_once_with( 154 | channel=mock_config.feed_channel_id, 155 | blocks=[ 156 | {"type": "section"}, 157 | { 158 | "type": "section", 159 | "block_id": "do_nothing", 160 | "text": { 161 | "type": "mrkdwn", 162 | "text": "<@user123> decided that no action was necessary :done_:", 163 | "verbatim": True, 164 | }, 165 | }, 166 | ], 167 | ts="12345", 168 | text="Do Nothing action selected", 169 | ) 170 | 171 | 172 | @pytest.mark.asyncio 173 | async def test_end_chat_handle(mock_slack_client, mock_config, mock_get_thread_summary): 174 | # Mock the Slack client and the database 175 | with patch( 176 | "incident_response_slackbot.handlers.Database", new_callable=AsyncMock 177 | ) as MockDatabase: 178 | # Instantiate the handler 179 | handler = InboundIncidentEndChatHandler(slack_client=mock_slack_client) 180 | 181 | # Define a namedtuple for args 182 | Args = namedtuple("Args", ["body"]) 183 | 184 | # Instantiate the args object 185 | args = Args(body={"user": {"id": "user_id"}, "message": {"ts": "message_ts"}}) 186 | 187 | # Mock the get_user_id method of the database to return a user id 188 | MockDatabase.get_user_id.return_value = "alert_user_id" 189 | 190 | # Call the handle method 191 | await handler.handle(args) 192 | 193 | # Assert that the correct methods were called with the expected arguments 194 | mock_slack_client.get_original_blocks.assert_called_once_with( 195 | "message_ts", mock_config.feed_channel_id 196 | ) 197 | mock_slack_client.update_message.assert_called() 198 | mock_slack_client.post_message.assert_called() 199 | -------------------------------------------------------------------------------- /bots/incident-response-slackbot/tests/test_openai.py: -------------------------------------------------------------------------------- 1 | # in tests/test_openai_utils.py 2 | from unittest.mock import patch 3 | 4 | import pytest 5 | from incident_response_slackbot.openai_utils import get_user_awareness 6 | 7 | 8 | @pytest.mark.asyncio 9 | @patch("openai.ChatCompletion.create") 10 | async def test_get_user_awareness(mock_create): 11 | # Arrange 12 | mock_create.return_value = { 13 | "choices": [ 14 | { 15 | "message": { 16 | "function_call": {"arguments": '{"has_answered": true, "is_aware": false}'} 17 | } 18 | } 19 | ] 20 | } 21 | inbound_direct_message = "mock_inbound_direct_message" 22 | 23 | # Act 24 | result = await get_user_awareness(inbound_direct_message) 25 | 26 | # Assert 27 | assert result == {"has_answered": True, "is_aware": False} 28 | -------------------------------------------------------------------------------- /bots/sdlc-slackbot/Makefile: -------------------------------------------------------------------------------- 1 | CWD := $(shell pwd) 2 | REPO_ROOT := $(shell git rev-parse --show-toplevel) 3 | ESCAPED_REPO_ROOT := $(shell echo $(REPO_ROOT) | sed 's/\//\\\//g') 4 | 5 | init-env-file: 6 | cp ./sdlc_slackbot/.env.template ./sdlc_slackbot/.env 7 | 8 | init-pyproject: 9 | cat $(CWD)/pyproject.template.toml | \ 10 | sed "s/\$$REPO_ROOT/$(ESCAPED_REPO_ROOT)/g" > $(CWD)/pyproject.toml 11 | -------------------------------------------------------------------------------- /bots/sdlc-slackbot/README.md: -------------------------------------------------------------------------------- 1 | 2 |

3 | sdlc-slackbot-logo 4 |

SDLC Slackbot

5 |

6 | 7 | SDLC Slackbot decides if a project merits a security review. 8 | 9 | ## Prerequisites 10 | 11 | You will need: 12 | 1. A Slack application (aka your sdlc bot) with Socket Mode enabled 13 | 2. OpenAI API key 14 | 15 | Generate an App-level token for your Slack app, by going to: 16 | ``` 17 | Your Slack App > Basic Information > App-Level Tokens > Generate Token and Scopes 18 | ``` 19 | Create a new token with `connections:write` scope. This is your `SOCKET_APP_TOKEN` token. 20 | 21 | Once you have them, from the current directory, run: 22 | ``` 23 | $ make init-env-file 24 | ``` 25 | and fill in the right values. 26 | 27 | Your Slack App needs the following scopes: 28 | 29 | - app\_mentions:read 30 | - channels:join 31 | - channels:read 32 | - channels:history 33 | - chat:write 34 | - groups:history 35 | - groups:read 36 | - groups:write 37 | - usergroups:read 38 | - users:read 39 | - users:read.email 40 | 41 | 42 | ## Setup 43 | 44 | From the current directory, run: 45 | ``` 46 | make init-pyproject 47 | ``` 48 | 49 | From the repo root, run: 50 | ``` 51 | make clean-venv 52 | source venv/bin/activate 53 | make build-bot BOT=sdlc-slackbot 54 | ``` 55 | 56 | ## Run bot with example configuration 57 | 58 | The example configuration is `config.toml`. Replace the configuration values as needed. 59 | You need to at least replace the `openai_organization_id` and `notification_channel_id`. 60 | 61 | For optional Google Docs integration you'll need a 'credentials.json' file: 62 | - Go to the Google Cloud Console. 63 | - Select your project. 64 | - Navigate to "APIs & Services" > "Credentials". 65 | - Under "OAuth 2.0 Client IDs", find your client ID and download the JSON file. 66 | - Save it in the `sdlc-slackbot/sdlc_slackbot` directory as `credentials.json`. 67 | 68 | 69 | 70 | ⚠️ *Make sure that the bot is added to the channels it needs to read from and post to.* ⚠️ 71 | 72 | From the repo root, run: 73 | 74 | ``` 75 | make run-bot BOT=sdlc-slackbot 76 | ``` 77 | 78 | -------------------------------------------------------------------------------- /bots/sdlc-slackbot/pyproject.template.toml: -------------------------------------------------------------------------------- 1 | [project] 2 | name = "openai-sdlc-slackbot" 3 | requires-python = ">=3.8" 4 | version = "1.0.0" 5 | dependencies = [ 6 | "toml", 7 | "openai_slackbot @ file://$REPO_ROOT/shared/openai-slackbot", 8 | "validators", 9 | "google-auth", 10 | "google-auth-httplib2", 11 | "google-auth-oauthlib", 12 | "google-api-python-client", 13 | "psycopg", 14 | "psycopg2-binary", 15 | "peewee", 16 | ] 17 | 18 | [build-system] 19 | requires = ["setuptools>=64.0"] 20 | build-backend = "setuptools.build_meta" 21 | 22 | [tool.pytest.ini_options] 23 | asyncio_mode = "auto" 24 | env = [ 25 | "SLACK_BOT_TOKEN=mock-token", 26 | "SOCKET_APP_TOKEN=mock-token", 27 | "OPENAI_API_KEY=mock-key", 28 | ] 29 | -------------------------------------------------------------------------------- /bots/sdlc-slackbot/requirements.txt: -------------------------------------------------------------------------------- 1 | openai 2 | python-dotenv 3 | slack-bolt 4 | validators 5 | google-auth 6 | google-auth-httplib2 7 | google-auth-oauthlib 8 | google-api-python-client 9 | psycopg 10 | psycopg2-binary 11 | peewee 12 | -------------------------------------------------------------------------------- /bots/sdlc-slackbot/sdlc_slackbot/.env.template: -------------------------------------------------------------------------------- 1 | SLACK_BOT_TOKEN=xoxb- # Go to your Slack app, Settings > Install App > Bot User OAuth Token 2 | SOCKET_APP_TOKEN=xapp- # Go to your Slack app, Settings > Basic Information > App-Level Tokens 3 | OPENAI_API_KEY= 4 | -------------------------------------------------------------------------------- /bots/sdlc-slackbot/sdlc_slackbot/bot.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | import hashlib 3 | import json 4 | import os 5 | import re 6 | import threading 7 | import time 8 | import traceback 9 | from logging import getLogger 10 | 11 | import validate 12 | import validators 13 | from database import * 14 | from gdoc import gdoc_get 15 | from openai_slackbot.bot import init_bot, start_app 16 | from openai_slackbot.utils.envvars import string 17 | from peewee import * 18 | from playhouse.db_url import * 19 | from playhouse.shortcuts import model_to_dict 20 | from sdlc_slackbot.config import get_config, load_config 21 | from slack_bolt import App 22 | from slack_bolt.adapter.socket_mode import SocketModeHandler 23 | from slack_sdk import WebClient 24 | from utils import * 25 | 26 | 27 | logger = getLogger(__name__) 28 | 29 | 30 | async def send_update_notification(input, response): 31 | risk_str, confidence_str = risk_and_confidence_to_string(response) 32 | risk_num = response["risk"] 33 | confidence_num = response["confidence"] 34 | 35 | msg = f""" 36 | Project {input['project_name']} has been updated and has a new decision: 37 | 38 | This new decision for the project is that it is: *{risk_str}({risk_num})* with *{confidence_str}({confidence_num})*. {response['justification']}." 39 | """ 40 | 41 | await app.client.chat_postMessage(channel=config.notification_channel_id, text=msg) 42 | 43 | 44 | def hash_content(content): 45 | return hashlib.sha256(content.encode("utf-8")).hexdigest() 46 | 47 | 48 | url_pat = re.compile( 49 | r"http[s]?://(?:[a-zA-Z]|[0-9]|[$-_@.&+]|[!*\\(\\),]|(?:%[0-9a-fA-F][0-9a-fA-F]))+\b(?!>)" 50 | ) 51 | 52 | 53 | def extract_urls(text): 54 | logger.info(f"extracting urls from {text}") 55 | urls = re.findall(url_pat, text) 56 | return [url for url in urls if validators.url(url)] 57 | 58 | 59 | async def async_fetch_slack(url): 60 | parts = url.split("/") 61 | channel = parts[-2] 62 | ts = parts[-1] 63 | ts = ts[1:] # trim p 64 | seconds = ts[:-6] 65 | nanoseconds = ts[-6:] 66 | result = await app.client.conversations_replies(channel=channel, ts=f"{seconds}.{nanoseconds}") 67 | return " ".join(message.get("text", "") for message in result.data.get("messages", [])) 68 | 69 | 70 | content_fetchers = [ 71 | ( 72 | lambda u: u.startswith(("https://docs.google.com/document", "docs.google.com/document")), 73 | gdoc_get, 74 | ), 75 | (lambda u: "slack.com/archives" in u, async_fetch_slack), 76 | ] 77 | 78 | 79 | async def fetch_content(url): 80 | for condition, fetcher in content_fetchers: 81 | if condition(url): 82 | if asyncio.iscoroutinefunction(fetcher): 83 | return await fetcher(url) # Await the result if it's a coroutine function 84 | else: 85 | return fetcher(url) # Call it directly if it's not a coroutine function 86 | 87 | 88 | form = [ 89 | input_block( 90 | "project_name", 91 | "Project Name", 92 | field("plain_text_input", "Enter the project name"), 93 | ), 94 | input_block( 95 | "project_description", 96 | "Project Description", 97 | field("plain_text_input", "Enter the project description", multiline=True), 98 | ), 99 | input_block( 100 | "links_to_resources", 101 | "Links to Resources", 102 | field("plain_text_input", "Enter links to resources", multiline=True), 103 | ), 104 | input_block("point_of_contact", "Point of Contact", field("users_select", "Select a user")), 105 | input_block( 106 | "estimated_go_live_date", 107 | "Estimated Go Live Date", 108 | field("datepicker", "Select a date"), 109 | ), 110 | submit_block("submit_form"), 111 | ] 112 | 113 | 114 | def risk_and_confidence_to_string(decision): 115 | # Lookup tables for risk and confidence 116 | risk_lookup = { 117 | (1, 2): "extremely low risk", 118 | (3, 3): "low risk", 119 | (4, 5): "medium risk", 120 | (6, 7): "medium-high risk", 121 | (8, 9): "high risk", 122 | (10, 10): "critical risk", 123 | } 124 | 125 | confidence_lookup = { 126 | (1, 2): "extremely low confidence", 127 | (3, 3): "low confidence", 128 | (4, 5): "medium confidence", 129 | (6, 7): "medium-high confidence", 130 | (8, 9): "high confidence", 131 | (10, 10): "extreme confidence", 132 | } 133 | 134 | # Function to find the appropriate string from a lookup table 135 | def find_in_lookup(value, lookup): 136 | for (min_val, max_val), descriptor in lookup.items(): 137 | if min_val <= value <= max_val: 138 | return descriptor 139 | return "unknown" 140 | 141 | # Convert risk and confidence using their respective lookup tables 142 | risk_str = find_in_lookup(decision["risk"], risk_lookup) 143 | confidence_str = find_in_lookup(decision["confidence"], confidence_lookup) 144 | 145 | return risk_str, confidence_str 146 | 147 | 148 | def decision_msg(response): 149 | risk_str, confidence_str = risk_and_confidence_to_string(response) 150 | risk_num = response["risk"] 151 | confidence_num = response["confidence"] 152 | 153 | 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']}." 154 | 155 | 156 | skip_params = set( 157 | [ 158 | "id", 159 | "project_name", 160 | "links_to_resources", 161 | "point_of_contact", 162 | "estimated_go_live_date", 163 | ] 164 | ) 165 | 166 | multiple_whitespace_pat = re.compile(r"\s+") 167 | 168 | 169 | def model_params_to_str(params): 170 | ss = (v for k, v in params.items() if k not in skip_params) 171 | return re.sub(multiple_whitespace_pat, " ", "\n".join(map(str, ss))).strip() 172 | 173 | 174 | def summarize_params(params): 175 | summary = {} 176 | for k, v in params.items(): 177 | if k not in skip_params: 178 | summary[k] = ask_ai( 179 | config.base_prompt + config.summary_prompt, v[: config.context_limit] 180 | ) 181 | else: 182 | summary[k] = v 183 | 184 | return summary 185 | 186 | 187 | async def handle_app_mention_events(say, event): 188 | logger.info("App mention event received:", event) 189 | await say(blocks=form, thread_ts=event["ts"]) 190 | 191 | 192 | async def handle_message_events(say, message): 193 | logger.info("message: ", message) 194 | if message["channel_type"] == "im": 195 | await say(blocks=form, thread_ts=message["ts"]) 196 | 197 | 198 | def get_response_with_retry(prompt, context, max_retries=1): 199 | prompt = prompt.strip().replace("\n", " ") 200 | retries = 0 201 | while retries <= max_retries: 202 | try: 203 | response = ask_ai(prompt, context) 204 | return response 205 | except json.JSONDecodeError as e: 206 | logger.error(f"JSON error on attempt {retries + 1}: {e}") 207 | retries += 1 208 | if retries > max_retries: 209 | return {} 210 | 211 | 212 | def normalize_response(response): 213 | if isinstance(response, list): 214 | return [json.loads(block.text) for block in response] 215 | elif isinstance(response, dict): 216 | return [response] 217 | else: 218 | raise TypeError("Unsupported response type") 219 | 220 | 221 | def clean_normalized_response(normalized_responses): 222 | """ 223 | Remove the 'decision' key from each dictionary in a list of dictionaries. 224 | Break it down into 'risk' and 'confidence' 225 | 226 | :param normalized_responses: A list of dictionaries. 227 | :return: The list of dictionaries with 'decision' key broken down. 228 | """ 229 | for response in normalized_responses: 230 | if "decision" in response: 231 | decision = response["decision"] 232 | response["risk"] = decision.get("risk") 233 | response["confidence"] = decision.get("confidence") 234 | response.pop("decision", None) 235 | 236 | return normalized_responses 237 | 238 | 239 | async def submit_form(ack, body, say): 240 | await ack() 241 | 242 | try: 243 | ts = body["container"]["message_ts"] 244 | values = body["state"]["values"] 245 | params = get_form_input( 246 | values, 247 | "project_name", 248 | "project_description", 249 | "links_to_resources", 250 | "point_of_contact", 251 | "estimated_go_live_date", 252 | ) 253 | 254 | validate.required(params, "project_name", "project_description", "point_of_contact") 255 | 256 | await say(text=config.reviewing_message, thread_ts=ts) 257 | 258 | try: 259 | assessment = Assessment.create(**params, user_id=body["user"]["id"]) 260 | except IntegrityError as e: 261 | raise validate.ValidationError("project_name", "must be unique") 262 | 263 | resources = [] 264 | for url in extract_urls(params.get("links_to_resources", "")): 265 | content = await fetch_content(url) 266 | if content: 267 | params[url] = content 268 | resources.append( 269 | dict( 270 | assessment=assessment, 271 | url=url, 272 | content_hash=hash_content(content), 273 | ) 274 | ) 275 | Resource.insert_many(resources).execute() 276 | 277 | context = model_params_to_str(params) 278 | if len(context) > config.context_limit: 279 | logger.info(f"context too long: {len(context)}. Summarizing...") 280 | summarized_context = summarize_params(params) 281 | context = model_params_to_str(summarized_context) 282 | # FIXME: is there a better way to handle this? currently, if the summary is still too long 283 | # we just give up and cut it off 284 | if len(context) > config.context_limit: 285 | logger.info(f"Summarized context too long: {len(context)}. Cutting off...") 286 | context = context[: config.context_limit] 287 | 288 | response = get_response_with_retry(config.base_prompt + config.initial_prompt, context) 289 | if not response: 290 | return 291 | 292 | normalized_response = normalize_response(response) 293 | clean_response = clean_normalized_response(normalized_response) 294 | 295 | for item in clean_response: 296 | if item["outcome"] == "decision": 297 | assessment.update(**item).execute() 298 | await say(text=decision_msg(item), thread_ts=ts) 299 | elif item["outcome"] == "followup": 300 | db_questions = [dict(assessment=assessment, question=q) for q in item["questions"]] 301 | Question.insert_many(db_questions).execute() 302 | 303 | form = [] 304 | for i, q in enumerate(item["questions"]): 305 | form.append( 306 | input_block( 307 | f"question_{i}", 308 | q, 309 | field("plain_text_input", "...", multiline=True), 310 | ) 311 | ) 312 | form.append(submit_block(f"submit_followup_questions_{assessment.id}")) 313 | 314 | await say(blocks=form, thread_ts=ts) 315 | except validate.ValidationError as e: 316 | await say(text=f"{e.field}: {e.issue}", thread_ts=ts) 317 | except Exception as e: 318 | import traceback 319 | 320 | traceback.print_exc() 321 | await say(text=config.irrecoverable_error_message, thread_ts=ts) 322 | 323 | 324 | async def submit_followup_questions(ack, body, say): 325 | await ack() 326 | 327 | try: 328 | assessment_id = int(body["actions"][0]["action_id"].split("_")[-1]) 329 | ts = body["container"]["message_ts"] 330 | assessment = Assessment.get(Assessment.id == assessment_id) 331 | params = model_to_dict(assessment) 332 | followup_questions = [q.question for q in assessment.questions] 333 | except Exception as e: 334 | logger.error(f"Failed to find params for user {body['user']['id']}", e) 335 | await say(text=config.recoverable_error_message, thread_ts=ts) 336 | return 337 | 338 | try: 339 | await say(text=config.reviewing_message, thread_ts=ts) 340 | 341 | values = body["state"]["values"] 342 | for i, q in enumerate(followup_questions): 343 | params[q] = values[f"question_{i}"][f"question_{i}_input"]["value"] 344 | 345 | for question in assessment.questions: 346 | question.answer = params[question.question] 347 | question.save() 348 | 349 | context = model_params_to_str(params) 350 | 351 | response = ask_ai(config.base_prompt, context) 352 | text_to_update = response 353 | if ( 354 | isinstance(response, dict) 355 | and "text" in response 356 | and "type" in response 357 | and response["type"] == "text" 358 | ): 359 | # Extract the text from the content block 360 | text_to_update = response.text 361 | 362 | normalized_response = normalize_response(text_to_update) 363 | clean_response = clean_normalized_response(normalized_response) 364 | 365 | for item in clean_response: 366 | if item["outcome"] == "decision": 367 | assessment.update(**item).execute() 368 | await say(text=decision_msg(item), thread_ts=ts) 369 | 370 | except Exception as e: 371 | logger.error(f"error: {e} processing followup questions: {json.dumps(body, indent=2)}") 372 | await say(text=config.irrecoverable_error_message, thread_ts=ts) 373 | 374 | 375 | def update_resources(): 376 | while True: 377 | time.sleep(monitor_thread_sleep_seconds) 378 | try: 379 | for assessment in Assessment.select(): 380 | logger.info(f"checking {assessment.project_name} for updates") 381 | 382 | assessment_params = model_to_dict(assessment) 383 | new_params = assessment_params.copy() 384 | 385 | changed = False 386 | 387 | previous_content = "" 388 | 389 | for resource in assessment.resources: 390 | new_content = asyncio.run(fetch_content(resource.url)) 391 | 392 | if resource.content_hash != hash_content(new_content): 393 | # just save previous content in memory temporarily 394 | previous_content = resource.content 395 | resource.content = new_content 396 | new_params[resource.url] = new_content 397 | changed = True 398 | 399 | if not changed: 400 | continue 401 | 402 | old_context = model_params_to_str(assessment_params) 403 | new_context = model_params_to_str(new_params) 404 | 405 | context = { 406 | "previous_context": previous_content, 407 | "previous_decision": { 408 | "risk": assessment.risk, 409 | "confidence": assessment.confidence, 410 | "justification": assessment.justification, 411 | }, 412 | "new_context": new_content, 413 | } 414 | 415 | context_json = json.dumps(context, indent=2) 416 | 417 | new_response = ask_ai(config.base_prompt + config.update_prompt, context_json) 418 | 419 | resource.content_hash = hash_content(new_content) 420 | resource.save() 421 | 422 | if new_response["outcome"] == "unchanged": 423 | continue 424 | 425 | normalized_response = normalize_response(new_response) 426 | clean_response = clean_normalized_response(normalized_response) 427 | 428 | for item in clean_response: 429 | assessment.update(**item).execute() 430 | 431 | asyncio.run(send_update_notification(assessment_params, new_response)) 432 | except Exception as e: 433 | logger.error(f"error: {e} updating resources") 434 | traceback.print_exc() 435 | 436 | 437 | monitor_thread_sleep_seconds = 6 438 | 439 | if __name__ == "__main__": 440 | current_dir = os.path.dirname(os.path.abspath(__file__)) 441 | load_config(os.path.join(current_dir, "config.toml")) 442 | 443 | template_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "templates") 444 | 445 | config = get_config() 446 | 447 | message_handler = [] 448 | action_handlers = [] 449 | view_submission_handlers = [] 450 | 451 | app = asyncio.run( 452 | init_bot( 453 | openai_organization_id=config.openai_organization_id, 454 | slack_message_handler=message_handler, 455 | slack_action_handlers=action_handlers, 456 | slack_template_path=template_path, 457 | ) 458 | ) 459 | 460 | # Register your custom event handlers 461 | app.event("app_mention")(handle_app_mention_events) 462 | app.message()(handle_message_events) 463 | 464 | app.action("submit_form")(submit_form) 465 | app.action(re.compile("submit_followup_questions.*"))(submit_followup_questions) 466 | 467 | t = threading.Thread(target=update_resources) 468 | t.start() 469 | 470 | # Start the app 471 | asyncio.run(start_app(app)) 472 | -------------------------------------------------------------------------------- /bots/sdlc-slackbot/sdlc_slackbot/config.py: -------------------------------------------------------------------------------- 1 | import os 2 | import typing as t 3 | 4 | import toml 5 | from dotenv import load_dotenv 6 | from pydantic import BaseModel, ValidationError, field_validator, model_validator 7 | from pydantic.functional_validators import AfterValidator, BeforeValidator 8 | 9 | _CONFIG = None 10 | 11 | 12 | def validate_channel(channel_id: str) -> str: 13 | if not channel_id.startswith("C"): 14 | raise ValueError("channel ID must start with 'C'") 15 | return channel_id 16 | 17 | 18 | class Config(BaseModel): 19 | # OpenAI organization ID associated with OpenAI API key. 20 | openai_organization_id: str 21 | 22 | context_limit: int 23 | 24 | # OpenAI prompts 25 | base_prompt: str 26 | initial_prompt: str 27 | update_prompt: str 28 | summary_prompt: str 29 | 30 | reviewing_message: str 31 | recoverable_error_message: str 32 | irrecoverable_error_message: str 33 | 34 | # Slack channel for notifications 35 | notification_channel_id: t.Annotated[str, AfterValidator(validate_channel)] 36 | 37 | 38 | def load_config(path: str): 39 | load_dotenv() 40 | 41 | with open(path) as f: 42 | cfg = toml.loads(f.read()) 43 | config = Config(**cfg) 44 | 45 | global _CONFIG 46 | _CONFIG = config 47 | return _CONFIG 48 | 49 | 50 | def get_config() -> Config: 51 | global _CONFIG 52 | if _CONFIG is None: 53 | raise Exception("config not initialized, call load_config() first") 54 | return _CONFIG 55 | -------------------------------------------------------------------------------- /bots/sdlc-slackbot/sdlc_slackbot/config.toml: -------------------------------------------------------------------------------- 1 | # Organization ID associated with OpenAI API key. 2 | openai_organization_id = "" 3 | 4 | notification_channel_id = "" 5 | 6 | context_limit = 31_500 7 | 8 | base_prompt = """ 9 | You'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. 10 | You 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 11 | the 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. 12 | Your 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. 13 | You should base your decision on a variety of factors including but not limited to: 14 | - if changes would affect any path to model weights or customer data 15 | - if changes are accessible from the internet 16 | - if changes affect end users 17 | - if changes affect security critical parts of the system, like authentication, authorization, encryption 18 | - if changes deal with historically risky technology like xml parsing 19 | - if changes will likely involve interpolating user input into a dynamic language like html, sql, or javascript 20 | 21 | If 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. 22 | 23 | Be 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. 24 | You've been asked to analyze a new project that is being developed by another team at your company 25 | and determine if and when it should be reviewed by your team. Your decision option should be two numeric scores: 26 | One 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. 27 | The second score is your confidence: how confident are you about your decision, with 1 meaning very low confidence, while 10 meaning super confident. 28 | Put both number in the "decision" as follows: 29 | 30 | decision: { "risk": 31 | "confidence": 32 | 33 | You should base your decision on how risky you think the project is to the company. 34 | You should also provide a brief justification for your decision. You should only respond with a json object. 35 | The 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..."}. 36 | 37 | Don't send any other responses. Our team has very limited resources and only wants to review the most important projects, so you 38 | should enforce a high bar for go live reviews. 39 | """ 40 | 41 | 42 | initial_prompt = """ 43 | You should ask as many questions as you need to make an informed, accurate decision. Don't hesitate at all to ask followup questions. 44 | Ask for clarification for any critical vague language in the fields below. If the project description doesn't contain information about 45 | factors that are critical to your decision, ask about them. 46 | If you need to ask a followup question, respond with {"outcome": "followup", "questions": ["What is the project's budget?", "What is the project's timeline?"]}. 47 | """ 48 | 49 | update_prompt = """ 50 | You've already reviewed this project before, but some information has changed. Below you'll find the previous project context 51 | your previous decision, a justification for your previous decision and the new content. If your decision still makes sense 52 | respond with a json object with a single property named "outcome" set to "unchanged". If your decision no longer makes sense 53 | respond 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. 54 | """ 55 | 56 | summary_prompt = """ 57 | You'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. 58 | You 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 59 | the greatest amount of security risk with your limited resources. 60 | Please 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" 61 | """ 62 | 63 | reviewing_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" 64 | 65 | recoverable_error_message = "Something went wrong. We've been notified and will fix it as soon as possible. Start a new conversation to try again" 66 | 67 | irrecoverable_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." 68 | -------------------------------------------------------------------------------- /bots/sdlc-slackbot/sdlc_slackbot/database.py: -------------------------------------------------------------------------------- 1 | import os 2 | 3 | from peewee import * 4 | from playhouse.db_url import * 5 | 6 | db_url = os.getenv("DATABASE_URL") or "postgres://postgres:postgres@localhost:5432/postgres" 7 | db = connect(db_url) 8 | 9 | 10 | class BaseModel(Model): 11 | class Meta: 12 | database = db 13 | 14 | 15 | class Assessment(BaseModel): 16 | project_name = CharField(unique=True) 17 | project_description = TextField() 18 | links_to_resources = TextField(null=True) 19 | point_of_contact = CharField() 20 | estimated_go_live_date = CharField(null=True) 21 | outcome = CharField(null=True) 22 | risk = IntegerField(null=True) # Storing risk as an integer 23 | confidence = IntegerField(null=True) # Storing confidence as an integer 24 | justification = TextField(null=True) 25 | 26 | 27 | class Question(Model): 28 | question = TextField() 29 | answer = TextField(null=True) 30 | assessment = ForeignKeyField(Assessment, backref="questions") 31 | 32 | class Meta: 33 | database = db 34 | indexes = ((("question", "assessment"), True),) 35 | 36 | 37 | class Resource(BaseModel): 38 | url = TextField() 39 | content_hash = CharField() 40 | content = TextField(null=True) 41 | assessment = ForeignKeyField(Assessment, backref="resources") 42 | 43 | 44 | db.connect() 45 | db.create_tables([Assessment, Question, Resource]) 46 | -------------------------------------------------------------------------------- /bots/sdlc-slackbot/sdlc_slackbot/gdoc.py: -------------------------------------------------------------------------------- 1 | from __future__ import print_function 2 | 3 | import os.path 4 | import re 5 | from logging import getLogger 6 | 7 | from google.auth.transport.requests import Request 8 | from google.oauth2.credentials import Credentials 9 | from google_auth_oauthlib.flow import InstalledAppFlow 10 | from googleapiclient.discovery import build 11 | from googleapiclient.errors import HttpError 12 | 13 | # If modifying these scopes, delete the file token.json. 14 | SCOPES = ["https://www.googleapis.com/auth/documents.readonly"] 15 | 16 | logger = getLogger(__name__) 17 | 18 | 19 | def read_paragraph_element(element): 20 | """Returns the text in the given ParagraphElement. 21 | 22 | Args: 23 | element: a ParagraphElement from a Google Doc. 24 | """ 25 | text_run = element.get("textRun") 26 | if not text_run: 27 | return "" 28 | return text_run.get("content") 29 | 30 | 31 | def read_structural_elements(elements): 32 | """Recurses through a list of Structural Elements to read a document's text where text may be 33 | in nested elements. 34 | 35 | Args: 36 | elements: a list of Structural Elements. 37 | """ 38 | text = "" 39 | for value in elements: 40 | if "paragraph" in value: 41 | elements = value.get("paragraph").get("elements") 42 | for elem in elements: 43 | text += read_paragraph_element(elem) 44 | elif "table" in value: 45 | # The text in table cells are in nested Structural Elements and tables may be 46 | # nested. 47 | table = value.get("table") 48 | for row in table.get("tableRows"): 49 | cells = row.get("tableCells") 50 | for cell in cells: 51 | text += read_structural_elements(cell.get("content")) 52 | elif "tableOfContents" in value: 53 | # The text in the TOC is also in a Structural Element. 54 | toc = value.get("tableOfContents") 55 | text += read_structural_elements(toc.get("content")) 56 | return text 57 | 58 | 59 | def gdoc_creds(): 60 | creds = None 61 | # The file token.json stores the user's access and refresh tokens, and is 62 | # created automatically when the authorization flow completes for the first 63 | # time. 64 | creds_path = "./bots/sdlc-slackbot/sdlc_slackbot/" 65 | 66 | if os.path.exists(creds_path + "token.json"): 67 | creds = Credentials.from_authorized_user_file(creds_path + "token.json", SCOPES) 68 | 69 | # If there are no (valid) credentials available, let the user log in. 70 | if not creds or not creds.valid: 71 | if creds and creds.expired and creds.refresh_token: 72 | creds.refresh(Request()) 73 | else: 74 | flow = InstalledAppFlow.from_client_secrets_file( 75 | creds_path + "credentials.json", SCOPES 76 | ) 77 | creds = flow.run_local_server(port=0) 78 | # Save the credentials for the next run 79 | with open(creds_path + "token.json", "w") as token: 80 | token.write(creds.to_json()) 81 | 82 | return creds 83 | 84 | 85 | def gdoc_get(gdoc_url): 86 | # https://docs.google.com/document/d//edit 87 | 88 | result = None 89 | logger.info(gdoc_url) 90 | if not gdoc_url.startswith("https://docs.google.com/document") and not gdoc_url.startswith( 91 | "docs.google.com/document" 92 | ): 93 | logger.error("Invalid google doc url") 94 | return result 95 | 96 | # This regex captures the ID after "/d/" and before an optional "/edit", "/" or the end of the string. 97 | pattern = r"/d/([^/]+)" 98 | match = re.search(pattern, gdoc_url) 99 | 100 | if match: 101 | document_id = match.group(1) 102 | logger.info(document_id) 103 | else: 104 | logger.error("No ID found in the URL") 105 | return result 106 | 107 | creds = gdoc_creds() 108 | try: 109 | service = build("docs", "v1", credentials=creds) 110 | 111 | # Retrieve the documents contents from the Docs service. 112 | document = service.documents().get(documentId=document_id).execute() 113 | 114 | logger.info("The title of the document is: {}".format(document.get("title"))) 115 | 116 | doc_content = document.get("body").get("content") 117 | result = read_structural_elements(doc_content) 118 | 119 | except HttpError as err: 120 | logger.error(err) 121 | 122 | return result 123 | -------------------------------------------------------------------------------- /bots/sdlc-slackbot/sdlc_slackbot/utils.py: -------------------------------------------------------------------------------- 1 | import json 2 | import os 3 | from logging import getLogger 4 | 5 | # import anthropic 6 | import openai 7 | 8 | logger = getLogger(__name__) 9 | 10 | 11 | def get_form_input(values, *fields): 12 | ret = {} 13 | for f in fields: 14 | container = values[f][f + "_input"] 15 | value = container.get("value") 16 | if value: 17 | ret[f] = container["value"] 18 | else: 19 | for key, item in container.items(): 20 | if key.startswith("selected_") and item: 21 | ret[f] = item 22 | break 23 | return ret 24 | 25 | 26 | def plain_text(text): 27 | return dict(type="plain_text", text=text) 28 | 29 | 30 | def field(type, placeholder, **kwargs): 31 | return dict(type=type, placeholder=plain_text(placeholder), **kwargs) 32 | 33 | 34 | def input_block(block_id, label, element): 35 | if "action_id" not in element: 36 | element["action_id"] = block_id + "_input" 37 | 38 | return dict( 39 | type="input", 40 | block_id=block_id, 41 | label=plain_text(label), 42 | element=element, 43 | ) 44 | 45 | 46 | def submit_block(action_id): 47 | return dict( 48 | type="actions", 49 | elements=[ 50 | dict( 51 | type="button", 52 | text=plain_text("Submit"), 53 | action_id=action_id, 54 | style="primary", 55 | ) 56 | ], 57 | ) 58 | 59 | 60 | def ask_ai(prompt, context): 61 | # return ask_claude(prompt, context) # YOU CAN USE CLAUDE HERE 62 | response = ask_gpt(prompt, context) 63 | 64 | # Removing leading and trailing backticks and whitespace 65 | clean_response = response.strip("`\n ") 66 | 67 | # Check if 'json' is at the beginning and remove it 68 | if clean_response.lower().startswith("json"): 69 | clean_response = clean_response[4:].strip() 70 | 71 | # Remove a trailing } if it exists 72 | if clean_response.endswith("}}"): 73 | clean_response = clean_response[:-1] # Remove the last character 74 | 75 | logger.info(clean_response) 76 | 77 | try: 78 | parsed_response = json.loads(clean_response) 79 | return parsed_response 80 | except json.JSONDecodeError as e: 81 | logger.error(f"Failed to parse JSON response from ask_gpt: {response}\nError: {e}") 82 | return None 83 | 84 | 85 | def ask_gpt(prompt, context): 86 | response = openai.chat.completions.create( 87 | model="gpt-4-32k", 88 | messages=[ 89 | {"role": "system", "content": prompt}, 90 | {"role": "user", "content": context}, 91 | ], 92 | ) 93 | return response.choices[0].message.content 94 | 95 | 96 | def ask_claude(prompt, context): 97 | client = anthropic.Anthropic(api_key=os.environ["CLAUDE_API_KEY"]) 98 | message = client.messages.create( 99 | model="claude-3-opus-20240229", 100 | max_tokens=4096, 101 | messages=[{"role": "user", "content": prompt}], 102 | ) 103 | return message.content 104 | -------------------------------------------------------------------------------- /bots/sdlc-slackbot/sdlc_slackbot/validate.py: -------------------------------------------------------------------------------- 1 | class ValidationError(Exception): 2 | def __init__(self, field, issue): 3 | self.field = field 4 | self.issue = issue 5 | super().__init__(f"{field} {issue}") 6 | 7 | 8 | def required(values, *fields): 9 | for f in fields: 10 | if f not in values: 11 | raise ValidationError(f, "required") 12 | if values[f] == "": 13 | raise ValidationError(f, "required") 14 | -------------------------------------------------------------------------------- /bots/sdlc-slackbot/setup.py: -------------------------------------------------------------------------------- 1 | from setuptools import setup, find_packages 2 | 3 | with open("requirements.txt", "r") as f: 4 | requirements = f.read().splitlines() 5 | 6 | setup( 7 | name="sdlc_bot", 8 | version="0.1", 9 | packages=find_packages(), 10 | install_requires=requirements, 11 | ) 12 | -------------------------------------------------------------------------------- /bots/triage-slackbot/Makefile: -------------------------------------------------------------------------------- 1 | CWD := $(shell pwd) 2 | REPO_ROOT := $(shell git rev-parse --show-toplevel) 3 | ESCAPED_REPO_ROOT := $(shell echo $(REPO_ROOT) | sed 's/\//\\\//g') 4 | 5 | init-env-file: 6 | cp ./triage_slackbot/.env.template ./triage_slackbot/.env 7 | 8 | init-pyproject: 9 | cat $(CWD)/pyproject.template.toml | \ 10 | sed "s/\$$REPO_ROOT/$(ESCAPED_REPO_ROOT)/g" > $(CWD)/pyproject.toml 11 | -------------------------------------------------------------------------------- /bots/triage-slackbot/README.md: -------------------------------------------------------------------------------- 1 |

2 | triage-slackbot-logo 3 |

Triage Slackbot

4 |

5 | 6 | Triage Slackbot triages inbound requests in a Slack channel to different sub-teams within your organization. 7 | 8 | ## Prerequisites 9 | 10 | You will need: 11 | 1. A Slack application (aka your triage bot) with Socket Mode enabled 12 | 2. OpenAI API key 13 | 14 | Generate an App-level token for your Slack app, by going to: 15 | ``` 16 | Your Slack App > Basic Information > App-Level Tokens > Generate Token and Scopes 17 | ``` 18 | Create a new token with `connections:write` scope. This is your `SOCKET_APP_TOKEN` token. 19 | 20 | Once you have them, from the current directory, run: 21 | ``` 22 | $ make init-env-file 23 | ``` 24 | and fill in the right values. 25 | 26 | Your Slack App needs the following scopes: 27 | 28 | - channels:history 29 | - chat:write 30 | - groups:history 31 | - reactions:read 32 | - reactions:write 33 | 34 | ## Setup 35 | 36 | From the current directory, run: 37 | ``` 38 | make init-pyproject 39 | ``` 40 | 41 | From the repo root, run: 42 | ``` 43 | make clean-venv 44 | source venv/bin/activate 45 | make build-bot BOT=triage-slackbot 46 | ``` 47 | 48 | ## Run bot with example configuration 49 | 50 | The example configuration is `config.toml`. Replace the configuration values as needed. 51 | 52 | ⚠️ *Make sure that the bot is added to the channels it needs to read from and post to.* ⚠️ 53 | 54 | From the repo root, run: 55 | 56 | ``` 57 | make run-bot BOT=triage-slackbot 58 | ``` 59 | 60 | ## Demo 61 | 62 | This demo is run with the provided `config.toml`. In this demo: 63 | 64 | ``` 65 | inbound_request_channel_id = ID of #inbound-security-requests channel 66 | feed_channel_id = ID of #inbound-security-requests-feed channel 67 | 68 | [[ categories ]] 69 | key = "appsec" 70 | ... 71 | oncall_slack_id = ID of #appsec-requests channel 72 | 73 | [[ categories ]] 74 | key = "privacy" 75 | ... 76 | oncall_slack_id = ID of @tiffany user 77 | ``` 78 | 79 | The following triage scenarios are supported: 80 | 81 | First, the bot categorizes the inbound requests accurately, and on-call acknowledges this prediction. 82 | 83 | https://github.com/openai/openai-security-bots/assets/10287796/2bb8b301-41b6-450f-a578-482e89a75050 84 | 85 | Secondly, the bot categorizes the request into a category that it can autorespond to, e.g. Physical Security, 86 | and there is no manual action from on-call required. 87 | 88 | https://github.com/openai/openai-security-bots/assets/10287796/e77bacf0-e16d-4ed3-9567-6f3caaab02ad 89 | 90 | Finally, on-call can re-route an inbound request to another category's on-call if the initial predicted 91 | category is not accurate. Additionally, if `other_category_enabled` is set to true, on-call can select any 92 | channels it can route the user to: 93 | 94 | https://github.com/openai/openai-security-bots/assets/10287796/04247a29-f904-42bc-82d8-12b7f2b7e170 95 | 96 | The bot will reply to the thread with this: 97 | 98 | autorespond 99 | -------------------------------------------------------------------------------- /bots/triage-slackbot/pyproject.template.toml: -------------------------------------------------------------------------------- 1 | [project] 2 | name = "openai-triage-slackbot" 3 | requires-python = ">=3.8" 4 | version = "1.0.0" 5 | dependencies = [ 6 | "toml", 7 | "openai_slackbot @ file://$REPO_ROOT/shared/openai-slackbot", 8 | ] 9 | 10 | [build-system] 11 | requires = ["setuptools>=64.0"] 12 | build-backend = "setuptools.build_meta" 13 | 14 | [tool.pytest.ini_options] 15 | asyncio_mode = "auto" 16 | env = [ 17 | "SLACK_BOT_TOKEN=mock-token", 18 | "SOCKET_APP_TOKEN=mock-token", 19 | "OPENAI_API_KEY=mock-key", 20 | ] -------------------------------------------------------------------------------- /bots/triage-slackbot/tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/openai/openai-security-bots/20fb09044cae19a6b25a192ccb605d24d3743d88/bots/triage-slackbot/tests/__init__.py -------------------------------------------------------------------------------- /bots/triage-slackbot/tests/conftest.py: -------------------------------------------------------------------------------- 1 | import os 2 | from unittest.mock import AsyncMock, MagicMock, patch 3 | 4 | import pytest 5 | from openai_slackbot.clients.slack import SlackClient 6 | from triage_slackbot.category import RequestCategory 7 | from triage_slackbot.config import load_config 8 | from triage_slackbot.handlers import MessageTemplatePath 9 | 10 | ########################## 11 | ##### HELPER METHODS ##### 12 | ########################## 13 | 14 | 15 | def bot_message_extra_data(): 16 | return { 17 | "bot_id": "bot_id", 18 | "bot_profile": {"id": "bot_profile_id"}, 19 | "team": "team", 20 | "type": "type", 21 | "user": "user", 22 | } 23 | 24 | 25 | def recategorize_message_data( 26 | *, 27 | ts, 28 | channel_id, 29 | user, 30 | message, 31 | category, 32 | conversation=None, 33 | ): 34 | return MagicMock( 35 | ack=AsyncMock(), 36 | body={ 37 | "actions": ["recategorize_submit_action"], 38 | "container": { 39 | "message_ts": ts, 40 | "channel_id": channel_id, 41 | }, 42 | "user": user, 43 | "message": message, 44 | "state": { 45 | "values": { 46 | "recategorize_select_category_block": { 47 | "recategorize_select_category_action": { 48 | "selected_option": { 49 | "value": category, 50 | } 51 | } 52 | }, 53 | "recategorize_select_conversation_block": { 54 | "recategorize_select_conversation_action": { 55 | "selected_conversation": conversation, 56 | } 57 | }, 58 | } 59 | }, 60 | }, 61 | ) 62 | 63 | 64 | #################### 65 | ##### FIXTURES ##### 66 | #################### 67 | 68 | 69 | @pytest.fixture(autouse=True) 70 | def mock_config(): 71 | current_dir = os.path.dirname(os.path.abspath(__file__)) 72 | config_path = os.path.join(current_dir, "test_config.toml") 73 | return load_config(config_path) 74 | 75 | 76 | @pytest.fixture() 77 | def mock_post_message_response(): 78 | return AsyncMock( 79 | return_value=MagicMock( 80 | ok=True, 81 | data={ 82 | "ok": True, 83 | "channel": "", 84 | "ts": "", 85 | "message": { 86 | "blocks": [], 87 | "text": "", 88 | "ts": "", 89 | **bot_message_extra_data(), 90 | }, 91 | }, 92 | ) 93 | ) 94 | 95 | 96 | @pytest.fixture() 97 | def mock_generic_slack_response(): 98 | return AsyncMock(return_value=MagicMock(ok=True, data={"ok": True})) 99 | 100 | 101 | @pytest.fixture() 102 | def mock_conversations_history_response(): 103 | return AsyncMock( 104 | return_value=MagicMock( 105 | ok=True, 106 | data={ 107 | "ok": True, 108 | "messages": [ 109 | { 110 | "blocks": [], 111 | "text": "", 112 | "ts": "", 113 | **bot_message_extra_data(), 114 | } 115 | ], 116 | }, 117 | ) 118 | ) 119 | 120 | 121 | @pytest.fixture() 122 | def mock_get_permalink_response(): 123 | return AsyncMock( 124 | return_value={ 125 | "ok": True, 126 | "permalink": "mockpermalink", 127 | }, 128 | ) 129 | 130 | 131 | @pytest.fixture 132 | def mock_slack_asyncwebclient( 133 | mock_conversations_history_response, 134 | mock_generic_slack_response, 135 | mock_post_message_response, 136 | mock_get_permalink_response, 137 | ): 138 | with patch("slack_sdk.web.async_client.AsyncWebClient", autospec=True) as mock_client: 139 | wc = mock_client.return_value 140 | wc.reactions_add = mock_generic_slack_response 141 | wc.chat_update = mock_generic_slack_response 142 | wc.conversations_history = mock_conversations_history_response 143 | wc.chat_postMessage = mock_post_message_response 144 | wc.chat_getPermalink = mock_get_permalink_response 145 | yield wc 146 | 147 | 148 | @pytest.fixture 149 | def mock_slack_client(mock_slack_asyncwebclient): 150 | template_path = os.path.join( 151 | os.path.dirname(os.path.abspath(__file__)), "../triage_slackbot/templates" 152 | ) 153 | return SlackClient(mock_slack_asyncwebclient, template_path) 154 | 155 | 156 | @pytest.fixture 157 | def mock_inbound_request_channel_id(mock_config): 158 | return mock_config.inbound_request_channel_id 159 | 160 | 161 | @pytest.fixture 162 | def mock_feed_channel_id(mock_config): 163 | return mock_config.feed_channel_id 164 | 165 | 166 | @pytest.fixture 167 | def mock_appsec_oncall_slack_channel_id(mock_config): 168 | return mock_config.categories["appsec"].oncall_slack_id 169 | 170 | 171 | @pytest.fixture 172 | def mock_privacy_oncall_slack_user_id(mock_config): 173 | return mock_config.categories["privacy"].oncall_slack_id 174 | 175 | 176 | @pytest.fixture 177 | def mock_appsec_oncall_slack_user(): 178 | return {"id": "U1234567890"} 179 | 180 | 181 | @pytest.fixture 182 | def mock_appsec_oncall_slack_user_id(mock_appsec_oncall_slack_user): 183 | return mock_appsec_oncall_slack_user["id"] 184 | 185 | 186 | @pytest.fixture 187 | def mock_inbound_request_ts(): 188 | return "t0" 189 | 190 | 191 | @pytest.fixture 192 | def mock_feed_message_ts(): 193 | return "t1" 194 | 195 | 196 | @pytest.fixture 197 | def mock_notify_appsec_oncall_message_ts(): 198 | return "t2" 199 | 200 | 201 | @pytest.fixture 202 | def mock_appsec_oncall_recategorize_ts(): 203 | return "t3" 204 | 205 | 206 | @pytest.fixture 207 | def mock_inbound_request(mock_inbound_request_channel_id, mock_inbound_request_ts): 208 | return MagicMock( 209 | ack=AsyncMock(), 210 | event={ 211 | "channel": mock_inbound_request_channel_id, 212 | "text": "sample inbound request", 213 | "thread_ts": None, 214 | "ts": mock_inbound_request_ts, 215 | }, 216 | ) 217 | 218 | 219 | @pytest.fixture 220 | def mock_inbound_request_permalink(mock_inbound_request_channel_id): 221 | return f"https://myorg.slack.com/archives/{mock_inbound_request_channel_id}/p1234567890" 222 | 223 | 224 | @pytest.fixture 225 | async def mock_notify_appsec_oncall_message_data( 226 | mock_slack_client, 227 | mock_config, 228 | mock_inbound_request_channel_id, 229 | mock_inbound_request_permalink, 230 | mock_inbound_request_ts, 231 | mock_feed_channel_id, 232 | mock_feed_message_ts, 233 | mock_appsec_oncall_slack_channel_id, 234 | mock_notify_appsec_oncall_message_ts, 235 | ): 236 | appsec_key = "appsec" 237 | remaining_categories = [c for c in mock_config.categories.values() if c.key != appsec_key] 238 | blocks = mock_slack_client.render_blocks_from_template( 239 | MessageTemplatePath.notify_oncall_channel.value, 240 | { 241 | "inbound_message_url": mock_inbound_request_permalink, 242 | "inbound_message_channel": mock_inbound_request_channel_id, 243 | "predicted_category": mock_config.categories[appsec_key].display_name, 244 | "options": RequestCategory.to_block_options(remaining_categories), 245 | }, 246 | ) 247 | 248 | return { 249 | "ok": True, 250 | "channel": mock_appsec_oncall_slack_channel_id, 251 | "ts": mock_notify_appsec_oncall_message_ts, 252 | "message": { 253 | "blocks": blocks, 254 | "text": "Notify on-call for new inbound request", 255 | "metadata": { 256 | "event_type": "notify_oncall", 257 | "event_payload": { 258 | "inbound_message_channel": mock_inbound_request_channel_id, 259 | "inbound_message_ts": mock_inbound_request_ts, 260 | "feed_message_channel": mock_feed_channel_id, 261 | "feed_message_ts": mock_feed_message_ts, 262 | "inbound_message_url": mock_inbound_request_permalink, 263 | "predicted_category": appsec_key, 264 | }, 265 | }, 266 | "ts": mock_notify_appsec_oncall_message_ts, 267 | **bot_message_extra_data(), 268 | }, 269 | } 270 | 271 | 272 | @pytest.fixture 273 | def mock_notify_appsec_oncall_message( 274 | mock_notify_appsec_oncall_message_data, 275 | mock_appsec_oncall_slack_channel_id, 276 | mock_appsec_oncall_slack_user, 277 | ): 278 | return MagicMock( 279 | ack=AsyncMock(), 280 | body={ 281 | "actions": ["acknowledge_submit_action"], 282 | "container": { 283 | "message_ts": mock_notify_appsec_oncall_message_data["ts"], 284 | "channel_id": mock_appsec_oncall_slack_channel_id, 285 | }, 286 | "user": mock_appsec_oncall_slack_user, 287 | "message": mock_notify_appsec_oncall_message_data["message"], 288 | }, 289 | ) 290 | 291 | 292 | @pytest.fixture 293 | def mock_appsec_oncall_recategorize_to_privacy_message( 294 | mock_appsec_oncall_recategorize_ts, 295 | mock_appsec_oncall_slack_channel_id, 296 | mock_appsec_oncall_slack_user, 297 | mock_notify_appsec_oncall_message_data, 298 | ): 299 | return recategorize_message_data( 300 | ts=mock_appsec_oncall_recategorize_ts, 301 | channel_id=mock_appsec_oncall_slack_channel_id, 302 | user=mock_appsec_oncall_slack_user, 303 | message=mock_notify_appsec_oncall_message_data["message"], 304 | category="privacy", 305 | ) 306 | 307 | 308 | @pytest.fixture 309 | def mock_appsec_oncall_recategorize_to_other_message( 310 | mock_appsec_oncall_recategorize_ts, 311 | mock_appsec_oncall_slack_channel_id, 312 | mock_appsec_oncall_slack_user, 313 | mock_notify_appsec_oncall_message_data, 314 | ): 315 | return recategorize_message_data( 316 | ts=mock_appsec_oncall_recategorize_ts, 317 | channel_id=mock_appsec_oncall_slack_channel_id, 318 | user=mock_appsec_oncall_slack_user, 319 | message=mock_notify_appsec_oncall_message_data["message"], 320 | category="other", 321 | conversation="C11111", 322 | ) 323 | -------------------------------------------------------------------------------- /bots/triage-slackbot/tests/test_config.toml: -------------------------------------------------------------------------------- 1 | # Organization ID associated with OpenAI API key. 2 | openai_organization_id = "org-1234" 3 | 4 | # Prompt to use for categorizing inbound requests. 5 | openai_prompt = """ 6 | You are currently an on-call engineer for a security team at a tech company. 7 | Your goal is to triage the following incoming Slack message into three categories: 8 | 1. Privacy, return "privacy" 9 | 2. Application security, return "appsec" 10 | 3. Physical security, return "physical_security" 11 | """ 12 | inbound_request_channel_id = "C12345" 13 | feed_channel_id = "C23456" 14 | other_category_enabled = true 15 | 16 | [[ categories ]] 17 | key = "appsec" 18 | display_name = "Application Security" 19 | oncall_slack_id = "C34567" 20 | autorespond = false 21 | 22 | [[ categories ]] 23 | key = "privacy" 24 | display_name = "Privacy" 25 | oncall_slack_id = "U12345" 26 | autorespond = false 27 | 28 | [[ categories ]] 29 | key = "physical_security" 30 | display_name = "Physical Security" 31 | autorespond = true 32 | autorespond_message = "Looking for Physical or Office Security? You can reach out to physical-security@company.com." -------------------------------------------------------------------------------- /bots/triage-slackbot/tests/test_handlers.py: -------------------------------------------------------------------------------- 1 | import json 2 | from unittest.mock import call, patch 3 | 4 | import pytest 5 | from triage_slackbot.handlers import ( 6 | InboundRequestAcknowledgeHandler, 7 | InboundRequestHandler, 8 | InboundRequestRecategorizeHandler, 9 | ) 10 | from triage_slackbot.openai_utils import openai 11 | 12 | 13 | def get_mock_chat_completion_response(category: str): 14 | category_args = json.dumps({"category": category}) 15 | return { 16 | "choices": [ 17 | { 18 | "message": { 19 | "function_call": { 20 | "arguments": category_args, 21 | } 22 | } 23 | } 24 | ] 25 | } 26 | 27 | 28 | def assert_chat_completion_called(mock_chat_completion, mock_config): 29 | mock_chat_completion.create.assert_called_once_with( 30 | model="gpt-4-32k-0613", 31 | messages=[ 32 | { 33 | "role": "system", 34 | "content": mock_config.openai_prompt, 35 | }, 36 | {"role": "user", "content": "sample inbound request"}, 37 | ], 38 | temperature=0, 39 | stream=False, 40 | functions=[ 41 | { 42 | "name": "get_predicted_category", 43 | "description": "Predicts the category of an inbound request.", 44 | "parameters": { 45 | "type": "object", 46 | "properties": { 47 | "category": { 48 | "type": "string", 49 | "enum": ["appsec", "privacy", "physical_security"], 50 | "description": "Predicted category of the inbound request", 51 | } 52 | }, 53 | "required": ["category"], 54 | }, 55 | } 56 | ], 57 | function_call={"name": "get_predicted_category"}, 58 | ) 59 | 60 | 61 | @patch.object(openai, "ChatCompletion") 62 | async def test_inbound_request_handler_handle( 63 | mock_chat_completion, 64 | mock_config, 65 | mock_slack_client, 66 | mock_inbound_request, 67 | ): 68 | # Setup mocks 69 | mock_chat_completion.create.return_value = get_mock_chat_completion_response("appsec") 70 | 71 | # Call handler 72 | handler = InboundRequestHandler(mock_slack_client) 73 | await handler.maybe_handle(mock_inbound_request) 74 | 75 | # Assert that handler calls OpenAI API 76 | assert_chat_completion_called(mock_chat_completion, mock_config) 77 | 78 | mock_slack_client._client.assert_has_calls( 79 | [ 80 | call.chat_getPermalink(channel="C12345", message_ts="t0"), 81 | call.chat_postMessage( 82 | channel="C23456", 83 | blocks=[ 84 | { 85 | "type": "section", 86 | "text": { 87 | "type": "mrkdwn", 88 | "text": "Received an in <#C12345>:", 89 | }, 90 | }, 91 | { 92 | "type": "context", 93 | "elements": [ 94 | { 95 | "type": "plain_text", 96 | "text": "Predicted category: Application Security", 97 | "emoji": True, 98 | }, 99 | {"type": "mrkdwn", "text": "Triaged to: <#C34567>"}, 100 | { 101 | "type": "plain_text", 102 | "text": "Triage updates in the :thread:", 103 | "emoji": True, 104 | }, 105 | ], 106 | }, 107 | ], 108 | text="New inbound request received", 109 | ), 110 | call.chat_postMessage( 111 | channel="C34567", 112 | thread_ts=None, 113 | blocks=[ 114 | { 115 | "type": "section", 116 | "text": { 117 | "type": "mrkdwn", 118 | "text": ":wave: Hi, we received an in <#C12345>, which was categorized as Application Security. Is this accurate?\n\n", 119 | }, 120 | }, 121 | { 122 | "type": "context", 123 | "elements": [ 124 | { 125 | "type": "plain_text", 126 | "text": ":thumbsup: Acknowledge this message and response directly to the inbound request.", 127 | "emoji": True, 128 | }, 129 | { 130 | "type": "plain_text", 131 | "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.", 132 | "emoji": True, 133 | }, 134 | ], 135 | }, 136 | { 137 | "type": "actions", 138 | "elements": [ 139 | { 140 | "type": "button", 141 | "text": { 142 | "type": "plain_text", 143 | "emoji": True, 144 | "text": "Acknowledge", 145 | }, 146 | "style": "primary", 147 | "value": "Application Security", 148 | "action_id": "acknowledge_submit_action", 149 | }, 150 | { 151 | "type": "button", 152 | "text": { 153 | "type": "plain_text", 154 | "emoji": True, 155 | "text": "Inaccurate, recategorize", 156 | }, 157 | "style": "danger", 158 | "value": "recategorize", 159 | "action_id": "recategorize_submit_action", 160 | }, 161 | ], 162 | }, 163 | { 164 | "type": "section", 165 | "block_id": "recategorize_select_category_block", 166 | "text": { 167 | "type": "mrkdwn", 168 | "text": "*Select a category from the dropdown list, or*", 169 | }, 170 | "accessory": { 171 | "type": "static_select", 172 | "placeholder": { 173 | "type": "plain_text", 174 | "text": "Select an item", 175 | "emoji": True, 176 | }, 177 | "options": [ 178 | { 179 | "text": { 180 | "type": "plain_text", 181 | "text": "Privacy", 182 | "emoji": True, 183 | }, 184 | "value": "privacy", 185 | }, 186 | { 187 | "text": { 188 | "type": "plain_text", 189 | "text": "Physical Security", 190 | "emoji": True, 191 | }, 192 | "value": "physical_security", 193 | }, 194 | { 195 | "text": { 196 | "type": "plain_text", 197 | "text": "Other", 198 | "emoji": True, 199 | }, 200 | "value": "other", 201 | }, 202 | ], 203 | "action_id": "recategorize_select_category_action", 204 | }, 205 | }, 206 | ], 207 | metadata={ 208 | "event_type": "notify_oncall", 209 | "event_payload": { 210 | "inbound_message_channel": "C12345", 211 | "inbound_message_ts": "t0", 212 | "feed_message_channel": "", 213 | "feed_message_ts": "", 214 | "inbound_message_url": "mockpermalink", 215 | "predicted_category": "appsec", 216 | }, 217 | }, 218 | text="Notify on-call for new inbound request", 219 | ), 220 | ] 221 | ) 222 | 223 | 224 | @patch.object(openai, "ChatCompletion") 225 | async def test_inbound_request_handler_handle_autorespond( 226 | mock_chat_completion, 227 | mock_config, 228 | mock_slack_client, 229 | mock_inbound_request, 230 | ): 231 | # Setup mocks 232 | mock_chat_completion.create.return_value = get_mock_chat_completion_response( 233 | "physical_security" 234 | ) 235 | 236 | # Call handler 237 | handler = InboundRequestHandler(mock_slack_client) 238 | await handler.maybe_handle(mock_inbound_request) 239 | 240 | # Assert that handler calls OpenAI API 241 | assert_chat_completion_called(mock_chat_completion, mock_config) 242 | 243 | mock_slack_client._client.assert_has_calls( 244 | [ 245 | call.chat_getPermalink(channel="C12345", message_ts="t0"), 246 | call.chat_postMessage( 247 | channel="C23456", 248 | blocks=[ 249 | { 250 | "type": "section", 251 | "text": { 252 | "type": "mrkdwn", 253 | "text": "Received an in <#C12345>:", 254 | }, 255 | }, 256 | { 257 | "type": "context", 258 | "elements": [ 259 | { 260 | "type": "plain_text", 261 | "text": "Predicted category: Physical Security", 262 | "emoji": True, 263 | }, 264 | { 265 | "type": "mrkdwn", 266 | "text": "Triaged to: No on-call assigned", 267 | }, 268 | { 269 | "type": "plain_text", 270 | "text": "Triage updates in the :thread:", 271 | "emoji": True, 272 | }, 273 | ], 274 | }, 275 | ], 276 | text="New inbound request received", 277 | ), 278 | call.chat_postMessage( 279 | channel="C12345", 280 | thread_ts="t0", 281 | text="Hi, thanks for reaching out! Looking for Physical or Office Security? You can reach out to physical-security@company.com.", 282 | blocks=[ 283 | { 284 | "type": "section", 285 | "text": { 286 | "type": "mrkdwn", 287 | "text": "Hi, thanks for reaching out! Looking for Physical or Office Security? You can reach out to physical-security@company.com.", 288 | }, 289 | }, 290 | { 291 | "type": "context", 292 | "elements": [ 293 | { 294 | "type": "plain_text", 295 | "text": "If you feel strongly this is a Security issue, respond to this thread and someone from our team will get back to you.", 296 | "emoji": True, 297 | } 298 | ], 299 | }, 300 | ], 301 | ), 302 | call.chat_getPermalink(channel="", message_ts=""), 303 | call.chat_postMessage( 304 | channel="", 305 | thread_ts="", 306 | text=" to inbound request.", 307 | ), 308 | ] 309 | ) 310 | 311 | 312 | async def test_inbound_request_acknowledge_handler( 313 | mock_slack_client, 314 | mock_notify_appsec_oncall_message, 315 | ): 316 | handler = InboundRequestAcknowledgeHandler(mock_slack_client) 317 | await handler.maybe_handle(mock_notify_appsec_oncall_message) 318 | mock_slack_client._client.assert_has_calls( 319 | [ 320 | call.reactions_add( 321 | blocks=[], 322 | channel="C34567", 323 | ts="t2", 324 | text=":thumbsup: <@U1234567890> acknowledged the triaged to Application Security.", 325 | ), 326 | call.chat_postMessage( 327 | blocks=[], 328 | channel="C23456", 329 | thread_ts="t1", 330 | text=":thumbsup: <@U1234567890> acknowledged the inbound message triaged to Application Security.", 331 | ), 332 | call.conversations_history(channel="C23456", inclusive=True, latest="t1", limit=1), 333 | call.reactions_add(channel="C23456", name="thumbsup", timestamp="t1"), 334 | ] 335 | ) 336 | 337 | 338 | async def test_inbound_request_recategorize_to_listed_category_handler( 339 | mock_slack_client, 340 | mock_appsec_oncall_recategorize_to_privacy_message, 341 | ): 342 | handler = InboundRequestRecategorizeHandler(mock_slack_client) 343 | await handler.maybe_handle(mock_appsec_oncall_recategorize_to_privacy_message) 344 | mock_slack_client._client.assert_has_calls( 345 | [ 346 | call.reactions_add( 347 | blocks=[], 348 | channel="C34567", 349 | ts="t3", 350 | text=":thumbsdown: <@U1234567890> reassigned the from Application Security to: Privacy.", 351 | ), 352 | call.reactions_add(channel="C23456", name="thumbsdown", timestamp="t1"), 353 | call.chat_postMessage( 354 | blocks=[], 355 | channel="C23456", 356 | thread_ts="t1", 357 | text=":thumbsdown: <@U1234567890> reassigned the inbound message from Application Security to: Privacy.", 358 | ), 359 | call.chat_postMessage( 360 | channel="C23456", 361 | thread_ts="t1", 362 | blocks=[ 363 | { 364 | "type": "section", 365 | "text": { 366 | "type": "mrkdwn", 367 | "text": ":wave: Hi <@U12345>, is this assignment accurate?\n\n", 368 | }, 369 | }, 370 | { 371 | "type": "context", 372 | "elements": [ 373 | { 374 | "type": "plain_text", 375 | "text": ":thumbsup: Acknowledge this message and response directly to the inbound request.", 376 | "emoji": True, 377 | }, 378 | { 379 | "type": "plain_text", 380 | "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.", 381 | "emoji": True, 382 | }, 383 | ], 384 | }, 385 | { 386 | "type": "actions", 387 | "elements": [ 388 | { 389 | "type": "button", 390 | "text": { 391 | "type": "plain_text", 392 | "emoji": True, 393 | "text": "Acknowledge", 394 | }, 395 | "style": "primary", 396 | "value": "Privacy", 397 | "action_id": "acknowledge_submit_action", 398 | }, 399 | { 400 | "type": "button", 401 | "text": { 402 | "type": "plain_text", 403 | "emoji": True, 404 | "text": "Inaccurate, recategorize", 405 | }, 406 | "style": "danger", 407 | "value": "recategorize", 408 | "action_id": "recategorize_submit_action", 409 | }, 410 | ], 411 | }, 412 | { 413 | "type": "section", 414 | "block_id": "recategorize_select_category_block", 415 | "text": { 416 | "type": "mrkdwn", 417 | "text": "*Select a category from the dropdown list, or*", 418 | }, 419 | "accessory": { 420 | "type": "static_select", 421 | "placeholder": { 422 | "type": "plain_text", 423 | "text": "Select an item", 424 | "emoji": True, 425 | }, 426 | "options": [ 427 | { 428 | "text": { 429 | "type": "plain_text", 430 | "text": "Physical Security", 431 | "emoji": True, 432 | }, 433 | "value": "physical_security", 434 | }, 435 | { 436 | "text": { 437 | "type": "plain_text", 438 | "text": "Other", 439 | "emoji": True, 440 | }, 441 | "value": "other", 442 | }, 443 | ], 444 | "action_id": "recategorize_select_category_action", 445 | }, 446 | }, 447 | ], 448 | metadata={ 449 | "event_type": "notify_oncall", 450 | "event_payload": { 451 | "inbound_message_channel": "C12345", 452 | "inbound_message_ts": "t0", 453 | "feed_message_channel": "C23456", 454 | "feed_message_ts": "t1", 455 | "inbound_message_url": "https://myorg.slack.com/archives/C12345/p1234567890", 456 | "predicted_category": "privacy", 457 | }, 458 | }, 459 | text="Notify on-call for new inbound request", 460 | ), 461 | ] 462 | ) 463 | 464 | 465 | async def test_inbound_request_recategorize_to_other_category_handler( 466 | mock_slack_client, 467 | mock_appsec_oncall_recategorize_to_other_message, 468 | ): 469 | handler = InboundRequestRecategorizeHandler(mock_slack_client) 470 | await handler.maybe_handle(mock_appsec_oncall_recategorize_to_other_message) 471 | mock_slack_client._client.assert_has_calls( 472 | [ 473 | call.reactions_add( 474 | blocks=[], 475 | channel="C34567", 476 | ts="t3", 477 | text=":thumbsdown: <@U1234567890> reassigned the from Application Security to: Other.", 478 | ), 479 | call.reactions_add(channel="C23456", name="thumbsdown", timestamp="t1"), 480 | call.chat_postMessage( 481 | blocks=[], 482 | channel="C23456", 483 | thread_ts="t1", 484 | text=":thumbsdown: <@U1234567890> reassigned the inbound message from Application Security to: Other.", 485 | ), 486 | call.chat_postMessage( 487 | channel="C12345", 488 | thread_ts="t0", 489 | 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.", 490 | blocks=[ 491 | { 492 | "type": "section", 493 | "text": { 494 | "type": "mrkdwn", 495 | "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.", 496 | }, 497 | }, 498 | { 499 | "type": "context", 500 | "elements": [ 501 | { 502 | "type": "plain_text", 503 | "text": "If you feel strongly this is a Security issue, respond to this thread and someone from our team will get back to you.", 504 | "emoji": True, 505 | } 506 | ], 507 | }, 508 | ], 509 | ), 510 | call.chat_getPermalink(channel="", message_ts=""), 511 | call.chat_postMessage( 512 | channel="C23456", 513 | thread_ts="t1", 514 | text=" to inbound request.", 515 | ), 516 | ] 517 | ) 518 | 519 | 520 | @pytest.mark.parametrize( 521 | "event_args_override", 522 | [ 523 | # Channel is not inbound request channel 524 | {"channel": "c0"}, 525 | # No text 526 | {"text": ""}, 527 | # Bot message 528 | {"subtype": "bot_message"}, 529 | # Thread response, not broadcasted 530 | {"thread_ts": "t0"}, 531 | ], 532 | ) 533 | @patch.object(openai, "ChatCompletion") 534 | async def test_inbound_request_handler_skip_handle( 535 | mock_chat_completion, event_args_override, mock_slack_client, mock_inbound_request 536 | ): 537 | mock_inbound_request.event = {**mock_inbound_request.event, **event_args_override} 538 | handler = InboundRequestHandler(mock_slack_client) 539 | 540 | await handler.maybe_handle(mock_inbound_request) 541 | mock_chat_completion.create.assert_not_called() 542 | -------------------------------------------------------------------------------- /bots/triage-slackbot/triage_slackbot/.env.template: -------------------------------------------------------------------------------- 1 | SLACK_BOT_TOKEN=xoxb- # Go to your Slack app, Settings > Install App > Bot User OAuth Token 2 | SOCKET_APP_TOKEN=xapp- # Go to your Slack app, Settings > Basic Information > App-Level Tokens 3 | OPENAI_API_KEY= -------------------------------------------------------------------------------- /bots/triage-slackbot/triage_slackbot/bot.py: -------------------------------------------------------------------------------- 1 | import asyncio 2 | import os 3 | 4 | from openai_slackbot.bot import start_bot 5 | from triage_slackbot.config import get_config, load_config 6 | from triage_slackbot.handlers import ( 7 | InboundRequestAcknowledgeHandler, 8 | InboundRequestHandler, 9 | InboundRequestRecategorizeHandler, 10 | InboundRequestRecategorizeSelectConversationHandler, 11 | InboundRequestRecategorizeSelectHandler, 12 | ) 13 | 14 | if __name__ == "__main__": 15 | current_dir = os.path.dirname(os.path.abspath(__file__)) 16 | load_config(os.path.join(current_dir, "config.toml")) 17 | 18 | message_handler = InboundRequestHandler 19 | action_handlers = [ 20 | InboundRequestAcknowledgeHandler, 21 | InboundRequestRecategorizeHandler, 22 | InboundRequestRecategorizeSelectHandler, 23 | InboundRequestRecategorizeSelectConversationHandler, 24 | ] 25 | 26 | template_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), "templates") 27 | 28 | config = get_config() 29 | asyncio.run( 30 | start_bot( 31 | openai_organization_id=config.openai_organization_id, 32 | slack_message_handler=message_handler, 33 | slack_action_handlers=action_handlers, 34 | slack_template_path=template_path, 35 | ) 36 | ) 37 | -------------------------------------------------------------------------------- /bots/triage-slackbot/triage_slackbot/category.py: -------------------------------------------------------------------------------- 1 | import typing as t 2 | 3 | from pydantic import BaseModel, ValidationError, model_validator 4 | 5 | OTHER_KEY = "other" 6 | 7 | 8 | class RequestCategory(BaseModel): 9 | # Key used to identify the category in the config. 10 | key: str 11 | 12 | # Display name of the category. 13 | display_name: str 14 | 15 | # Slack ID of the user or channel to route the request to. 16 | # If user is specified, user will be tagged on the message 17 | # in the feed channel. 18 | oncall_slack_id: t.Optional[str] = None 19 | 20 | # If true, no manual triage is required for this category 21 | # and that the bot will autorespond to the inbound request. 22 | autorespond: bool = False 23 | 24 | # Message to send when autoresponding to the inbound request. 25 | autorespond_message: t.Optional[str] = None 26 | 27 | @model_validator(mode="after") 28 | def check_autorespond(self) -> "RequestCategory": 29 | if self.autorespond and not self.autorespond_message: 30 | raise ValidationError("autorespond_message must be set if autorespond is True") 31 | return self 32 | 33 | @property 34 | def route_to_channel(self) -> bool: 35 | return (self.oncall_slack_id or "").startswith("C") 36 | 37 | @classmethod 38 | def to_block_options(cls, categories: t.List["RequestCategory"]) -> t.Dict[str, str]: 39 | return dict((c.key, c.display_name) for c in categories) 40 | 41 | def is_other(self) -> bool: 42 | return self.key == OTHER_KEY 43 | -------------------------------------------------------------------------------- /bots/triage-slackbot/triage_slackbot/config.py: -------------------------------------------------------------------------------- 1 | import os 2 | import typing as t 3 | 4 | import toml 5 | from dotenv import load_dotenv 6 | from pydantic import BaseModel, ValidationError, field_validator, model_validator 7 | from pydantic.functional_validators import AfterValidator, BeforeValidator 8 | from triage_slackbot.category import OTHER_KEY, RequestCategory 9 | 10 | _CONFIG = None 11 | 12 | 13 | def convert_categories(v: t.List[t.Dict]): 14 | categories = {} 15 | 16 | for category in v: 17 | categories[category["key"]] = category 18 | return categories 19 | 20 | 21 | def validate_channel(channel_id: str) -> str: 22 | if not channel_id.startswith("C"): 23 | raise ValueError("channel ID must start with 'C'") 24 | return channel_id 25 | 26 | 27 | class Config(BaseModel): 28 | # OpenAI organization ID associated with OpenAI API key. 29 | openai_organization_id: str 30 | 31 | # OpenAI prompt to categorize the request. 32 | openai_prompt: str 33 | 34 | # Slack channel where inbound requests are received. 35 | inbound_request_channel_id: t.Annotated[str, AfterValidator(validate_channel)] 36 | 37 | # Slack channel where triage updates are posted. 38 | feed_channel_id: t.Annotated[str, AfterValidator(validate_channel)] 39 | 40 | # Valid categories for inbound requests to be triaged into. 41 | categories: t.Annotated[t.Dict[str, RequestCategory], BeforeValidator(convert_categories)] 42 | 43 | # Enables "Other" category, which will allow triager to 44 | # route the request to a specific conversation. 45 | other_category_enabled: bool 46 | 47 | @model_validator(mode="after") 48 | def check_category_keys(config: "Config") -> "Config": 49 | if config.other_category_enabled: 50 | if OTHER_KEY in config.categories: 51 | raise ValidationError("other category is reserved and cannot be used") 52 | 53 | category_keys = set(config.categories.keys()) 54 | if len(category_keys) != len(config.categories): 55 | raise ValidationError("category keys must be unique") 56 | 57 | return config 58 | 59 | 60 | def load_config(path: str): 61 | load_dotenv() 62 | 63 | with open(path) as f: 64 | cfg = toml.loads(f.read()) 65 | config = Config(**cfg) 66 | 67 | if config.other_category_enabled: 68 | other_category = RequestCategory( 69 | key=OTHER_KEY, 70 | display_name=OTHER_KEY.capitalize(), 71 | oncall_slack_id=None, 72 | autorespond=True, 73 | autorespond_message="Our team looked at your request, and this is actually something that we don't own. We recommend reaching out to {} instead.", 74 | ) 75 | config.categories[other_category.key] = other_category 76 | 77 | global _CONFIG 78 | _CONFIG = config 79 | return _CONFIG 80 | 81 | 82 | def get_config() -> Config: 83 | global _CONFIG 84 | if _CONFIG is None: 85 | raise Exception("config not initialized, call load_config() first") 86 | return _CONFIG 87 | -------------------------------------------------------------------------------- /bots/triage-slackbot/triage_slackbot/config.toml: -------------------------------------------------------------------------------- 1 | # Organization ID associated with OpenAI API key. 2 | openai_organization_id = "" 3 | 4 | # Prompt to use for categorizing inbound requests. 5 | openai_prompt = """ 6 | You are currently an on-call engineer for a security team at a tech company. 7 | Your goal is to triage the following incoming Slack message into three categories: 8 | 1. Privacy, return "privacy" 9 | 2. Application security, return "appsec" 10 | 3. Physical security, return "physical_security" 11 | """ 12 | inbound_request_channel_id = "" 13 | feed_channel_id = "" 14 | other_category_enabled = true 15 | 16 | [[ categories ]] 17 | key = "appsec" 18 | display_name = "Application Security" 19 | oncall_slack_id = "" 20 | autorespond = false 21 | 22 | [[ categories ]] 23 | key = "privacy" 24 | display_name = "Privacy" 25 | oncall_slack_id = "" 26 | autorespond = false 27 | 28 | [[ categories ]] 29 | key = "physical_security" 30 | display_name = "Physical Security" 31 | autorespond = true 32 | autorespond_message = "Looking for Physical or Office Security? You can reach out to physical-security@company.com." 33 | -------------------------------------------------------------------------------- /bots/triage-slackbot/triage_slackbot/handlers.py: -------------------------------------------------------------------------------- 1 | import typing as t 2 | from enum import Enum 3 | from logging import getLogger 4 | 5 | from openai_slackbot.clients.slack import CreateSlackMessageResponse, SlackClient 6 | from openai_slackbot.handlers import BaseActionHandler, BaseHandler, BaseMessageHandler 7 | from openai_slackbot.utils.slack import ( 8 | RenderedSlackBlock, 9 | block_id_exists, 10 | extract_text_from_event, 11 | get_block_by_id, 12 | remove_block_id_if_exists, 13 | render_slack_id_to_mention, 14 | render_slack_url, 15 | ) 16 | from triage_slackbot.category import RequestCategory 17 | from triage_slackbot.config import get_config 18 | from triage_slackbot.openai_utils import get_predicted_category 19 | 20 | logger = getLogger(__name__) 21 | 22 | 23 | class BlockId(str, Enum): 24 | # Block that will be rendered if on-call recagorizes inbound request but doesn't select a new category. 25 | empty_category_warning = "empty_category_warning_block" 26 | 27 | # Block that will be rendreed if on-call recategorizes inbound request to "Other" but doesn't select a conversation. 28 | empty_conversation_warning = "empty_conversation_warning_block" 29 | 30 | # Block that will be rendered to show on-call all the remaining categories to route to. 31 | recategorize_select_category = "recategorize_select_category_block" 32 | 33 | # Block that will be rendered to show on-call all the conversations they can reroute the user 34 | # to, if they select "Other" as the category. 35 | recategorize_select_conversation = "recategorize_select_conversation_block" 36 | 37 | 38 | class MessageTemplatePath(str, Enum): 39 | # Template for feed channel message that summarizes triage updates. 40 | feed = "messages/feed.j2" 41 | 42 | # Template for message that notifies oncall about inbound request in the same channel as the feed channel. 43 | notify_oncall_in_feed = "messages/notify_oncall_in_feed.j2" 44 | 45 | # Template for message that notifies oncall about inbound request in a different channel from the feed channel. 46 | notify_oncall_channel = "messages/notify_oncall_channel.j2" 47 | 48 | # Template for message that will autorespond to inbound requests. 49 | autorespond = "messages/autorespond.j2" 50 | 51 | 52 | BlockIdToTemplatePath: t.Dict[BlockId, str] = { 53 | BlockId.empty_category_warning: "blocks/empty_category_warning.j2", 54 | BlockId.empty_conversation_warning: "blocks/empty_conversation_warning.j2", 55 | BlockId.recategorize_select_conversation: "blocks/select_conversation.j2", 56 | } 57 | 58 | 59 | class InboundRequestHandlerMixin(BaseHandler): 60 | def __init__(self, slack_client: SlackClient) -> None: 61 | super().__init__(slack_client) 62 | self.config = get_config() 63 | 64 | def render_block_if_not_exists( 65 | self, *, block_id: BlockId, blocks: t.List[RenderedSlackBlock] 66 | ) -> t.List[RenderedSlackBlock]: 67 | if not block_id_exists(blocks, block_id): 68 | template_path = BlockIdToTemplatePath[block_id] 69 | block = self._slack_client.render_blocks_from_template(template_path) 70 | blocks.append(block) 71 | return blocks 72 | 73 | def get_selected_category(self, body: t.Dict[str, t.Any]) -> t.Optional[RequestCategory]: 74 | category = ( 75 | body["state"] 76 | .get("values", {}) 77 | .get(BlockId.recategorize_select_category, {}) 78 | .get("recategorize_select_category_action", {}) 79 | .get("selected_option", {}) 80 | or {} 81 | ).get("value") 82 | 83 | if not category: 84 | return None 85 | 86 | return self.config.categories[category] 87 | 88 | def get_selected_conversation(self, body: t.Dict[str, t.Any]) -> t.Optional[str]: 89 | return ( 90 | body["state"] 91 | .get("values", {}) 92 | .get(BlockId.recategorize_select_conversation, {}) 93 | .get("recategorize_select_conversation_action", {}) 94 | .get("selected_conversation") 95 | ) 96 | 97 | async def notify_oncall( 98 | self, 99 | *, 100 | predicted_category: RequestCategory, 101 | selected_conversation: t.Optional[str], 102 | remaining_categories: t.List[RequestCategory], 103 | inbound_message_channel: str, 104 | inbound_message_ts: str, 105 | feed_message_channel: str, 106 | feed_message_ts: str, 107 | inbound_message_url: str, 108 | ) -> None: 109 | autoresponded = await self._maybe_autorespond( 110 | predicted_category, 111 | selected_conversation, 112 | inbound_message_channel, 113 | inbound_message_ts, 114 | feed_message_channel, 115 | feed_message_ts, 116 | ) 117 | 118 | if autoresponded: 119 | logger.info(f"Autoresponded to inbound request: {inbound_message_url}") 120 | return 121 | 122 | # This metadata will continue to be passed along to the subsequent 123 | # notify on-call messages. 124 | metadata = { 125 | "event_type": "notify_oncall", 126 | "event_payload": { 127 | "inbound_message_channel": inbound_message_channel, 128 | "inbound_message_ts": inbound_message_ts, 129 | "feed_message_channel": feed_message_channel, 130 | "feed_message_ts": feed_message_ts, 131 | "inbound_message_url": inbound_message_url, 132 | "predicted_category": predicted_category.key, 133 | }, 134 | } 135 | 136 | block_args = { 137 | "predicted_category": predicted_category, 138 | "remaining_categories": remaining_categories, 139 | "inbound_message_channel": inbound_message_channel, 140 | } 141 | 142 | if predicted_category.route_to_channel: 143 | channel = predicted_category.oncall_slack_id 144 | thread_ts = None # This will be a new message, not a thread. 145 | blocks = await self._get_notify_oncall_channel_blocks( 146 | **block_args, 147 | inbound_message_url=inbound_message_url, 148 | ) 149 | else: 150 | channel = feed_message_channel 151 | thread_ts = feed_message_ts # Post this as a thread reply to the original feed message. 152 | blocks = await self._get_notify_oncall_in_feed_blocks(**block_args) 153 | 154 | await self._slack_client.post_message( 155 | channel=channel, 156 | thread_ts=thread_ts, 157 | blocks=blocks, 158 | metadata=metadata, 159 | text="Notify on-call for new inbound request", 160 | ) 161 | 162 | async def _get_notify_oncall_in_feed_blocks( 163 | self, 164 | *, 165 | predicted_category: RequestCategory, 166 | remaining_categories: t.List[RequestCategory], 167 | inbound_message_channel: str, 168 | ): 169 | oncall_mention = self._get_oncall_mention(predicted_category) 170 | predicted_category_display_name = predicted_category.display_name 171 | oncall_greeting = ( 172 | f":wave: Hi {oncall_mention}" 173 | if oncall_mention 174 | else f"No on-call defined for {predicted_category_display_name}" 175 | ) 176 | 177 | return self._slack_client.render_blocks_from_template( 178 | MessageTemplatePath.notify_oncall_in_feed.value, 179 | { 180 | "predicted_category": predicted_category_display_name, 181 | "oncall_greeting": oncall_greeting, 182 | "options": RequestCategory.to_block_options(remaining_categories), 183 | "inbound_message_channel": inbound_message_channel, 184 | }, 185 | ) 186 | 187 | async def _get_notify_oncall_channel_blocks( 188 | self, 189 | *, 190 | predicted_category: RequestCategory, 191 | remaining_categories: t.List[RequestCategory], 192 | inbound_message_channel: str, 193 | inbound_message_url: str, 194 | ): 195 | return self._slack_client.render_blocks_from_template( 196 | MessageTemplatePath.notify_oncall_channel.value, 197 | { 198 | "inbound_message_url": inbound_message_url, 199 | "inbound_message_channel": inbound_message_channel, 200 | "predicted_category": predicted_category.display_name, 201 | "options": RequestCategory.to_block_options(remaining_categories), 202 | }, 203 | ) 204 | 205 | def _get_oncall_mention(self, predicted_category: RequestCategory) -> t.Optional[str]: 206 | oncall_slack_id = predicted_category.oncall_slack_id 207 | return render_slack_id_to_mention(oncall_slack_id) if oncall_slack_id else None 208 | 209 | async def _maybe_autorespond( 210 | self, 211 | predicted_category: RequestCategory, 212 | selected_conversation: t.Optional[str], 213 | inbound_message_channel: str, 214 | inbound_message_ts: str, 215 | feed_message_channel: str, 216 | feed_message_ts: str, 217 | ) -> bool: 218 | if not predicted_category.autorespond: 219 | return False 220 | 221 | text = "Hi, thanks for reaching out!" 222 | if predicted_category.autorespond_message: 223 | rendered_selected_conversation = ( 224 | render_slack_id_to_mention(selected_conversation) if selected_conversation else None 225 | ) 226 | text += ( 227 | f" {predicted_category.autorespond_message.format(rendered_selected_conversation)}" 228 | ) 229 | 230 | blocks = self._slack_client.render_blocks_from_template( 231 | MessageTemplatePath.autorespond.value, {"text": text} 232 | ) 233 | message = await self._slack_client.post_message( 234 | channel=inbound_message_channel, 235 | thread_ts=inbound_message_ts, 236 | text=text, 237 | blocks=blocks, 238 | ) 239 | message_link = await self._slack_client.get_message_link( 240 | channel=message.channel, message_ts=message.ts 241 | ) 242 | 243 | # Post an update to the feed channel. 244 | feed_message = ( 245 | f"{render_slack_url(url=message_link, text='Autoresponded')} to inbound request." 246 | ) 247 | await self._slack_client.post_message( 248 | channel=feed_message_channel, thread_ts=feed_message_ts, text=feed_message 249 | ) 250 | 251 | return True 252 | 253 | 254 | class InboundRequestHandler(BaseMessageHandler, InboundRequestHandlerMixin): 255 | """ 256 | Handles inbound requests in inbound request channel. 257 | """ 258 | 259 | async def handle(self, args): 260 | event = args.event 261 | 262 | channel = event.get("channel") 263 | ts = event.get("ts") 264 | 265 | logging_extra = self.logging_extra(args) 266 | 267 | text = extract_text_from_event(event) 268 | if not text: 269 | logger.info("No text in event, done processing", extra=logging_extra) 270 | return 271 | 272 | predicted_category = await self._predict_category(text) 273 | logger.info(f"Predicted category: {predicted_category}", extra=logging_extra) 274 | 275 | message_link = await self._slack_client.get_message_link(channel=channel, message_ts=ts) 276 | feed_message = await self._update_feed( 277 | predicted_category=predicted_category, 278 | message_channel=channel, 279 | message_link=message_link, 280 | ) 281 | logger.info( 282 | f"Updated feed channel for inbound message link: {message_link}", 283 | extra=logging_extra, 284 | ) 285 | 286 | remaining_categories = [ 287 | r for r in self.config.categories.values() if r != predicted_category 288 | ] 289 | await self.notify_oncall( 290 | predicted_category=predicted_category, 291 | selected_conversation=None, 292 | remaining_categories=remaining_categories, 293 | inbound_message_channel=channel, 294 | inbound_message_ts=ts, 295 | feed_message_channel=feed_message.channel, 296 | feed_message_ts=feed_message.ts, 297 | inbound_message_url=message_link, 298 | ) 299 | logger.info("Notified on-call", extra=logging_extra) 300 | 301 | async def should_handle(self, args): 302 | event = args.event 303 | 304 | return ( 305 | event["channel"] == self.config.inbound_request_channel_id 306 | and 307 | # Don't respond to messages in threads (with the exception of thread replies 308 | # that are also sent to the channel) 309 | ( 310 | ( 311 | event.get("thread_ts") is None 312 | and (not event.get("subtype") or event.get("subtype") == "file_share") 313 | ) 314 | or event.get("subtype") == "thread_broadcast" 315 | ) 316 | ) 317 | 318 | async def _predict_category(self, body) -> RequestCategory: 319 | predicted_category = await get_predicted_category(body) 320 | return self.config.categories[predicted_category] 321 | 322 | async def _update_feed( 323 | self, 324 | *, 325 | predicted_category: RequestCategory, 326 | message_channel: str, 327 | message_link: str, 328 | ) -> CreateSlackMessageResponse: 329 | oncall_mention = self._get_oncall_mention(predicted_category) or "No on-call assigned" 330 | blocks = self._slack_client.render_blocks_from_template( 331 | MessageTemplatePath.feed.value, 332 | { 333 | "predicted_category": predicted_category.display_name, 334 | "inbound_message_channel": message_channel, 335 | "inbound_message_url": message_link, 336 | "oncall_mention": oncall_mention, 337 | }, 338 | ) 339 | 340 | message = await self._slack_client.post_message( 341 | channel=self.config.feed_channel_id, 342 | blocks=blocks, 343 | text="New inbound request received", 344 | ) 345 | return message 346 | 347 | 348 | class InboundRequestAcknowledgeHandler(BaseActionHandler, InboundRequestHandlerMixin): 349 | """ 350 | Once InboundRequestHandler has predicted the category of an inbound request 351 | and notifies the corresponding on-call, this handler will be called if on-call 352 | acknowledges the prediction, i.e. they think the prediction is accurate. 353 | """ 354 | 355 | @property 356 | def action_id(self): 357 | return "acknowledge_submit_action" 358 | 359 | async def handle(self, args): 360 | body = args.body 361 | 362 | notify_oncall_msg = body["container"] 363 | notify_oncall_msg_ts = notify_oncall_msg["message_ts"] 364 | notify_oncall_msg_channel = notify_oncall_msg["channel_id"] 365 | 366 | feed_message_metadata = body["message"].get("metadata", {}).get("event_payload", {}) 367 | feed_message_ts = feed_message_metadata["feed_message_ts"] 368 | feed_message_channel = feed_message_metadata["feed_message_channel"] 369 | inbound_message_url = feed_message_metadata["inbound_message_url"] 370 | predicted_category = feed_message_metadata["predicted_category"] 371 | 372 | # Oncall that was notified. 373 | user = body["user"] 374 | 375 | await self._slack_client.update_message( 376 | blocks=[], 377 | channel=notify_oncall_msg_channel, 378 | ts=notify_oncall_msg_ts, 379 | # If oncall is notified in the feed channel, don't need to include 380 | # the inbound message URL since oncall will be notified in the feed 381 | # message thread, and the URL is already in the original message. 382 | text=self._get_message( 383 | user=user, 384 | category=predicted_category, 385 | inbound_message_url=inbound_message_url, 386 | with_url=notify_oncall_msg_channel != feed_message_channel, 387 | ), 388 | ) 389 | 390 | # If oncall gets notified in a separate channel and not the feed channel, 391 | # update the feed thread with the acknowledgment. 392 | if notify_oncall_msg_channel != feed_message_channel: 393 | await self._slack_client.post_message( 394 | blocks=[], 395 | channel=feed_message_channel, 396 | thread_ts=feed_message_ts, 397 | text=self._get_message( 398 | user=user, 399 | category=predicted_category, 400 | inbound_message_url=inbound_message_url, 401 | with_url=False, 402 | ), 403 | ) 404 | 405 | feed_message = await self._slack_client.get_message( 406 | channel=feed_message_channel, ts=feed_message_ts 407 | ) 408 | if feed_message: 409 | # If the original message has been thumbs-downed, this means 410 | # that the bot's original prediction is wrong, so don't thumbs 411 | # up the feed message. 412 | wrong_original_prediction = any( 413 | [r["name"] == "-1" for r in feed_message.get("reactions", [])] 414 | ) 415 | 416 | if not wrong_original_prediction: 417 | await self._slack_client.add_reaction( 418 | channel=feed_message_channel, 419 | name="thumbsup", 420 | timestamp=feed_message_ts, 421 | ) 422 | 423 | def _get_message( 424 | self, user: t.Dict, category: str, inbound_message_url: str, with_url: bool 425 | ) -> str: 426 | message = f":thumbsup: {render_slack_id_to_mention(user['id'])} acknowledged the " 427 | if with_url: 428 | message += render_slack_url(url=inbound_message_url, text="inbound message") 429 | else: 430 | message += "inbound message" 431 | 432 | return f"{message} triaged to {self.config.categories[category].display_name}." 433 | 434 | 435 | class InboundRequestRecategorizeHandler(BaseActionHandler, InboundRequestHandlerMixin): 436 | """ 437 | This handler will be called if on-call wants to recategorize the request 438 | that they get notified about. 439 | """ 440 | 441 | @property 442 | def action_id(self): 443 | return "recategorize_submit_action" 444 | 445 | async def handle(self, args): 446 | body = args.body 447 | 448 | notify_oncall_msg = body["container"] 449 | notify_oncall_msg_ts = notify_oncall_msg["message_ts"] 450 | notify_oncall_msg_channel = notify_oncall_msg["channel_id"] 451 | 452 | msg_metadata = body["message"].get("metadata", {}).get("event_payload", {}) 453 | feed_message_ts = msg_metadata["feed_message_ts"] 454 | feed_message_channel = msg_metadata["feed_message_channel"] 455 | inbound_message_url = msg_metadata["inbound_message_url"] 456 | 457 | # Predicted category that turned out to be incorrect 458 | # and wanted to be recategorized. 459 | predicted_category = self.config.categories[msg_metadata.pop("predicted_category")] 460 | assert predicted_category 461 | 462 | user: t.Dict = body["user"] 463 | 464 | notify_oncall_msg_blocks = body["message"]["blocks"] 465 | selection_block = get_block_by_id( 466 | notify_oncall_msg_blocks, BlockId.recategorize_select_category 467 | ) 468 | remaining_category_keys: t.List[str] = [ 469 | o["value"] for o in selection_block["accessory"]["options"] 470 | ] 471 | 472 | selected_category: t.Optional[RequestCategory] = self.get_selected_category(body) 473 | selected_conversation: t.Optional[str] = self.get_selected_conversation(body) 474 | valid, notify_oncall_msg_blocks = await self._validate_selection( 475 | selected_category, selected_conversation, notify_oncall_msg_blocks 476 | ) 477 | if valid: 478 | assert selected_category, "selected_category should be set if valid" 479 | message_kwargs = { 480 | "user": user, 481 | "predicted_category": predicted_category, 482 | "selected_category": selected_category, 483 | "selected_conversation": selected_conversation, 484 | "inbound_message_url": inbound_message_url, 485 | } 486 | 487 | await self._slack_client.update_message( 488 | blocks=[], 489 | channel=notify_oncall_msg_channel, 490 | ts=notify_oncall_msg_ts, 491 | # If the feed message is in the same channel as the notify on-call message, don't need to include 492 | # the URL since it's already in the original feed message. 493 | text=self._get_message( 494 | **message_kwargs, 495 | with_url=notify_oncall_msg_channel != feed_message_channel, 496 | ), 497 | ) 498 | 499 | # Indicate that the previous predicted category is not accurate. 500 | await self._slack_client.add_reaction( 501 | channel=feed_message_channel, 502 | name="thumbsdown", 503 | timestamp=feed_message_ts, 504 | ) 505 | 506 | # If the feed message is in a different channel than the notify on-call message, 507 | # post recategorization update to the feed channel. 508 | if notify_oncall_msg_channel != feed_message_channel: 509 | await self._slack_client.post_message( 510 | blocks=[], 511 | channel=feed_message_channel, 512 | thread_ts=feed_message_ts, 513 | text=self._get_message(**message_kwargs, with_url=False), 514 | ) 515 | 516 | remaining_categories = [ 517 | self.config.categories[category_key] 518 | for category_key in remaining_category_keys 519 | if category_key != selected_category.key 520 | ] 521 | 522 | # Route this to the next oncall. 523 | await self.notify_oncall( 524 | predicted_category=selected_category, 525 | selected_conversation=selected_conversation, 526 | remaining_categories=remaining_categories, 527 | **msg_metadata, 528 | ) 529 | else: 530 | # Display warning. 531 | await self._slack_client.update_message( 532 | blocks=notify_oncall_msg_blocks, 533 | channel=notify_oncall_msg_channel, 534 | ts=notify_oncall_msg_ts, 535 | text="", 536 | ) 537 | 538 | def _get_message( 539 | self, 540 | *, 541 | user: t.Dict, 542 | predicted_category: RequestCategory, 543 | selected_category: RequestCategory, 544 | selected_conversation: t.Optional[str], 545 | inbound_message_url: str, 546 | with_url: bool, 547 | ) -> str: 548 | rendered_selected_conversation = ( 549 | render_slack_id_to_mention(selected_conversation) if selected_conversation else None 550 | ) 551 | selected_category_display_name = selected_category.display_name.format( 552 | rendered_selected_conversation 553 | ) 554 | 555 | message_text = f"<{inbound_message_url}|inbound message>" if with_url else "inbound message" 556 | return f":thumbsdown: {render_slack_id_to_mention(user['id'])} reassigned the {message_text} from {predicted_category.display_name} to: {selected_category_display_name}." 557 | 558 | async def _validate_selection( 559 | self, 560 | selected_category: t.Optional[RequestCategory], 561 | selected_conversation: t.Optional[str], 562 | blocks: t.List[RenderedSlackBlock], 563 | ) -> t.Tuple[bool, t.List[RenderedSlackBlock]]: 564 | if not selected_category: 565 | return False, self.render_block_if_not_exists( 566 | block_id=BlockId.empty_category_warning, blocks=blocks 567 | ) 568 | elif selected_category.is_other() and not selected_conversation: 569 | return False, self.render_block_if_not_exists( 570 | block_id=BlockId.empty_conversation_warning, blocks=blocks 571 | ) 572 | 573 | return True, blocks 574 | 575 | 576 | class InboundRequestRecategorizeSelectHandler(BaseActionHandler, InboundRequestHandlerMixin): 577 | """ 578 | This handler will be called if on-call selects a new category for a request they 579 | get notififed about. 580 | """ 581 | 582 | @property 583 | def action_id(self): 584 | return "recategorize_select_category_action" 585 | 586 | async def handle(self, args): 587 | body = args.body 588 | 589 | notify_oncall_msg = body["container"] 590 | notify_oncall_msg_ts = notify_oncall_msg["message_ts"] 591 | notify_oncall_msg_channel = notify_oncall_msg["channel_id"] 592 | 593 | notify_oncall_msg_blocks = body["message"]["blocks"] 594 | notify_oncall_msg_blocks = remove_block_id_if_exists( 595 | notify_oncall_msg_blocks, BlockId.empty_category_warning 596 | ) 597 | 598 | selected_category = self.get_selected_category(body) 599 | if selected_category.is_other(): 600 | # Prompt on-call to select a conversation if Other category is selected. 601 | notify_oncall_msg_blocks = self.render_block_if_not_exists( 602 | block_id=BlockId.recategorize_select_conversation, 603 | blocks=notify_oncall_msg_blocks, 604 | ) 605 | else: 606 | # Remove warning if on-call updates their selection from Other to non-Other. 607 | notify_oncall_msg_blocks = remove_block_id_if_exists( 608 | notify_oncall_msg_blocks, BlockId.recategorize_select_conversation 609 | ) 610 | 611 | # Update message with warnings, if any. 612 | await self._slack_client.update_message( 613 | blocks=notify_oncall_msg_blocks, 614 | channel=notify_oncall_msg_channel, 615 | ts=notify_oncall_msg_ts, 616 | ) 617 | 618 | 619 | class InboundRequestRecategorizeSelectConversationHandler(BaseActionHandler): 620 | """ 621 | This handler will be called if on-call selects a conversation to route the request to. 622 | """ 623 | 624 | @property 625 | def action_id(self): 626 | return "recategorize_select_conversation_action" 627 | 628 | async def handle(self, args): 629 | pass 630 | -------------------------------------------------------------------------------- /bots/triage-slackbot/triage_slackbot/openai_utils.py: -------------------------------------------------------------------------------- 1 | import json 2 | from functools import cache 3 | 4 | import openai 5 | from triage_slackbot.category import OTHER_KEY, RequestCategory 6 | from triage_slackbot.config import get_config 7 | 8 | 9 | @cache 10 | def predict_category_functions(categories: list[RequestCategory]) -> list[dict]: 11 | return [ 12 | { 13 | "name": "get_predicted_category", 14 | "description": "Predicts the category of an inbound request.", 15 | "parameters": { 16 | "type": "object", 17 | "properties": { 18 | "category": { 19 | "type": "string", 20 | "enum": [ 21 | category.key for category in categories if category.key != OTHER_KEY 22 | ], 23 | "description": "Predicted category of the inbound request", 24 | }, 25 | }, 26 | "required": ["category"], 27 | }, 28 | } 29 | ] 30 | 31 | 32 | async def get_predicted_category(inbound_request_content: str) -> str: 33 | """ 34 | This function uses the OpenAI Chat Completion API to predict the category of an inbound request. 35 | """ 36 | config = get_config() 37 | 38 | # Define the prompt 39 | messages = [ 40 | {"role": "system", "content": config.openai_prompt}, 41 | {"role": "user", "content": inbound_request_content}, 42 | ] 43 | 44 | # Call the API 45 | response = openai.chat.completions.create( 46 | model="gpt-4-32k", 47 | messages=messages, 48 | temperature=0, 49 | stream=False, 50 | functions=predict_category_functions(config.categories.values()), 51 | function_call={"name": "get_predicted_category"}, 52 | ) 53 | 54 | function_args = json.loads(response.choices[0].message.function_call.arguments) # type: ignore 55 | return function_args["category"] 56 | -------------------------------------------------------------------------------- /bots/triage-slackbot/triage_slackbot/templates/blocks/empty_category_warning.j2: -------------------------------------------------------------------------------- 1 | { 2 | "type": "context", 3 | "block_id": "empty_category_warning_block", 4 | "elements": [{"type": "plain_text", "text": "Category is required."}] 5 | } -------------------------------------------------------------------------------- /bots/triage-slackbot/triage_slackbot/templates/blocks/empty_conversation_warning.j2: -------------------------------------------------------------------------------- 1 | { 2 | "type": "context", 3 | "block_id": "empty_conversation_warning_block", 4 | "elements": [{"type": "plain_text", "text": "Conversation is required."}] 5 | } 6 | -------------------------------------------------------------------------------- /bots/triage-slackbot/triage_slackbot/templates/blocks/select_conversation.j2: -------------------------------------------------------------------------------- 1 | { 2 | "type": "section", 3 | "text": {"type": "mrkdwn", "text": "*Select a channel*"}, 4 | "accessory": { 5 | "type": "conversations_select", 6 | "placeholder": { 7 | "type": "plain_text", 8 | "text": "Select conversations", 9 | "emoji": true 10 | }, 11 | "action_id": "recategorize_select_conversation_action" 12 | }, 13 | "block_id": "recategorize_select_conversation_block" 14 | } -------------------------------------------------------------------------------- /bots/triage-slackbot/triage_slackbot/templates/messages/_notify_oncall_body.j2: -------------------------------------------------------------------------------- 1 | { 2 | "type": "context", 3 | "elements": [ 4 | { 5 | "type": "plain_text", 6 | "text": ":thumbsup: Acknowledge this message and response directly to the inbound request.", 7 | "emoji": true 8 | }, 9 | { 10 | "type": "plain_text", 11 | "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.", 12 | "emoji": true 13 | } 14 | ] 15 | }, 16 | { 17 | "type": "actions", 18 | "elements": [ 19 | { 20 | "type": "button", 21 | "text": { 22 | "type": "plain_text", 23 | "emoji": true, 24 | "text": "Acknowledge" 25 | }, 26 | "style": "primary", 27 | "value": "{{ predicted_category }}", 28 | "action_id": "acknowledge_submit_action" 29 | }, 30 | { 31 | "type": "button", 32 | "text": { 33 | "type": "plain_text", 34 | "emoji": true, 35 | "text": "Inaccurate, recategorize" 36 | }, 37 | "style": "danger", 38 | "value": "recategorize", 39 | "action_id": "recategorize_submit_action" 40 | } 41 | ] 42 | }, 43 | { 44 | "type": "section", 45 | "block_id": "recategorize_select_category_block", 46 | "text": { 47 | "type": "mrkdwn", 48 | "text": "*Select a category from the dropdown list, or*" 49 | }, 50 | "accessory": { 51 | "type": "static_select", 52 | "placeholder": { 53 | "type": "plain_text", 54 | "text": "Select an item", 55 | "emoji": true 56 | }, 57 | "options": [ 58 | {% for value, text in options.items() %} 59 | { 60 | "text": { 61 | "type": "plain_text", 62 | "text": "{{ text }}", 63 | "emoji": true 64 | }, 65 | "value": "{{ value }}" 66 | }{% if not loop.last %},{% endif %} 67 | {% endfor %} 68 | ], 69 | "action_id": "recategorize_select_category_action" 70 | } 71 | } 72 | -------------------------------------------------------------------------------- /bots/triage-slackbot/triage_slackbot/templates/messages/autorespond.j2: -------------------------------------------------------------------------------- 1 | [ 2 | { 3 | "type": "section", 4 | "text": { 5 | "type": "mrkdwn", 6 | "text": "{{ text }}" 7 | } 8 | }, 9 | { 10 | "type": "context", 11 | "elements": [ 12 | { 13 | "type": "plain_text", 14 | "text": "If you feel strongly this is a Security issue, respond to this thread and someone from our team will get back to you.", 15 | "emoji": true 16 | } 17 | ] 18 | } 19 | ] 20 | -------------------------------------------------------------------------------- /bots/triage-slackbot/triage_slackbot/templates/messages/feed.j2: -------------------------------------------------------------------------------- 1 | [ 2 | { 3 | "type": "section", 4 | "text": { 5 | "type": "mrkdwn", 6 | "text": "Received an <{{ inbound_message_url }}|inbound message> in <#{{ inbound_message_channel }}>:" 7 | } 8 | }, 9 | { 10 | "type": "context", 11 | "elements": [ 12 | { 13 | "type": "plain_text", 14 | "text": "Predicted category: {{ predicted_category }}", 15 | "emoji": true 16 | }, 17 | { 18 | "type": "mrkdwn", 19 | "text": "Triaged to: {{ oncall_mention }}" 20 | }, 21 | { 22 | "type": "plain_text", 23 | "text": "Triage updates in the :thread:", 24 | "emoji": true 25 | } 26 | ] 27 | } 28 | ] 29 | -------------------------------------------------------------------------------- /bots/triage-slackbot/triage_slackbot/templates/messages/notify_oncall_channel.j2: -------------------------------------------------------------------------------- 1 | [ 2 | { 3 | "type": "section", 4 | "text": { 5 | "type": "mrkdwn", 6 | "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" 7 | } 8 | }, 9 | {% include 'messages/_notify_oncall_body.j2' %} 10 | ] 11 | -------------------------------------------------------------------------------- /bots/triage-slackbot/triage_slackbot/templates/messages/notify_oncall_in_feed.j2: -------------------------------------------------------------------------------- 1 | [ 2 | { 3 | "type": "section", 4 | "text": { 5 | "type": "mrkdwn", 6 | "text": "{{ oncall_greeting }}, is this assignment accurate?\n\n" 7 | } 8 | }, 9 | {% include 'messages/_notify_oncall_body.j2' %} 10 | ] 11 | -------------------------------------------------------------------------------- /shared/openai-slackbot/openai_slackbot/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/openai/openai-security-bots/20fb09044cae19a6b25a192ccb605d24d3743d88/shared/openai-slackbot/openai_slackbot/__init__.py -------------------------------------------------------------------------------- /shared/openai-slackbot/openai_slackbot/bot.py: -------------------------------------------------------------------------------- 1 | import typing as t 2 | from logging import getLogger 3 | 4 | import openai 5 | from openai_slackbot.clients.slack import SlackClient 6 | from openai_slackbot.handlers import BaseActionHandler, BaseMessageHandler 7 | from openai_slackbot.utils.envvars import string 8 | from slack_bolt.adapter.socket_mode.async_handler import AsyncSocketModeHandler 9 | from slack_bolt.app.async_app import AsyncApp 10 | 11 | logger = getLogger(__name__) 12 | 13 | 14 | async def register_app_handlers( 15 | *, 16 | app: AsyncApp, 17 | message_handler: t.Type[BaseMessageHandler], 18 | action_handlers: t.List[t.Type[BaseActionHandler]], 19 | slack_client: SlackClient, 20 | ): 21 | if message_handler: 22 | app.event("message")(message_handler(slack_client).maybe_handle) 23 | 24 | if action_handlers: 25 | for action_handler in action_handlers: 26 | handler = action_handler(slack_client) 27 | app.action(handler.action_id)(handler.maybe_handle) 28 | 29 | 30 | async def init_bot( 31 | *, 32 | openai_organization_id: str, 33 | slack_message_handler: t.Type[BaseMessageHandler], 34 | slack_action_handlers: t.List[t.Type[BaseActionHandler]], 35 | slack_template_path: str, 36 | ): 37 | slack_bot_token = string("SLACK_BOT_TOKEN") 38 | openai_api_key = string("OPENAI_API_KEY") 39 | 40 | # Init OpenAI API 41 | openai.organization = openai_organization_id 42 | openai.api_key = openai_api_key 43 | 44 | # Init slack bot 45 | app = AsyncApp(token=slack_bot_token) 46 | slack_client = SlackClient(app.client, slack_template_path) 47 | await register_app_handlers( 48 | app=app, 49 | message_handler=slack_message_handler, 50 | action_handlers=slack_action_handlers, 51 | slack_client=slack_client, 52 | ) 53 | 54 | return app 55 | 56 | 57 | async def start_app(app): 58 | socket_app_token = string("SOCKET_APP_TOKEN") 59 | handler = AsyncSocketModeHandler(app, socket_app_token) 60 | await handler.start_async() 61 | 62 | 63 | async def start_bot( 64 | *, 65 | openai_organization_id: str, 66 | slack_message_handler: t.Type[BaseMessageHandler], 67 | slack_action_handlers: t.List[t.Type[BaseActionHandler]], 68 | slack_template_path: str, 69 | ): 70 | app = await init_bot( 71 | openai_organization_id=openai_organization_id, 72 | slack_message_handler=slack_message_handler, 73 | slack_action_handlers=slack_action_handlers, 74 | slack_template_path=slack_template_path, 75 | ) 76 | 77 | await start_app(app) 78 | -------------------------------------------------------------------------------- /shared/openai-slackbot/openai_slackbot/clients/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/openai/openai-security-bots/20fb09044cae19a6b25a192ccb605d24d3743d88/shared/openai-slackbot/openai_slackbot/clients/__init__.py -------------------------------------------------------------------------------- /shared/openai-slackbot/openai_slackbot/clients/slack.py: -------------------------------------------------------------------------------- 1 | import json 2 | import os 3 | import typing as t 4 | from logging import getLogger 5 | 6 | from jinja2 import Environment, FileSystemLoader 7 | from pydantic import BaseModel 8 | from slack_sdk.errors import SlackApiError 9 | from slack_sdk.web.async_client import AsyncWebClient 10 | 11 | logger = getLogger(__name__) 12 | 13 | 14 | class SlackMessage(BaseModel): 15 | app_id: t.Optional[str] = None 16 | blocks: t.Optional[t.List[t.Any]] = None 17 | bot_id: t.Optional[str] = None 18 | bot_profile: t.Optional[t.Dict[str, t.Any]] = None 19 | team: str 20 | text: str 21 | ts: str 22 | type: str 23 | user: t.Optional[str] = None 24 | 25 | 26 | class CreateSlackMessageResponse(BaseModel): 27 | ok: bool 28 | channel: str 29 | ts: str 30 | message: SlackMessage 31 | 32 | 33 | class SlackClient: 34 | """ 35 | SlackClient wraps the Slack AsyncWebClient implementation and 36 | provides some additional functionality specific to the Slackbot 37 | implementation. 38 | """ 39 | 40 | def __init__(self, client: AsyncWebClient, template_path: str) -> None: 41 | self._client = client 42 | self._jinja = self._init_jinja(template_path) 43 | 44 | async def get_message_link(self, **kwargs) -> str: 45 | response = await self._client.chat_getPermalink(**kwargs) 46 | if not response["ok"]: 47 | raise Exception(f"Failed to get Slack message link: {response['error']}") 48 | return response["permalink"] 49 | 50 | async def get_message(self, channel: str, ts: str) -> t.Optional[t.Dict[str, t.Any]]: 51 | """Follows: https://api.slack.com/messaging/retrieving.""" 52 | result = await self._client.conversations_history( 53 | channel=channel, 54 | inclusive=True, 55 | latest=ts, 56 | limit=1, 57 | ) 58 | return result["messages"][0] if result["messages"] else None 59 | 60 | async def post_message(self, **kwargs) -> CreateSlackMessageResponse: 61 | response = await self._client.chat_postMessage(**kwargs) 62 | if not response["ok"]: 63 | raise Exception(f"Failed to post Slack message: {response['error']}") 64 | 65 | assert isinstance(response.data, dict) 66 | return CreateSlackMessageResponse(**response.data) 67 | 68 | async def update_message(self, **kwargs) -> t.Dict[str, t.Any]: 69 | response = await self._client.chat_update(**kwargs) 70 | if not response["ok"]: 71 | raise Exception(f"Failed to update Slack message: {response['error']}") 72 | 73 | assert isinstance(response.data, dict) 74 | return response.data 75 | 76 | async def add_reaction(self, **kwargs) -> t.Dict[str, t.Any]: 77 | try: 78 | response = await self._client.reactions_add(**kwargs) 79 | except SlackApiError as e: 80 | if e.response["error"] == "already_reacted": 81 | return {} 82 | raise e 83 | 84 | assert isinstance(response.data, dict) 85 | return response.data 86 | 87 | async def get_thread_messages(self, channel: str, thread_ts: str) -> t.List[t.Dict[str, t.Any]]: 88 | response = await self._client.conversations_replies(channel=channel, ts=thread_ts) 89 | if not response["ok"]: 90 | raise Exception(f"Failed to get thread messages: {response['error']}") 91 | 92 | assert isinstance(response.data, dict) 93 | return response.data["messages"] 94 | 95 | async def get_user_display_name(self, user_id: str) -> str: 96 | response = await self._client.users_info(user=user_id) 97 | if not response["ok"]: 98 | raise Exception(f"Failed to get user info: {response['error']}") 99 | return response["user"]["profile"]["display_name"] 100 | 101 | async def get_original_blocks(self, thread_ts: str, channel: str) -> None: 102 | """Given a thread_ts, get original message block""" 103 | response = await self._client.conversations_replies( 104 | channel=channel, 105 | ts=thread_ts, 106 | ) 107 | try: 108 | messages = response.get("messages", []) 109 | if not messages: 110 | raise ValueError(f"Error fetching original message for thread_ts {thread_ts}") 111 | blocks = messages[0].get("blocks") 112 | if not blocks: 113 | raise ValueError(f"Error fetching original message for thread_ts {thread_ts}") 114 | return blocks 115 | except Exception as e: 116 | logger.exception(f"Error fetching original message for thread_ts {thread_ts}: {e}") 117 | 118 | def render_blocks_from_template(self, template_filename: str, context: t.Dict = {}) -> t.Any: 119 | rendered_template = self._jinja.get_template(template_filename).render(context) 120 | return json.loads(rendered_template) 121 | 122 | def _init_jinja(self, template_path: str): 123 | templates_dir = os.path.join(template_path) 124 | return Environment(loader=FileSystemLoader(templates_dir)) 125 | -------------------------------------------------------------------------------- /shared/openai-slackbot/openai_slackbot/handlers.py: -------------------------------------------------------------------------------- 1 | import abc 2 | import typing as t 3 | from logging import getLogger 4 | 5 | from openai_slackbot.clients.slack import SlackClient 6 | 7 | logger = getLogger(__name__) 8 | 9 | 10 | class BaseHandler(abc.ABC): 11 | def __init__(self, slack_client: SlackClient) -> None: 12 | self._slack_client = slack_client 13 | 14 | async def maybe_handle(self, args): 15 | await args.ack() 16 | 17 | logging_extra = self.logging_extra(args) 18 | try: 19 | should_handle = await self.should_handle(args) 20 | logger.info( 21 | f"Handler: {self.__class__.__name__}, should handle: {should_handle}", 22 | extra=logging_extra, 23 | ) 24 | if should_handle: 25 | await self.handle(args) 26 | except Exception: 27 | logger.exception("Failed to handle event", extra=logging_extra) 28 | 29 | @abc.abstractmethod 30 | async def should_handle(self, args) -> bool: 31 | ... 32 | 33 | @abc.abstractmethod 34 | async def handle(self, args): 35 | ... 36 | 37 | @abc.abstractmethod 38 | def logging_extra(self, args) -> t.Dict[str, t.Any]: 39 | ... 40 | 41 | 42 | class BaseMessageHandler(BaseHandler): 43 | def logging_extra(self, args) -> t.Dict[str, t.Any]: 44 | fields = {} 45 | for field in ["type", "subtype", "channel", "ts"]: 46 | fields[field] = args.event.get(field) 47 | return fields 48 | 49 | 50 | class BaseActionHandler(BaseHandler): 51 | @abc.abstractproperty 52 | def action_id(self) -> str: 53 | ... 54 | 55 | async def should_handle(self, args) -> bool: 56 | return True 57 | 58 | def logging_extra(self, args) -> t.Dict[str, t.Any]: 59 | return { 60 | "action_type": args.body.get("type"), 61 | "action": args.body.get("actions", [])[0], 62 | } 63 | -------------------------------------------------------------------------------- /shared/openai-slackbot/openai_slackbot/utils/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/openai/openai-security-bots/20fb09044cae19a6b25a192ccb605d24d3743d88/shared/openai-slackbot/openai_slackbot/utils/__init__.py -------------------------------------------------------------------------------- /shared/openai-slackbot/openai_slackbot/utils/envvars.py: -------------------------------------------------------------------------------- 1 | import os 2 | import typing as t 3 | 4 | 5 | def string(key: str, default: t.Optional[str] = None) -> str: 6 | val = os.environ.get(key) 7 | if not val: 8 | if default is None: 9 | raise ValueError(f"Missing required environment variable: {key}") 10 | return default 11 | return val 12 | -------------------------------------------------------------------------------- /shared/openai-slackbot/openai_slackbot/utils/slack.py: -------------------------------------------------------------------------------- 1 | import typing as t 2 | 3 | RenderedSlackBlock = t.NewType("RenderedSlackBlock", t.Dict[str, t.Any]) 4 | 5 | 6 | def block_id_exists(blocks: t.List[RenderedSlackBlock], block_id: str) -> bool: 7 | return any([block.get("block_id") == block_id for block in blocks]) 8 | 9 | 10 | def remove_block_id_if_exists(blocks: t.List[RenderedSlackBlock], block_id: str) -> t.List: 11 | return [block for block in blocks if block.get("block_id") != block_id] 12 | 13 | 14 | def get_block_by_id(blocks: t.Dict, block_id: str) -> t.Dict: 15 | for block in blocks: 16 | if block.get("block_id") == block_id: 17 | return block 18 | return {} 19 | 20 | 21 | def extract_text_from_event(event) -> str: 22 | """Extracts text from either plaintext and block message.""" 23 | 24 | # Extract text from plaintext message. 25 | text = event.get("text") 26 | if text: 27 | return text 28 | 29 | # Extract text from message blocks. 30 | texts = [] 31 | attachments = event.get("attachments", []) 32 | for attachment in attachments: 33 | attachment_message_blocks = attachment.get("message_blocks", []) 34 | for amb in attachment_message_blocks: 35 | message_blocks = amb.get("message", {}).get("blocks", []) 36 | for mb in message_blocks: 37 | mb_elements = mb.get("elements", []) 38 | for mbe in mb_elements: 39 | mbe_elements = mbe.get("elements", []) 40 | for mbee in mbe_elements: 41 | if mbee.get("type") == "text": 42 | texts.append(mbee["text"]) 43 | 44 | return " ".join(texts).strip() 45 | 46 | 47 | def render_slack_id_to_mention(id: str): 48 | """Render a usergroup or user ID to a mention.""" 49 | 50 | if not id: 51 | return "" 52 | elif id.startswith("U"): 53 | return f"<@{id}>" 54 | elif id.startswith("S"): 55 | return f"" 56 | elif id.startswith("C"): 57 | return f"<#{id}>" 58 | else: 59 | raise ValueError(f"Unsupported/invalid ID type: {id}") 60 | 61 | 62 | def render_slack_url(*, url: str, text: str) -> str: 63 | """Render a URL to a clickable link.""" 64 | return f"<{url}|{text}>" 65 | -------------------------------------------------------------------------------- /shared/openai-slackbot/pyproject.toml: -------------------------------------------------------------------------------- 1 | [project] 2 | name = "openai-slackbot" 3 | requires-python = ">=3.8" 4 | version = "1.0.0" 5 | dependencies = [ 6 | "aiohttp", 7 | "Jinja2", 8 | "openai", 9 | "pydantic", 10 | "python-dotenv", 11 | "slack-bolt", 12 | "slack-sdk", 13 | "pytest", 14 | "pytest-env", 15 | "pytest-asyncio", 16 | "aiohttp", 17 | ] 18 | 19 | [build-system] 20 | requires = ["setuptools>=64.0"] 21 | build-backend = "setuptools.build_meta" 22 | 23 | [tool.pytest.ini_options] 24 | asyncio_mode = "auto" 25 | env = [ 26 | "SLACK_BOT_TOKEN=mock-token", 27 | "SOCKET_APP_TOKEN=mock-token", 28 | "OPENAI_API_KEY=mock-key", 29 | ] 30 | -------------------------------------------------------------------------------- /shared/openai-slackbot/setup.cfg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/openai/openai-security-bots/20fb09044cae19a6b25a192ccb605d24d3743d88/shared/openai-slackbot/setup.cfg -------------------------------------------------------------------------------- /shared/openai-slackbot/tests/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/openai/openai-security-bots/20fb09044cae19a6b25a192ccb605d24d3743d88/shared/openai-slackbot/tests/__init__.py -------------------------------------------------------------------------------- /shared/openai-slackbot/tests/clients/__init__.py: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/openai/openai-security-bots/20fb09044cae19a6b25a192ccb605d24d3743d88/shared/openai-slackbot/tests/clients/__init__.py -------------------------------------------------------------------------------- /shared/openai-slackbot/tests/clients/test_slack.py: -------------------------------------------------------------------------------- 1 | from unittest.mock import AsyncMock, MagicMock 2 | 3 | import pytest 4 | from openai_slackbot.clients.slack import CreateSlackMessageResponse 5 | from slack_sdk.errors import SlackApiError 6 | 7 | 8 | async def test_get_message_link_success(mock_slack_client): 9 | mock_slack_client._client.chat_getPermalink = AsyncMock( 10 | return_value={ 11 | "ok": True, 12 | "channel": "C123456", 13 | "permalink": "https://myorg.slack.com/archives/C123456/p1234567890", 14 | } 15 | ) 16 | link = await mock_slack_client.get_message_link(channel="channel", message_ts="message_ts") 17 | mock_slack_client._client.chat_getPermalink.assert_called_once_with( 18 | channel="channel", message_ts="message_ts" 19 | ) 20 | assert link == "https://myorg.slack.com/archives/C123456/p1234567890" 21 | 22 | 23 | async def test_get_message_link_failed(mock_slack_client): 24 | mock_slack_client._client.chat_getPermalink = AsyncMock( 25 | return_value={"ok": False, "error": "failed"} 26 | ) 27 | with pytest.raises(Exception): 28 | await mock_slack_client.get_message_link(channel="channel", message_ts="message_ts") 29 | mock_slack_client._client.chat_getPermalink.assert_called_once_with( 30 | channel="channel", message_ts="message_ts" 31 | ) 32 | 33 | 34 | async def test_post_message_success(mock_slack_client): 35 | mock_message_data = { 36 | "ok": True, 37 | "channel": "C234567", 38 | "ts": "ts", 39 | "message": { 40 | "bot_id": "bot_id", 41 | "bot_profile": {"id": "bot_profile_id"}, 42 | "team": "team", 43 | "text": "text", 44 | "ts": "ts", 45 | "type": "type", 46 | "user": "user", 47 | }, 48 | } 49 | mock_response = MagicMock(data=mock_message_data) 50 | mock_response.__getitem__.side_effect = mock_message_data.__getitem__ 51 | mock_slack_client._client.chat_postMessage = AsyncMock(return_value=mock_response) 52 | 53 | response = await mock_slack_client.post_message(channel="C234567", text="text") 54 | assert response == CreateSlackMessageResponse(**mock_message_data) 55 | 56 | 57 | async def test_post_message_failed(mock_slack_client): 58 | mock_slack_client._client.chat_postMessage = AsyncMock( 59 | return_value={"ok": False, "error": "failed"} 60 | ) 61 | with pytest.raises(Exception): 62 | await mock_slack_client.post_message(channel="channel", text="text") 63 | mock_slack_client._client.chat_postMessage.assert_called_once_with( 64 | channel="channel", text="text" 65 | ) 66 | 67 | 68 | async def test_update_message_success(mock_slack_client): 69 | mock_message_data = { 70 | "ok": True, 71 | "channel": "C234567", 72 | "ts": "ts", 73 | "message": { 74 | "bot_id": "bot_id", 75 | "bot_profile": {"id": "bot_profile_id"}, 76 | "team": "team", 77 | "text": "text", 78 | "ts": "ts", 79 | "type": "type", 80 | "user": "user", 81 | }, 82 | } 83 | mock_response = MagicMock(data=mock_message_data) 84 | mock_response.__getitem__.side_effect = mock_message_data.__getitem__ 85 | mock_slack_client._client.chat_update = AsyncMock(return_value=mock_response) 86 | 87 | response = await mock_slack_client.update_message(channel="C234567", ts="ts", text="text") 88 | assert response == mock_message_data 89 | 90 | 91 | async def test_update_message_failed(mock_slack_client): 92 | mock_slack_client._client.chat_update = AsyncMock(return_value={"ok": False, "error": "failed"}) 93 | with pytest.raises(Exception): 94 | await mock_slack_client.update_message(channel="channel", ts="ts", text="text") 95 | mock_slack_client._client.chat_update.assert_called_once_with( 96 | channel="channel", ts="ts", text="text" 97 | ) 98 | 99 | 100 | async def test_add_reaction_success(mock_slack_client): 101 | mock_response_data = {"ok": True} 102 | mock_response = MagicMock(data=mock_response_data) 103 | mock_response.__getitem__.side_effect = mock_response_data.__getitem__ 104 | mock_slack_client._client.reactions_add = AsyncMock(return_value=mock_response) 105 | await mock_slack_client.add_reaction(channel="channel", name="thumbsup", timestamp="timestamp") 106 | 107 | 108 | async def test_add_reaction_already_reacted(mock_slack_client): 109 | mock_slack_client._client.reactions_add = AsyncMock( 110 | side_effect=SlackApiError("already_reacted", {"error": "already_reacted"}) 111 | ) 112 | response = await mock_slack_client.add_reaction( 113 | channel="channel", name="thumbsup", timestamp="timestamp" 114 | ) 115 | assert response == {} 116 | 117 | 118 | async def test_add_reaction_failed(mock_slack_client): 119 | mock_slack_client._client.reactions_add = AsyncMock( 120 | side_effect=SlackApiError("failed", {"error": "invalid_reaction"}) 121 | ) 122 | with pytest.raises(Exception): 123 | await mock_slack_client.add_reaction( 124 | channel="channel", name="thumbsup", timestamp="timestamp" 125 | ) 126 | -------------------------------------------------------------------------------- /shared/openai-slackbot/tests/conftest.py: -------------------------------------------------------------------------------- 1 | from unittest.mock import AsyncMock, MagicMock, patch 2 | 3 | import pytest 4 | from openai_slackbot.clients.slack import SlackClient 5 | from openai_slackbot.handlers import BaseActionHandler, BaseMessageHandler 6 | 7 | 8 | @pytest.fixture 9 | def mock_slack_app(): 10 | with patch("slack_bolt.app.async_app.AsyncApp") as mock_app: 11 | yield mock_app.return_value 12 | 13 | 14 | @pytest.fixture 15 | def mock_socket_mode_handler(): 16 | with patch( 17 | "slack_bolt.adapter.socket_mode.async_handler.AsyncSocketModeHandler" 18 | ) as mock_handler: 19 | mock_handler_object = mock_handler.return_value 20 | mock_handler_object.start_async = AsyncMock() 21 | yield mock_handler_object 22 | 23 | 24 | @pytest.fixture 25 | def mock_openai(): 26 | mock_openai = MagicMock() 27 | with patch.dict("sys.modules", openai=mock_openai): 28 | yield mock_openai 29 | 30 | 31 | @pytest.fixture 32 | def mock_slack_asyncwebclient(): 33 | with patch("slack_sdk.web.async_client.AsyncWebClient") as mock_client: 34 | yield mock_client.return_value 35 | 36 | 37 | @pytest.fixture 38 | def mock_slack_client(mock_slack_asyncwebclient): 39 | return SlackClient(mock_slack_asyncwebclient, "template_path") 40 | 41 | 42 | @pytest.fixture 43 | def mock_message_handler(mock_slack_client): 44 | return MockMessageHandler(mock_slack_client) 45 | 46 | 47 | @pytest.fixture 48 | def mock_action_handler(mock_slack_client): 49 | return MockActionHandler(mock_slack_client) 50 | 51 | 52 | class MockMessageHandler(BaseMessageHandler): 53 | def __init__(self, slack_client): 54 | super().__init__(slack_client) 55 | self.mock_handler = AsyncMock() 56 | 57 | async def should_handle(self, args): 58 | return args.event["subtype"] != "bot_message" 59 | 60 | async def handle(self, args): 61 | await self.mock_handler(args) 62 | 63 | 64 | class MockActionHandler(BaseActionHandler): 65 | def __init__(self, slack_client): 66 | super().__init__(slack_client) 67 | self.mock_handler = AsyncMock() 68 | 69 | async def handle(self, args): 70 | await self.mock_handler(args) 71 | 72 | @property 73 | def action_id(self): 74 | return "mock_action" 75 | -------------------------------------------------------------------------------- /shared/openai-slackbot/tests/test_bot.py: -------------------------------------------------------------------------------- 1 | import pytest 2 | 3 | 4 | async def test_start_bot( 5 | mock_slack_app, mock_socket_mode_handler, mock_message_handler, mock_action_handler 6 | ): 7 | from openai_slackbot.bot import start_bot 8 | 9 | await start_bot( 10 | openai_organization_id="org-id", 11 | slack_message_handler=mock_message_handler.__class__, 12 | slack_action_handlers=[mock_action_handler.__class__], 13 | slack_template_path="/path/to/templates", 14 | ) 15 | 16 | mock_slack_app.event.assert_called_once_with("message") 17 | mock_slack_app.action.assert_called_once_with("mock_action") 18 | mock_socket_mode_handler.start_async.assert_called_once() 19 | -------------------------------------------------------------------------------- /shared/openai-slackbot/tests/test_handlers.py: -------------------------------------------------------------------------------- 1 | from unittest.mock import AsyncMock, MagicMock 2 | 3 | import pytest 4 | 5 | 6 | @pytest.mark.parametrize("subtype, should_handle", [("message", True), ("bot_message", False)]) 7 | async def test_message_handler(mock_message_handler, subtype, should_handle): 8 | args = MagicMock( 9 | ack=AsyncMock(), 10 | event={"type": "message", "subtype": subtype, "channel": "channel", "ts": "ts"}, 11 | ) 12 | 13 | await mock_message_handler.maybe_handle(args) 14 | args.ack.assert_awaited_once() 15 | if should_handle: 16 | mock_message_handler.mock_handler.assert_awaited_once_with(args) 17 | else: 18 | mock_message_handler.mock_handler.assert_not_awaited() 19 | 20 | assert mock_message_handler.logging_extra(args) == { 21 | "type": "message", 22 | "subtype": subtype, 23 | "channel": "channel", 24 | "ts": "ts", 25 | } 26 | 27 | 28 | async def test_action_handler(mock_action_handler): 29 | args = MagicMock( 30 | ack=AsyncMock(), 31 | body={ 32 | "type": "type", 33 | "actions": ["action"], 34 | }, 35 | ) 36 | 37 | await mock_action_handler.maybe_handle(args) 38 | args.ack.assert_awaited_once() 39 | mock_action_handler.mock_handler.assert_awaited_once_with(args) 40 | 41 | assert mock_action_handler.logging_extra(args) == { 42 | "action_type": "type", 43 | "action": "action", 44 | } 45 | --------------------------------------------------------------------------------