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:
 | 
				
			||||||
        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)
 | 
				
			||||||
| 
						 | 
					
 | 
				
			||||||
| 
						 | 
					@ -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…
	
	Add table
		Add a link
		
	
		Reference in a new issue