257 lines
8.6 KiB
Python
257 lines
8.6 KiB
Python
import asyncio
|
|
import contextlib
|
|
import signal
|
|
from collections.abc import Callable, Coroutine
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
import asyncpg
|
|
import uvloop
|
|
|
|
from dependencies.container import container
|
|
from userbot.modules.capture import CaptureContext, build_capture_context
|
|
from userbot.modules.catchup import SavedState, catch_up, read_saved_state
|
|
from userbot.modules.client import PyroClient
|
|
from userbot.modules.jobs import JobConsumer
|
|
from userbot.modules.uptime import repository as uptime
|
|
from utils.env import env
|
|
from utils.jobs import enqueue
|
|
from utils.listen import Listener
|
|
from utils.logging import logger, setup_logging
|
|
from utils.read.accounts import (
|
|
ACCOUNTS_CHANGED_CHANNEL,
|
|
inactive_session_names,
|
|
session_device_models,
|
|
sync_account,
|
|
)
|
|
from utils.read.watches import WATCHES_CHANGED_CHANNEL
|
|
from utils.storage import ContentAddressedStorage
|
|
|
|
setup_logging()
|
|
|
|
SESSION_FLUSH_INTERVAL = 60
|
|
SESSION_RETRY_INTERVAL = 60
|
|
UPTIME_INTERVAL = 60
|
|
|
|
|
|
@dataclass
|
|
class RunningAccount:
|
|
client: PyroClient
|
|
tasks: list[asyncio.Task]
|
|
device_model: str | None
|
|
|
|
|
|
def _sessions_dir() -> Path:
|
|
sessions_dir = Path(env.tg.sessions_dir)
|
|
sessions_dir.mkdir(parents=True, exist_ok=True)
|
|
return sessions_dir
|
|
|
|
|
|
async def _cancel(task: asyncio.Task) -> None:
|
|
task.cancel()
|
|
with contextlib.suppress(asyncio.CancelledError):
|
|
await task
|
|
|
|
|
|
class AccountRegistry:
|
|
def __init__(self, pool: asyncpg.Pool, storage: ContentAddressedStorage) -> None:
|
|
self._pool = pool
|
|
self._storage = storage
|
|
self._running: dict[str, RunningAccount] = {}
|
|
self._saved: dict[str, SavedState] = {}
|
|
self._lock = asyncio.Lock()
|
|
|
|
@property
|
|
def clients(self) -> list[PyroClient]:
|
|
return [account.client for account in self._running.values()]
|
|
|
|
async def sync(self) -> None:
|
|
async with self._lock:
|
|
inactive = await inactive_session_names(self._pool)
|
|
models = await session_device_models(self._pool)
|
|
present: set[str] = set()
|
|
for path in sorted(_sessions_dir().glob("*.session")):
|
|
if path.stem in inactive:
|
|
await self._log_out(path)
|
|
continue
|
|
present.add(path.stem)
|
|
device_model = models.get(path.stem)
|
|
running = self._running.get(path.stem)
|
|
if running is not None and running.device_model != device_model:
|
|
await self._stop(path.stem)
|
|
running = None
|
|
if running is None:
|
|
await self._start(path, device_model)
|
|
for session_name in set(self._running) - present:
|
|
await self._stop(session_name)
|
|
|
|
async def close(self) -> None:
|
|
for session_name in list(self._running):
|
|
await self._stop(session_name)
|
|
|
|
async def _start(self, path: Path, device_model: str | None) -> None:
|
|
session_name = path.stem
|
|
if session_name not in self._saved:
|
|
self._saved[session_name] = await asyncio.to_thread(read_saved_state, path)
|
|
client = PyroClient(
|
|
session_name, workdir=str(path.parent), device_model=device_model
|
|
)
|
|
try:
|
|
await client.start()
|
|
me = client.me
|
|
if me is None:
|
|
msg = f"session {session_name} is not authorized"
|
|
raise RuntimeError(msg)
|
|
account_id = await sync_account(self._pool, me, session_name)
|
|
client.capture = await build_capture_context(
|
|
client, self._pool, self._storage, account_id
|
|
)
|
|
except Exception:
|
|
logger.exception(
|
|
f"[red]Failed to start session:[/] {session_name}, "
|
|
f"retrying in {SESSION_RETRY_INTERVAL}s"
|
|
)
|
|
with contextlib.suppress(Exception):
|
|
await client.stop()
|
|
return
|
|
consumer = JobConsumer(client, self._pool, account_id)
|
|
saved = self._saved.pop(session_name)
|
|
tasks = [
|
|
asyncio.create_task(consumer.run()),
|
|
asyncio.create_task(catch_up(client, saved)),
|
|
]
|
|
self._running[session_name] = RunningAccount(client, tasks, device_model)
|
|
logger.info(f"[green]Client started:[/] {me.full_name} ({me.id})")
|
|
await enqueue(self._pool, account_id, "sync_dialogs", {})
|
|
await enqueue(self._pool, account_id, "sync_contacts", {})
|
|
|
|
async def _stop(self, session_name: str) -> None:
|
|
account = self._running.pop(session_name, None)
|
|
if account is None:
|
|
return
|
|
for task in account.tasks:
|
|
await _cancel(task)
|
|
with contextlib.suppress(Exception):
|
|
await account.client.stop()
|
|
logger.info(f"[yellow]Client stopped:[/] {session_name}")
|
|
|
|
async def _log_out(self, path: Path) -> None:
|
|
account = self._running.pop(path.stem, None)
|
|
if account is not None:
|
|
for task in account.tasks:
|
|
await _cancel(task)
|
|
client = account.client
|
|
else:
|
|
client = PyroClient(
|
|
path.stem, workdir=str(path.parent), load_handlers=False
|
|
)
|
|
with contextlib.suppress(Exception):
|
|
if account is None:
|
|
await client.start()
|
|
await client.log_out()
|
|
with contextlib.suppress(Exception):
|
|
await client.stop()
|
|
path.unlink(missing_ok=True) # noqa: ASYNC240
|
|
logger.info(f"[yellow]Account logged out:[/] {path.stem}")
|
|
|
|
|
|
async def _flush_sessions(registry: AccountRegistry) -> None:
|
|
while True:
|
|
await asyncio.sleep(SESSION_FLUSH_INTERVAL)
|
|
for client in registry.clients:
|
|
with contextlib.suppress(Exception):
|
|
await client.storage.save()
|
|
|
|
|
|
async def _retry_sessions(registry: AccountRegistry) -> None:
|
|
while True:
|
|
await asyncio.sleep(SESSION_RETRY_INTERVAL)
|
|
try:
|
|
await registry.sync()
|
|
except Exception:
|
|
logger.exception("[red]Session sync failed[/]")
|
|
|
|
|
|
async def _keep_uptime(registry: AccountRegistry, pool: asyncpg.Pool) -> None:
|
|
while True:
|
|
for client in registry.clients:
|
|
if client.capture is not None and client.is_connected:
|
|
with contextlib.suppress(Exception):
|
|
await uptime.heartbeat(pool, client.capture.account_id)
|
|
await asyncio.sleep(UPTIME_INTERVAL)
|
|
|
|
|
|
def _change_listener(registry: AccountRegistry, tasks: set[asyncio.Task]) -> Listener:
|
|
def spawn(coro: Coroutine[Any, Any, None]) -> None:
|
|
task = asyncio.create_task(coro)
|
|
tasks.add(task)
|
|
task.add_done_callback(tasks.discard)
|
|
|
|
def reload(
|
|
make_coro: Callable[[CaptureContext], Coroutine[Any, Any, None]],
|
|
) -> None:
|
|
for client in registry.clients:
|
|
if client.capture is not None:
|
|
spawn(make_coro(client.capture))
|
|
|
|
def on_policy(_payload: str) -> None:
|
|
reload(lambda capture: capture.reload_policies())
|
|
|
|
def on_watch(_payload: str) -> None:
|
|
reload(lambda capture: capture.watches.refresh())
|
|
|
|
def on_accounts(_payload: str) -> None:
|
|
spawn(registry.sync())
|
|
|
|
def on_reconnect() -> None:
|
|
on_policy("")
|
|
on_watch("")
|
|
on_accounts("")
|
|
|
|
return Listener(
|
|
"userbot",
|
|
{
|
|
"policy_changed": on_policy,
|
|
WATCHES_CHANGED_CHANNEL: on_watch,
|
|
ACCOUNTS_CHANGED_CHANNEL: on_accounts,
|
|
},
|
|
on_reconnect,
|
|
)
|
|
|
|
|
|
async def runner() -> None:
|
|
pool = await container.get(asyncpg.Pool)
|
|
storage = await container.get(ContentAddressedStorage)
|
|
|
|
registry = AccountRegistry(pool, storage)
|
|
tasks: set[asyncio.Task] = set()
|
|
listener = _change_listener(registry, tasks)
|
|
stop = asyncio.Event()
|
|
asyncio.get_running_loop().add_signal_handler(signal.SIGTERM, stop.set)
|
|
flush = asyncio.create_task(_flush_sessions(registry))
|
|
retry = asyncio.create_task(_retry_sessions(registry))
|
|
beat = asyncio.create_task(_keep_uptime(registry, pool))
|
|
try:
|
|
await registry.sync()
|
|
if not registry.clients:
|
|
logger.warning("[yellow]No sessions yet. Add an account in the web UI.[/]")
|
|
await listener.start()
|
|
logger.info("[green]Userbot running.[/]")
|
|
await stop.wait()
|
|
finally:
|
|
await _cancel(flush)
|
|
await _cancel(retry)
|
|
await _cancel(beat)
|
|
await listener.stop()
|
|
await registry.close()
|
|
await container.close()
|
|
|
|
|
|
def main() -> None:
|
|
uvloop.install()
|
|
logger.info("Starting userbot...")
|
|
with contextlib.suppress(KeyboardInterrupt):
|
|
asyncio.run(runner())
|
|
logger.info("[red]Userbot stopped.[/]")
|