diff --git a/docker-compose.yaml b/docker-compose.yaml index 35f50f2..0452494 100644 --- a/docker-compose.yaml +++ b/docker-compose.yaml @@ -9,6 +9,7 @@ services: image: localhost:5001/tg-newsmaker:${APP_VERSION} environment: - TG_OWNER_ID=258985362 + - TG_ADMIN_IDS=${TG_ADMIN_IDS:-} - CONFIG_FILE=/mnt/config.json - SESSION_FILE=/mnt/newsmaker.session - TG_APP_ID=${TG_APP_ID} diff --git a/src/bot/bot.py b/src/bot/bot.py index 14da03f..7da8d02 100644 --- a/src/bot/bot.py +++ b/src/bot/bot.py @@ -11,16 +11,18 @@ logger.setLevel(logging.INFO) class Bot(): - def __init__(self, + def __init__(self, session_file: str, config_file: str, api_id: int, api_hash: str, - owner_id: int) -> None: + owner_id: int, + admin_ids: set = None) -> None: self.__config_file = config_file self.__config = None self.__me = None self.__owner_id = owner_id + self.__admin_ids = set(admin_ids) if admin_ids else set() self.__reload() self.__bot = TelegramClient(session=session_file, api_id=api_id, api_hash=api_hash) self.__weather_queue = Queue() @@ -32,7 +34,9 @@ class Bot(): raise events.StopPropagation @self.__bot.on(events.NewMessage(incoming=True, pattern="!config")) - async def onMessageReload(event): + async def onMessageConfig(event): + if not self.__is_admin(event.sender_id): + raise events.StopPropagation self.__reload() await event.reply("```\n{0}\n```".format(self.__config.read_as_text())) await self.__bot.send_read_acknowledge(event.chat_id, event.message) @@ -40,6 +44,8 @@ class Bot(): @self.__bot.on(events.NewMessage(incoming=True, pattern="!reload")) async def onMessageReload(event): + if not self.__is_admin(event.sender_id): + raise events.StopPropagation self.__reload() await event.reply("Reloaded") await self.__bot.send_read_acknowledge(event.chat_id, event.message) @@ -47,10 +53,21 @@ class Bot(): @self.__bot.on(events.NewMessage(incoming=True, pattern="!chats")) async def onMessageChats(event): + if not self.__is_admin(event.sender_id): + raise events.StopPropagation await self.__send_all_chats_to_owner() await self.__bot.send_read_acknowledge(event.chat_id, event.message) raise events.StopPropagation + @self.__bot.on(events.NewMessage(incoming=True, pattern="!sub")) + async def onMessageSub(event): + if not self.__is_admin(event.sender_id): + raise events.StopPropagation + reply = self.__handle_sub_command(event.message.text or "") + await event.reply(reply) + await self.__bot.send_read_acknowledge(event.chat_id, event.message) + raise events.StopPropagation + @self.__bot.on(events.ChatAction) async def onChatAction(event): if event.user_left and event.user.id == self.__me: @@ -96,6 +113,70 @@ class Bot(): def __reload(self) -> None: self.__config = config.Config(filename=self.__config_file) + def __is_admin(self, user_id) -> bool: + return user_id == self.__owner_id or user_id in self.__admin_ids + + def __handle_sub_command(self, text: str) -> str: + parts = text.strip().split() + if len(parts) < 2: + return self.__sub_help() + cmd = parts[1].lower() + try: + if cmd == 'list': + return self.__sub_list() + if cmd == 'add': + # !sub add + tokens = text.split(None, 5) + if len(tokens) < 6: + return "Usage: !sub add " + src = int(tokens[2]) + dst = int(tokens[3]) + ftype = tokens[4].lower() + if ftype not in ('contain', 'regexp'): + return "filter type must be 'contain' or 'regexp'" + value = tokens[5] + rule = self.__config.add_subscription(src, ftype, value, dst) + return "Added #{0}: {1} [{2}] '{3}' -> {4}".format( + rule['id'], src, ftype, value, dst) + if cmd in ('del', 'rm', 'delete'): + sub_id = int(parts[2]) + ok = self.__config.remove_subscription(sub_id) + return "Removed #{0}".format(sub_id) if ok else "Not found: #{0}".format(sub_id) + if cmd in ('on', 'enable'): + sub_id = int(parts[2]) + ok = self.__config.set_enabled(sub_id, True) + return "Enabled #{0}".format(sub_id) if ok else "Not found: #{0}".format(sub_id) + if cmd in ('off', 'disable'): + sub_id = int(parts[2]) + ok = self.__config.set_enabled(sub_id, False) + return "Disabled #{0}".format(sub_id) if ok else "Not found: #{0}".format(sub_id) + except (ValueError, IndexError): + return self.__sub_help() + return self.__sub_help() + + def __sub_list(self) -> str: + subs = self.__config.list_subscriptions() + if not subs: + return "No subscriptions" + lines = [] + for s in subs: + flt = s.get('filter', {}) + act = s.get('action', {}) + status = 'on' if s.get('enabled', True) else 'off' + name = s.get('name', '') + lines.append("#{0} [{1}] {2} {3}:'{4}' -> {5}{6}".format( + s.get('id'), status, s.get('srcChatId'), + flt.get('type'), flt.get('value'), act.get('chatId'), + " ({0})".format(name) if name else "")) + return "\n".join(lines) + + def __sub_help(self) -> str: + return ("Commands:\n" + "!sub list\n" + "!sub add \n" + "!sub del \n" + "!sub on | !sub off ") + async def __send_all_chats_to_owner(self) -> None: text = "" async for dialog in self.__bot.iter_dialogs(): @@ -113,6 +194,9 @@ class Bot(): i = 0 for config in configs: i = i + 1 + if config.get('enabled') is False: + logger.info("(chat_id, msg_id, cfg_id)=({0}, {1}, {2}): Disabled".format(chat_id, message.id, i)) + continue filter = config['filter'] if 'filter' in config else None action = config['action'] if 'action' in config else None action_type = action['type'] if 'type' in action else None diff --git a/src/main.py b/src/main.py index aea827b..763adc9 100644 --- a/src/main.py +++ b/src/main.py @@ -23,13 +23,16 @@ if __name__ == "__main__": api_hash = os.environ['TG_APP_HASH'] owner_id = int(os.environ['TG_OWNER_ID']) config_file = os.environ['CONFIG_FILE'] + admin_ids = {int(x) for x in os.environ.get('TG_ADMIN_IDS', '').replace(',', ' ').split() if x.strip()} logger.info(f"SESSION_FILE: {session_file}") logger.info(f"CONFIG_FILE: {config_file}") + logger.info(f"ADMIN_IDS: {admin_ids}") bot = Bot(session_file=session_file, api_id=app_id, api_hash=api_hash, config_file=config_file, - owner_id=owner_id) + owner_id=owner_id, + admin_ids=admin_ids) def start_bot(): logger.info("Starting bot") diff --git a/src/settings/config.py b/src/settings/config.py index 6429dfe..ea09fee 100644 --- a/src/settings/config.py +++ b/src/settings/config.py @@ -1,11 +1,16 @@ import json +import os +import tempfile +import threading from typing import Optional + class Config(): def __init__(self, filename: str) -> None: self.__filename = filename self.__config = {} + self.__lock = threading.RLock() self.__read() def read_as_text(self) -> str: @@ -14,20 +19,114 @@ class Config(): def get_on_message_for_chat_id(self, chat_id: int) -> Optional[dict]: chat_id_str = str(chat_id) - if 'onMessage' in self.__config: - on_message = self.__config['onMessage'] - if chat_id_str in on_message: - return on_message[chat_id_str] - return None + on_message = self.__config.get('onMessage', {}) + return on_message.get(chat_id_str) def get_weather_destination_chat_id(self, key: str) -> Optional[int]: key_str = str(key) - if 'weather' in self.__config: - weather = self.__config['weather'] - if key_str in weather: - return weather[key_str] - return None + weather = self.__config.get('weather', {}) + return weather.get(key_str) + + # --- subscription management --- + + def list_subscriptions(self) -> list: + result = [] + on_message = self.__config.get('onMessage', {}) + for src_chat_id, rules in on_message.items(): + if not isinstance(rules, list): + continue + for rule in rules: + item = dict(rule) + item['srcChatId'] = int(src_chat_id) + result.append(item) + result.sort(key=lambda r: r.get('id', 0)) + return result + + def add_subscription(self, src_chat_id: int, filter_type: str, filter_value: str, + dst_chat_id: int, name: Optional[str] = None) -> dict: + with self.__lock: + on_message = self.__config.setdefault('onMessage', {}) + rules = on_message.setdefault(str(src_chat_id), []) + rule = { + 'id': self.__next_id(), + 'enabled': True, + 'filter': {'type': filter_type, 'value': filter_value}, + 'action': {'type': 'forward', 'chatId': dst_chat_id}, + } + if name: + rule['name'] = name + rules.append(rule) + self.__save() + return rule + + def remove_subscription(self, sub_id: int) -> bool: + with self.__lock: + on_message = self.__config.get('onMessage', {}) + for key in list(on_message.keys()): + rules = on_message[key] + if not isinstance(rules, list): + continue + for i, rule in enumerate(rules): + if rule.get('id') == sub_id: + rules.pop(i) + if not rules: + del on_message[key] + self.__save() + return True + return False + + def set_enabled(self, sub_id: int, enabled: bool) -> bool: + with self.__lock: + for rules in self.__config.get('onMessage', {}).values(): + if not isinstance(rules, list): + continue + for rule in rules: + if rule.get('id') == sub_id: + rule['enabled'] = enabled + self.__save() + return True + return False + + def __next_id(self) -> int: + max_id = 0 + for rules in self.__config.get('onMessage', {}).values(): + if not isinstance(rules, list): + continue + for rule in rules: + rid = rule.get('id') + if isinstance(rid, int) and rid > max_id: + max_id = rid + return max_id + 1 + + def __backfill(self) -> None: + changed = False + for rules in self.__config.get('onMessage', {}).values(): + if not isinstance(rules, list): + continue + for rule in rules: + if 'id' not in rule: + rule['id'] = self.__next_id() + changed = True + if 'enabled' not in rule: + rule['enabled'] = True + changed = True + if changed: + self.__save() def __read(self) -> None: with open(self.__filename, 'r') as f: self.__config = json.load(f) + self.__backfill() + + def __save(self) -> None: + with self.__lock: + directory = os.path.dirname(os.path.abspath(self.__filename)) + fd, tmp_path = tempfile.mkstemp(dir=directory, suffix='.tmp') + try: + with os.fdopen(fd, 'w') as f: + json.dump(self.__config, f, ensure_ascii=False, indent=4) + os.replace(tmp_path, self.__filename) + except Exception: + if os.path.exists(tmp_path): + os.remove(tmp_path) + raise