From 84931ba0d30f4ff22cd66d7e486939702bf70451 Mon Sep 17 00:00:00 2001 From: Von Random Date: Tue, 31 Oct 2023 02:28:16 +0200 Subject: [PATCH 01/10] docs, gitignore --- .gitignore | 1 + README.md | 16 ++++++++++++++++ 2 files changed, 17 insertions(+) diff --git a/.gitignore b/.gitignore index b4794e7..3f140f0 100644 --- a/.gitignore +++ b/.gitignore @@ -5,3 +5,4 @@ __pycache__ config.yml conf.d .venv +bot_session.session diff --git a/README.md b/README.md index 20111a7..88ee136 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,19 @@ # pgbot Use `config.yml` to set it up. It needs a list of regex with tokens and a database to match them. This doc is probably going to be abandoned right away, but at least I have added this line I dunno. + +Dependencies: +* bs4 - for parsing html +* fake_headers - for tricking search engines +* psycopg - for pgsql +* pyyaml - for config parsing +* requests - for http requests +* schedule - for scheduling, duh +* telethon - for interacting with bot api + +Initial setup: +``` +python -m venv .venv +source .venv/bin/activate +pip install bs4 fake_headers psycopg pyyaml requests schedule telethon +``` From 7d48ffa35c362ac578328cb05aedfa7fe372e095 Mon Sep 17 00:00:00 2001 From: Von Random Date: Wed, 1 Nov 2023 00:26:09 +0000 Subject: [PATCH 02/10] refactor some shit, enable adding user aliases --- pgbot | 10 ++---- pgbotlib/api.py | 22 ++++++++------ pgbotlib/commands.py | 72 ++++++++++++++++++++++++++++++++++---------- pgbotlib/misc.py | 1 - pgbotlib/response.py | 15 ++++++--- 5 files changed, 81 insertions(+), 39 deletions(-) diff --git a/pgbot b/pgbot index 40f9b82..67cd595 100755 --- a/pgbot +++ b/pgbot @@ -25,8 +25,6 @@ def init(args: list) -> tuple: 'bot_session', config['api_id'], config['api_hash']).start(bot_token=config['bot_token']) - # db_conn = pgbotlib.dbstuff.DBConn( - # f'dbname={config['db_name']} user={config['db_user']}') db_conn = pgbotlib.dbstuff.DBConn(config['db_spec']) return config, db_conn, client @@ -35,9 +33,10 @@ def init(args: list) -> tuple: def main(): config, db_conn, client = init(sys.argv[1:]) - responder = pgbotlib.response.Responder(config, client, db_conn) + namegen = pgbotlib.misc.NameGenerator(config, db_conn) + responder = pgbotlib.response.Responder(config, client, db_conn, namegen) commander = pgbotlib.commands.Commander(config, client, config['admins'], - db_conn, responder) + db_conn, namegen, responder) sched_thread = threading.Thread( target=pgbotlib.sched.spawn_scheduler, @@ -47,9 +46,6 @@ def main(): @client.on(telethon.events.NewMessage()) async def handle_new_message(event): - chat = await event.get_chat() - result = await client.get_messages(chat.id, ids=[event.message.reply_to.reply_to_msg_id]) - print(result) if event.message.text.startswith('/'): await commander.action(event) else: diff --git a/pgbotlib/api.py b/pgbotlib/api.py index 848f3c0..20ea1ec 100644 --- a/pgbotlib/api.py +++ b/pgbotlib/api.py @@ -3,6 +3,7 @@ import json import random import re +import typing import requests import bs4 @@ -27,16 +28,17 @@ class ApiWrapper: # this is the entry point for the api calls # if you add another api, make sure there is a match here - def call(self, api: str, data: str | None, message: str) -> str: - match api: - case 'img_url': return self.format_img(data) - case 'gif': return self.get_gif() - case 'kmp': return self.get_kmp() - case 'fga': return self.get_fga() - case 'fakenews': return self.get_fakenews() - case 'anek': return self.get_anek() - case 'y_search': return self.y_search(message) - case _: return self.FAILED + # this could have used match - case statement, but python 3.9 + def call(self, api: str, data: typing.Union[str, None], + message: str) -> str: + if api == 'img_url': return self.format_img(data) + elif api == 'gif': return self.get_gif() + elif api == 'kmp': return self.get_kmp() + elif api == 'fga': return self.get_fga() + elif api == 'fakenews': return self.get_fakenews() + elif api == 'anek': return self.get_anek() + elif api == 'y_search': return self.y_search(message) + return self.FAILED def __sanitize_search(self, message: str) -> str: """Removes one of each of the search tokens from the query diff --git a/pgbotlib/commands.py b/pgbotlib/commands.py index 360b938..e55b2aa 100644 --- a/pgbotlib/commands.py +++ b/pgbotlib/commands.py @@ -4,22 +4,31 @@ import telethon import pgbotlib.api import pgbotlib.dbstuff +import pgbotlib.misc import pgbotlib.response +# TODO: quote via response? +# chat = await event.get_chat() +# result = await client.get_messages(chat.id, ids=[event.message.reply_to.reply_to_msg_id]) +# print(result) class Commander: - T_START = frozenset(['start_cmd']) - T_STOP = frozenset(['stop_cmd']) + T_START = frozenset(['cmd_start']) + T_START_E = frozenset(['cmd_start_enabled']) + T_STOP = frozenset(['cmd_stop']) + T_STOP_D = frozenset(['cmd_stop_d']) def __init__(self, config: dict, client: telethon.TelegramClient, admins: list, db_conn: pgbotlib.dbstuff.DBConn, + namegen: pgbotlib.misc.NameGenerator, responder: pgbotlib.response.Responder) -> None: self.config = config self.client = client self.admins = admins self.db_conn = db_conn + self.namegen = namegen self.responder = responder self.available_tokens = [ str(token) for token, _ in self.responder.tokens] @@ -37,25 +46,56 @@ class Commander: values = (','.join(sorted(input_tokenset)), phrase.strip()) return self.db_conn.update(query, values) + def __add_user(self, caller: int, userspec: str) -> bool: + if caller not in self.admins: + print('fuck off!') + return None + 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 + + def __start_response(self) -> str: + if self.responder.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(): + return self.responder.get_response(self.T_STOP) + return self.responder.get_response(self.T_STOP_D) + + def __list_users(self, users: list) -> str: + userlist = [f'{user.id}: {self.namegen.get_tg_name(user)}' + for user in users] + return '\n'.join(userlist) + async def action(self, event: telethon.events.common.EventBuilder) -> None: command = event.message.text sender = await event.get_sender() response = None - match command: - case command if command.startswith('/add '): - if self.__add_entry(sender.id, command[5:]): - response = 'success' - else: - response = 'failure' - case '/list': - response = ', '.join(self.available_tokens) - case '/start': - self.responder.enable() - response = self.responder.get_response(self.T_START) - case '/stop': - self.responder.disable() - response = self.responder.get_response(self.T_STOP) + 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 == '/list': + response = ', '.join(self.available_tokens) + elif command == '/users': + users = await self.client.get_participants( + entity=event.message.peer_id) + response = self.__list_users(users) + elif command == '/start': + response = self.__start_response() + self.responder.enable() + elif command == '/stop': + response = self.__stop_response() + self.responder.disable() if response: await self.client.send_message(event.message.peer_id, response) return None diff --git a/pgbotlib/misc.py b/pgbotlib/misc.py index cc51a94..47eda4a 100644 --- a/pgbotlib/misc.py +++ b/pgbotlib/misc.py @@ -1,6 +1,5 @@ import telethon import pgbotlib.dbstuff -import pgbotlib.response class NameGenerator: diff --git a/pgbotlib/response.py b/pgbotlib/response.py index a46559f..df0960b 100644 --- a/pgbotlib/response.py +++ b/pgbotlib/response.py @@ -4,6 +4,7 @@ import telethon import yaml import pgbotlib.api import pgbotlib.dbstuff +import pgbotlib.misc def get_token(token_name: str, token_regex: list) -> tuple: @@ -22,10 +23,11 @@ def get_tokens(path: str) -> list: class Responder: def __init__(self, config: dict, client: telethon.TelegramClient, - db_connection: pgbotlib.dbstuff.DBConn) -> None: + db_connection: pgbotlib.dbstuff.DBConn, + namegen: pgbotlib.misc.NameGenerator) -> None: # apiregex matches "{apiname}optional data" # message itself is also passed to the api call method - self.started = True + self.enabled = True self.apiregex = re.compile(r'^\{(\w+)\}(.+)?$') self.namegen = pgbotlib.misc.NameGenerator(config, db_connection) self.tokens = get_tokens(config['response_tokens']) @@ -54,10 +56,13 @@ class Responder: "SELECT response FROM responses WHERE tokens = %s", (key,)) def enable(self) -> None: - self.started = True + self.enabled = True def disable(self) -> None: - self.started = False + self.enabled = False + + def enabled(self) -> bool: + return self.enabled def get_response(self, tokens: frozenset) -> str: counter = 0 @@ -93,7 +98,7 @@ class Responder: async def respond(self, event: telethon.events.common.EventBuilder) -> None: - if not self.started: + if not self.enabled: return None message = event.message.text.lower() tokens = self.__tokenize(message) From 54be2516fe0b508df445c5f6388fd377a7b3d1f1 Mon Sep 17 00:00:00 2001 From: Von Random Date: Wed, 1 Nov 2023 02:01:19 +0000 Subject: [PATCH 03/10] a script to populate databases, and a command to see chat id --- pgbotlib/commands.py | 2 ++ populate | 37 +++++++++++++++++++++++++++++++++++++ 2 files changed, 39 insertions(+) create mode 100755 populate diff --git a/pgbotlib/commands.py b/pgbotlib/commands.py index e55b2aa..63df18c 100644 --- a/pgbotlib/commands.py +++ b/pgbotlib/commands.py @@ -84,6 +84,8 @@ class Commander: 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': response = ', '.join(self.available_tokens) elif command == '/users': diff --git a/populate b/populate new file mode 100755 index 0000000..6710f9d --- /dev/null +++ b/populate @@ -0,0 +1,37 @@ +#!/usr/bin/env python3 +import psycopg +import yaml +import sys + +with open('config.yml', 'r', encoding='UTF-8') as data: + config = yaml.safe_load(data.read()) +with open(config['response_tokens'], 'r', encoding='UTF-8') as data: + valid_tokens = {item for item in yaml.safe_load(data.read())} +with open(sys.argv[1], 'r', encoding='UTF-8') as data: + phrases = yaml.safe_load(data.read()) + +with psycopg.connect(config['db_spec']) as conn: + query_phrases = 'INSERT INTO responses (tokens, response) VALUES (%s, %s)' + with conn.cursor() as cur: + for regexref, responses in phrases.items(): + tokens = set(regexref.split(',')) + if tokens != tokens & valid_tokens: + print(f'{str(tokens)} failed to add!') + continue + token_string = ','.join(sorted(tokens)) + for response in responses: + cur.execute('SAVEPOINT sp1') + try: + cur.execute( query_phrases, (token_string, response)) + except psycopg.errors.UniqueViolation as err: + cur.execute('ROLLBACK TO SAVEPOINT sp1') + # print(err) + continue + cur.execute('RELEASE SAVEPOINT sp1') + conn.commit() + #for item in names: + # usernames = names[item] + # for username in usernames: + # cur.execute('INSERT INTO names (tg_id, name) VALUES (%s, %s)', + # (item, username)) + #conn.commit() From bee10a2e8965774b71dd02c6db81c087dcb63812 Mon Sep 17 00:00:00 2001 From: Von Random Date: Sun, 5 Nov 2023 01:01:33 +0000 Subject: [PATCH 04/10] implement a functioning scheduler, get rid of threads --- pgbot | 14 ++++----- pgbotlib/cron.py | 32 ++++++++++++++++++++ pgbotlib/response.py | 4 +-- pgbotlib/sched.py | 70 -------------------------------------------- sched.dist.yml | 13 +++----- 5 files changed, 44 insertions(+), 89 deletions(-) create mode 100644 pgbotlib/cron.py delete mode 100644 pgbotlib/sched.py diff --git a/pgbot b/pgbot index 67cd595..7f0774c 100755 --- a/pgbot +++ b/pgbot @@ -1,5 +1,6 @@ #!/usr/bin/env python3 +import asyncio import sys import threading @@ -8,9 +9,9 @@ import yaml import pgbotlib.dbstuff import pgbotlib.commands +import pgbotlib.cron import pgbotlib.misc import pgbotlib.response -import pgbotlib.sched def init(args: list) -> tuple: @@ -38,12 +39,6 @@ def main(): commander = pgbotlib.commands.Commander(config, client, config['admins'], db_conn, namegen, responder) - sched_thread = threading.Thread( - target=pgbotlib.sched.spawn_scheduler, - args=(config, client, responder), - daemon=True) - sched_thread.start() - @client.on(telethon.events.NewMessage()) async def handle_new_message(event): if event.message.text.startswith('/'): @@ -51,7 +46,10 @@ def main(): else: await responder.respond(event) - client.run_until_disconnected() + cron = pgbotlib.cron.Cron(config, client, responder) + cron.plan() + loop = asyncio.get_event_loop() + loop.run_forever() if __name__ == '__main__': diff --git a/pgbotlib/cron.py b/pgbotlib/cron.py new file mode 100644 index 0000000..f731731 --- /dev/null +++ b/pgbotlib/cron.py @@ -0,0 +1,32 @@ +import time +import random + +import yaml +import aiocron +import telethon +import pgbotlib.response + + +class Cron: + def __init__(self, + config: dict, + client: telethon.TelegramClient, + responder: pgbotlib.response.Responder) -> None: + with open(config['schedule'], 'r', encoding='utf-8') as data: + self.sched = yaml.safe_load(data.read()) + self.responder = responder + self.client = client + + def __mkjob(self, job: dict) -> callable: + tokens = frozenset(job['tokens'].split(',')) + async def send_message() -> None: + if 'rand' in job: + time.sleep(random.randint(0, job['rand']) * 60) + message = self.responder.get_response(tokens) + message = self.responder.api_match(message, '') + await self.client.send_message(job['chat'], message) + return send_message + + def plan(self) -> None: + for job in self.sched: + aiocron.crontab(job['cron'], func=self.__mkjob(job)) diff --git a/pgbotlib/response.py b/pgbotlib/response.py index df0960b..a4914a8 100644 --- a/pgbotlib/response.py +++ b/pgbotlib/response.py @@ -35,7 +35,7 @@ class Responder: self.db_connection = db_connection self.client = client - def __tokenize(self, message: str) -> frozenset: + def tokenize(self, message: str) -> frozenset: tokens = set() for token, regexi in self.tokens: for regex in regexi: @@ -101,7 +101,7 @@ class Responder: if not self.enabled: return None message = event.message.text.lower() - tokens = self.__tokenize(message) + tokens = self.tokenize(message) response = self.get_response(tokens) if not response: return None diff --git a/pgbotlib/sched.py b/pgbotlib/sched.py deleted file mode 100644 index 9b3ada2..0000000 --- a/pgbotlib/sched.py +++ /dev/null @@ -1,70 +0,0 @@ -import asyncio -import time -import random - -import yaml -import schedule -import telethon -import pgbotlib.response - - -class Scheduler: - def __init__(self, - config: dict, - client: telethon.TelegramClient, - responder: pgbotlib.response.Responder) -> None: - self.responder = responder - self.client = client - with open(config['schedule'], 'r', encoding='utf-8') as data: - self.sched = yaml.safe_load(data.read()) - self.days = ( - schedule.every().day, - schedule.every().monday, - schedule.every().tuesday, - schedule.every().wednesday, - schedule.every().thursday, - schedule.every().friday, - schedule.every().saturday, - schedule.every().sunday - ) - - def __get_job(self, tokens: frozenset, - chat_id: int, rand: int) -> callable: - async def send_message(): - if rand: - time.sleep(random.randint(0, rand) * 60) - message = self.responder.get_response(tokens) - message = self.responder.api_match(message, '') - await self.client.send_message(chat_id, message) - - def job(): - loop = asyncio.get_event_loop() - coroutine = send_message() - loop.run_until_complete(coroutine) - return job - - def __schedule_job(self, tokens: str, chat: int, - day: int, t: str, rand: int) -> None: - job_tokens = frozenset(tokens.split(',')) - job = self.__get_job(job_tokens, chat, rand) - self.days[day].at(t).do(job) - - def build(self) -> None: - for i in self.sched: - for day in i.get('days', [0]): - for timespec in i['time']: - self.__schedule_job(i['tokens'], i['chat'], - day, timespec, i.get('rand', 0)) - - def run(self) -> None: - while True: - schedule.run_pending() - time.sleep(1) - - -def spawn_scheduler(config: dict, client: telethon.TelegramClient, - responder: pgbotlib.response.Responder) -> Scheduler: - asyncio.set_event_loop(asyncio.new_event_loop()) - scheduler = Scheduler(config, client, responder) - scheduler.build() - scheduler.run() diff --git a/sched.dist.yml b/sched.dist.yml index 997d0d5..f39ba36 100644 --- a/sched.dist.yml +++ b/sched.dist.yml @@ -1,15 +1,10 @@ # schedule things here, see examples - tokens: botname,praise - chat: 00000000 - days: [1, 5] - time: - - "19:59" + cron: 59 19 * * 1-5 rand: 5 + chat: 00000000 - tokens: greeting - chat: 00000000 - days: [1, 2, 3] - time: - - "13:05" - - "13:10" + cron: 5,10 13 * * 1-3 rand: 3 + chat: 00000000 From c9331151f818d8ab7bb11c0686d77c7ecd058381 Mon Sep 17 00:00:00 2001 From: Von Random Date: Sun, 5 Nov 2023 01:07:09 +0000 Subject: [PATCH 05/10] update readme --- README.md | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 88ee136..810a32d 100644 --- a/README.md +++ b/README.md @@ -3,17 +3,17 @@ Use `config.yml` to set it up. It needs a list of regex with tokens and a database to match them. This doc is probably going to be abandoned right away, but at least I have added this line I dunno. Dependencies: +* aiocron - for the scheduler * bs4 - for parsing html * fake_headers - for tricking search engines * psycopg - for pgsql * pyyaml - for config parsing * requests - for http requests -* schedule - for scheduling, duh * telethon - for interacting with bot api Initial setup: ``` python -m venv .venv source .venv/bin/activate -pip install bs4 fake_headers psycopg pyyaml requests schedule telethon +pip install aiocron bs4 fake_headers psycopg pyyaml requests telethon ``` From 696bf5492bfdec38d33cd13cd5e569407cbbfa5c Mon Sep 17 00:00:00 2001 From: Von Random Date: Sun, 5 Nov 2023 11:31:42 +0000 Subject: [PATCH 06/10] some refactoring, fixed db stuff and help section --- pgbot | 2 +- pgbotlib/commands.py | 76 ++++++++++++++++++++++++++++++-------------- pgbotlib/dbstuff.py | 16 ++++++++-- pgbotlib/response.py | 5 ++- 4 files changed, 70 insertions(+), 29 deletions(-) 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) From 6da6769c25313a28f7de0f519aca692faa0bbed3 Mon Sep 17 00:00:00 2001 From: Von Random Date: Sun, 5 Nov 2023 15:05:06 +0000 Subject: [PATCH 07/10] fix check for chats in response module --- pgbotlib/response.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pgbotlib/response.py b/pgbotlib/response.py index cbeac8a..9c02643 100644 --- a/pgbotlib/response.py +++ b/pgbotlib/response.py @@ -101,7 +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: + chat_id = telethon.utils.get_peer_id(event.message.peer_id) + if chat_id not in self.chats: return None message = event.message.text.lower() tokens = self.tokenize(message) From 29d8ff41c54981e75deba093ebef97bebc6f8917 Mon Sep 17 00:00:00 2001 From: Von Random Date: Mon, 6 Nov 2023 10:55:39 +0200 Subject: [PATCH 08/10] use asyncio.sleep instead of time.sleep in cron tasks --- pgbotlib/cron.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/pgbotlib/cron.py b/pgbotlib/cron.py index f731731..1959293 100644 --- a/pgbotlib/cron.py +++ b/pgbotlib/cron.py @@ -1,4 +1,4 @@ -import time +import asyncio import random import yaml @@ -21,7 +21,8 @@ class Cron: tokens = frozenset(job['tokens'].split(',')) async def send_message() -> None: if 'rand' in job: - time.sleep(random.randint(0, job['rand']) * 60) + wait_seconds = random.randint(0, job['rand']) * 60 + await asyncio.sleep(wait_seconds) message = self.responder.get_response(tokens) message = self.responder.api_match(message, '') await self.client.send_message(job['chat'], message) From 75c8727097678ba580983a76d22ec6434a469b41 Mon Sep 17 00:00:00 2001 From: Von Random Date: Mon, 6 Nov 2023 23:21:00 +0200 Subject: [PATCH 09/10] add .token command --- pgbotlib/commands.py | 21 +++++++++++++++++---- 1 file changed, 17 insertions(+), 4 deletions(-) diff --git a/pgbotlib/commands.py b/pgbotlib/commands.py index 8c9d514..dd792fe 100644 --- a/pgbotlib/commands.py +++ b/pgbotlib/commands.py @@ -18,6 +18,8 @@ class Commander: T_START_E = frozenset(['cmd_start_enabled']) T_STOP = frozenset(['cmd_stop']) T_STOP_D = frozenset(['cmd_stop_disabled']) + NOPE = "а ты что ещё за хуй с горы?" + YEP = "да, господин!" DOC = """ Команды: __.start__ @@ -26,6 +28,8 @@ class Commander: остановить бота __.list__ перечислить доступные токены + __.regex token__ + перечислить регулярные выражения, относящиеся к токену __.chat__ получить id текущего чата __.users__ @@ -55,7 +59,7 @@ class Commander: def __add_response(self, caller: int, command: str) -> bool: if caller not in self.admins: - return 'а ты что ещё за хуй с горы?' + return self.NOPE input_tokens, phrase = command.strip().split(' ', 1) input_tokenset = frozenset(input_tokens.split(',')) for token in input_tokenset: @@ -64,17 +68,17 @@ class Commander: query = 'INSERT INTO responses (tokens, response) values (%s,%s)' values = (','.join(sorted(input_tokenset)), phrase.strip()) self.db_conn.update(query, values) - return 'да, господин!' + return self.YEP def __add_user(self, caller: int, userspec: str) -> bool: if caller not in self.admins: - return 'а ты что ещё за хуй с горы?' + return self.NOPE 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 'да, господин!' + return self.YEP def __start_response(self) -> str: @@ -92,6 +96,13 @@ class Commander: for user in users] return '\n'.join(userlist) + def __list_regex(self, token: str) -> str: + for t, r in self.responder.tokens: + if token == t: + regexlist = [i.pattern for i in r] + return '\n'.join(regexlist) + return 'not found!' + async def action(self, event: telethon.events.common.EventBuilder) -> None: chat_id = telethon.utils.get_peer_id(event.message.peer_id) @@ -114,6 +125,8 @@ class Commander: response = str(chat_id) elif command == '.list': response = ', '.join(self.available_tokens) + elif command.startswith('.regex '): + response = self.__list_regex(command[7:].strip()) elif command == '.users': users = await self.client.get_participants( entity=event.message.peer_id) From c65787b6bd3157e0979ffbcea8d5176a013bf073 Mon Sep 17 00:00:00 2001 From: Von Random Date: Fri, 10 Nov 2023 01:02:27 +0200 Subject: [PATCH 10/10] fix y_search --- pgbotlib/api.py | 14 +++++++------- pgbotlib/commands.py | 1 - 2 files changed, 7 insertions(+), 8 deletions(-) diff --git a/pgbotlib/api.py b/pgbotlib/api.py index 20ea1ec..7172c6d 100644 --- a/pgbotlib/api.py +++ b/pgbotlib/api.py @@ -5,9 +5,9 @@ import random import re import typing -import requests import bs4 import fake_headers +import requests import pgbotlib.dbstuff @@ -69,14 +69,14 @@ class ApiWrapper: 'isize': 'medium'}, headers=self.headers.generate()) parser = bs4.BeautifulSoup(request.text, 'html.parser') - items_place = parser.find('div', {'class': 'serp-list'}) - items = items_place.find_all('div', {'class': 'serp-item'}) + items_tag = parser.find('div', {'role': 'main'}) + items_full = json.loads(items_tag.find('div')['data-state']) + items = items_full['initialState']['serpList']['items']['entities'] images = [] - for item in items: - data = json.loads(item.get('data-bem')) - images.append(data['serp-item']['img_href']) + for item in items.values(): + images.append(item.get('origUrl')) if not images: - return None + return self.FAILED result = random.choice(images) return f'[url]({result})' diff --git a/pgbotlib/commands.py b/pgbotlib/commands.py index dd792fe..1dfe22b 100644 --- a/pgbotlib/commands.py +++ b/pgbotlib/commands.py @@ -3,7 +3,6 @@ import telethon import telethon.utils -import pgbotlib.api import pgbotlib.dbstuff import pgbotlib.misc import pgbotlib.response