bot.py

import os
import shutil
import subprocess
import sys
from queue import Queue, Empty
from threading import Thread
from time import sleep

import psutil
from loguru import logger

from config import Config
from database import BotDBUtil, session, DBVersion

encode = 'UTF-8'


class RestartBot(Exception):
    pass


def get_pid(name):
    return [p.pid for p in psutil.process_iter() if p.name().find(name) != -1]


def enqueue_output(out, queue):
    for line in iter(out.readline, b''):
        queue.put(line)
    out.close()


def init_bot():
    base_superuser = Config('base_superuser')
    if base_superuser is not None:
        BotDBUtil.SenderInfo(base_superuser).edit('isSuperUser', True)


pidlst = []


def run_bot():
    cache_path = os.path.abspath(Config('cache_path'))
    if os.path.exists(cache_path):
        shutil.rmtree(cache_path)
        os.mkdir(cache_path)
    else:
        os.mkdir(cache_path)

    pid_cache = os.path.abspath('.pid_last')
    if os.path.exists(pid_cache):
        with open(pid_cache, 'r') as f:
            pid_last = f.read().split('\n')
            running_pids = get_pid('python')
            for pid in pid_last:
                if int(pid) in running_pids:
                    try:
                        os.kill(int(pid), 9)
                    except (PermissionError, ProcessLookupError):
                        pass
        os.remove(pid_cache)
    envs = os.environ.copy()
    envs['PYTHONIOENCODING'] = 'UTF-8'
    envs['PYTHONPATH'] = os.path.abspath('.')
    botdir = './bots/'
    lst = os.listdir(botdir)
    runlst = []
    for x in lst:
        bot = os.path.abspath(f'{botdir}{x}/bot.py')
        if os.path.exists(bot):
            p = subprocess.Popen([sys.executable, bot], shell=False, stdout=subprocess.PIPE, stderr=subprocess.STDOUT,
                                 cwd=os.path.abspath('.'), env=envs)
            runlst.append(p)
            pidlst.append(p.pid)

    with open(pid_cache, 'w') as c:
        c.write('\n'.join(str(p) for p in pidlst))

    q = Queue()
    threads = []
    for p in runlst:
        threads.append(Thread(target=enqueue_output, args=(p.stdout, q)))

    for t in threads:
        t.daemon = True
        t.start()

    while True:
        try:
            line = q.get_nowait()
        except Empty:
            sleep(1)
        else:
            try:
                logger.info(line.decode(encode)[:-1])
            except UnicodeDecodeError:
                encode_list = ['GBK']
                for e in encode_list:
                    try:
                        logger.warning(f'Cannot decode string from UTF-8, decode with {e}: '
                                       + line.decode(e)[:-1])
                        break
                    except Exception:
                        if encode_list[-1] != e:
                            logger.warning(f'Cannot decode string from {e}, '
                                           f'attempting with {encode_list[encode_list.index(e) + 1]}.')
                        else:
                            logger.error(f'Cannot decode string from {e}, no more attempts.')

        # break when all processes are done.
        if all(p.poll() is not None for p in runlst):
            break

        for p in runlst:
            if p.poll() == 233:
                logger.warning(f'{p.pid} exited with code 233, restart all bots.')
                pidlst.remove(p.pid)
                raise RestartBot
        sleep(0.0001)


if __name__ == '__main__':
    init_bot()
    logger.remove()
    logger.add(sys.stderr, format='{message}', level="INFO")
    query_dbver = session.query(DBVersion).first()
    if query_dbver is None:
        session.add_all([DBVersion(value=str(BotDBUtil.database_version))])
        session.commit()
        query_dbver = session.query(DBVersion).first()
    if (current_ver := int(query_dbver.value)) < (target_ver := BotDBUtil.database_version):
        logger.info(f'Updating database from {current_ver} to {target_ver}...')
        from database.update import update_database

        update_database()
        logger.info('Database updated successfully!')
    try:
        while True:
            try:
                run_bot()  # Process will block here so
                logger.error('All bots exited unexpectedly, please check the output')
                break
            except RestartBot:
                for x in pidlst:
                    try:
                        os.kill(x, 9)
                    except (PermissionError, ProcessLookupError):
                        pass
                pidlst.clear()
                sleep(5)
                continue
    except KeyboardInterrupt:
        for x in pidlst:
            try:
                os.kill(x, 9)
            except (PermissionError, ProcessLookupError):
                pass