diff --git a/TEx/core/mapper/telethon_message_mapper.py b/TEx/core/mapper/telethon_message_mapper.py index 6b0a021..98bf1f2 100644 --- a/TEx/core/mapper/telethon_message_mapper.py +++ b/TEx/core/mapper/telethon_message_mapper.py @@ -51,7 +51,7 @@ async def to_finder_notification_facade_entity(message: Message, downloaded_medi group_name=mapped_chat_props.chat_title, group_id=mapped_chat_props.chat_id, from_id=message.from_id.user_id if isinstance(message.from_id, PeerUser) else None, - to_id=message.to_id.channel_id if message.to_id is not None else None, + to_id=message.to_id.channel_id if message.to_id is not None and hasattr(message.to_id, 'channel_id') else None, reply_to_msg_id=message.reply_to.reply_to_msg_id if message.is_reply and message.reply_to else None, message_id=message.id, is_reply=message.is_reply, diff --git a/TEx/core/ocr/tesseract_ocr_engine.py b/TEx/core/ocr/tesseract_ocr_engine.py index fabcec6..2b1ea11 100644 --- a/TEx/core/ocr/tesseract_ocr_engine.py +++ b/TEx/core/ocr/tesseract_ocr_engine.py @@ -47,6 +47,10 @@ def configure(self, config: Optional[SectionProxy]) -> None: def run(self, file_path: str) -> Optional[str]: """Run Tesseract Engine and Return Detected Text.""" try: + + if not os.path.exists(file_path): + return '' + return cast(str, tesseract.image_to_string(file_path, lang=self.language)) except Exception as ex: diff --git a/TEx/models/facade/media_handler_facade_entity.py b/TEx/models/facade/media_handler_facade_entity.py index 3cba68c..9b4ac41 100644 --- a/TEx/models/facade/media_handler_facade_entity.py +++ b/TEx/models/facade/media_handler_facade_entity.py @@ -12,3 +12,12 @@ class MediaHandlingEntity(BaseModel): size_bytes: int disk_file_path: str is_ocr_supported: bool + + def is_image(self) -> bool: + """Return if Downloaded Image are an Image.""" + return self.content_type in ['image/gif', 'image/jpeg', 'image/png', 'image/webp', 'application/gif'] + + def is_video(self) -> bool: + """Return if Downloaded Image are a Video.""" + return self.content_type in ['application/ogg', 'video/mp4', 'video/quicktime', 'video/webm'] + diff --git a/TEx/modules/telegram_connection_manager.py b/TEx/modules/telegram_connection_manager.py index 89b197a..2cb3445 100644 --- a/TEx/modules/telegram_connection_manager.py +++ b/TEx/modules/telegram_connection_manager.py @@ -91,6 +91,7 @@ async def __get_telegram_client(self, session_dir: str, config: ConfigParser, ap catch_up=True, device_model=device_model, proxy=proxy_settings, + timeout=int(config['CONFIGURATION'].get('timeout', fallback='10')), ) def __get_device_model_name(self, config: ConfigParser) -> str: diff --git a/TEx/modules/telegram_messages_listener.py b/TEx/modules/telegram_messages_listener.py index 4589e69..6fd3f08 100644 --- a/TEx/modules/telegram_messages_listener.py +++ b/TEx/modules/telegram_messages_listener.py @@ -83,8 +83,9 @@ async def __handler(self, event: NewMessage.Event) -> None: await self.__ensure_group_exists(event=event) # Download Media - downloaded_media: Optional[MediaHandlingEntity] = await self.media_handler.handle_medias(message, event.chat.id, - self.data_path) if self.download_media else None + downloaded_media: Optional[MediaHandlingEntity] = await self.media_handler.handle_medias( + message, event.chat.id, self.data_path, + ) if self.download_media else None # Process OCR ocr_content: Optional[str] = None @@ -100,7 +101,7 @@ async def __handler(self, event: NewMessage.Event) -> None: 'date_time': message.date.astimezone(tz=pytz.utc), 'message': self.__build_final_message(message.message, ocr_content), 'raw': self.__build_final_message(message.raw_text, ocr_content), - 'to_id': message.to_id.channel_id if message.to_id is not None else None, + 'to_id': message.to_id.channel_id if message.to_id is not None and hasattr(message.to_id, 'channel_id') else None, 'media_id': downloaded_media.media_id if downloaded_media else None, 'is_reply': message.is_reply, 'reply_to_msg_id': message.reply_to.reply_to_msg_id if message.is_reply else None, diff --git a/TEx/notifier/discord_notifier.py b/TEx/notifier/discord_notifier.py index b5d4632..610cc39 100644 --- a/TEx/notifier/discord_notifier.py +++ b/TEx/notifier/discord_notifier.py @@ -1,10 +1,12 @@ """Discord Notifier.""" from __future__ import annotations +import os from configparser import SectionProxy from typing import Union -from discord_webhook import DiscordEmbed, DiscordWebhook +import aiofiles +from discord_webhook import AsyncDiscordWebhook, DiscordEmbed from TEx.models.facade.finder_notification_facade_entity import FinderNotificationMessageEntity from TEx.models.facade.signal_notification_model import SignalNotificationEntityModel @@ -26,6 +28,13 @@ def configure(self, url: str, config: SectionProxy) -> None: async def run(self, entity: Union[FinderNotificationMessageEntity, SignalNotificationEntityModel], rule_id: str, source: str) -> None: """Run Discord Notifier.""" + # Run the Notification Process + webhook = AsyncDiscordWebhook( # type: ignore + url=self.url, + rate_limit_retry=True, + timeout=self.timeout_seconds, + ) + embed: DiscordEmbed if isinstance(entity, FinderNotificationMessageEntity): is_duplicated, duplication_tag = self.check_is_duplicated(message=entity.raw_text) @@ -39,16 +48,45 @@ async def run(self, entity: Union[FinderNotificationMessageEntity, SignalNotific duplication_tag=duplication_tag, ) + # Handle Attachments + await self.__handle_attachment( + entity=entity, + webhook=webhook, + embed=embed, + ) + else: embed = await self.__get_signal_notification_embed( entity=entity, source=source, ) - # Run the Notification Process - webhook = DiscordWebhook(url=self.url, rate_limit_retry=True) webhook.add_embed(embed) - webhook.execute() + await webhook.execute(remove_embeds=True) + + async def __handle_attachment(self, entity: FinderNotificationMessageEntity, webhook: AsyncDiscordWebhook, embed: DiscordEmbed) -> None: + """Handle the Attachment Upload.""" + if not entity.downloaded_media_info or not self.media_attachments_enabled: + return + + # Check Max Size + if entity.downloaded_media_info.size_bytes > self.media_attachments_max_size_bytes: + return + + # Upload File + if os.path.exists(entity.downloaded_media_info.disk_file_path): + + # Open and Upload + async with aiofiles.open(entity.downloaded_media_info.disk_file_path, 'rb') as f: + webhook.add_file(file=await f.read(), filename=f'{entity.downloaded_media_info.file_name}') + await f.close() + + # Add on Embed + if entity.downloaded_media_info.is_image(): + embed.set_image(url=f'attachment://{entity.downloaded_media_info.file_name}') + + elif entity.downloaded_media_info.is_video(): + embed.set_video(url=f'attachment://{entity.downloaded_media_info.file_name}') async def __get_signal_notification_embed(self, entity: SignalNotificationEntityModel, source: str) -> DiscordEmbed: """Return the Embed Object for Signals.""" diff --git a/TEx/notifier/elastic_search_notifier.py b/TEx/notifier/elastic_search_notifier.py index 8415d97..f0a15bf 100644 --- a/TEx/notifier/elastic_search_notifier.py +++ b/TEx/notifier/elastic_search_notifier.py @@ -32,7 +32,7 @@ def configure(self, config: SectionProxy) -> None: api_key=config.get('api_key', fallback=None), verify_certs=config.get('verify_ssl_cert', fallback='True') == 'True', cloud_id=config.get('cloud_id', fallback=None), - request_timeout=20, + request_timeout=30, max_retries=10, ssl_show_warn=False, ) diff --git a/TEx/notifier/notifier_base.py b/TEx/notifier/notifier_base.py index 97dfc7d..f9ee8b2 100644 --- a/TEx/notifier/notifier_base.py +++ b/TEx/notifier/notifier_base.py @@ -18,10 +18,16 @@ class BaseNotifier: def __init__(self) -> None: """Initialize the Base Notifier.""" self.cache: Optional[TTLCache] = None + self.timeout_seconds: int + self.media_attachments_enabled: bool + self.media_attachments_max_size_bytes: int def configure_base(self, config: SectionProxy) -> None: """Configure Base Notifier.""" - self.cache = TTLCache(maxsize=4096, ttl=int(config['prevent_duplication_for_minutes']) * 60) + self.cache = TTLCache(maxsize=4096, ttl=int(config.get('prevent_duplication_for_minutes', fallback='240')) * 60) + self.timeout_seconds = int(config.get('timeout_seconds', fallback='30')) + self.media_attachments_enabled = config.get('media_attachments_enabled', fallback='false') == 'true' + self.media_attachments_max_size_bytes = int(config.get('media_attachments_max_size_bytes', fallback='10000000')) def check_is_duplicated(self, message: str) -> Tuple[bool, str]: """Check if Message is Duplicated on Notifier.""" diff --git a/docs/changelog/v030.md b/docs/changelog/v030.md index 8165f34..c872298 100644 --- a/docs/changelog/v030.md +++ b/docs/changelog/v030.md @@ -2,7 +2,7 @@ !!! warning "Python Version" - This are the latest version os Telegram Explorer that supports Python 3.8 and 3.9. + This are the latest version of Telegram Explorer that supports Python 3.8 and 3.9. Please, consider upgrate to Python 3.10+ as possible. @@ -10,6 +10,8 @@ - Proxy (HTTP, SOCKS4, SOCKS5) support ([#26](https://github.com/guibacellar/TEx/issues/26)) - Discord Notifications now have a source information with account/phone number +- It is now possible to set the connection timeout for the Telegram servers connectors +- Discord Notifications now allow to send downloaded files as attachments ([#41](https://github.com/guibacellar/TEx/issues/41)) - New Message Finder Rule to Catch All Messages - New Notification connector for Elastic Search ([#12](https://github.com/guibacellar/TEx/issues/12)) - Fine Control on Media Download Settings ([#37](https://github.com/guibacellar/TEx/issues/37)) diff --git a/docs/configuration/basic.md b/docs/configuration/basic.md index 0a99640..682372e 100644 --- a/docs/configuration/basic.md +++ b/docs/configuration/basic.md @@ -8,6 +8,7 @@ api_hash=my_api_hash phone_number=my_phone_number data_path=my_data_path device_model=device_model_name +timeout=30 ``` * **api_id** > Required - Telegram API ID. From https://my.telegram.org/ > login > API development tools @@ -17,7 +18,10 @@ device_model=device_model_name * **device_model** > Optional - Defines which device model is passed to Telegram Servers. * If Blank or Absent - Uses 'TeX' for backwards compatibility * If set as 'AUTO' - Uses the computer/system device model +* **timeout** > Optional - Defines the Timeout in seconds for Telegram Client. + * Default: 10 + !!! warning "Note about 'device_model'" If you are using versions prior to 0.2.15 or have already connected to Telegram and have not configured the 'device_model' parameter, do not make the change, as Telegram may interpret this operation as an attack on your account. @@ -32,4 +36,5 @@ api_hash=dead1f29db5d1fa56cc42757acbabeef phone_number=15552809753 data_path=/usr/home/tex_data/ device_model=AMD64 +timeout=15 ``` diff --git a/docs/configuration/complete_configuration_file_example.md b/docs/configuration/complete_configuration_file_example.md index 28c281c..d88dee8 100644 --- a/docs/configuration/complete_configuration_file_example.md +++ b/docs/configuration/complete_configuration_file_example.md @@ -9,6 +9,7 @@ api_hash=dead1f29db5d1fa56cc42757acbabeef phone_number=15552809753 data_path=/usr/home/tex_data/ device_model=AMD64 +timeout=30 [PROXY] type=HTTP @@ -47,14 +48,19 @@ notifier=NOTIFIER.ELASTIC_SEARCH.GENERAL [NOTIFIER.DISCORD.MY_HOOK_1] webhook=https://discord.com/api/webhooks/1157896186751897357/o7foobar4txvAvKSdeadHiI-9XYeXaGlQtd-5PtrrX_eCE0XElWktpPqjrZ0KbeefPtQC prevent_duplication_for_minutes=240 +timeout_seconds=30 +media_attachments_enabled=false [NOTIFIER.DISCORD.MY_HOOK_2] webhook=https://discord.com/api/webhooks/1128765187657681875/foobarqOMFp_4tM2ic2mbeefNPOZqJnBZZdfaubQv2vJgbYzfdeadZd5aqGX6FmCmbNjX prevent_duplication_for_minutes=240 +media_attachments_enabled=false [NOTIFIER.DISCORD.SIGNALS_HOOK] webhook=https://discord.com/api/webhooks/1128765187657681875/foobarqOMFp_457EDs2mbeefNPPeqJnBZZdfaubQvOKIUHYzfdeadZd5aqGX6FmCmbNjv prevent_duplication_for_minutes=0 +media_attachments_enabled=true +media_attachments_max_size_bytes=10000000 [NOTIFIER.ELASTIC_SEARCH.GENERAL] address=https://localhost:9200 diff --git a/docs/notification/notification_discord.md b/docs/notification/notification_discord.md index b59dc66..d1bc06e 100644 --- a/docs/notification/notification_discord.md +++ b/docs/notification/notification_discord.md @@ -14,22 +14,39 @@ For each notification hook you must set a configuration using the default name s * **webhook** > Required - Discord Webhook URI * **prevent_duplication_for_minutes** > Required - Time (in minutes) that the system keep track of messages sent to Discord servers to prevent others message with same content to be sent to the webhook. If you don't want to use this feature, just set the parameter to 0. - + * **timeout_seconds** > Optional - Timeout (in seconds) that waits to send the message. If the message sent take more that time, the message will be ignored. + * Default: 30 + * **media_attachments_enabled** > Optional - Enable/Disable the behavior for sending downloaded medias on messages that have been reported. + * Default: false + * **media_attachments_max_size_bytes** > Optional - Set the max size in bytes to send the medias on the notifications. + * Default: 10000000 + +=true +media_attachments_max_size_bytes=10000000 **Changes on Configuration File** ```ini [NOTIFIER.DISCORD.MY_HOOK_1] webhook=https://discord.com/api/webhooks/1157896186751897357/o7foobar4txvAvKSdeadHiI-9XYeXaGlQtd-5PtrrX_eCE0XElWktpPqjrZ0KbeefPtQC prevent_duplication_for_minutes=240 +timeout_seconds=30 +media_attachments_enabled=true +media_attachments_max_size_bytes=10000000 [NOTIFIER.DISCORD.MY_HOOK_2] webhook=https://discord.com/api/webhooks/1128765187657681875/foobarqOMFp_4tM2ic2mbeefNPOZqJnBZZdfaubQv2vJgbYzfdeadZd5aqGX6FmCmbNjX prevent_duplication_for_minutes=240 +media_attachments_enabled=false +media_attachments_max_size_bytes=10000000 [NOTIFIER.DISCORD.MY_HOOK_3] webhook=https://discord.com/api/webhooks/1256789875462124045/bQ9TZqOzgA05PLVu8E2LU3N5foobarFU8-0nQbeefP5oIgAUOlydeadf7Uc19Hs00OJQ prevent_duplication_for_minutes=60 +timeout_seconds=30 +media_attachments_enabled=true +media_attachments_max_size_bytes=25000000 [NOTIFIER.DISCORD.MY_HOOK_4] webhook=https://discord.com/api/webhooks/1487651987651004895/mR0v3zOywH3Z5HvdeadrGEqqndkcYepgCM-Q6foobardjAMXAEbeefuA_F7-h5JcBM4RT prevent_duplication_for_minutes=240 +media_attachments_enabled=true ``` diff --git a/pyproject.toml b/pyproject.toml index a1804c4..5196080 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -66,7 +66,7 @@ requests = ">=2.31.0,<3" cachetools = ">=5.3.1,<6" toml = ">=0.10.2" tox = "^4.10.0" -discord-webhook = ">=1.3.0,<2" +discord-webhook = {extras = ["async"], version=">=1.3.0,<2"} aiofiles = "23.2.1" types-aiofiles = "23.2.0.0" python-socks = "2.4.3" diff --git a/tests/core/ocr/test_tesseract_ocr_engine.py b/tests/core/ocr/test_tesseract_ocr_engine.py index 19827a9..c6caa8f 100644 --- a/tests/core/ocr/test_tesseract_ocr_engine.py +++ b/tests/core/ocr/test_tesseract_ocr_engine.py @@ -113,8 +113,8 @@ def test_configure_tesseract_cmd_not_found(self): self.assertEqual(f'Tesseract command cannot be found at "/folder/to/cmd/file"', context.exception.args[0]) @mock.patch('TEx.core.ocr.tesseract_ocr_engine.tesseract') - def test_run_ocr_error(self, mocked_tesseract): - """Test Tesseract Engine 'run' method returning a Exception.""" + def test_run_ocr_file_not_found(self, mocked_tesseract): + """Test Tesseract Engine 'run' method returning Empty Value due a File not Found.""" # Configure Mock mocked_tesseract.image_to_string = mock.MagicMock(side_effect=Exception()) @@ -123,8 +123,4 @@ def test_run_ocr_error(self, mocked_tesseract): target: OcrEngineBase = TesseractOcrEngine() # Call Run - with self.assertLogs() as captured: - target.run(file_path='/path/to/target/image') - - # Check Log Message - self.assertEqual('OCR Fail', captured.records[0].message) \ No newline at end of file + self.assertEqual('', target.run(file_path='/path/to/target/image')) diff --git a/tests/modules/test_telegram_connection_manager.py b/tests/modules/test_telegram_connection_manager.py index f287624..2735c74 100644 --- a/tests/modules/test_telegram_connection_manager.py +++ b/tests/modules/test_telegram_connection_manager.py @@ -65,7 +65,8 @@ def test_run_connect(self, name: str, use_proxy: bool): 'deff1f2587358746548deadbeef58ddd', catch_up=True, 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 + proxy={'proxy_type': 'HTTP', 'addr': '1.2.3.4', 'port': 4444, 'username': 'ut_username', 'password': 'ut_password', 'rdns': True} if use_proxy else None, + timeout=20 ) # Check Logs @@ -149,7 +150,8 @@ def test_run_reuse(self, name: str, use_proxy: bool): 'MyTestApiHash2', catch_up=True, device_model='TeX', - proxy={'proxy_type': 'HTTP', 'addr': '1.2.3.4', 'port': 4444} if use_proxy else None + proxy={'proxy_type': 'HTTP', 'addr': '1.2.3.4', 'port': 4444} if use_proxy else None, + timeout=20 ) # Check Logs @@ -253,7 +255,8 @@ def test_constructor_call_with_auto_device_model(self): 'deff1f2587358746548deadbeef58ddd', catch_up=True, device_model=platform.uname().machine, - proxy={'proxy_type': 'HTTP', 'addr': '1.2.3.4', 'port': 4444, 'username': 'ut_username', 'password': 'ut_password', 'rdns': True} + proxy={'proxy_type': 'HTTP', 'addr': '1.2.3.4', 'port': 4444, 'username': 'ut_username', 'password': 'ut_password', 'rdns': True}, + timeout=20 ) diff --git a/tests/notifier/test_discord_notifier.py b/tests/notifier/test_discord_notifier.py index f9f95a3..57b852c 100644 --- a/tests/notifier/test_discord_notifier.py +++ b/tests/notifier/test_discord_notifier.py @@ -6,6 +6,8 @@ from unittest import mock from TEx.models.facade.finder_notification_facade_entity import FinderNotificationMessageEntity +from TEx.models.facade.media_handler_facade_entity import MediaHandlingEntity +from TEx.models.facade.signal_notification_model import SignalNotificationEntityModel from TEx.notifier.discord_notifier import DiscordNotifier from tests.modules.common import TestsCommon from tests.modules.mockups_groups_mockup_data import channel_1_mocked @@ -44,7 +46,7 @@ def test_run_no_duplication(self): data: Dict = {} TestsCommon.execute_basic_pipeline_steps_for_initialization(config=self.config, args=args, data=data) - with mock.patch('TEx.notifier.discord_notifier.DiscordWebhook', return_value=discord_webhook_mock): + with mock.patch('TEx.notifier.discord_notifier.AsyncDiscordWebhook', return_value=discord_webhook_mock): # Execute Discord Notifier Configure Method target.configure( config=self.config['NOTIFIER.DISCORD.NOT_001'], @@ -78,7 +80,7 @@ def test_run_no_duplication(self): self.assertEqual(call_arg.fields[6], {'inline': False, 'name': 'Tag', 'value': 'de33f5dda9c686c64d23b8aec2eebfc7'}) # Check if Webhook was Executed - discord_webhook_mock.execute.assert_called_once() + discord_webhook_mock.execute.assert_awaited_once() def test_run_duplication_control(self): """Test Run Method First Time - With Duplication Detection.""" @@ -107,7 +109,7 @@ def test_run_duplication_control(self): data: Dict = {} TestsCommon.execute_basic_pipeline_steps_for_initialization(config=self.config, args=args, data=data) - with mock.patch('TEx.notifier.discord_notifier.DiscordWebhook', return_value=discord_webhook_mock): + with mock.patch('TEx.notifier.discord_notifier.AsyncDiscordWebhook', return_value=discord_webhook_mock): # Execute Discord Notifier Configure Method target.configure( config=self.config['NOTIFIER.DISCORD.NOT_001'], @@ -139,4 +141,212 @@ def test_run_duplication_control(self): discord_webhook_mock.add_embed.assert_called_once() # Check if Webhook was Executed Exact 1 Time - discord_webhook_mock.execute.assert_called_once() + discord_webhook_mock.execute.assert_awaited_once() + + def test_run_with_downloaded_media_image(self): + """Test Run Method With Downloaded Media as Image.""" + + # Setup Mock + discord_webhook_mock = mock.AsyncMock() + discord_webhook_mock.add_embed = mock.MagicMock() + discord_webhook_mock.add_file = mock.MagicMock() + + message_entity: FinderNotificationMessageEntity = FinderNotificationMessageEntity( + date_time=datetime.datetime(2023, 10, 1, 9, 58, 22), + raw_text="Mocked Raw Text", + group_name="Channel 1972142108", + group_id=1972142108, + from_id="1234", + to_id=9876, + reply_to_msg_id=5544, + message_id=5975883, + is_reply=False, + downloaded_media_info=MediaHandlingEntity( + media_id=555, + file_name='122761750_387013276008970_8208112669996447119_n.jpg', + content_type='image/png', + size_bytes=1520, + disk_file_path='resources/122761750_387013276008970_8208112669996447119_n.jpg', + is_ocr_supported=True + ), + ) + + target: DiscordNotifier = DiscordNotifier() + args: Dict = { + 'config': 'unittest_configfile.config' + } + data: Dict = {} + TestsCommon.execute_basic_pipeline_steps_for_initialization(config=self.config, args=args, data=data) + + with mock.patch('TEx.notifier.discord_notifier.AsyncDiscordWebhook', return_value=discord_webhook_mock): + # Execute Discord Notifier Configure Method + target.configure( + config=self.config['NOTIFIER.DISCORD.NOT_001'], + url='url.domain/path' + ) + + loop = asyncio.get_event_loop() + loop.run_until_complete( + + # Invoke Test Target + target.run( + entity=message_entity, + rule_id='RULE_UT_01', + source='+15558987453' + ) + ) + + # Check is Embed was Added into Webhook + discord_webhook_mock.add_embed.assert_called_once() + discord_webhook_mock.add_file.assert_called_once() + embed_call_arg = discord_webhook_mock.add_embed.call_args[0][0] + add_file_arg = discord_webhook_mock.add_file.call_args[1] + + self.assertEqual(embed_call_arg.title, '**Channel 1972142108** (1972142108)') + self.assertEqual(embed_call_arg.description, 'Mocked Raw Text') + + self.assertEqual(len(embed_call_arg.fields), 7) + self.assertEqual(embed_call_arg.fields[0], {'inline': True, 'name': 'Source', 'value': '+15558987453'}) + self.assertEqual(embed_call_arg.fields[1], {'inline': True, 'name': 'Rule', 'value': 'RULE_UT_01'}) + self.assertEqual(embed_call_arg.fields[2], {'inline': False, 'name': 'Message ID', 'value': '5975883'}) + self.assertEqual(embed_call_arg.fields[3], {'inline': True, 'name': 'Group Name', 'value': 'Channel 1972142108'}) + self.assertEqual(embed_call_arg.fields[4], {'inline': True, 'name': 'Group ID', 'value': '1972142108'}) + self.assertEqual(embed_call_arg.fields[6], {'inline': False, 'name': 'Tag', 'value': 'de33f5dda9c686c64d23b8aec2eebfc7'}) + + self.assertEqual(embed_call_arg.image['url'], 'attachment://122761750_387013276008970_8208112669996447119_n.jpg') + + self.assertEqual(add_file_arg['filename'], '122761750_387013276008970_8208112669996447119_n.jpg') + self.assertIsNotNone(add_file_arg['file']) + + # Check if Webhook was Executed + discord_webhook_mock.execute.assert_awaited_once() + + def test_run_with_downloaded_media_video(self): + """Test Run Method With Downloaded Media as Video.""" + + # Setup Mock + discord_webhook_mock = mock.AsyncMock() + discord_webhook_mock.add_embed = mock.MagicMock() + discord_webhook_mock.add_file = mock.MagicMock() + + message_entity: FinderNotificationMessageEntity = FinderNotificationMessageEntity( + date_time=datetime.datetime(2023, 10, 1, 9, 58, 22), + raw_text="Mocked Raw Text", + group_name="Channel 1972142108", + group_id=1972142108, + from_id="1234", + to_id=9876, + reply_to_msg_id=5544, + message_id=5975883, + is_reply=False, + downloaded_media_info=MediaHandlingEntity( + media_id=555, + file_name='unknow.mp4', + content_type='video/mp4', + size_bytes=1520, + disk_file_path='resources/unknow.mp4', + is_ocr_supported=False + ), + ) + + target: DiscordNotifier = DiscordNotifier() + args: Dict = { + 'config': 'unittest_configfile.config' + } + data: Dict = {} + TestsCommon.execute_basic_pipeline_steps_for_initialization(config=self.config, args=args, data=data) + + with mock.patch('TEx.notifier.discord_notifier.AsyncDiscordWebhook', return_value=discord_webhook_mock): + # Execute Discord Notifier Configure Method + target.configure( + config=self.config['NOTIFIER.DISCORD.NOT_001'], + url='url.domain/path' + ) + + loop = asyncio.get_event_loop() + loop.run_until_complete( + + # Invoke Test Target + target.run( + entity=message_entity, + rule_id='RULE_UT_01', + source='+15558987453' + ) + ) + + # Check is Embed was Added into Webhook + discord_webhook_mock.add_embed.assert_called_once() + discord_webhook_mock.add_file.assert_called_once() + embed_call_arg = discord_webhook_mock.add_embed.call_args[0][0] + add_file_arg = discord_webhook_mock.add_file.call_args[1] + + self.assertEqual(embed_call_arg.title, '**Channel 1972142108** (1972142108)') + self.assertEqual(embed_call_arg.description, 'Mocked Raw Text') + + self.assertEqual(len(embed_call_arg.fields), 7) + self.assertEqual(embed_call_arg.fields[0], {'inline': True, 'name': 'Source', 'value': '+15558987453'}) + self.assertEqual(embed_call_arg.fields[1], {'inline': True, 'name': 'Rule', 'value': 'RULE_UT_01'}) + self.assertEqual(embed_call_arg.fields[2], {'inline': False, 'name': 'Message ID', 'value': '5975883'}) + self.assertEqual(embed_call_arg.fields[3], {'inline': True, 'name': 'Group Name', 'value': 'Channel 1972142108'}) + self.assertEqual(embed_call_arg.fields[4], {'inline': True, 'name': 'Group ID', 'value': '1972142108'}) + self.assertEqual(embed_call_arg.fields[6], {'inline': False, 'name': 'Tag', 'value': 'de33f5dda9c686c64d23b8aec2eebfc7'}) + + self.assertEqual(embed_call_arg.video['url'], 'attachment://unknow.mp4') + + self.assertEqual(add_file_arg['filename'], 'unknow.mp4') + self.assertIsNotNone(add_file_arg['file']) + + # Check if Webhook was Executed + discord_webhook_mock.execute.assert_awaited_once() + + def test_run_with_signal(self): + """Test Run Method With Signal Input.""" + + # Setup Mock + discord_webhook_mock = mock.AsyncMock() + discord_webhook_mock.add_embed = mock.MagicMock() + discord_webhook_mock.add_file = mock.MagicMock() + + message_entity: SignalNotificationEntityModel = SignalNotificationEntityModel( + signal='INITIALIZATION', + date_time=datetime.datetime(2023, 10, 1, 9, 58, 22), + content='Signal Content' + ) + + target: DiscordNotifier = DiscordNotifier() + args: Dict = { + 'config': 'unittest_configfile.config' + } + data: Dict = {} + TestsCommon.execute_basic_pipeline_steps_for_initialization(config=self.config, args=args, data=data) + + with mock.patch('TEx.notifier.discord_notifier.AsyncDiscordWebhook', return_value=discord_webhook_mock): + # Execute Discord Notifier Configure Method + target.configure( + config=self.config['NOTIFIER.DISCORD.NOT_001'], + url='url.domain/path' + ) + + loop = asyncio.get_event_loop() + loop.run_until_complete( + + # Invoke Test Target + target.run( + entity=message_entity, + rule_id='RULE_UT_01', + source='+15558987453' + ) + ) + + # Check is Embed was Added into Webhook + discord_webhook_mock.add_embed.assert_called_once() + embed_call_arg = discord_webhook_mock.add_embed.call_args[0][0] + + self.assertEqual(embed_call_arg.title, 'INITIALIZATION') + self.assertEqual(embed_call_arg.description, 'Signal Content') + + self.assertEqual(len(embed_call_arg.fields), 2) + self.assertEqual(embed_call_arg.fields[0], {'inline': True, 'name': 'Source', 'value': '+15558987453'}) + + # Check if Webhook was Executed + discord_webhook_mock.execute.assert_awaited_once() diff --git a/tests/notifier/test_elastic_search_notifier.py b/tests/notifier/test_elastic_search_notifier.py index 78c118d..27aae7d 100644 --- a/tests/notifier/test_elastic_search_notifier.py +++ b/tests/notifier/test_elastic_search_notifier.py @@ -9,6 +9,7 @@ from TEx.models.facade.finder_notification_facade_entity import FinderNotificationMessageEntity from TEx.models.facade.media_handler_facade_entity import MediaHandlingEntity +from TEx.models.facade.signal_notification_model import SignalNotificationEntityModel from TEx.notifier.elastic_search_notifier import ElasticSearchNotifier from tests.modules.common import TestsCommon @@ -237,3 +238,59 @@ def test_run_with_downloaded_file(self): } self.assertEqual(submited_document, expected_document) + + def test_run_with_message_signal(self): + """Test Run Method With Message as Signal.""" + + # Setup Mock + elastic_search_api_mock = mock.AsyncMock() + elastic_search_api_mock.index = mock.AsyncMock() + + target: ElasticSearchNotifier = ElasticSearchNotifier() + args: Dict = { + 'config': 'unittest_configfile.config' + } + data: Dict = {} + TestsCommon.execute_basic_pipeline_steps_for_initialization(config=self.config, args=args, data=data) + + # Set Message + message_entity: SignalNotificationEntityModel = SignalNotificationEntityModel( + signal='INITIALIZATION', + date_time=datetime.datetime(2023, 10, 1, 9, 58, 22, tzinfo=pytz.UTC), + content='Signal Content' + ) + + with mock.patch('TEx.notifier.elastic_search_notifier.AsyncElasticsearch', return_value=elastic_search_api_mock): + # Execute Discord Notifier Configure Method + target.configure( + config=self.config['NOTIFIER.ELASTIC_SEARCH.UT_01'] + ) + + loop = asyncio.get_event_loop() + loop.run_until_complete( + + # Invoke Test Target + target.run( + entity=message_entity, + rule_id='RULE_UT_01', + source='+15558987453' + ) + ) + + # Check .index call + elastic_search_api_mock.index.assert_called_once() + call_arg = elastic_search_api_mock.index.call_args[1] + + self.assertEqual(call_arg['index'], 'test_index_name') + self.assertEqual(call_arg['pipeline'], 'test_pipeline_name') + + submited_document = call_arg['document'] + expected_document = { + 'time': datetime.datetime(2023, 10, 1, 9, 58, 22, tzinfo=pytz.UTC), + 'source': '+15558987453', + 'signal': 'INITIALIZATION', + 'content': 'Signal Content', + } + + self.assertEqual(submited_document, expected_document) + diff --git a/tests/unittest_configfile.config b/tests/unittest_configfile.config index 0bd9cd8..fcdcb81 100644 --- a/tests/unittest_configfile.config +++ b/tests/unittest_configfile.config @@ -4,6 +4,7 @@ api_hash=deff1f2587358746548deadbeef58ddd phone_number=5526986587745 data_path=_data device_model=UT_DEVICE_01 +timeout=20 [OCR] enabled=true @@ -73,10 +74,14 @@ notifier=NOTIFIER.DISCORD.NOT_002 [NOTIFIER.DISCORD.NOT_001] webhook=https://uri.domain.com/webhook/001 prevent_duplication_for_minutes=240 +media_attachments_enabled=true +media_attachments_max_size_bytes=10000000 [NOTIFIER.DISCORD.NOT_002] webhook=https://uri.domain.com/webhook/002 prevent_duplication_for_minutes=240 +media_attachments_enabled=true +media_attachments_max_size_bytes=10000000 [NOTIFIER.ELASTIC_SEARCH.UT_01] address=https://localhost:666