diff --git a/TEx/modules/telegram_connection_manager.py b/TEx/modules/telegram_connection_manager.py index 5d67fba..89b197a 100644 --- a/TEx/modules/telegram_connection_manager.py +++ b/TEx/modules/telegram_connection_manager.py @@ -4,8 +4,8 @@ import logging import os.path import platform -from configparser import ConfigParser -from typing import Dict, cast +from configparser import ConfigParser, SectionProxy +from typing import Dict, Optional, cast from telethon import TelegramClient @@ -35,20 +35,17 @@ async def run(self, config: ConfigParser, args: Dict, data: Dict) -> None: if not os.path.exists(session_dir): os.mkdir(session_dir) - device_model: str = self.__get_device_model_name(config=config) - # Check Activation Command if args['connect']: # New Connection logger.info('\t\tAuthorizing on Telegram...') # Connect - client = TelegramClient( - os.path.join(session_dir, config['CONFIGURATION']['phone_number']), - config['CONFIGURATION']['api_id'], - config['CONFIGURATION']['api_hash'], - catch_up=True, - device_model=device_model, - ) + client = await self.__get_telegram_client( + session_dir=session_dir, + config=config, + api_id=config['CONFIGURATION']['api_id'], + api_hash=config['CONFIGURATION']['api_hash'], + ) await client.start(phone=config['CONFIGURATION']['phone_number']) client.session.save() @@ -60,27 +57,42 @@ async def run(self, config: ConfigParser, args: Dict, data: Dict) -> None: } else: # Reuse Previous Connection - # Check if Contains the Required Data if 'telegram_connection' not in data or \ 'api_id' not in data['telegram_connection'] or \ 'api_hash' not in data['telegram_connection'] or \ 'target_phone_number' not in data['telegram_connection']: logger.warning('\t\tNot Authenticated on Telegram. Please use the "connect" command.') + data['internals']['panic'] = True return - client = TelegramClient( - os.path.join(session_dir, config['CONFIGURATION']['phone_number']), - data['telegram_connection']['api_id'], - data['telegram_connection']['api_hash'], - catch_up=True, - device_model=device_model, - ) + client = await self.__get_telegram_client( + session_dir=session_dir, + config=config, + api_id=data['telegram_connection']['api_id'], + api_hash=data['telegram_connection']['api_hash'], + ) await client.start(phone=data['telegram_connection']['target_phone_number']) data['telegram_client'] = client logger.info(f'\t\tUser Authorized on Telegram: {await client.is_user_authorized()}') + async def __get_telegram_client(self, session_dir: str, config: ConfigParser, api_id: str, api_hash: str) -> TelegramClient: + """Return a Telegram Client.""" + device_model: str = self.__get_device_model_name(config=config) + + proxy_settings: Optional[Dict] = await self.__get_proxy_settings(config=config) + if proxy_settings: + logger.info(f'\t\tUsing {proxy_settings["proxy_type"]} Proxy') + + return TelegramClient( + os.path.join(session_dir, config['CONFIGURATION']['phone_number']), + api_id, api_hash, + catch_up=True, + device_model=device_model, + proxy=proxy_settings, + ) + def __get_device_model_name(self, config: ConfigParser) -> str: """ Compute Device Model Name for Telegram API. @@ -100,6 +112,37 @@ def __get_device_model_name(self, config: ConfigParser) -> str: return device_model_name + async def __get_proxy_settings(self, config: ConfigParser) -> Optional[Dict]: + """Return Proxy Setting.""" + # Check if Config Contains Proxy + if not config.has_section('PROXY'): + return None + + proxy_section: SectionProxy = config['PROXY'] + + # Check Minimum Proxy Settings + if 'type' not in proxy_section or 'address' not in proxy_section or 'port' not in proxy_section: + return None + + # Set Basic Result + h_result: Dict = { + 'proxy_type': proxy_section['type'], + 'addr': proxy_section['address'], + 'port': int(proxy_section['port']), + } + + # Check Proxy Auth and Password + if 'username' in proxy_section and proxy_section['username'] != '': + h_result['username'] = proxy_section['username'] + + if 'password' in proxy_section and proxy_section['password'] != '': + h_result['password'] = proxy_section['password'] + + if 'rdns' in proxy_section and proxy_section['rdns'] != '': + h_result['rdns'] = bool(proxy_section['rdns']) + + return h_result + class TelegramDisconnector(BaseModule): """Telegram Connection Manager - Connect.""" diff --git a/docs/changelog/v030.md b/docs/changelog/v030.md new file mode 100644 index 0000000..5dba494 --- /dev/null +++ b/docs/changelog/v030.md @@ -0,0 +1,11 @@ +# Changelog - V0.3.0 + +**🚀 Features** + +- Proxy (HTTP, SOCKS4, SOCKS5) support (#26) + +**🐛 Bug Fixes** + +**⚙️ Internal Improvements** + +- Replace Pylint, PyDocStyle and Flake8 code quality tools for Ruff (#22) \ No newline at end of file diff --git a/docs/configuration/complete_configuration_file_example.md b/docs/configuration/complete_configuration_file_example.md index 231f5ca..7ea68ad 100644 --- a/docs/configuration/complete_configuration_file_example.md +++ b/docs/configuration/complete_configuration_file_example.md @@ -10,6 +10,14 @@ phone_number=15552809753 data_path=/usr/home/tex_data/ device_model=AMD64 +[PROXY] +type=HTTP +address=127.0.0.1 +port=3128 +username=proxy username +password=proxy password +rdns=true + [FINDER] enabled=true diff --git a/docs/configuration/proxy.md b/docs/configuration/proxy.md new file mode 100644 index 0000000..40e44c8 --- /dev/null +++ b/docs/configuration/proxy.md @@ -0,0 +1,19 @@ +# Proxy +If you need to use a proxy server, you can configure this behavior within the configuration file. If not, just omit this section from your file. + +```ini +[PROXY] +type=HTTP +address=127.0.0.1 +port=3128 +username=proxy username +password=proxy password +rdns=true +``` + +* **type** > Required - Protocol to use (HTTP, SOCKS5 or SOCKS4) +* **address** > Required - Proxy Address +* **port** > Required - Proxy IP Port +* **username** > Optional - Username if the proxy requires auth +* **password** > Optional - Password if the proxy requires auth +* **rdns** > Optional - Whether to use remote or local resolve, default remote diff --git a/mkdocs.yml b/mkdocs.yml index 35b0e86..0f3fc15 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -21,6 +21,7 @@ nav: - 'Contact': 'contact.md' - 'Configuration': - 'Basic Configuration': 'configuration/basic.md' + - 'Proxy': 'configuration/proxy.md' - 'Complete Configuration File Example': 'configuration/complete_configuration_file_example.md' - 'How to Use': - 'Basic Usage': 'how_use/how_to_use_basic.md' @@ -40,6 +41,8 @@ nav: - 'Text Report': 'report/report_text.md' - 'Maintenance': - 'Purging Old Data': 'maintenance/purge_old_data.md' + - 'Changelog': + - 'V0.3.0': 'changelog/v030.md' site_author: Th3 0bservator diff --git a/pyproject.toml b/pyproject.toml index e3c4c6b..2e35809 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -67,8 +67,10 @@ cachetools = ">=5.3.1,<6" toml = ">=0.10.2" tox = "^4.10.0" discord-webhook = ">=1.3.0,<2" -aiofiles = "23.2.1" +aiofiles = "23.2.1" types-aiofiles = "23.2.0.0" +python-socks = "2.4.3" +async-timeout = "4.0.3" [tool.poetry.dev-dependencies] pytest = ">=7.4.0" @@ -84,6 +86,7 @@ types-toml = ">=0.10.8.7" mkdocs = ">=1.5.3,<2" mkdocs-material = ">=9.4.2,<10" ruff = "0.0.292" +parameterized = "0.9.0" [build-system] requires = ["poetry-core"] diff --git a/tests/modules/test_telegram_connection_manager.py b/tests/modules/test_telegram_connection_manager.py index aeb7697..f287624 100644 --- a/tests/modules/test_telegram_connection_manager.py +++ b/tests/modules/test_telegram_connection_manager.py @@ -7,6 +7,7 @@ from configparser import ConfigParser from typing import Dict from unittest import mock +from parameterized import parameterized, parameterized_class from TEx.modules.telegram_connection_manager import TelegramConnector, TelegramDisconnector from tests.modules.common import TestsCommon @@ -19,8 +20,12 @@ def setUp(self) -> None: self.config = ConfigParser() self.config.read('../../config.ini') - def test_run_connect(self): - """Test Run Method with Telegram Server Connection.""" + @parameterized.expand([ + ('With Proxy', True), + ('Without Proxy', False) + ]) + def test_run_connect(self, name: str, use_proxy: bool): + """Test Run Method with Telegram Server Connection without Proxy Setting.""" # Setup Mock telegram_client_mockup = mock.Mock() @@ -38,6 +43,10 @@ def test_run_connect(self): data: Dict = {} TestsCommon.execute_basic_pipeline_steps_for_initialization(config=self.config, args=args, data=data) + # If not Use Proxy, remove + if not use_proxy: + self.config.remove_section('PROXY') + with mock.patch('TEx.modules.telegram_connection_manager.TelegramClient', return_value=telegram_client_mockup) as client_base: with self.assertLogs() as captured: loop = asyncio.get_event_loop() @@ -55,12 +64,18 @@ def test_run_connect(self): '12345678', 'deff1f2587358746548deadbeef58ddd', catch_up=True, - device_model='UT_DEVICE_01' + device_model='UT_DEVICE_01', + proxy={'proxy_type': 'HTTP', 'addr': '1.2.3.4', 'port': 4444, 'username': 'ut_username', 'password': 'ut_password', 'rdns': True} if use_proxy else None ) # Check Logs - self.assertEqual(2, len(captured.records)) - self.assertEqual(' User Authorized on Telegram: True', captured.records[1].message) + if use_proxy: + self.assertEqual(3, len(captured.records)) + self.assertEqual('\t\tUsing HTTP Proxy', captured.records[1].message) + self.assertEqual(' User Authorized on Telegram: True', captured.records[2].message) + else: + self.assertEqual(2, len(captured.records)) + self.assertEqual(' User Authorized on Telegram: True', captured.records[1].message) # Validate Mock Calls telegram_client_mockup.start.assert_awaited_once_with(phone='5526986587745') @@ -72,7 +87,11 @@ def test_run_connect(self): self.assertEqual('5526986587745', data['telegram_connection']['target_phone_number']) self.assertEqual(telegram_client_mockup, data['telegram_client']) - def test_run_reuse(self): + @parameterized.expand([ + ('With Proxy', True), + ('Without Proxy', False) + ]) + def test_run_reuse(self, name: str, use_proxy: bool): """Test Run Method with Reused Connection.""" # Setup Mock @@ -103,6 +122,15 @@ def test_run_reuse(self): # Force Delete device_model from Configuration, to Ensure the FallBack System Works del self.config['CONFIGURATION']['device_model'] + # If not Use Proxy, remove + if not use_proxy: + self.config.remove_section('PROXY') + else: + # Remove Password, Username and RDNS + del self.config['PROXY']['username'] + del self.config['PROXY']['password'] + del self.config['PROXY']['rdns'] + with mock.patch('TEx.modules.telegram_connection_manager.TelegramClient', return_value=telegram_client_mockup) as client_base: with self.assertLogs() as captured: loop = asyncio.get_event_loop() @@ -120,12 +148,18 @@ def test_run_reuse(self): 'MyTestApiID2', 'MyTestApiHash2', catch_up=True, - device_model='TeX' + device_model='TeX', + proxy={'proxy_type': 'HTTP', 'addr': '1.2.3.4', 'port': 4444} if use_proxy else None ) # Check Logs - self.assertEqual(1, len(captured.records)) - self.assertEqual(' User Authorized on Telegram: False', captured.records[0].message) + if use_proxy: + self.assertEqual(2, len(captured.records)) + self.assertEqual('\t\tUsing HTTP Proxy', captured.records[0].message) + self.assertEqual(' User Authorized on Telegram: False', captured.records[1].message) + else: + self.assertEqual(1, len(captured.records)) + self.assertEqual(' User Authorized on Telegram: False', captured.records[0].message) # Validate Mock Calls telegram_client_mockup.start.assert_awaited_once_with(phone='MyTestPhoneNumber2') @@ -156,7 +190,8 @@ def test_run_reuse_without_authentication(self): } data: Dict = { 'telegram_connection': { - } + }, + 'internals': {'panic': False} } TestsCommon.execute_basic_pipeline_steps_for_initialization(config=self.config, args=args, data=data) @@ -175,6 +210,9 @@ def test_run_reuse_without_authentication(self): self.assertEqual(1, len(captured.records)) self.assertEqual('\t\tNot Authenticated on Telegram. Please use the "connect" command.', captured.records[0].message) + # Check Panic Control + self.assertTrue(data['internals']['panic']) + def test_constructor_call_with_auto_device_model(self): """Test If Auto Configuration for device_model works.""" @@ -214,7 +252,8 @@ def test_constructor_call_with_auto_device_model(self): '12345678', 'deff1f2587358746548deadbeef58ddd', catch_up=True, - device_model=platform.uname().machine + device_model=platform.uname().machine, + proxy={'proxy_type': 'HTTP', 'addr': '1.2.3.4', 'port': 4444, 'username': 'ut_username', 'password': 'ut_password', 'rdns': True} ) diff --git a/tests/unittest_configfile.config b/tests/unittest_configfile.config index 997f1a7..818edb5 100644 --- a/tests/unittest_configfile.config +++ b/tests/unittest_configfile.config @@ -5,6 +5,14 @@ phone_number=5526986587745 data_path=_data device_model=UT_DEVICE_01 +[PROXY] +type=HTTP +address=1.2.3.4 +port=4444 +username=ut_username +password=ut_password +rdns=true + [FINDER] enabled=true