diff --git a/pastepwn/actions/__init__.py b/pastepwn/actions/__init__.py index cad74ae..700aba1 100644 --- a/pastepwn/actions/__init__.py +++ b/pastepwn/actions/__init__.py @@ -8,6 +8,7 @@ from .databaseaction import DatabaseAction from .savejsonaction import SaveJSONAction from .twitteraction import TwitterAction +from .discordaction import DiscordAction __all__ = ( "BasicAction", @@ -18,4 +19,5 @@ "DatabaseAction", "SaveJSONAction", "TwitterAction", + "DiscordAction", ) diff --git a/pastepwn/actions/discordaction.py b/pastepwn/actions/discordaction.py new file mode 100644 index 0000000..0ee418d --- /dev/null +++ b/pastepwn/actions/discordaction.py @@ -0,0 +1,114 @@ +# -*- coding: utf-8 -*- +import asyncio +import json +import logging +import sys +import websockets +from string import Template + +from pastepwn.util import Request, DictWrapper +from .basicaction import BasicAction + + +class DiscordAction(BasicAction): + """Action to send a Discord message to a certain webhook or channel.""" + name = "DiscordAction" + + def __init__(self, webhook=None, token=None, channel_id=None, custom_payload=None, template=None): + super().__init__() + self.logger = logging.getLogger(__name__) + + self.webhook = webhook + if webhook is None: + if token is None or channel_id is None: + raise ValueError('Invalid arguments: requires either webhook or token+channel_id arguments') + self.token = token + self.channel_id = channel_id + self.identified = False + self.custom_payload = custom_payload + if template is not None: + self.template = Template(template) + else: + self.template = None + + @asyncio.coroutine + def _identify(self, ws_url): + """Connect to the Discord Gateway to identify the bot.""" + # Docs: https://discordapp.com/developers/docs/topics/gateway#connecting-to-the-gateway + # Open connection to the Discord Gateway + socket = yield from websockets.connect(ws_url + '/?v=6&encoding=json') + try: + # Receive Hello + hello_str = yield from socket.recv() + hello = json.loads(hello_str) + if hello.get('op') != 10: + self.logger.warning('[ws] Expected Hello payload but received %s', hello_str) + + # Send heartbeat and receive ACK + yield from socket.send(json.dumps({"op": 1, "d": {}})) + ack_str = yield from socket.recv() + ack = json.loads(ack_str) + if ack.get('op') != 11: + self.logger.warning('[ws] Expected Heartbeat ACK payload but received %s', ack_str) + + # Identify + payload = { + "token": self.token, + "properties": { + "$os": sys.platform, + "$browser": "pastepwn", + "$device": "pastepwn" + } + } + yield from socket.send(json.dumps({"op": 2, "d": payload})) + + # Receive READY event + ready_str = yield from socket.recv() + ready = json.loads(ready_str) + if ready.get('t') != 'READY': + self.logger.warning('[ws] Expected READY event but received %s', ready_str) + finally: + # Close websocket connection + yield from socket.close() + + def initialize_gateway(self): + """Initialize the bot token so Discord identifies it properly.""" + if self.webhook is not None: + raise NotImplementedError('Gateway initialization is only necessary for bot accounts.') + + # Call Get Gateway Bot to get the websocket URL + r = Request() + r.headers = {'Authorization': 'Bot {}'.format(self.token)} + res = json.loads(r.get('https://discordapp.com/api/gateway/bot')) + ws_url = res.get('url') + + # Start websocket client + asyncio.get_event_loop().run_until_complete(self._identify(ws_url)) + self.identified = True + + def perform(self, paste, analyzer_name=None): + """Send a message via Discord to a specified channel, without checking for errors""" + r = Request() + if self.template is None: + text = "New paste matched by analyzer '{0}' - Link: {1}".format(analyzer_name, paste.full_url) + else: + paste_dict = paste.to_dict() + paste_dict["analyzer_name"] = analyzer_name + text = self.template.safe_substitute(DictWrapper(paste_dict)) + + if self.webhook is not None: + # Send to a webhook (no authentication) + url = self.webhook + else: + # Send through Discord bot API (header-based authentication) + url = 'https://discordapp.com/api/channels/{0}/messages'.format(self.channel_id) + r.headers = {'Authorization': 'Bot {}'.format(self.token)} + + res = json.loads(r.post(url, {"content": text})) + + if res.get('code') == 40001 and self.webhook is None and not self.identified: + # Unauthorized access, bot token hasn't been identified to Discord Gateway + self.logger.info('Accessing Discord Gateway to initialize token') + self.initialize_gateway() + # Retry action + self.perform(paste, analyzer_name=analyzer_name) diff --git a/requirements.txt b/requirements.txt index dab93a2..c6120db 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ pymongo mysql-connector-python requests -python-twitter \ No newline at end of file +python-twitter +websockets>=7.0,<8 \ No newline at end of file