initial commit
This commit is contained in:
		
							parent
							
								
									1307dbcbee
								
							
						
					
					
						commit
						4302290482
					
				
					 8 changed files with 481 additions and 0 deletions
				
			
		| 
						 | 
				
			
			@ -1,2 +1,3 @@
 | 
			
		|||
# 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.
 | 
			
		||||
| 
						 | 
				
			
			
 | 
			
		|||
							
								
								
									
										62
									
								
								pgbot
									
										
									
									
									
										Executable file
									
								
							
							
						
						
									
										62
									
								
								pgbot
									
										
									
									
									
										Executable file
									
								
							| 
						 | 
				
			
			@ -0,0 +1,62 @@
 | 
			
		|||
#!/usr/bin/env python3
 | 
			
		||||
 | 
			
		||||
import sys
 | 
			
		||||
import threading
 | 
			
		||||
 | 
			
		||||
import telethon
 | 
			
		||||
import yaml
 | 
			
		||||
 | 
			
		||||
import pgbotlib.dbstuff
 | 
			
		||||
import pgbotlib.commands
 | 
			
		||||
import pgbotlib.misc
 | 
			
		||||
import pgbotlib.response
 | 
			
		||||
import pgbotlib.sched
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
def init(args: list) -> tuple:
 | 
			
		||||
    conf_path = args[0] if args else 'config.yml'
 | 
			
		||||
    try:
 | 
			
		||||
        with open(conf_path, 'r', encoding='utf-8') as data:
 | 
			
		||||
            config = yaml.safe_load(data.read())
 | 
			
		||||
    except FileNotFoundError as err:
 | 
			
		||||
        sys.exit(err)
 | 
			
		||||
 | 
			
		||||
    client = telethon.TelegramClient(
 | 
			
		||||
            '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
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
def main():
 | 
			
		||||
    config, db_conn, client = init(sys.argv[1:])
 | 
			
		||||
 | 
			
		||||
    responder = pgbotlib.response.Responder(config, client, db_conn)
 | 
			
		||||
    commander = pgbotlib.commands.Commander(config, client, config['admins'],
 | 
			
		||||
                                            db_conn, 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):
 | 
			
		||||
        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:
 | 
			
		||||
            await responder.respond(event)
 | 
			
		||||
 | 
			
		||||
    client.run_until_disconnected()
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
if __name__ == '__main__':
 | 
			
		||||
    main()
 | 
			
		||||
							
								
								
									
										118
									
								
								pgbotlib/api.py
									
										
									
									
									
										Normal file
									
								
							
							
						
						
									
										118
									
								
								pgbotlib/api.py
									
										
									
									
									
										Normal file
									
								
							| 
						 | 
				
			
			@ -0,0 +1,118 @@
 | 
			
		|||
""" Some functions for api calls """
 | 
			
		||||
 | 
			
		||||
import json
 | 
			
		||||
import random
 | 
			
		||||
import re
 | 
			
		||||
 | 
			
		||||
import requests
 | 
			
		||||
import bs4
 | 
			
		||||
import fake_headers
 | 
			
		||||
 | 
			
		||||
import pgbotlib.dbstuff
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
class ApiWrapper:
 | 
			
		||||
    FAILED = 'я обосрался :<'
 | 
			
		||||
    GIF_REGEX = {
 | 
			
		||||
        'part': re.compile(r'(?<=\<center\>).*(?=\<\/center\>)'),
 | 
			
		||||
        'gif': re.compile(r'(?<=src=").*(?="\s)')
 | 
			
		||||
    }
 | 
			
		||||
    SEARCH_TOKENS = ['botname', '!find']
 | 
			
		||||
 | 
			
		||||
    def __init__(self, tokens: dict, db_conn: pgbotlib.dbstuff.DBConn) -> None:
 | 
			
		||||
        self.tokens = tokens
 | 
			
		||||
        self.db_conn = db_conn
 | 
			
		||||
        self.nonw = re.compile(r'\W')
 | 
			
		||||
        self.headers = fake_headers.Headers(headers=True)
 | 
			
		||||
 | 
			
		||||
    # 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
 | 
			
		||||
 | 
			
		||||
    def __sanitize_search(self, message: str) -> str:
 | 
			
		||||
        """Removes one of each of the search tokens from the query
 | 
			
		||||
           so that "bot find" phrase does not poison the search query
 | 
			
		||||
 | 
			
		||||
           It's not guaranteed it will delete the first match though,
 | 
			
		||||
           and I see no point in implementing that"""
 | 
			
		||||
        keywords = self.nonw.sub(' ', message)
 | 
			
		||||
        for token_spec in self.tokens:
 | 
			
		||||
            if token_spec[0] not in self.SEARCH_TOKENS:
 | 
			
		||||
                continue
 | 
			
		||||
            for regex in token_spec[1]:
 | 
			
		||||
                sub_spec = regex.subn('', keywords, count=1)
 | 
			
		||||
                if sub_spec[1]:
 | 
			
		||||
                    keywords = sub_spec[0]
 | 
			
		||||
                    break
 | 
			
		||||
        return keywords
 | 
			
		||||
 | 
			
		||||
    def y_search(self, message: str) -> str:
 | 
			
		||||
        """Pretty much copy & paste from the original bot
 | 
			
		||||
           I have no fucking clue how this black magic works"""
 | 
			
		||||
        query = self.__sanitize_search(message)
 | 
			
		||||
        request = requests.get('https://yandex.ru/images/search',
 | 
			
		||||
                               timeout=30,
 | 
			
		||||
                               params={'text': query,
 | 
			
		||||
                                       'nomisspell': 1,
 | 
			
		||||
                                       'noreask': 1,
 | 
			
		||||
                                       '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'})
 | 
			
		||||
        images = []
 | 
			
		||||
        for item in items:
 | 
			
		||||
            data = json.loads(item.get('data-bem'))
 | 
			
		||||
            images.append(data['serp-item']['img_href'])
 | 
			
		||||
        if not images:
 | 
			
		||||
            return None
 | 
			
		||||
        result = random.choice(images)
 | 
			
		||||
        return f'[url]({result})'
 | 
			
		||||
 | 
			
		||||
    def get_gif(self) -> str:
 | 
			
		||||
        resp = requests.get("http://xdgif.ru/random/", timeout=30)
 | 
			
		||||
        part = self.GIF_REGEX['part'].search(resp.text).group(0)
 | 
			
		||||
        gif = self.GIF_REGEX['gif'].search(part).group(0)
 | 
			
		||||
        return gif
 | 
			
		||||
 | 
			
		||||
    @staticmethod
 | 
			
		||||
    def get_kmp() -> str:
 | 
			
		||||
        request = requests.get("https://killpls.me/random/", timeout=30)
 | 
			
		||||
        parser = bs4.BeautifulSoup(request.text, features="html.parser")
 | 
			
		||||
        result = parser.find("div", attrs={
 | 
			
		||||
            "style": "margin:0.5em 0;line-height:1.785em"})
 | 
			
		||||
        return result.text.strip()
 | 
			
		||||
 | 
			
		||||
    @staticmethod
 | 
			
		||||
    def get_fga() -> str:
 | 
			
		||||
        request = requests.get("http://fucking-great-advice.ru/api/random",
 | 
			
		||||
                               timeout=30)
 | 
			
		||||
        return json.loads(request.text)["text"]
 | 
			
		||||
 | 
			
		||||
    @staticmethod
 | 
			
		||||
    def get_fakenews() -> str:
 | 
			
		||||
        request = requests.get("http://news.olegmakarenko.ru/news", timeout=30)
 | 
			
		||||
        parser = bs4.BeautifulSoup(request.text, features="html.parser")
 | 
			
		||||
        news = [item.text.strip() for item in parser.find_all(
 | 
			
		||||
                    "span", attrs={"class": "headlinetext"})]
 | 
			
		||||
        return random.choice(news)
 | 
			
		||||
 | 
			
		||||
    @staticmethod
 | 
			
		||||
    def get_anek() -> str:
 | 
			
		||||
        request = requests.get("http://rzhunemogu.ru/Rand.aspx?CType=11",
 | 
			
		||||
                               timeout=30)
 | 
			
		||||
        result = request.text.split('<content>')[1].split('</content>')[0]
 | 
			
		||||
        return result.strip()
 | 
			
		||||
 | 
			
		||||
    @staticmethod
 | 
			
		||||
    def format_img(data: str) -> str:
 | 
			
		||||
        return f'[url]({data})'
 | 
			
		||||
							
								
								
									
										61
									
								
								pgbotlib/commands.py
									
										
									
									
									
										Normal file
									
								
							
							
						
						
									
										61
									
								
								pgbotlib/commands.py
									
										
									
									
									
										Normal file
									
								
							| 
						 | 
				
			
			@ -0,0 +1,61 @@
 | 
			
		|||
""" Respond to commands """
 | 
			
		||||
 | 
			
		||||
import telethon
 | 
			
		||||
 | 
			
		||||
import pgbotlib.api
 | 
			
		||||
import pgbotlib.dbstuff
 | 
			
		||||
import pgbotlib.response
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
class Commander:
 | 
			
		||||
    T_START = frozenset(['start_cmd'])
 | 
			
		||||
    T_STOP = frozenset(['stop_cmd'])
 | 
			
		||||
 | 
			
		||||
    def __init__(self, config: dict,
 | 
			
		||||
                 client: telethon.TelegramClient,
 | 
			
		||||
                 admins: list,
 | 
			
		||||
                 db_conn: pgbotlib.dbstuff.DBConn,
 | 
			
		||||
                 responder: pgbotlib.response.Responder) -> None:
 | 
			
		||||
        self.config = config
 | 
			
		||||
        self.client = client
 | 
			
		||||
        self.admins = admins
 | 
			
		||||
        self.db_conn = db_conn
 | 
			
		||||
        self.responder = responder
 | 
			
		||||
        self.available_tokens = [
 | 
			
		||||
                str(token) for token, _ in self.responder.tokens]
 | 
			
		||||
 | 
			
		||||
    def __add_entry(self, caller: int, command: str) -> bool:
 | 
			
		||||
        if caller not in self.admins:
 | 
			
		||||
            print('fuck off!')
 | 
			
		||||
            return None
 | 
			
		||||
        input_tokens, phrase = command.strip().split(' ', 1)
 | 
			
		||||
        input_tokenset = frozenset(input_tokens.split(','))
 | 
			
		||||
        for token in input_tokenset:
 | 
			
		||||
            if token not in self.available_tokens:
 | 
			
		||||
                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)
 | 
			
		||||
 | 
			
		||||
    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 response:
 | 
			
		||||
            await self.client.send_message(event.message.peer_id, response)
 | 
			
		||||
        return None
 | 
			
		||||
							
								
								
									
										22
									
								
								pgbotlib/dbstuff.py
									
										
									
									
									
										Normal file
									
								
							
							
						
						
									
										22
									
								
								pgbotlib/dbstuff.py
									
										
									
									
									
										Normal file
									
								
							| 
						 | 
				
			
			@ -0,0 +1,22 @@
 | 
			
		|||
import random
 | 
			
		||||
import psycopg
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
class DBConn:
 | 
			
		||||
    def __init__(self, *args, **kwargs) -> None:
 | 
			
		||||
        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 query_raw(self, query: str, values: tuple) -> list:
 | 
			
		||||
        self.cursor.execute(query, values)
 | 
			
		||||
        return self.cursor.fetchall()
 | 
			
		||||
 | 
			
		||||
    def query_random(self, query: str, values: tuple) -> str:
 | 
			
		||||
        result = self.query_raw(query, values)
 | 
			
		||||
        if not result:
 | 
			
		||||
            return None
 | 
			
		||||
        return random.choice(result)[0]
 | 
			
		||||
							
								
								
									
										42
									
								
								pgbotlib/misc.py
									
										
									
									
									
										Normal file
									
								
							
							
						
						
									
										42
									
								
								pgbotlib/misc.py
									
										
									
									
									
										Normal file
									
								
							| 
						 | 
				
			
			@ -0,0 +1,42 @@
 | 
			
		|||
import telethon
 | 
			
		||||
import pgbotlib.dbstuff
 | 
			
		||||
import pgbotlib.response
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
class NameGenerator:
 | 
			
		||||
    def __init__(self, config: dict, db: pgbotlib.dbstuff.DBConn) -> None:
 | 
			
		||||
        self.db = db
 | 
			
		||||
 | 
			
		||||
    def get_name(self, sender: telethon.tl.types.User) -> str:
 | 
			
		||||
        query = 'SELECT name FROM names WHERE tg_id = %s'
 | 
			
		||||
        result = self.db.query_random(query, (sender.id,))
 | 
			
		||||
        if result:
 | 
			
		||||
            return result
 | 
			
		||||
        return self.get_tg_name(sender)
 | 
			
		||||
 | 
			
		||||
    @staticmethod
 | 
			
		||||
    def get_tg_name(sender: telethon.tl.types.User) -> str:
 | 
			
		||||
        result = [sender.first_name, sender.last_name]
 | 
			
		||||
        while None in result:
 | 
			
		||||
            result.remove(None)
 | 
			
		||||
        if result:
 | 
			
		||||
            return ' '.join(result)
 | 
			
		||||
        if sender.username:
 | 
			
		||||
            return sender.username
 | 
			
		||||
        return sender.id
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
class MiscReactor:
 | 
			
		||||
    def __init__(self, config: dict, db: pgbotlib.dbstuff.DBConn) -> None:
 | 
			
		||||
        self.namegen = NameGenerator(config, db)
 | 
			
		||||
        self.db = db
 | 
			
		||||
 | 
			
		||||
    def spawn_edited_handler(self,
 | 
			
		||||
                             client: telethon.TelegramClient,
 | 
			
		||||
                             trigger: telethon.events.common.EventBuilder) -> None:
 | 
			
		||||
        @client.on(trigger)
 | 
			
		||||
        async def handle_edited_message(event):
 | 
			
		||||
            sender = await event.get_sender()
 | 
			
		||||
            sender_name = self.namegen.get_name(sender)
 | 
			
		||||
            chat_id = event.message.peer_id
 | 
			
		||||
            await client.send_message(chat_id, f'Я всё видел! {sender_name} совсем охуел, сообщения правит!')
 | 
			
		||||
							
								
								
									
										105
									
								
								pgbotlib/response.py
									
										
									
									
									
										Normal file
									
								
							
							
						
						
									
										105
									
								
								pgbotlib/response.py
									
										
									
									
									
										Normal file
									
								
							| 
						 | 
				
			
			@ -0,0 +1,105 @@
 | 
			
		|||
import re
 | 
			
		||||
 | 
			
		||||
import telethon
 | 
			
		||||
import yaml
 | 
			
		||||
import pgbotlib.api
 | 
			
		||||
import pgbotlib.dbstuff
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
def get_token(token_name: str, token_regex: list) -> tuple:
 | 
			
		||||
    regex = []
 | 
			
		||||
    for i in token_regex:
 | 
			
		||||
        regex.append(re.compile(i))
 | 
			
		||||
    return token_name, regex
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
def get_tokens(path: str) -> list:
 | 
			
		||||
    with open(path, 'r', encoding='utf-8') as data:
 | 
			
		||||
        tokens = yaml.safe_load(data.read())
 | 
			
		||||
    return [get_token(i, tokens[i]) for i in tokens]
 | 
			
		||||
 | 
			
		||||
 | 
			
		||||
class Responder:
 | 
			
		||||
    def __init__(self, config: dict,
 | 
			
		||||
                 client: telethon.TelegramClient,
 | 
			
		||||
                 db_connection: pgbotlib.dbstuff.DBConn) -> None:
 | 
			
		||||
        # apiregex matches "{apiname}optional data"
 | 
			
		||||
        # message itself is also passed to the api call method
 | 
			
		||||
        self.started = True
 | 
			
		||||
        self.apiregex = re.compile(r'^\{(\w+)\}(.+)?$')
 | 
			
		||||
        self.namegen = pgbotlib.misc.NameGenerator(config, db_connection)
 | 
			
		||||
        self.tokens = get_tokens(config['response_tokens'])
 | 
			
		||||
        self.api = pgbotlib.api.ApiWrapper(self.tokens, db_connection)
 | 
			
		||||
        self.db_connection = db_connection
 | 
			
		||||
        self.client = client
 | 
			
		||||
 | 
			
		||||
    def __tokenize(self, message: str) -> frozenset:
 | 
			
		||||
        tokens = set()
 | 
			
		||||
        for token, regexi in self.tokens:
 | 
			
		||||
            for regex in regexi:
 | 
			
		||||
                if regex.search(message):
 | 
			
		||||
                    tokens.add(token)
 | 
			
		||||
                    break
 | 
			
		||||
        return frozenset(tokens)
 | 
			
		||||
 | 
			
		||||
    def __get_keys(self) -> dict:
 | 
			
		||||
        result = {}
 | 
			
		||||
        query = 'SELECT DISTINCT tokens FROM responses'
 | 
			
		||||
        for i in self.db_connection.query_raw(query, tuple()):
 | 
			
		||||
            result[frozenset(i[0].split(','))] = i[0]
 | 
			
		||||
        return result
 | 
			
		||||
 | 
			
		||||
    def __response_choice(self, key: str) -> str:
 | 
			
		||||
        return self.db_connection.query_random(
 | 
			
		||||
                "SELECT response FROM responses WHERE tokens = %s", (key,))
 | 
			
		||||
 | 
			
		||||
    def enable(self) -> None:
 | 
			
		||||
        self.started = True
 | 
			
		||||
 | 
			
		||||
    def disable(self) -> None:
 | 
			
		||||
        self.started = False
 | 
			
		||||
 | 
			
		||||
    def get_response(self, tokens: frozenset) -> str:
 | 
			
		||||
        counter = 0
 | 
			
		||||
        keys = self.__get_keys()
 | 
			
		||||
        for items, string in keys.items():
 | 
			
		||||
            if items <= tokens:
 | 
			
		||||
                # check for priority tokens
 | 
			
		||||
                for token in items:
 | 
			
		||||
                    if token.startswith('!'):
 | 
			
		||||
                        return self.__response_choice(string)
 | 
			
		||||
                match_length = len(items & tokens)
 | 
			
		||||
                if match_length > counter:
 | 
			
		||||
                    counter = match_length
 | 
			
		||||
                    key = string
 | 
			
		||||
        if not counter:
 | 
			
		||||
            return None
 | 
			
		||||
        return self.__response_choice(key)
 | 
			
		||||
 | 
			
		||||
    def api_match(self, response: str, message: str) -> str:
 | 
			
		||||
        match = self.apiregex.search(response)
 | 
			
		||||
        if not match:
 | 
			
		||||
            return response
 | 
			
		||||
        api_spec = match.groups()
 | 
			
		||||
        return self.api.call(*api_spec, message)
 | 
			
		||||
 | 
			
		||||
    async def username(self, response: str,
 | 
			
		||||
                       event: telethon.events.common.EventBuilder) -> str:
 | 
			
		||||
        if '<username>' not in response:
 | 
			
		||||
            return response
 | 
			
		||||
        sender = await event.get_sender()
 | 
			
		||||
        username = self.namegen.get_name(sender)
 | 
			
		||||
        return response.replace('<username>', username)
 | 
			
		||||
 | 
			
		||||
    async def respond(self,
 | 
			
		||||
                      event: telethon.events.common.EventBuilder) -> None:
 | 
			
		||||
        if not self.started:
 | 
			
		||||
            return None
 | 
			
		||||
        message = event.message.text.lower()
 | 
			
		||||
        tokens = self.__tokenize(message)
 | 
			
		||||
        response = self.get_response(tokens)
 | 
			
		||||
        if not response:
 | 
			
		||||
            return None
 | 
			
		||||
        response = self.api_match(response, message)
 | 
			
		||||
        response = await self.username(response, event)
 | 
			
		||||
        await self.client.send_message(event.message.peer_id, response)
 | 
			
		||||
							
								
								
									
										70
									
								
								pgbotlib/sched.py
									
										
									
									
									
										Normal file
									
								
							
							
						
						
									
										70
									
								
								pgbotlib/sched.py
									
										
									
									
									
										Normal file
									
								
							| 
						 | 
				
			
			@ -0,0 +1,70 @@
 | 
			
		|||
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()
 | 
			
		||||
		Loading…
	
	Add table
		Add a link
		
	
		Reference in a new issue