feat: add telegram bots healthcheck (#42)

* feat: add telegram bots healthcheck

* chore: add tests

.gitignore

@@ -129,4 +129,6 @@ dmypy.json
 .pyre/
 
 # Custom
-config/settings.yaml
+config/settings.yaml
+src/healthcheck_user_client.session
+.vscode

README.md

@@ -12,11 +12,27 @@ A simple Python script to verify if a service is up. Whenever the service falls,
 **Make a copy** of the file `config/settings.yaml.dist` in the same directory and rename it to `settings.yaml`:
 - Set `bot_token`, your bot token
 - Set `chat_ids`, the ID(s) the bot will use to communicate any downtime. It's possible to set multiple IDs, semicolon separated without any space
+- Set the urls you'd like to check in `urls.json`
+
+### Telegram bots
+This is bot is now able to contact other bots and check if their running correctly!
+In order to setup this module, obtain an `api_id` and `api_hash` following the [Telegram guide](https://core.telegram.org/api/obtaining_api_id), now add them your `settings.yaml`. 
+The last step is to add the bot(s) to check in the same settings file, completing the value for each key:
+- Set `username`, the bot to be checked username without the prepending '@'
+- Set `command`, the command that will be sent to your bot
+- Set `expected_response`, the expected response from your bot. Any special character should be replaced with their equivalent character (a new line should be replaced with `\n`)
+Multiple bots could be added by added just by repeating the same structure of a sequence of objects in yaml. 
 
 ### Example
 ```yaml
 bot_token: "123456:ABC-DEF1234ghIkl-zyx57W2v1u123ew11" # Your bot token
 chat_ids: ["10000000", "10000001", "10000002"]  # Single ID or Multiple IDs
+api_id: "123456"
+api_hash: "123456:ABC-DEF1234ghIkl-zyx57W2v"
+bots_to_check:
+  - username: "@WebpageBot"
+    command: "/start"
+    expected_response: "Hello, I'm the Webpage Bot!\nPlease send me up to 10 links and I will update previews for them."
 ```
 
 ### Run it every 5 minutes using crontab

config/settings.yaml.dist

@@ -1,2 +1,8 @@
-bot_token: ""
-chat_ids: []
+token: ""
+chat_ids: []
+api_id: ""
+api_hash: ""
+bots_to_check:
+  - username: ""
+    command: ""
+    expected_response: ""

requirements.txt

@@ -1,3 +1,4 @@
 httpx == 0.23.0
 httpx[http2] == 0.23.0
 pyyaml == 6.0.1
+pyrogram == 2.0.106

src/main.py

@@ -1,7 +1,7 @@
 #!/usr/bin/python3
 
 from httpx import AsyncClient, codes # Yes, we could use requests but httpx supports async tasks and HTTP/2!
-from os import getenv
+from src.telegram_checker import bot_checker
 from asyncio import run
 import platform    # For getting the operating system name
 import subprocess  # For executing a shell command
@@ -19,10 +19,10 @@
         config_map = yaml.load(yaml_config, Loader=yaml.SafeLoader)
 
 def get_token() -> str:
-    return config_map["token"]
+    return config_map['token']
 
 def get_users() -> list[str]:
-    return config_map["chat_ids"]
+    return config_map['chat_ids']
 
 
 async def check_ok(url: str) -> bool:
@@ -87,12 +87,14 @@ def handle_communication(url: str, method: str) -> None:
 
 
 def main() -> None:
-    with open("src/urls.json", "r") as f:
+    with open('src/urls.json', 'r') as f:
         urls_list = load(f)
 
     for method, urls in urls_list.items():
         for url in urls:
             handle_urls(url, method)
+
+    bot_checker(config_map)
 
 
 

src/telegram_checker.py

@@ -0,0 +1,65 @@
+from pyrogram import Client, filters
+from pyrogram.errors import RPCError
+from pyrogram.types import Message
+from os.path import exists
+from asyncio import sleep
+
+config_map = None
+app = Client('healthcheck_user_client')
+received_answer = []
+
+@app.on_message(filters.text & filters.bot)
+async def check_welcome(client: Client, message: Message) -> None:
+    global received_answer
+
+    for bots_to_check in config_map['bots_to_check']:
+        if message.from_user.username == bots_to_check['username']:
+            if message.text == bots_to_check['expected_response']:
+                received_answer.insert(0, message.from_user.username)
+            return
+
+
+def bot_checker(config: dict) -> None:
+    global config_map
+    config_map = config
+
+    # No bots to check
+    if 'bots_to_check' not in config_map or not config_map['bots_to_check']:
+        return
+
+    # Checking if already exists a session
+    if not exists('healthcheck_user_client.session'):
+        # We need to check that all the needed parameters are available
+        for key in ['api_id', 'api_hash']:
+            if key not in config_map:
+                print(f'Missing required parameter: {key}')
+                return
+
+        app.api_id=config_map['api_id']
+        app.api_hash=config_map['api_hash']
+
+        print('Starting login, after the client is connected you can exit with CTRL+C')
+
+    app.run(main())
+
+
+async def main() -> None:
+    for bot in config_map['bots_to_check']:
+        if not all(key in bot for key in ['username', 'command', 'expected_response']):
+            print(f'Skipping {bot["username"] if "id" in bot else "a bot without a specified username c:"}')
+            continue
+
+        try:
+            await app.send_message(chat_id=bot['username'], text=bot['command'])
+        except RPCError as e:
+            print(f'There was a problem communicating with {bot["username"]}:', e)
+
+    await sleep(10) # Just to be sure the bot is not busy
+    # To prevent circular imports
+    # TODO: in a future refactor we could split the main code in another file only for webpages, 
+    #       would improve modularity and reuse of code with an helper file
+    from src.main import make_request_to_telegram
+    for bot in config_map['bots_to_check']:
+        if bot['username'] not in received_answer:
+            for user_to_notify in config_map['chat_ids']:
+                await make_request_to_telegram(f'@{bot["username"]}', 'Telegram', user_to_notify)

tests/main_test.py

@@ -109,12 +109,23 @@
         'func': main.main,
         'expected_res': None,
         'arg': tuple(),
-        'mock_obj': [main],
-        'mock_func': ['handle_urls'],
-        'mock_ret': [None]
+        'mock_obj': [main, main],
+        'mock_func': ['handle_urls', 'bot_checker'],
+        'mock_ret': [None, None]
+    },
+    {
+        'func': main.get_token,
+        'expected_res': '',
+        'arg': tuple(),
+    },
+    {
+        'func': main.get_users,
+        'expected_res': [],
+        'arg': tuple(),  
     }
 ]
 
+@pytest.mark.asyncio
 @pytest.mark.parametrize('test', tests)
 async def test_generic(mocker: MockerFixture, test: dict) -> None:
     spyed_objs = []

tests/telegram_checker_test.py

@@ -0,0 +1,82 @@
+import pytest
+from pytest_mock import MockerFixture
+from unittest.mock import AsyncMock
+import src.telegram_checker as tg
+import src.main as main
+from pyrogram.client import Client
+import asyncio
+
+tests = [
+    {
+        'func': tg.bot_checker,
+        'expected_res': None,
+        'arg': ({'bots_to_check': [{'username':'WebpageBot', 'command':'/start', 'expected_response':'Hello, I\'m the Webpage Bot!\nPlease send me up to 10 links and I will update previews for them.'}]},),
+        'is_async': False
+    },
+    {
+        'func': tg.bot_checker,
+        'expected_res': None,
+        'arg': ({},),
+        'is_async': False
+    },
+    {
+        'func': tg.bot_checker,
+        'expected_res': None,
+        'arg': ({'api_id':'12345', 'api_hash':'0123456789abcdef0123456789abcdef','bots_to_check': [{'username':'WebpageBot', 'command':'/start', 'expected_response':'Hello, I\'m the Webpage Bot!\nPlease send me up to 10 links and I will update previews for them.'}]},),
+        'mock_obj': [Client],
+        'mock_func': ['run'],
+        'mock_ret': [None],
+        'is_async': False
+    },
+    {
+        'func': tg.main,
+        'expected_res': None,
+        'arg': tuple(),
+        'mock_obj': [main, Client],
+        'mock_func': ['make_request_to_telegram', 'send_message'],
+        'mock_ret': [None, None],
+        'mock_dict': {'bots_to_check': [{'username':'WebpageBot', 'command':'/start', 'expected_response':'Hello, I\'m the Webpage Bot!\nPlease send me up to 10 links and I will update previews for them.'}]},
+        'is_async': True
+    },
+    {
+        'func': tg.main,
+        'expected_res': None,
+        'arg': tuple(),
+        'mock_obj': [main, Client],
+        'mock_func': ['make_request_to_telegram', 'send_message'],
+        'mock_ret': [None, None],
+        'mock_dict': {'bots_to_check': [{'command':'/start'}]},
+        'is_async': True
+    },
+    {
+        'func': tg.check_welcome,
+        'expected_res': None,
+        'arg': (None, {'from_user': {'username': 'username'}, 'text': 'A response'},),
+        'mock_dict': {'bots_to_check': [{'username':'WebpageBot', 'command':'/start', 'expected_response':'Hello, I\'m the Webpage Bot!\nPlease send me up to 10 links and I will update previews for them.'}]},
+        'is_async': True
+    }
+]
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize('test', tests)
+async def test_generic(mocker: MockerFixture, test: dict) -> None:
+    spyed_objs = []
+    spyed_dicts = []
+
+    if test.get('mock_obj'):
+        for index, obj in enumerate(test['mock_obj']):
+            mocker.patch.object(obj, test['mock_func'][index], return_value=test['mock_ret'][index])
+            spyed_objs.append(mocker.spy(obj, test['mock_func'][index]))
+
+    if test.get('mock_dict'):
+        mocker.patch.dict(tg.config_map, test['mock_dict'])
+        spyed_dicts.append(mocker.spy(tg, 'config_map'))
+
+    if test.get('is_async'):
+        res = await test['func'](*test['arg'])
+    else:
+        res = test['func'](*test['arg'])       
+
+    assert res == test['expected_res']
+    for index, spy in enumerate(spyed_objs):
+        assert spy.spy_return == test['mock_ret'][index]