Merge pull request #39841 from rafaelcaricio/telegram-beacon

Telegram beacon

AUTHORS

@@ -97,6 +97,7 @@ Pedro Algarvio <pedro@algarvio.me>
 Peter Baumgartner
 Pierre Carrier <pierre@spotify.com>
 Rhys Elsmore <me@rhys.io>
+Rafael Caricio <rafael@caricio.com>
 Robert Fielding
 Sean Channel <pentabular@gmail.com>
 Seth House <seth@eseth.com>

doc/ref/beacons/all/index.rst

@@ -32,5 +32,6 @@ beacon modules
     service
     sh
     status
+    telegram_bot_msg
     twilio_txt_msg
     wtmp

doc/ref/beacons/all/salt.beacons.telegram_bot_msg.rst

@@ -0,0 +1,6 @@
+=============================
+salt.beacons.telegram_bot_msg
+=============================
+
+.. automodule:: salt.beacons.telegram_bot_msg
+    :members:

salt/beacons/telegram_bot_msg.py

@@ -0,0 +1,95 @@
+# -*- coding: utf-8 -*-
+'''
+Beacon to emit Telegram messages
+'''
+
+# Import Python libs
+from __future__ import absolute_import
+import logging
+
+# Import 3rd Party libs
+try:
+    import telegram
+    logging.getLogger('telegram').setLevel(logging.CRITICAL)
+    HAS_TELEGRAM = True
+except ImportError:
+    HAS_TELEGRAM = False
+
+log = logging.getLogger(__name__)
+
+
+__virtualname__ = 'telegram_bot_msg'
+
+
+def __virtual__():
+    if HAS_TELEGRAM:
+        return __virtualname__
+    else:
+        return False
+
+
+def __validate__(config):
+    '''
+    Validate the beacon configuration
+    '''
+    if not isinstance(config, dict):
+        return False, ('Configuration for telegram_bot_msg '
+                       'beacon must be a dictionary.')
+
+    if not all(config.get(required_config)
+               for required_config in ['token', 'accept_from']):
+        return False, ('Not all required configuration for '
+                       'telegram_bot_msg are set.')
+
+    if not isinstance(config.get('accept_from'), list):
+        return False, ('Configuration for telegram_bot_msg, '
+                       'accept_from must be a list of usernames.')
+
+    return True, 'Valid beacon configuration.'
+
+
+def beacon(config):
+    '''
+    Emit a dict with a key "msgs" whose value is a list of messages
+    sent to the configured bot by one of the allowed usernames.
+
+    .. code-block:: yaml
+
+        beacons:
+          telegram_bot_msg:
+            token: "<bot access token>"
+            accept_from:
+              - "<valid username>"
+            interval: 10
+
+    '''
+    log.debug('telegram_bot_msg beacon starting')
+    ret = []
+    output = {}
+    output['msgs'] = []
+
+    bot = telegram.Bot(config['token'])
+    updates = bot.get_updates(limit=100, timeout=0, network_delay=10)
+
+    log.debug('Num updates: {0}'.format(len(updates)))
+    if not updates:
+        log.debug('Telegram Bot beacon has no new messages')
+        return ret
+
+    latest_update_id = 0
+    for update in updates:
+        message = update.message
+
+        if update.update_id > latest_update_id:
+            latest_update_id = update.update_id
+
+        if message.chat.username in config['accept_from']:
+            output['msgs'].append(message.to_dict())
+
+    # mark in the server that previous messages are processed
+    bot.get_updates(offset=latest_update_id + 1)
+
+    log.debug('Emitting {0} messages.'.format(len(output['msgs'])))
+    if output['msgs']:
+        ret.append(output)
+    return ret

tests/unit/beacons/test_telegram_bot_msg_beacon.py

@@ -0,0 +1,134 @@
+# coding: utf-8
+
+# Python libs
+from __future__ import absolute_import
+import datetime
+
+# Salt libs
+from salt.beacons import telegram_bot_msg
+
+# Salt testing libs
+from tests.support.unit import TestCase, skipIf
+from tests.support.mock import NO_MOCK, NO_MOCK_REASON, MagicMock, patch
+
+# Third-party libs
+try:
+    import telegram
+    HAS_TELEGRAM = True
+except ImportError:
+    HAS_TELEGRAM = False
+
+
+@skipIf(not HAS_TELEGRAM, 'telegram is not available')
+@skipIf(NO_MOCK, NO_MOCK_REASON)
+class TelegramBotMsgBeaconTestCase(TestCase):
+    '''
+    Test case for salt.beacons.telegram_bot
+    '''
+    def setUp(self):
+        telegram_bot_msg.__context__ = {}
+
+    def test_validate_empty_config(self, *args, **kwargs):
+        ret = telegram_bot_msg.__validate__(None)
+        self.assertEqual(ret, (False, ('Configuration for telegram_bot_msg '
+                                       'beacon must be a dictionary.')))
+
+    def test_validate_missing_accept_from_config(self, *args, **kwargs):
+        ret = telegram_bot_msg.__validate__({
+            'token': 'bcd'
+        })
+        self.assertEqual(ret, (False, ('Not all required configuration for '
+                                       'telegram_bot_msg are set.')))
+
+    def test_validate_missing_token_config(self, *args, **kwargs):
+        ret = telegram_bot_msg.__validate__({
+            'accept_from': []
+        })
+        self.assertEqual(ret, (False, ('Not all required configuration for '
+                                       'telegram_bot_msg are set.')))
+
+    def test_validate_config_not_list_in_accept_from(self, *args, **kwargs):
+        ret = telegram_bot_msg.__validate__({
+            'token': 'bcd',
+            'accept_from': {'nodict': "1"}
+        })
+        self.assertEqual(ret, (False, ('Configuration for telegram_bot_msg, '
+                                       'accept_from must be a list of '
+                                       'usernames.')))
+
+    def test_validate_valid_config(self, *args, **kwargs):
+        ret = telegram_bot_msg.__validate__({
+            'token': 'bcd',
+            'accept_from': [
+                'username'
+            ]
+        })
+        self.assertEqual(ret, (True, 'Valid beacon configuration.'))
+
+    @patch("salt.beacons.telegram_bot_msg.telegram")
+    def test_call_no_updates(self, telegram_api, *args, **kwargs):
+        token = 'abc'
+        config = {
+            'token': token,
+            'accept_from': ['tester']
+        }
+        inst = MagicMock(name='telegram.Bot()')
+        telegram_api.Bot = MagicMock(name='telegram', return_value=inst)
+        inst.get_updates.return_value = []
+
+        ret = telegram_bot_msg.beacon(config)
+
+        telegram_api.Bot.assert_called_once_with(token)
+        self.assertEqual(ret, [])
+
+    @patch("salt.beacons.telegram_bot_msg.telegram")
+    def test_call_telegram_return_no_updates_for_user(
+            self, telegram_api, *args, **kwargs):
+        token = 'abc'
+        username = 'tester'
+        config = {
+            'token': token,
+            'accept_from': [username]
+        }
+        inst = MagicMock(name='telegram.Bot()')
+        telegram_api.Bot = MagicMock(name='telegram', return_value=inst)
+
+        username = 'different_user'
+        user = telegram.User(id=1, first_name='', username=username)
+        chat = telegram.Chat(1, 'private', username=username)
+        date = datetime.datetime(2016, 12, 18, 0, 0)
+        message = telegram.Message(1, user, date=date, chat=chat)
+        update = telegram.update.Update(update_id=1, message=message)
+
+        inst.get_updates.return_value = [update]
+
+        ret = telegram_bot_msg.beacon(config)
+
+        telegram_api.Bot.assert_called_once_with(token)
+        self.assertEqual(ret, [])
+
+    @patch("salt.beacons.telegram_bot_msg.telegram")
+    def test_call_telegram_returning_updates(self, telegram_api,
+                                             *args, **kwargs):
+        token = 'abc'
+        username = 'tester'
+        config = {
+            'token': token,
+            'accept_from': [username]
+        }
+        inst = MagicMock(name='telegram.Bot()')
+        telegram_api.Bot = MagicMock(name='telegram', return_value=inst)
+
+        user = telegram.User(id=1, first_name='', username=username)
+        chat = telegram.Chat(1, 'private', username=username)
+        date = datetime.datetime(2016, 12, 18, 0, 0)
+        message = telegram.Message(1, user, date=date, chat=chat)
+        update = telegram.update.Update(update_id=1, message=message)
+
+        inst.get_updates.return_value = [update]
+
+        ret = telegram_bot_msg.beacon(config)
+
+        telegram_api.Bot.assert_called_once_with(token)
+        self.assertTrue(ret)
+        self.assertEqual(ret[0]['msgs'][0], message.to_dict())