Merge pull request #30 from guibacellar/26-feature-request-add-proxy-…

…support

New Feature - Proxy Support for Telegram Connection

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."""

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) 

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
 

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

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
 

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"]

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}
                 )
 
 

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