2024-09-15 23:10:27 +01:00
|
|
|
import asyncio
|
2024-09-18 17:22:55 +01:00
|
|
|
import copy
|
2024-09-18 17:18:36 +01:00
|
|
|
import time
|
2024-09-15 23:10:27 +01:00
|
|
|
|
2024-07-28 03:59:19 +01:00
|
|
|
import click
|
|
|
|
import logging
|
2024-09-18 17:22:55 +01:00
|
|
|
|
|
|
|
import niobot
|
2024-07-28 03:59:19 +01:00
|
|
|
from tortoise import run_async
|
|
|
|
|
|
|
|
|
|
|
|
log = logging.getLogger(__name__)
|
|
|
|
|
|
|
|
|
|
|
|
@click.group()
|
2024-07-31 23:17:09 +01:00
|
|
|
@click.option(
|
|
|
|
"--log-level",
|
|
|
|
default="INFO",
|
2024-09-15 23:03:24 +01:00
|
|
|
type=click.Choice(
|
|
|
|
["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"], case_sensitive=False
|
|
|
|
),
|
|
|
|
help="Set the log level.",
|
|
|
|
)
|
|
|
|
@click.option(
|
|
|
|
"--log-file",
|
|
|
|
default=None,
|
|
|
|
type=click.Path(),
|
|
|
|
help="Log to a file as well as stdout",
|
2024-07-31 23:17:09 +01:00
|
|
|
)
|
2024-08-04 17:21:00 +01:00
|
|
|
def cli(log_level: str, log_file: str | None):
|
|
|
|
logging.basicConfig(
|
|
|
|
level=log_level.upper(),
|
|
|
|
format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
|
|
|
|
datefmt="%Y-%m-%d %H:%M:%S",
|
|
|
|
)
|
|
|
|
if log_file is not None:
|
|
|
|
file_handler = logging.FileHandler(log_file)
|
2024-08-04 18:03:28 +01:00
|
|
|
file_handler.setLevel(logging.DEBUG)
|
2024-09-15 23:03:24 +01:00
|
|
|
file_handler.setFormatter(
|
|
|
|
logging.Formatter("%(asctime)s - %(name)s - %(levelname)s - %(message)s")
|
|
|
|
)
|
2024-08-04 17:21:00 +01:00
|
|
|
logging.getLogger().addHandler(file_handler)
|
2024-07-28 03:59:19 +01:00
|
|
|
|
|
|
|
|
|
|
|
@cli.command()
|
2024-09-18 17:18:36 +01:00
|
|
|
@click.option(
|
|
|
|
"--resume/--no-resume",
|
2024-09-18 17:19:08 +01:00
|
|
|
default=True,
|
2024-09-18 17:18:36 +01:00
|
|
|
help="Whether to do a full sync, or resume from previous"
|
|
|
|
)
|
|
|
|
def run(resume: bool):
|
2024-07-28 03:59:19 +01:00
|
|
|
"""Runs the bot"""
|
|
|
|
log.info("Starting bot.")
|
2024-09-18 17:22:55 +01:00
|
|
|
log.debug("Importing instance...")
|
2024-07-28 03:59:19 +01:00
|
|
|
from .main import bot
|
2024-09-18 17:22:55 +01:00
|
|
|
log.debug("Import complete.")
|
2024-07-28 23:33:15 +01:00
|
|
|
|
2024-09-18 17:18:36 +01:00
|
|
|
if resume is False:
|
2024-09-18 17:22:55 +01:00
|
|
|
log.debug("Resume is false, hijacking config...")
|
|
|
|
original_config = copy.copy(bot.config)
|
|
|
|
new_config = niobot.AsyncClientConfig(
|
|
|
|
store=original_config.store,
|
|
|
|
encryption_enabled=original_config.encryption_enabled,
|
|
|
|
store_name=original_config.store_name,
|
|
|
|
pickle_key=original_config.pickle_key,
|
|
|
|
store_sync_tokens=False,
|
|
|
|
custom_headers=original_config.custom_headers,
|
|
|
|
max_limit_exceeded=original_config.max_limit_exceeded,
|
|
|
|
max_timeouts=original_config.max_timeouts,
|
|
|
|
backoff_factor=original_config.backoff_factor,
|
|
|
|
max_timeout_retry_wait_time=original_config.max_timeout_retry_wait_time,
|
|
|
|
request_timeout=original_config.request_timeout,
|
|
|
|
io_chunk_size=original_config.io_chunk_size
|
|
|
|
)
|
|
|
|
bot.config = new_config
|
|
|
|
bot.load_store()
|
2024-09-18 17:18:36 +01:00
|
|
|
bot.loaded_sync_token = ""
|
2024-09-18 17:22:55 +01:00
|
|
|
log.debug("Bot is now prepared for a complete sync")
|
2024-09-18 17:18:36 +01:00
|
|
|
log.critical("Not resuming from previous sync. This may be slow!")
|
|
|
|
time.sleep(3)
|
|
|
|
|
2024-09-08 01:33:31 +01:00
|
|
|
# run_async(bot.start(access_token=bot.cfg["bot"]["access_token"]))
|
2024-09-18 17:22:55 +01:00
|
|
|
log.debug("All pistons firing")
|
2024-09-15 23:10:27 +01:00
|
|
|
asyncio.run(bot.start(access_token=bot.cfg["bot"]["access_token"]))
|
2024-07-28 03:59:19 +01:00
|
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
cli()
|