"""Command-line interface.""" from __future__ import annotations import argparse import asyncio import os import sys from collections.abc import Callable, Sequence from pathlib import Path import aiohttp from dotenv import load_dotenv from telegram import Bot from telegram.error import InvalidToken, TelegramError from sticker_downloader import __version__ from sticker_downloader.config import ( ALL_FILE_TYPES, DEFAULT_CONCURRENCY, DEFAULT_OUTPUT_DIR, DEFAULT_RETRIES, DEFAULT_URL_FILE, FILE_TYPE_HELP, DownloadConfig, parse_file_types, ) from sticker_downloader.downloader import StickerDownloader from sticker_downloader.errors import StickerDownloaderError from sticker_downloader.fetcher import AiohttpFetcher from sticker_downloader.progress import ConsoleProgress, NullProgress, ProgressReporter from sticker_downloader.results import PackResult from sticker_downloader.urls import parse_pack_reference_lines, read_pack_references EXIT_OK = 0 EXIT_FAILURES = 1 EXIT_USAGE = 2 TOKEN_ENV_VAR = "TELEGRAM_BOT_TOKEN" _TOKEN_INSTRUCTIONS = f"""\ Create a bot with @BotFather on Telegram, then either: * put {TOKEN_ENV_VAR}= in a .env file next to this project, or * export {TOKEN_ENV_VAR}=, or * pass --token .""" _TOKEN_MISSING = f"No Telegram bot token found.\n\n{_TOKEN_INSTRUCTIONS}" _TOKEN_REJECTED = f"Telegram rejected the bot token.\n\n{_TOKEN_INSTRUCTIONS}" class UsageError(StickerDownloaderError): """Raised for bad invocations; reported without a traceback.""" def build_parser() -> argparse.ArgumentParser: types_list = ", ".join(f"{name} ({FILE_TYPE_HELP[name]})" for name in ALL_FILE_TYPES) parser = argparse.ArgumentParser( prog="tg-stickers", description="Download Telegram sticker packs into per-pack folders.", epilog=( "examples:\n" " tg-stickers https://t.me/addstickers/SomePack\n" " tg-stickers SomePack AnotherPack --types webp,png\n" f" tg-stickers --from-file {DEFAULT_URL_FILE} --concurrency 16\n" " tg-stickers --from-file - < my-packs.txt\n" " tg-stickers # interactive prompts\n" ), formatter_class=argparse.RawDescriptionHelpFormatter, ) parser.add_argument( "packs", nargs="*", metavar="PACK", help="sticker pack URLs, tg:// links or bare pack names", ) parser.add_argument( "-f", "--from-file", metavar="PATH", help=f"read references from a file, one per line ('-' for stdin); " f"blank lines and # comments are ignored (default: {DEFAULT_URL_FILE} " f"when chosen interactively)", ) parser.add_argument( "-t", "--types", metavar="LIST", help=f"comma-separated file types to save, or 'all'. Supported: {types_list}", ) parser.add_argument( "-o", "--output", metavar="DIR", type=Path, default=DEFAULT_OUTPUT_DIR, help=f"directory to write packs into (default: {DEFAULT_OUTPUT_DIR})", ) parser.add_argument( "-c", "--concurrency", type=int, default=DEFAULT_CONCURRENCY, metavar="N", help=f"files to download in parallel per pack (default: {DEFAULT_CONCURRENCY})", ) parser.add_argument( "--retries", type=int, default=DEFAULT_RETRIES, metavar="N", help=f"attempts per network call (default: {DEFAULT_RETRIES})", ) parser.add_argument( "--overwrite", action="store_true", help="re-download files that are already on disk (default: skip them)", ) parser.add_argument( "--emoji-names", action="store_true", help="include each sticker's emoji in its file name", ) parser.add_argument( "--no-metadata", dest="write_metadata", action="store_false", help="do not write pack.json alongside the stickers", ) parser.add_argument( "--no-pack-thumbnail", dest="pack_thumbnail", action="store_false", help="do not download the pack's own cover thumbnail", ) parser.add_argument( "--dry-run", action="store_true", help="list what would be downloaded without writing anything", ) parser.add_argument( "--token", metavar="TOKEN", help=f"bot token (default: ${TOKEN_ENV_VAR}, also read from --env-file)", ) parser.add_argument( "--env-file", metavar="PATH", default=".env", help="dotenv file to load the token from (default: .env)", ) verbosity = parser.add_mutually_exclusive_group() verbosity.add_argument( "-v", "--verbose", action="store_true", help="log every file individually" ) verbosity.add_argument( "-q", "--quiet", action="store_true", help="only report fatal errors" ) parser.add_argument("--version", action="version", version=f"%(prog)s {__version__}") return parser def _resolve_references( args: argparse.Namespace, progress: ProgressReporter ) -> list[str]: """Work out which packs to download, prompting only when nothing was given.""" if args.packs and args.from_file: raise UsageError("pass either PACK arguments or --from-file, not both") if args.packs: references, problems = parse_pack_reference_lines(args.packs, source="argument") for problem in problems: progress.warn(problem.split(": ", 1)[-1]) if not references: raise UsageError("no valid sticker pack references given") return references if args.from_file: return _references_from_file(args.from_file, progress) return [] def _references_from_file(source: str, progress: ProgressReporter) -> list[str]: if source == "-": references, problems = parse_pack_reference_lines( sys.stdin.read().splitlines(), source="" ) else: path = Path(source) if not path.is_file(): raise UsageError(f"{path} not found") references, problems = read_pack_references(path) for problem in problems: progress.warn(problem) if not references: raise UsageError(f"no valid sticker pack references in {source}") return references def collect_interactively( *, types_given: str | None, url_file: Path = DEFAULT_URL_FILE, ask: Callable[[str], str] = input, say: Callable[[str], None] = print, ) -> tuple[str | None, list[str]]: """Ask for the file types and pack references. Returns ``(types_text, references)``; ``types_text`` is ``None`` when the caller already supplied ``--types``. """ say(f"--- Telegram Sticker Pack Downloader {__version__} ---") if types_given is None: say("") say("Available file types:") for name in ALL_FILE_TYPES: say(f" {name:<5} {FILE_TYPE_HELP[name]}") types_given = ask("File types (comma-separated, or 'all') [all]: ").strip() # Validate now so a typo does not surface after the bot has connected. parse_file_types(types_given) say("") say("Choose an option:") say(" 1: download one or more packs from URLs you type") say(f" 2: download every pack listed in {url_file}") choice = ask("Your choice (1 or 2) [1]: ").strip() or "1" if choice == "1": answer = ask("Sticker pack URL(s), separated by spaces: ").strip() references = answer.split() if not references: raise UsageError("no sticker pack URL given") return types_given, references if choice == "2": if not url_file.is_file(): raise UsageError(f"{url_file} not found") references, problems = read_pack_references(url_file) for problem in problems: say(f"warning: {problem}") if not references: raise UsageError(f"no valid sticker pack references in {url_file}") say(f"Found {len(references)} pack(s) in {url_file}.") return types_given, references raise UsageError(f"invalid choice {choice!r}; expected 1 or 2") def _make_progress(args: argparse.Namespace) -> ProgressReporter: if args.quiet: return NullProgress() return ConsoleProgress(sys.stderr, verbose=args.verbose) def _resolve_token(args: argparse.Namespace) -> str: env_file = Path(args.env_file) if env_file.is_file(): load_dotenv(env_file) else: load_dotenv() token = (args.token or os.getenv(TOKEN_ENV_VAR) or "").strip() if not token: raise UsageError(_TOKEN_MISSING) return token async def _download( token: str, references: Sequence[str], config: DownloadConfig, progress: ProgressReporter, ) -> list[PackResult]: timeout = aiohttp.ClientTimeout(total=300, sock_connect=30, sock_read=60) async with Bot(token) as bot: me = await bot.get_me() progress.info(f"Connected as @{me.username}.") async with aiohttp.ClientSession(timeout=timeout) as session: downloader = StickerDownloader( bot=bot, fetcher=AiohttpFetcher(session), config=config, progress=progress, ) return await downloader.download_all(references) def main(argv: Sequence[str] | None = None) -> int: """Run the CLI. Returns the process exit code.""" args = build_parser().parse_args(argv) progress = _make_progress(args) try: references = _resolve_references(args, progress) types_text = args.types if not references: if not sys.stdin.isatty(): raise UsageError( "no packs given. Pass PACK arguments, --from-file PATH, " "or run interactively from a terminal." ) types_text, references = collect_interactively(types_given=args.types) config = DownloadConfig( file_types=parse_file_types(types_text), output_dir=args.output, concurrency=args.concurrency, retries=args.retries, overwrite=args.overwrite, write_metadata=args.write_metadata, pack_thumbnail=args.pack_thumbnail, emoji_names=args.emoji_names, dry_run=args.dry_run, ) token = _resolve_token(args) except (UsageError, StickerDownloaderError, ValueError) as error: print(f"error: {error}", file=sys.stderr) return EXIT_USAGE except (EOFError, KeyboardInterrupt): print("\nAborted.", file=sys.stderr) return EXIT_USAGE if config.dry_run: progress.info("Dry run: nothing will be written to disk.") try: results = asyncio.run(_download(token, references, config, progress)) except InvalidToken: print(f"error: {_TOKEN_REJECTED}", file=sys.stderr) return EXIT_USAGE except TelegramError as error: print(f"error: could not reach Telegram: {error}", file=sys.stderr) return EXIT_FAILURES except KeyboardInterrupt: print("\nInterrupted.", file=sys.stderr) return EXIT_FAILURES return EXIT_OK if all(result.ok for result in results) else EXIT_FAILURES def run() -> int: # pragma: no cover - thin console-script wrapper return main() if __name__ == "__main__": # pragma: no cover raise SystemExit(main())