diff --git a/.gitignore b/.gitignore index 912b1d7..31eaba4 100644 --- a/.gitignore +++ b/.gitignore @@ -129,4 +129,6 @@ dmypy.json .pyre/ # Custom -config/settings.yaml \ No newline at end of file +config/settings.yaml +src/healthcheck_user_client.session +.vscode \ No newline at end of file diff --git a/README.md b/README.md index 6965b5e..d1f9c8c 100644 --- a/README.md +++ b/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 diff --git a/config/settings.yaml.dist b/config/settings.yaml.dist index 6e8463b..a63439b 100644 --- a/config/settings.yaml.dist +++ b/config/settings.yaml.dist @@ -1,2 +1,8 @@ -bot_token: "" -chat_ids: [] \ No newline at end of file +token: "" +chat_ids: [] +api_id: "" +api_hash: "" +bots_to_check: + - username: "" + command: "" + expected_response: "" \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index cb0566f..e286f72 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,3 +1,4 @@ httpx == 0.23.0 httpx[http2] == 0.23.0 pyyaml == 6.0.1 +pyrogram == 2.0.106 diff --git a/src/main.py b/src/main.py index e5568ec..113f4b0 100644 --- a/src/main.py +++ b/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) diff --git a/src/telegram_checker.py b/src/telegram_checker.py new file mode 100644 index 0000000..90a0e69 --- /dev/null +++ b/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) \ No newline at end of file diff --git a/tests/main_test.py b/tests/main_test.py index 9629778..90b91c3 100644 --- a/tests/main_test.py +++ b/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 = [] diff --git a/tests/telegram_checker_test.py b/tests/telegram_checker_test.py new file mode 100644 index 0000000..397d960 --- /dev/null +++ b/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]