From cafc0d13ce0e24e50d1eda7258fb2094d07e0bfa Mon Sep 17 00:00:00 2001 From: gmankab Date: Wed, 24 Jul 2024 00:51:03 +0400 Subject: [PATCH] [test] --- reposter/core/common.py | 4 +- reposter/db/init.py | 12 +++ reposter/db/models.py | 25 ++++++ reposter/funcs/other.py | 4 +- reposter/funcs/parse_conf.py | 2 + reposter/handlers/forward_unrestricted.py | 15 +++- reposter/handlers/on_message.py | 46 ----------- reposter/handlers/on_msg.py | 98 +++++++++++++++++++++++ reposter/handlers/resend_restricted.py | 14 +++- reposter/handlers/set.py | 22 ++--- 10 files changed, 179 insertions(+), 63 deletions(-) create mode 100644 reposter/db/init.py create mode 100644 reposter/db/models.py delete mode 100644 reposter/handlers/on_message.py create mode 100644 reposter/handlers/on_msg.py diff --git a/reposter/core/common.py b/reposter/core/common.py index e4a5ed3..9ed740a 100644 --- a/reposter/core/common.py +++ b/reposter/core/common.py @@ -11,9 +11,10 @@ class path: app_dir = src_dir.parent.resolve() pyproject_toml: Path = app_dir / 'pyproject.toml' config_json: Path - data_dir: Path errors_dir: Path + data_dir: Path session: Path + db_path: Path class app: @@ -28,6 +29,7 @@ class app: console: rich.console.Console = rich.console.Console() progress: rich.progress.Progress exit_code: int = 0 + db_url: str class tg: diff --git a/reposter/db/init.py b/reposter/db/init.py new file mode 100644 index 0000000..8b54554 --- /dev/null +++ b/reposter/db/init.py @@ -0,0 +1,12 @@ +import reposter.core.common +import tortoise + + +async def init(): + await tortoise.Tortoise.init( + #db_url=reposter.core.common.app.db_url, + db_url='sqlite://:memory:', + modules={'models': ['reposter.db.models']} + ) + await tortoise.Tortoise.generate_schemas() + diff --git a/reposter/db/models.py b/reposter/db/models.py new file mode 100644 index 0000000..50d251a --- /dev/null +++ b/reposter/db/models.py @@ -0,0 +1,25 @@ +import reposter.core.common +from tortoise.models import Model +from tortoise import fields + + +class Msg(Model): + src_msg = fields.IntField() + src_chat = fields.IntField() + target_msg = fields.IntField() + target_chat = fields.IntField() + + +async def create(): + await Msg.create( + src_msg=1, + src_chat=2, + target_msg=3, + target_chat=4, + ) + reposter.core.common.log( + await Msg.all().first().values() + ) + import os + os._exit(0) + diff --git a/reposter/funcs/other.py b/reposter/funcs/other.py index 760ed5f..bde3a02 100644 --- a/reposter/funcs/other.py +++ b/reposter/funcs/other.py @@ -1,8 +1,9 @@ from reposter.core import common, config import reposter.funcs.parse_conf import reposter.core.config -import reposter.tg.save_file import pyrogram.client +import reposter.tg.save_file +import reposter.db.init import types import sys import io @@ -16,6 +17,7 @@ async def init(): reposter.funcs.parse_conf.check_env() reposter.funcs.parse_conf.read_config() reposter.funcs.parse_conf.check_config() + await reposter.db.init.init() get_client() await common.tg.client.start() await start_log() diff --git a/reposter/funcs/parse_conf.py b/reposter/funcs/parse_conf.py index 010a723..47d8f98 100644 --- a/reposter/funcs/parse_conf.py +++ b/reposter/funcs/parse_conf.py @@ -63,6 +63,8 @@ def read_env() -> None: if not config.env.session_name: config.env.session_name = 'tg_bot' common.path.session = common.path.data_dir / f'{config.env.session_name}.session' + common.path.db_path = common.path.data_dir / 'db.sqlite' + common.app.db_url = f'sqlite://{common.path.db_path}' def check_env(): diff --git a/reposter/handlers/forward_unrestricted.py b/reposter/handlers/forward_unrestricted.py index 6f8e482..72e6978 100644 --- a/reposter/handlers/forward_unrestricted.py +++ b/reposter/handlers/forward_unrestricted.py @@ -3,6 +3,7 @@ import reposter.funcs.handle import reposter.core.config import reposter.core.types +import reposter.db.models import pyrogram.types @@ -10,15 +11,17 @@ class ForwardUnrestricted: def __init__( self, target_any: reposter.core.types.target, - src_msg: pyrogram.types.Message, + src_to_forward: pyrogram.types.Message, + src_in_db: pyrogram.types.Message, ) -> None: self.target_any: reposter.core.types.target = target_any - self.src_msg = src_msg + self.src_msg = src_to_forward + self.src_in_db: pyrogram.types.Message = src_in_db assert isinstance(self.target_any, (str, int, list)) async def forward_all( self, - ): + ) -> None: await reposter.handlers.other.parse_targets( target_any=self.target_any, to_call=self.forward_one, @@ -34,6 +37,12 @@ async def forward_one( drop_author=reposter.core.config.json.drop_author, ) assert isinstance(target_msg, pyrogram.types.Message) + await reposter.db.models.Msg.create( + src_msg=self.src_in_db.id, + src_chat=self.src_in_db.chat.id, + target_msg=target_msg.id, + target_chat=target_msg.chat.id + ) reposter.funcs.logging.log_msg( to_log='[green]\\[forward][/]', source_msg=self.src_msg, diff --git a/reposter/handlers/on_message.py b/reposter/handlers/on_message.py deleted file mode 100644 index 1322b93..0000000 --- a/reposter/handlers/on_message.py +++ /dev/null @@ -1,46 +0,0 @@ -import reposter.handlers.forward_unrestricted -import reposter.handlers.resend_restricted -import reposter.handlers.stream_notify -import reposter.handlers.service -import reposter.core.types -import pyrogram.types - - -class OnMessage: - def __init__( - self, - target_any: reposter.core.types.target - ) -> None: - self.target_any: reposter.core.types.target = target_any - assert isinstance(self.target_any, (str, int, list)) - - async def handler( - self, - _, - src_msg: pyrogram.types.Message, - ) -> None: - if src_msg.service: - service = reposter.handlers.service.Service( - target_any=self.target_any, - src_msg=src_msg, - ) - await service.service_all() - if src_msg.service == pyrogram.enums.MessageServiceType.VIDEO_CHAT_STARTED: - stream_notify = reposter.handlers.stream_notify.StreamNotify( - target_any=self.target_any, - ) - await stream_notify.notify_all() - return - if src_msg.has_protected_content or src_msg.chat.has_protected_content: - real_time_resend = reposter.handlers.resend_restricted.ResendRestricted( - src_msg=src_msg, - target_any=self.target_any, - ) - await real_time_resend.resend_all() - else: - real_time_forward = reposter.handlers.forward_unrestricted.ForwardUnrestricted( - src_msg=src_msg, - target_any=self.target_any, - ) - await real_time_forward.forward_all() - diff --git a/reposter/handlers/on_msg.py b/reposter/handlers/on_msg.py new file mode 100644 index 0000000..f216659 --- /dev/null +++ b/reposter/handlers/on_msg.py @@ -0,0 +1,98 @@ +import reposter.handlers.forward_unrestricted +import reposter.handlers.resend_restricted +import reposter.handlers.stream_notify +import reposter.handlers.service +import reposter.funcs.other +import reposter.core.common +import reposter.core.types +import reposter.db.models +import pyrogram.types + + +class OnMsg: + def __init__( + self, + target_any: reposter.core.types.target + ) -> None: + self.target_any: reposter.core.types.target = target_any + assert isinstance(self.target_any, (str, int, list)) + + async def on_new_msg( + self, + _, + src_msg: pyrogram.types.Message, + ) -> None: + link = reposter.funcs.other.single_link(src_msg) + reposter.core.common.log( + f'[green]\\[new msg] [blue]{link}' + ) + if src_msg.service: + service = reposter.handlers.service.Service( + target_any=self.target_any, + src_msg=src_msg, + ) + await service.service_all() + if src_msg.service == pyrogram.enums.MessageServiceType.VIDEO_CHAT_STARTED: + stream_notify = reposter.handlers.stream_notify.StreamNotify( + target_any=self.target_any, + ) + await stream_notify.notify_all() + return + if src_msg.has_protected_content or src_msg.chat.has_protected_content: + real_time_resend = reposter.handlers.resend_restricted.ResendRestricted( + src_msg=src_msg, + target_any=self.target_any, + ) + await real_time_resend.resend_all() + else: + real_time_forward = reposter.handlers.forward_unrestricted.ForwardUnrestricted( + target_any=self.target_any, + src_to_forward=src_msg, + src_in_db=src_msg, + ) + await real_time_forward.forward_all() + + async def on_edited_msg( + self, + _, + src_msg: pyrogram.types.Message, + ) -> None: + link = reposter.funcs.other.single_link(src_msg) + reposter.core.common.log( + f'[bright_cyan]\\[edited msg] [blue]{link}' + ) + loaded = await reposter.db.models.Msg.get_or_none( + src_msg=src_msg.id, + src_chat=src_msg.chat.id + ) + if loaded: + reposter.core.common.log( + loaded + ) + else: + reposter.core.common.log( + f'[yellow]\\[warn] [blue]{link} edited but was never saved in db' + ) + + async def on_deleted_msg( + self, + _, + src_msg: pyrogram.types.Message, + ) -> None: + link = reposter.funcs.other.single_link(src_msg) + reposter.core.common.log( + f'[red]\\[deleted msg] [blue]{link}' + ) + loaded = await reposter.db.models.Msg.get_or_none( + src_msg=src_msg.id, + src_chat=src_msg.chat.id + ) + if loaded: + reposter.core.common.log( + loaded + ) + else: + reposter.core.common.log( + f'[yellow]\\[warn] [blue]{link} deleted but was never saved in db' + ) + diff --git a/reposter/handlers/resend_restricted.py b/reposter/handlers/resend_restricted.py index be1b4d9..2979288 100644 --- a/reposter/handlers/resend_restricted.py +++ b/reposter/handlers/resend_restricted.py @@ -4,6 +4,7 @@ import reposter.core.config import reposter.core.types import reposter.tg.restricted +import reposter.db.models import pyrogram.types @@ -26,15 +27,22 @@ async def resend_all(self) -> None: target=reposter.core.config.json.logs_chat, ) real_time_forward = reposter.handlers.forward_unrestricted.ForwardUnrestricted( - src_msg=resent_to_log_chat, target_any=self.target_any, + src_to_forward=resent_to_log_chat, + src_in_db=self.src_msg, ) - await real_time_forward.forward_all() + return await real_time_forward.forward_all() elif isinstance(self.target_any, (str, int)): - await self.resend_one( + target_msg = await self.resend_one( src_msg=self.src_msg, target=self.target_any, ) + await reposter.db.models.Msg.create( + src_msg=self.src_msg.id, + src_chat=self.src_msg.chat, + target_msg=target_msg.id, + target_chat=target_msg.chat.id + ) else: raise AssertionError diff --git a/reposter/handlers/set.py b/reposter/handlers/set.py index b05a8a3..6d665f6 100644 --- a/reposter/handlers/set.py +++ b/reposter/handlers/set.py @@ -1,9 +1,10 @@ -import reposter.handlers.on_message +import reposter.handlers.on_msg import reposter.funcs.handle import reposter.funcs.other import reposter.core.common import reposter.core.config import pyrogram.raw.functions +import pyrogram.handlers import asyncio import sys @@ -58,21 +59,24 @@ def set_handlers(): reposter.core.config.json.stream_notify_chats = notify for source_to_fix, target_to_fix in reposter.core.config.json.chats.items(): source = chat_str_fix(source_to_fix) + source_filter = pyrogram.filters.chat(source) if isinstance(target_to_fix, list): target = [] for i in target_to_fix: target.append(chat_str_fix(i)) else: target = chat_str_fix(target_to_fix) - on_message = reposter.handlers.on_message.OnMessage(target_any=target) - reposter.core.common.tg.client.add_handler( - pyrogram.handlers.message_handler.MessageHandler( - callback=on_message.handler, - filters=pyrogram.filters.chat(source), - ) - ) + on_msg = reposter.handlers.on_msg.OnMsg(target_any=target) + for handler, to_run in { + pyrogram.handlers.edited_message_handler.EditedMessageHandler: on_msg.on_edited_msg, + pyrogram.handlers.message_handler.MessageHandler: on_msg.on_new_msg, + pyrogram.handlers.deleted_messages_handler.DeletedMessagesHandler: on_msg.on_deleted_msg, + }.items(): + reposter.core.common.tg.client.add_handler(handler( + callback=to_run, + filters=source_filter, + )) reposter.core.common.log( f'[green]\\[added handler][/] [blue]{source}[/] -> [yellow]{target}' ) -