Spaces:
Running
Running
| """Entry point: ``python -m bot``.""" | |
| from __future__ import annotations | |
| import asyncio | |
| import logging | |
| import shutil | |
| import signal | |
| import sys | |
| from pathlib import Path | |
| from typing import Awaitable, Callable | |
| from . import __version__, errorlog, sysinfo | |
| from .config import Config, ConfigError, load_config | |
| from .db import Database | |
| log = logging.getLogger("bot") | |
| TEMP_PREFIXES = ("job", "tracks", "ethumb", "minfo") # the folders the bot makes in TEMP_DIR | |
| def clean_temp_dir(path: str) -> None: | |
| """Remove leftovers of a crashed run; only touches our own job folders.""" | |
| root = Path(path) | |
| root.mkdir(parents=True, exist_ok=True) | |
| for child in root.iterdir(): | |
| if child.is_dir() and child.name.startswith(TEMP_PREFIXES): | |
| shutil.rmtree(child, ignore_errors=True) | |
| async def build_client(cfg: Config, db: Database): | |
| from telethon import TelegramClient | |
| from telethon.sessions import StringSession | |
| token_id = cfg.bot_token.split(":")[0] | |
| session = cfg.session_string | |
| if not session and await db.kv_get("session_bot") == token_id: | |
| session = await db.kv_get("session") or "" # keeps the login across Space restarts | |
| client = TelegramClient( | |
| StringSession(session or None), cfg.api_id, cfg.api_hash, | |
| connection_retries=-1, retry_delay=2, auto_reconnect=True, flood_sleep_threshold=60, | |
| ) | |
| await client.start(bot_token=cfg.bot_token) | |
| await db.kv_set("session", client.session.save()) | |
| await db.kv_set("session_bot", token_id) | |
| return client | |
| async def amain() -> int: | |
| from .app import App | |
| from .web import start_web | |
| cfg = load_config() | |
| logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(name)s: %(message)s") | |
| logging.getLogger("telethon").setLevel(logging.WARNING) | |
| logging.getLogger("sqlalchemy.engine").setLevel(logging.WARNING) | |
| log.info("Auto Rename Bot v%s starting", __version__) | |
| clean_temp_dir(cfg.temp_dir) | |
| db = Database(cfg.database_url) | |
| await db.init() | |
| if db.url.startswith("sqlite"): | |
| log.warning("DATABASE_URL is not set: using local SQLite, so settings are lost when the Space restarts") | |
| runner = await start_web(cfg.port) # up first: the Space must answer its health check while we log in | |
| client = await build_client(cfg, db) | |
| me = await client.get_me() | |
| app = App(client, cfg, db) | |
| username = getattr(me, "username", "") or "" | |
| app.telegraph.set_author(getattr(me, "first_name", "") or username, f"https://t.me/{username}" if username else "") | |
| await app.setup() | |
| channel_log = errorlog.install(app._log_send, errorlog.secrets_from(cfg)) if cfg.log_channel else None # errors -> LOG_CHANNEL | |
| log.info("@%s is online (DC %s, %d connection(s) per transfer, fast transfer %s)", | |
| me.username, client.session.dc_id, cfg.connections, "on" if cfg.fast_transfer else "off") | |
| speed = sysinfo.speed_summary(cfg.fast_transfer, cfg.connections, cfg.stream_transfer) | |
| log.info("speed: %s", speed) | |
| for _short, detail, fine in sysinfo.speed_checks(cfg.fast_transfer, cfg.connections, cfg.stream_transfer): | |
| if not fine: | |
| log.warning("speed: %s", detail) | |
| if cfg.log_channel: | |
| try: | |
| await app._log_send(f"✅ <b>Bot started</b> · v{__version__} · DC {client.session.dc_id}\n⚡ {speed}") | |
| except Exception as exc: # noqa: BLE001 | |
| log.warning("could not post to LOG_CHANNEL: %s", exc) | |
| stop = asyncio.Event() | |
| loop = asyncio.get_running_loop() | |
| for sig in (signal.SIGINT, signal.SIGTERM): | |
| try: | |
| loop.add_signal_handler(sig, stop.set) | |
| except NotImplementedError: # Windows | |
| pass | |
| disconnected = asyncio.ensure_future(client.run_until_disconnected()) | |
| stopper = asyncio.ensure_future(stop.wait()) | |
| await asyncio.wait({disconnected, stopper}, return_when=asyncio.FIRST_COMPLETED) | |
| crashed = disconnected.done() and not stop.is_set() | |
| log.info("shutting down") | |
| await app.shutdown() | |
| if channel_log is not None: | |
| await channel_log.aclose() # sends the reports still on their way, while the client is connected | |
| await runner.cleanup() | |
| try: | |
| await client.disconnect() | |
| except Exception: # noqa: BLE001 | |
| pass | |
| await db.close() | |
| return 1 if crashed else 0 | |
| def _run(main: Callable[[], Awaitable[int]]) -> int: | |
| try: | |
| import uvloop | |
| except ImportError: | |
| return asyncio.run(main()) | |
| if sys.version_info >= (3, 12): | |
| return asyncio.run(main(), loop_factory=uvloop.new_event_loop) | |
| uvloop.install() | |
| return asyncio.run(main()) | |
| def main() -> None: | |
| try: | |
| code = _run(amain) | |
| except ConfigError as exc: | |
| print(f"Configuration error: {exc}", file=sys.stderr) | |
| sys.exit(2) | |
| except KeyboardInterrupt: | |
| code = 0 | |
| sys.exit(code) | |
| if __name__ == "__main__": | |
| main() | |