Skip to content

Commit

Permalink
sync
Browse files Browse the repository at this point in the history
  • Loading branch information
DoroWolf committed Dec 13, 2024
1 parent ff0d9ce commit 040c374
Show file tree
Hide file tree
Showing 27 changed files with 855 additions and 766 deletions.
135 changes: 80 additions & 55 deletions bot.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,16 +7,11 @@
from datetime import datetime
from time import sleep

import core.scripts.config_generate # noqa
from core.config import Config, CFGManager
from core.constants.default import base_superuser_default
from core.constants.path import cache_path
from core.database import BotDBUtil, session, DBVersion
from core.logger import Logger
from core.utils.info import Info
from loguru import logger as loggerFallback


ascii_art = r'''
._. _ .____ ._.
_ _ ____ _
/\ | | (_) | _ \ | |
/ \ | | ____ _ _ __ _ | |_) | ___ | |_
/ /\ \ | |/ / _` | '__| | | _ < / _ \| __|
Expand Down Expand Up @@ -48,20 +43,65 @@ class RestartBot(Exception):


failed_to_start_attempts = {}
disabled_bots = []
processes = []


def init_bot():
import core.scripts.config_generate # noqa
from core.config import Config, CFGManager # noqa
from core.constants.default import base_superuser_default # noqa
from core.database import BotDBUtil, session, DBVersion # noqa
from core.logger import Logger # noqa

query_dbver = session.query(DBVersion).first()
if not query_dbver:
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 core.database.update import update_database

update_database()
Logger.info('Database updated successfully!')
print(ascii_art)
base_superuser = Config('base_superuser', base_superuser_default, cfg_type=(str, list))
if base_superuser:
if isinstance(base_superuser, str):
base_superuser = [base_superuser]
for bu in base_superuser:
BotDBUtil.SenderInfo(bu).init()
BotDBUtil.SenderInfo(bu).edit('isSuperUser', True)
print(ascii_art)
else:
Logger.warning("The base superuser was not found, please setup it in the config file.")

disabled_bots.clear()
for t in CFGManager.values:
if t.startswith('bot_') and not t.endswith('_secret'):
if 'enable' in CFGManager.values[t][t]:
if not CFGManager.values[t][t]['enable']:
disabled_bots.append(t[4:])


def multiprocess_run_until_complete(func):
p = multiprocessing.Process(
target=func,)
p.start()

while True:
if not p.is_alive():
break
sleep(1)
p.terminate()
p.join()
p.close()


def go(bot_name: str = None, subprocess: bool = False, binary_mode: bool = False):
from core.logger import Logger # noqa
from core.utils.info import Info # noqa

Logger.info(f"[{bot_name}] Here we go!")
Info.subprocess = subprocess
Info.binary_mode = binary_mode
Expand All @@ -75,41 +115,34 @@ def go(bot_name: str = None, subprocess: bool = False, binary_mode: bool = False
sys.exit(1)


disabled_bots = []
processes = []

for t in CFGManager.values:
if t.startswith('bot_') and not t.endswith('_secret'):
if 'enable' in CFGManager.values[t][t]:
if not CFGManager.values[t][t]['enable']:
disabled_bots.append(t[4:])


def restart_process(bot_name: str):
if bot_name not in failed_to_start_attempts or datetime.now(
).timestamp() - failed_to_start_attempts[bot_name]['timestamp'] > 60:
failed_to_start_attempts[bot_name] = {}
failed_to_start_attempts[bot_name]['count'] = 0
def run_bot():
from core.constants.path import cache_path # noqa
from core.config import Config # noqa
from core.logger import Logger # noqa

def restart_process(bot_name: str):
if bot_name not in failed_to_start_attempts or datetime.now(
).timestamp() - failed_to_start_attempts[bot_name]['timestamp'] > 60:
failed_to_start_attempts[bot_name] = {}
failed_to_start_attempts[bot_name]['count'] = 0
failed_to_start_attempts[bot_name]['timestamp'] = datetime.now().timestamp()
failed_to_start_attempts[bot_name]['count'] += 1
failed_to_start_attempts[bot_name]['timestamp'] = datetime.now().timestamp()
failed_to_start_attempts[bot_name]['count'] += 1
failed_to_start_attempts[bot_name]['timestamp'] = datetime.now().timestamp()
if failed_to_start_attempts[bot_name]['count'] >= 3:
Logger.error(f'Bot {bot_name} failed to start 3 times, abort to restart, please check the log.')
return

Logger.warning(f'Restarting bot {bot_name}...')
p = multiprocessing.Process(
target=go,
args=(
bot_name,
True,
bool(not sys.argv[0].endswith('.py'))),
name=bot_name)
p.start()
processes.append(p)
if failed_to_start_attempts[bot_name]['count'] >= 3:
Logger.error(f'Bot {bot_name} failed to start 3 times, abort to restart, please check the log.')
return

Logger.warning(f'Restarting bot {bot_name}...')
p = multiprocessing.Process(
target=go,
args=(
bot_name,
True,
bool(not sys.argv[0].endswith('.py'))),
name=bot_name)
p.start()
processes.append(p)

def run_bot():
if os.path.exists(cache_path):
shutil.rmtree(cache_path)
os.makedirs(cache_path, exist_ok=True)
Expand Down Expand Up @@ -152,45 +185,37 @@ def run_bot():
processes.remove(p)
p.terminate()
p.join()
p.close()
restart_process(p.name)
break

if not processes:
break
sleep(1)
Logger.critical('All bots exited unexpectedly, please check the output.')


if __name__ == '__main__':
query_dbver = session.query(DBVersion).first()
if not query_dbver:
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 core.database.update import update_database

update_database()
Logger.info('Database updated successfully!')
init_bot()
try:
while True:
try:
multiprocess_run_until_complete(init_bot)
run_bot() # Process will block here so
Logger.critical('All bots exited unexpectedly, please check the output.')
break
except RestartBot:
for ps in processes:
ps.terminate()
ps.join()
ps.close()
processes.clear()
continue
except Exception:
Logger.critical('An error occurred, please check the output.')
loggerFallback.critical('An error occurred, please check the output.')
traceback.print_exc()
break
except (KeyboardInterrupt, SystemExit):
for ps in processes:
ps.terminate()
ps.join()
ps.close()
processes.clear()
6 changes: 5 additions & 1 deletion core/builtins/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ class Bot:
ModuleHookContext = ModuleHookContext
ExecutionLockList = ExecutionLockList
Info = Info
Temp = Temp

@staticmethod
async def send_message(target: Union[FetchedSession, MessageSession, str],
Expand All @@ -37,7 +38,10 @@ async def send_message(target: Union[FetchedSession, MessageSession, str],
if isinstance(msg, list):
msg = MessageChain(msg)
Logger.info(target.__dict__)
await target.send_direct_message(msg, disable_secret_check, enable_split_image)
await target.send_direct_message(message_chain=msg,
disable_secret_check=disable_secret_check,
enable_parse_message=enable_parse_message,
enable_split_image=enable_split_image)

@staticmethod
async def fetch_target(target: str):
Expand Down
Loading

0 comments on commit 040c374

Please sign in to comment.