some refactoring, fixed db stuff and help section

This commit is contained in:
Von Random 2023-11-05 11:31:42 +00:00
parent c9331151f8
commit 696bf5492b
4 changed files with 70 additions and 29 deletions

2
pgbot
View file

@ -41,7 +41,7 @@ def main():
@client.on(telethon.events.NewMessage()) @client.on(telethon.events.NewMessage())
async def handle_new_message(event): async def handle_new_message(event):
if event.message.text.startswith('/'): if event.message.text.startswith('.'):
await commander.action(event) await commander.action(event)
else: else:
await responder.respond(event) await responder.respond(event)

View file

@ -1,6 +1,7 @@
""" Respond to commands """ """ Respond to commands """
import telethon import telethon
import telethon.utils
import pgbotlib.api import pgbotlib.api
import pgbotlib.dbstuff import pgbotlib.dbstuff
@ -16,7 +17,26 @@ class Commander:
T_START = frozenset(['cmd_start']) T_START = frozenset(['cmd_start'])
T_START_E = frozenset(['cmd_start_enabled']) T_START_E = frozenset(['cmd_start_enabled'])
T_STOP = frozenset(['cmd_stop']) 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, def __init__(self, config: dict,
client: telethon.TelegramClient, client: telethon.TelegramClient,
@ -24,7 +44,7 @@ class Commander:
db_conn: pgbotlib.dbstuff.DBConn, db_conn: pgbotlib.dbstuff.DBConn,
namegen: pgbotlib.misc.NameGenerator, namegen: pgbotlib.misc.NameGenerator,
responder: pgbotlib.response.Responder) -> None: responder: pgbotlib.response.Responder) -> None:
self.config = config self.chats = config['chats']
self.client = client self.client = client
self.admins = admins self.admins = admins
self.db_conn = db_conn self.db_conn = db_conn
@ -33,10 +53,9 @@ class Commander:
self.available_tokens = [ self.available_tokens = [
str(token) for token, _ in self.responder.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: if caller not in self.admins:
print('fuck off!') return 'а ты что ещё за хуй с горы?'
return None
input_tokens, phrase = command.strip().split(' ', 1) input_tokens, phrase = command.strip().split(' ', 1)
input_tokenset = frozenset(input_tokens.split(',')) input_tokenset = frozenset(input_tokens.split(','))
for token in input_tokenset: for token in input_tokenset:
@ -44,26 +63,27 @@ class Commander:
return False return False
query = 'INSERT INTO responses (tokens, response) values (%s,%s)' query = 'INSERT INTO responses (tokens, response) values (%s,%s)'
values = (','.join(sorted(input_tokenset)), phrase.strip()) 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: def __add_user(self, caller: int, userspec: str) -> bool:
if caller not in self.admins: if caller not in self.admins:
print('fuck off!') return 'а ты что ещё за хуй с горы?'
return None
user_id, names = userspec.strip().split(' ', 1) user_id, names = userspec.strip().split(' ', 1)
for name in names.strip().split(','): for name in names.strip().split(','):
query = 'INSERT INTO names (tg_id, name) values(%s,%s)' query = 'INSERT INTO names (tg_id, name) values(%s,%s)'
values = (user_id, name) values = (user_id, name)
self.db_conn.update(query, values) self.db_conn.update(query, values)
return True return 'да, господин!'
def __start_response(self) -> str: 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_E)
return self.responder.get_response(self.T_START) return self.responder.get_response(self.T_START)
def __stop_response(self) -> str: 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)
return self.responder.get_response(self.T_STOP_D) return self.responder.get_response(self.T_STOP_D)
@ -74,30 +94,38 @@ class Commander:
async def action(self, async def action(self,
event: telethon.events.common.EventBuilder) -> None: 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 command = event.message.text
sender = await event.get_sender() sender = await event.get_sender()
response = None response = None
if command.startswith('/add '): if command.startswith('.add '):
if self.__add_entry(sender.id, command[5:]): try:
response = 'success' response = self.__add_response(sender.id, command[5:])
else: except Exception as e:
response = 'failure' response = str(e)
elif command.startswith('/adduser '): elif command.startswith('.adduser '):
self.__add_user(sender.id, command[9:]) try:
elif command == '/chat': response = self.__add_user(sender.id, command[9:])
response = str(event.message.peer_id) except Exception as e:
elif command == '/list': response = str(e)
elif command == '.chat':
response = str(chat_id)
elif command == '.list':
response = ', '.join(self.available_tokens) response = ', '.join(self.available_tokens)
elif command == '/users': elif command == '.users':
users = await self.client.get_participants( users = await self.client.get_participants(
entity=event.message.peer_id) entity=event.message.peer_id)
response = self.__list_users(users) response = self.__list_users(users)
elif command == '/start': elif command == '.start':
response = self.__start_response() response = self.__start_response()
self.responder.enable() self.responder.enable()
elif command == '/stop': elif command == '.stop':
response = self.__stop_response() response = self.__stop_response()
self.responder.disable() self.responder.disable()
elif command == '.help':
response = self.DOC
if response: if response:
await self.client.send_message(event.message.peer_id, response) await self.client.send_message(event.message.peer_id, response)
return None return None

View file

@ -7,9 +7,19 @@ class DBConn:
self.connection = psycopg.connect(*args, **kwargs) self.connection = psycopg.connect(*args, **kwargs)
self.cursor = self.connection.cursor() self.cursor = self.connection.cursor()
def update(self, query: str, values: tuple) -> list: def update(self, query: str, values: tuple) -> None:
self.cursor.execute(query, values) failure = None
return self.connection.commit() 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: def query_raw(self, query: str, values: tuple) -> list:
self.cursor.execute(query, values) self.cursor.execute(query, values)

View file

@ -31,6 +31,7 @@ class Responder:
self.apiregex = re.compile(r'^\{(\w+)\}(.+)?$') self.apiregex = re.compile(r'^\{(\w+)\}(.+)?$')
self.namegen = pgbotlib.misc.NameGenerator(config, db_connection) self.namegen = pgbotlib.misc.NameGenerator(config, db_connection)
self.tokens = get_tokens(config['response_tokens']) self.tokens = get_tokens(config['response_tokens'])
self.chats = config['chats']
self.api = pgbotlib.api.ApiWrapper(self.tokens, db_connection) self.api = pgbotlib.api.ApiWrapper(self.tokens, db_connection)
self.db_connection = db_connection self.db_connection = db_connection
self.client = client self.client = client
@ -61,7 +62,7 @@ class Responder:
def disable(self) -> None: def disable(self) -> None:
self.enabled = False self.enabled = False
def enabled(self) -> bool: def is_enabled(self) -> bool:
return self.enabled return self.enabled
def get_response(self, tokens: frozenset) -> str: def get_response(self, tokens: frozenset) -> str:
@ -100,6 +101,8 @@ class Responder:
event: telethon.events.common.EventBuilder) -> None: event: telethon.events.common.EventBuilder) -> None:
if not self.enabled: if not self.enabled:
return None return None
if event.message.peer_id not in self.chats:
return None
message = event.message.text.lower() message = event.message.text.lower()
tokens = self.tokenize(message) tokens = self.tokenize(message)
response = self.get_response(tokens) response = self.get_response(tokens)