329 lines
14 KiB
Python
329 lines
14 KiB
Python
"""Entry point.
|
|
|
|
python main.py run # long-running worker (default; used by the Heroku Procfile)
|
|
python main.py process # headless: sync, process the queue once, exit
|
|
python main.py sync [--rescan]
|
|
python main.py status
|
|
python main.py check-config
|
|
python main.py login [--session-string]
|
|
python main.py drive-auth [--no-browser]
|
|
python main.py drive-init --name "Telegram Videos"
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import asyncio
|
|
import logging
|
|
import os
|
|
import signal
|
|
import sys
|
|
from dataclasses import dataclass, field
|
|
|
|
from config import Config, ConfigError
|
|
from database import Database, InstanceLock, InstanceLockError
|
|
from errors import BotError
|
|
from logging_config import setup_logging
|
|
|
|
log = logging.getLogger("main")
|
|
|
|
|
|
@dataclass
|
|
class App:
|
|
config: Config
|
|
db: Database
|
|
lock: InstanceLock
|
|
queue: object
|
|
source: object
|
|
clients: list = field(default_factory=list)
|
|
|
|
async def close(self) -> None:
|
|
for client in self.clients:
|
|
try:
|
|
await client.disconnect()
|
|
except Exception as exc: # pragma: no cover
|
|
log.debug("disconnect failed: %s", exc)
|
|
self.db.close()
|
|
self.lock.release()
|
|
|
|
|
|
async def build_app(config: Config, *, control: bool) -> App:
|
|
"""Wire every component. Fails fast with actionable errors."""
|
|
from telethon import events, utils
|
|
|
|
from bot_handlers import CommandRouter, TelethonNotifier, register_command_handlers
|
|
from downloader import Downloader
|
|
from drive_client import DriveClient, FolderResolver
|
|
from progress import LogNotifier, ProgressReporter
|
|
from queue_manager import CURSOR_KEY, QueueManager
|
|
from telegram_client import ACCESS_HASH_KEY, TelethonSource, build_client, resolve_channel
|
|
|
|
lock = InstanceLock(config, wait_seconds=config.instance_lock_wait)
|
|
await asyncio.to_thread(lock.acquire)
|
|
db = Database.from_config(config)
|
|
log.info("Database: %s", db.describe())
|
|
if config.is_heroku and not config.database_url:
|
|
log.warning("ALLOW_EPHEMERAL_DATABASE=true: progress is lost whenever the dyno restarts!")
|
|
db.recover_stale()
|
|
|
|
drive = await asyncio.to_thread(DriveClient.from_config, config)
|
|
await asyncio.to_thread(drive.check_access)
|
|
root_id = config.drive_root_folder_id or "root"
|
|
if config.drive_root_folder_id:
|
|
await asyncio.to_thread(drive.check_root_folder, root_id)
|
|
else:
|
|
log.warning("GOOGLE_DRIVE_ROOT_FOLDER_ID not set; topic folders are created in the top level of My Drive")
|
|
log.info("Google Drive credentials validated")
|
|
|
|
clients = []
|
|
if config.telegram_mode == "bot":
|
|
source_client = build_client(config, "bot")
|
|
await source_client.start(bot_token=config.bot_token)
|
|
control_client = source_client
|
|
clients.append(source_client)
|
|
else:
|
|
source_client = build_client(config, "user")
|
|
await source_client.connect()
|
|
clients.append(source_client)
|
|
if not await source_client.is_user_authorized():
|
|
raise ConfigError(["Telegram user session is not authorized. Run `python main.py login` (see README)."])
|
|
control_client = None
|
|
if config.bot_token and control:
|
|
control_client = build_client(config, "bot")
|
|
await control_client.start(bot_token=config.bot_token)
|
|
clients.append(control_client)
|
|
|
|
is_bot = config.telegram_mode == "bot"
|
|
try:
|
|
entity = await resolve_channel(source_client, config.source_channel, is_bot=is_bot, db=db)
|
|
channel_id = utils.get_peer_id(entity)
|
|
except BotError as exc:
|
|
if not is_bot:
|
|
raise
|
|
# A bot learns a private channel from its first live post (handled below).
|
|
log.warning("%s", exc.public)
|
|
entity = None
|
|
channel_id = int(config.source_channel) if config.source_channel.lstrip("-").isdigit() else None
|
|
if channel_id is None:
|
|
raise
|
|
log.info("Source channel resolved (id %s, mode %s)", channel_id, config.telegram_mode)
|
|
|
|
source = TelethonSource(source_client, entity, channel_id=channel_id, mode=config.telegram_mode, config=config, db=db)
|
|
downloader = Downloader(
|
|
source,
|
|
config.download_dir,
|
|
min_free_bytes=config.min_free_disk_space,
|
|
max_temp_bytes=config.max_temp_storage,
|
|
compute_sha256=config.compute_sha256,
|
|
)
|
|
notifier = TelethonNotifier(control_client) if control_client else LogNotifier()
|
|
reporter = ProgressReporter(notifier, config.notify_chat_ids, config.progress_interval)
|
|
queue = QueueManager(
|
|
config, db, drive, downloader, FolderResolver(drive, db, root_id), source=source, reporter=reporter
|
|
)
|
|
removed = downloader.cleanup_orphans(db, keep_failed=config.keep_failed_downloads)
|
|
log.info("Startup cleanup removed %d abandoned temporary file(s)", removed)
|
|
|
|
if config.process_mode in ("new", "both"):
|
|
queue.listening = True
|
|
|
|
@source_client.on(events.NewMessage(func=lambda e: e.chat_id == channel_id))
|
|
async def _on_channel_post(event) -> None: # pragma: no cover - needs live Telegram
|
|
if source.entity is None:
|
|
source.entity = await event.get_input_chat()
|
|
db.set_kv(ACCESS_HASH_KEY.format(source.entity.channel_id), str(source.entity.access_hash))
|
|
outcome = queue.register(source.convert(event.message))
|
|
log.info("Live post %s: %s", event.message.id, outcome)
|
|
|
|
if config.process_mode == "new" and db.get_int(CURSOR_KEY, 0) == 0 and entity is not None:
|
|
latest = await source.latest_message_id()
|
|
if latest:
|
|
db.raise_int(CURSOR_KEY, latest)
|
|
|
|
if control_client:
|
|
register_command_handlers(control_client, CommandRouter(config, db, queue))
|
|
return App(config=config, db=db, lock=lock, queue=queue, source=source, clients=clients)
|
|
|
|
|
|
async def run_service(config: Config) -> None:
|
|
app = await build_app(config, control=True)
|
|
queue = app.queue
|
|
stop_event = asyncio.Event()
|
|
loop = asyncio.get_running_loop()
|
|
for sig in (signal.SIGINT, signal.SIGTERM):
|
|
try:
|
|
loop.add_signal_handler(sig, stop_event.set)
|
|
except NotImplementedError: # pragma: no cover - Windows
|
|
pass
|
|
|
|
from models import ACTIVE_TASK_STATUSES
|
|
|
|
desired = queue.desired_state()
|
|
if desired != "paused" and (config.auto_start or desired == "running"):
|
|
if config.process_mode != "new" and app.source.entity is not None:
|
|
try:
|
|
result = await queue.sync()
|
|
log.info("Startup sync: %s", result.summary().replace("\n", "; "))
|
|
except BotError as exc:
|
|
log.error("Startup sync failed: %s", exc)
|
|
queue.start()
|
|
booted = queue.boot() # also resumes RUNNING/QUEUED /upload tasks from the database
|
|
active_tasks = app.db.tasks_with_status(*ACTIVE_TASK_STATUSES)
|
|
if booted == "paused":
|
|
state_text = "queue is paused; send /resume to continue"
|
|
elif booted == "running":
|
|
state_text = "processing resumed automatically"
|
|
else:
|
|
state_text = "send /upload to process a message range, or /sync then /run for the whole channel"
|
|
if active_tasks:
|
|
state_text += "; upload tasks: " + ", ".join(f"#{t.id} {t.status.value}" for t in active_tasks)
|
|
if queue.reporter:
|
|
await queue.reporter.notify(f"Bot online ({app.db.describe()}); {state_text}.")
|
|
|
|
waiters = [asyncio.ensure_future(stop_event.wait())]
|
|
waiters += [asyncio.ensure_future(client.disconnected) for client in app.clients]
|
|
await asyncio.wait(waiters, return_when=asyncio.FIRST_COMPLETED)
|
|
log.info("Shutting down: checkpointing current work")
|
|
# Heroku sends SIGTERM and allows 30 seconds before SIGKILL.
|
|
await queue.shutdown(timeout=20)
|
|
for waiter in waiters:
|
|
waiter.cancel()
|
|
await app.close()
|
|
log.info("Shutdown complete")
|
|
|
|
|
|
async def run_headless(config: Config, *, sync_only: bool = False, rescan: bool = False) -> None:
|
|
app = await build_app(config, control=False)
|
|
try:
|
|
if app.source.entity is not None:
|
|
result = await app.queue.sync(rescan=rescan)
|
|
print(result.summary())
|
|
if not sync_only:
|
|
app.queue.listening = False
|
|
app.queue.start()
|
|
await app.queue.wait()
|
|
finally:
|
|
await app.close()
|
|
|
|
|
|
async def login(config: Config, *, session_string: bool) -> None:
|
|
from telethon import TelegramClient
|
|
from telethon.sessions import StringSession
|
|
|
|
if session_string:
|
|
client = TelegramClient(StringSession(), config.api_id, config.api_hash)
|
|
else:
|
|
from telegram_client import build_client
|
|
|
|
client = build_client(config, "user")
|
|
# Telethon asks for the phone number, the login code and (only if enabled) the
|
|
# 2FA password interactively; none of them is stored by this project.
|
|
await client.start()
|
|
me = await client.get_me()
|
|
print(f"Logged in as user id {me.id}.")
|
|
if session_string:
|
|
path = config.session_dir / "session_string.txt"
|
|
config.session_dir.mkdir(parents=True, exist_ok=True)
|
|
fd = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_TRUNC, 0o600)
|
|
with os.fdopen(fd, "w") as fh:
|
|
fh.write(client.session.save())
|
|
print(
|
|
f"Session string written to {path} (permissions 600). It grants full access to your Telegram "
|
|
"account: never share or commit it.\nSet it on Heroku with:\n"
|
|
f' heroku config:set TELEGRAM_SESSION_STRING="$(cat {path})" --app YOUR_APP_NAME\n'
|
|
"then delete the file."
|
|
)
|
|
await client.disconnect()
|
|
|
|
|
|
def drive_init(config: Config, name: str) -> None:
|
|
from drive_client import DriveClient
|
|
|
|
drive = DriveClient.from_config(config)
|
|
folder = drive.create_folder("root", name)
|
|
print(f"Created Drive folder {name!r}. Set GOOGLE_DRIVE_ROOT_FOLDER_ID={folder.id}")
|
|
|
|
|
|
def print_status(config: Config) -> None:
|
|
db = Database.from_config(config)
|
|
print(f"Database: {db.describe()}")
|
|
for state, count in db.counts().items():
|
|
print(f" {state:<12} {count}")
|
|
paused = db.get_kv("queue_paused", "0") == "1"
|
|
print(f"Queue: {'paused' if paused else ('running' if db.get_kv('queue_global_run', '0') == '1' else 'stopped')}")
|
|
for task in db.recent_tasks(5):
|
|
print(f" task #{task.id} {task.status.value} range {task.range_text} uploaded {task.uploaded_count}")
|
|
db.close()
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> int:
|
|
parser = argparse.ArgumentParser(description="Telegram channel -> Google Drive organizer")
|
|
sub = parser.add_subparsers(dest="command")
|
|
sub.add_parser("run", help="run the bot worker (default)")
|
|
sub.add_parser("process", help="headless: sync and process the queue once, then exit")
|
|
p_sync = sub.add_parser("sync", help="discover channel videos and exit")
|
|
p_sync.add_argument("--rescan", action="store_true", help="scan again from START_MESSAGE_ID")
|
|
sub.add_parser("status", help="print queue counts")
|
|
sub.add_parser("check-config", help="validate configuration and exit")
|
|
p_login = sub.add_parser("login", help="log in the Telegram user account (user mode)")
|
|
p_login.add_argument("--session-string", action="store_true", help="create a TELEGRAM_SESSION_STRING for Heroku")
|
|
p_auth = sub.add_parser("drive-auth", help="authorize Google Drive (OAuth) and save the token")
|
|
p_auth.add_argument("--no-browser", action="store_true")
|
|
p_init = sub.add_parser("drive-init", help="create a Drive root folder owned by this app")
|
|
p_init.add_argument("--name", required=True)
|
|
args = parser.parse_args(argv)
|
|
command = args.command or "run"
|
|
|
|
try:
|
|
config = Config.from_env()
|
|
except ConfigError as exc:
|
|
print(exc, file=sys.stderr)
|
|
return 2
|
|
setup_logging(config.log_level, log_dir=None if config.is_heroku else config.log_dir, secrets=config.secrets())
|
|
|
|
try:
|
|
if command in ("run", "process", "sync", "check-config"):
|
|
config.validate("telegram", "drive", "database")
|
|
if command == "check-config":
|
|
print("Configuration OK.")
|
|
elif command == "run":
|
|
asyncio.run(run_service(config))
|
|
elif command == "process":
|
|
asyncio.run(run_headless(config))
|
|
elif command == "sync":
|
|
asyncio.run(run_headless(config, sync_only=True, rescan=args.rescan))
|
|
elif command == "status":
|
|
config.validate("database")
|
|
print_status(config)
|
|
elif command == "login":
|
|
if not config.api_id or not config.api_hash:
|
|
raise ConfigError(["TELEGRAM_API_ID and TELEGRAM_API_HASH are required for login"])
|
|
asyncio.run(login(config, session_string=args.session_string))
|
|
elif command == "drive-auth":
|
|
from drive_client import run_oauth_flow
|
|
|
|
path = run_oauth_flow(config, open_browser=not args.no_browser)
|
|
print(
|
|
f"Token saved to {path} (permissions 600).\nFor Heroku run:\n"
|
|
f' heroku config:set GOOGLE_OAUTH_TOKEN_JSON="$(cat {path})" --app YOUR_APP_NAME'
|
|
)
|
|
elif command == "drive-init":
|
|
config.validate("drive")
|
|
drive_init(config, args.name)
|
|
except ConfigError as exc:
|
|
log.error("%s", exc)
|
|
return 2
|
|
except InstanceLockError as exc:
|
|
log.error("%s", exc)
|
|
return 3
|
|
except BotError as exc:
|
|
log.error("%s", exc.public)
|
|
log.debug("details: %s", exc)
|
|
return 1
|
|
except KeyboardInterrupt:
|
|
return 130
|
|
return 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|