some refactoring, fixed db stuff and help section
This commit is contained in:
parent
c9331151f8
commit
696bf5492b
4 changed files with 70 additions and 29 deletions
2
pgbot
2
pgbot
|
@ -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)
|
||||||
|
|
|
@ -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
|
||||||
|
|
|
@ -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:
|
||||||
|
failure = None
|
||||||
|
try:
|
||||||
|
self.cursor.execute('SAVEPOINT sp1')
|
||||||
self.cursor.execute(query, values)
|
self.cursor.execute(query, values)
|
||||||
return self.connection.commit()
|
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)
|
||||||
|
|
|
@ -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)
|
||||||
|
|
Loading…
Reference in a new issue