diff --git a/pgbot b/pgbot index 7f0774c..c61e283 100755 --- a/pgbot +++ b/pgbot @@ -41,7 +41,7 @@ def main(): @client.on(telethon.events.NewMessage()) async def handle_new_message(event): - if event.message.text.startswith('/'): + if event.message.text.startswith('.'): await commander.action(event) else: await responder.respond(event) diff --git a/pgbotlib/commands.py b/pgbotlib/commands.py index 63df18c..8c9d514 100644 --- a/pgbotlib/commands.py +++ b/pgbotlib/commands.py @@ -1,6 +1,7 @@ """ Respond to commands """ import telethon +import telethon.utils import pgbotlib.api import pgbotlib.dbstuff @@ -16,7 +17,26 @@ class Commander: T_START = frozenset(['cmd_start']) T_START_E = frozenset(['cmd_start_enabled']) T_STOP = frozenset(['cmd_stop']) - T_STOP_D = frozenset(['cmd_stop_d']) + T_STOP_D = frozenset(['cmd_stop_disabled']) + DOC = """ + Команды: + __.start__ + запустить бота + __.stop__ + остановить бота + __.list__ + перечислить доступные токены + __.chat__ + получить id текущего чата + __.users__ + перечислить id пользователей + __.add token1[,token2,...] your phrase here__ + добавить фразу your phrase here для реакции на токены + __.adduser id имя__ + добавить пользователю имя + __.help__ + вывести этот текст + """ def __init__(self, config: dict, client: telethon.TelegramClient, @@ -24,7 +44,7 @@ class Commander: db_conn: pgbotlib.dbstuff.DBConn, namegen: pgbotlib.misc.NameGenerator, responder: pgbotlib.response.Responder) -> None: - self.config = config + self.chats = config['chats'] self.client = client self.admins = admins self.db_conn = db_conn @@ -33,10 +53,9 @@ class Commander: self.available_tokens = [ str(token) for token, _ in self.responder.tokens] - def __add_entry(self, caller: int, command: str) -> bool: + def __add_response(self, caller: int, command: str) -> bool: if caller not in self.admins: - print('fuck off!') - return None + return 'а ты что ещё за хуй с горы?' input_tokens, phrase = command.strip().split(' ', 1) input_tokenset = frozenset(input_tokens.split(',')) for token in input_tokenset: @@ -44,26 +63,27 @@ class Commander: return False query = 'INSERT INTO responses (tokens, response) values (%s,%s)' values = (','.join(sorted(input_tokenset)), phrase.strip()) - return self.db_conn.update(query, values) + self.db_conn.update(query, values) + return 'да, господин!' def __add_user(self, caller: int, userspec: str) -> bool: if caller not in self.admins: - print('fuck off!') - return None + return 'а ты что ещё за хуй с горы?' user_id, names = userspec.strip().split(' ', 1) for name in names.strip().split(','): query = 'INSERT INTO names (tg_id, name) values(%s,%s)' values = (user_id, name) self.db_conn.update(query, values) - return True + return 'да, господин!' + def __start_response(self) -> str: - if self.responder.enabled(): + if self.responder.is_enabled(): return self.responder.get_response(self.T_START_E) return self.responder.get_response(self.T_START) def __stop_response(self) -> str: - if self.responder.enabled(): + if self.responder.is_enabled(): return self.responder.get_response(self.T_STOP) return self.responder.get_response(self.T_STOP_D) @@ -74,30 +94,38 @@ class Commander: async def action(self, event: telethon.events.common.EventBuilder) -> None: + chat_id = telethon.utils.get_peer_id(event.message.peer_id) + if chat_id not in self.chats: + return None command = event.message.text sender = await event.get_sender() response = None - if command.startswith('/add '): - if self.__add_entry(sender.id, command[5:]): - response = 'success' - else: - response = 'failure' - elif command.startswith('/adduser '): - self.__add_user(sender.id, command[9:]) - elif command == '/chat': - response = str(event.message.peer_id) - elif command == '/list': + if command.startswith('.add '): + try: + response = self.__add_response(sender.id, command[5:]) + except Exception as e: + response = str(e) + elif command.startswith('.adduser '): + try: + response = self.__add_user(sender.id, command[9:]) + except Exception as e: + response = str(e) + elif command == '.chat': + response = str(chat_id) + elif command == '.list': response = ', '.join(self.available_tokens) - elif command == '/users': + elif command == '.users': users = await self.client.get_participants( entity=event.message.peer_id) response = self.__list_users(users) - elif command == '/start': + elif command == '.start': response = self.__start_response() self.responder.enable() - elif command == '/stop': + elif command == '.stop': response = self.__stop_response() self.responder.disable() + elif command == '.help': + response = self.DOC if response: await self.client.send_message(event.message.peer_id, response) return None diff --git a/pgbotlib/dbstuff.py b/pgbotlib/dbstuff.py index e207fae..510ec7c 100644 --- a/pgbotlib/dbstuff.py +++ b/pgbotlib/dbstuff.py @@ -7,9 +7,19 @@ class DBConn: self.connection = psycopg.connect(*args, **kwargs) self.cursor = self.connection.cursor() - def update(self, query: str, values: tuple) -> list: - self.cursor.execute(query, values) - return self.connection.commit() + def update(self, query: str, values: tuple) -> None: + failure = None + try: + self.cursor.execute('SAVEPOINT sp1') + self.cursor.execute(query, values) + except Exception as e: + failure = e + self.cursor.execute('ROLLBACK TO SAVEPOINT sp1') + else: + self.cursor.execute('RELEASE SAVEPOINT sp1') + self.connection.commit() + if failure: + raise failure def query_raw(self, query: str, values: tuple) -> list: self.cursor.execute(query, values) diff --git a/pgbotlib/response.py b/pgbotlib/response.py index a4914a8..cbeac8a 100644 --- a/pgbotlib/response.py +++ b/pgbotlib/response.py @@ -31,6 +31,7 @@ class Responder: self.apiregex = re.compile(r'^\{(\w+)\}(.+)?$') self.namegen = pgbotlib.misc.NameGenerator(config, db_connection) self.tokens = get_tokens(config['response_tokens']) + self.chats = config['chats'] self.api = pgbotlib.api.ApiWrapper(self.tokens, db_connection) self.db_connection = db_connection self.client = client @@ -61,7 +62,7 @@ class Responder: def disable(self) -> None: self.enabled = False - def enabled(self) -> bool: + def is_enabled(self) -> bool: return self.enabled def get_response(self, tokens: frozenset) -> str: @@ -100,6 +101,8 @@ class Responder: event: telethon.events.common.EventBuilder) -> None: if not self.enabled: return None + if event.message.peer_id not in self.chats: + return None message = event.message.text.lower() tokens = self.tokenize(message) response = self.get_response(tokens)