script for sending files from stdin to Telegram using mtproto

tg_uploader.py

#!/usr/bin/env -S uv run --script
# /// script
# requires-python = ">=3.12"
# dependencies = [ "pyrofork", "tgcrypto-pyrofork", "rich", "pysocks" ]
# ///

import argparse
import asyncio
import io
import math
import os
import signal
import sys
import types
import typing
import time
from urllib.parse import urlparse, unquote

from rich.console import Console
from rich.progress import (
    Progress,
    SpinnerColumn,
    TextColumn,
    BarColumn,
    MofNCompleteColumn,
    DownloadColumn,
    TransferSpeedColumn,
    TimeRemainingColumn,
)

import pyrogram
import pyrogram.client
import pyrogram.methods.advanced.save_file
import pyrogram.raw
import pyrogram.session.session
from pyrogram.enums import ChatAction
from pyrogram.errors import FloodWait

MAX_CHUNK_SIZE = 512 * 1024

MIN_BIG_FILE_SIZE = 10 * 1024 * 1024

DEFAULT_WORKERS = 4
DEFAULT_BUFFER_PARTS = 16

console = Console()


async def gen_typing(client, chat_id, typing_task):
    async def cycle():
        while True:
            await client.send_chat_action(chat_id, ChatAction.UPLOAD_DOCUMENT)
            await asyncio.sleep(4)

    if typing_task is True:
        typing_task = asyncio.create_task(cycle())
        return typing_task
    else:
        typing_task.cancel()
        await client.send_chat_action(chat_id, ChatAction.CANCEL)
        try:
            await typing_task
        except asyncio.CancelledError:
            pass


def make_progress(transient: bool = True) -> Progress:
    return Progress(
        SpinnerColumn(),
        TextColumn("[bold]Upload[/]"),
        BarColumn(),
        MofNCompleteColumn(),
        DownloadColumn(),
        TransferSpeedColumn(),
        TimeRemainingColumn(),
        transient=transient,
        console=console,
    )

def human_size(n: float) -> str:
    units = ("B", "KB", "MB", "GB", "TB", "PB")
    i = 0
    while n >= 1024 and i < len(units) - 1:
        n /= 1024.0
        i += 1
    return f"{n:.1f} {units[i]}"


class LineProgress:
    def __init__(self, console: Console, interval: float = 1.0):
        self.console = console
        self.interval = float(interval)
        self._tasks: dict[int, dict[str, typing.Any]] = {}
        self._next_id = 0

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc, tb):
        return False

    def add_task(self, name: str, total: int = 0) -> int:
        t_id = self._next_id
        self._next_id += 1
        now = time.monotonic()
        self._tasks[t_id] = {
            "name": name,
            "total": int(total),
            "current": 0,
            "start": now,
            "last_print": 0.0,
        }
        self._print(t_id, force=True)
        return t_id

    def update(self, task_id: int, advance: int = 0):
        task = self._tasks.get(task_id)
        if not task:
            return
        task["current"] += int(advance)
        now = time.monotonic()
        done = task["current"] >= task["total"] > 0
        if done or (now - task["last_print"] >= self.interval):
            self._print(task_id, force=True)

    def _print(self, task_id: int, force: bool = False):
        task = self._tasks[task_id]
        now = time.monotonic()
        elapsed = max(1e-6, now - task["start"])
        cur = task["current"]
        total = task["total"]
        speed = cur / elapsed
        pct = (cur / total) * 100 if total else 0.0
        remain = (total - cur) / speed if speed > 0 else float("inf")

        bar_w = 30
        filled = int(bar_w * pct / 100)
        bar = "█" * filled + "░" * (bar_w - filled)

        msg = (
            f"[{bar}] {pct:6.2f}%  "
            f"{human_size(cur)}/{human_size(total)}  "
            f"{human_size(speed)}/s  ETA {format_eta(remain)}"
        )
        self.console.log(msg)
        task["last_print"] = now


def format_eta(seconds: float | None) -> str:
    if seconds is None or seconds == float("inf"):
        return "--:--"
    seconds = int(max(0, seconds))
    m, s = divmod(seconds, 60)
    h, m = divmod(m, 60)
    if h:
        return f"{h:02d}:{m:02d}:{s:02d}"
    return f"{m:02d}:{s:02d}"

class StdinChunkGen:
    def __init__(
        self,
        total_size: int,
        name: str,
        read_fd=sys.stdin.buffer,
        progress: Progress | None = None,
        task_id: int | None = None,
        cancel_event: asyncio.Event | None = None,
    ):
        self.total = int(total_size)
        self.name = name
        self._fd = read_fd
        self.progress = progress
        self.task_id = task_id
        self.cancel_event = cancel_event or asyncio.Event()

    def __aiter__(self):
        return self._stream()

    async def _read(self, n: int) -> bytes:
        loop = asyncio.get_running_loop()
        return await loop.run_in_executor(None, self._fd.read, n)

    async def _stream(self) -> typing.AsyncGenerator[bytes, None]:
        remaining = self.total
        while remaining > 0:
            if self.cancel_event.is_set():
                raise asyncio.CancelledError()
            to_read = min(MAX_CHUNK_SIZE, remaining)
            chunk = await self._read(to_read)
            if not chunk:
                break
            remaining -= len(chunk)
            if self.progress and self.task_id is not None:
                self.progress.update(self.task_id, advance=len(chunk))
            yield chunk


async def save_file_custom_wrapper(
    self: pyrogram.client.Client,
    path: str | typing.BinaryIO | StdinChunkGen,
    file_id: int | None = None,
    file_part: int = 0,
    progress: typing.Callable | None = None,
    progress_args: tuple = (),
):
    if isinstance(path, StdinChunkGen):
        return await save_file_from_bytes_gen(self=self, chunk_gen=path)
    return await pyrogram.methods.advanced.save_file.SaveFile.save_file(
        self=self,
        path=path,
        file_id=file_id,
        file_part=file_part,
        progress=progress,
        progress_args=progress_args,
    )


async def save_file_from_bytes_gen(
    self: pyrogram.client.Client,
    chunk_gen: StdinChunkGen,
):
    assert self.me

    file_size_limit_mib = 4000 if self.me.is_premium else 2000

    workers = int(getattr(chunk_gen, "workers", DEFAULT_WORKERS))
    buffer_parts = int(getattr(chunk_gen, "buffer_parts", DEFAULT_BUFFER_PARTS))
    workers = max(1, workers)
    buffer_parts = max(workers, buffer_parts)

    dc_id = await self.storage.dc_id()
    auth_key = await self.storage.auth_key()
    test_mode = await self.storage.test_mode()

    sessions: list[pyrogram.session.session.Session] = []

    async def make_session() -> pyrogram.session.session.Session:
        s = pyrogram.session.session.Session(
            self,
            dc_id,
            auth_key,
            test_mode,
            is_media=True,
        )
        await s.start()
        return s

    try:
        for _ in range(workers):
            sessions.append(await make_session())

        file_id = self.rnd_id()
        file_total_parts = math.ceil(chunk_gen.total / MAX_CHUNK_SIZE)

        queue: asyncio.Queue[typing.Any] = asyncio.Queue(maxsize=buffer_parts)

        cancel_event = chunk_gen.cancel_event

        uploaded_size = 0
        uploaded_lock = asyncio.Lock()

        async def producer():
            nonlocal uploaded_size
            try:
                idx = 0
                async for chunk in chunk_gen:
                    if cancel_event.is_set():
                        raise asyncio.CancelledError()

                    assert len(chunk) <= MAX_CHUNK_SIZE
                    if idx + 1 != file_total_parts:
                        assert len(chunk) % 1024 == 0

                    async with uploaded_lock:
                        uploaded_size_local = uploaded_size + len(chunk)
                        if uploaded_size_local > file_size_limit_mib * 1024 * 1024:
                            raise ValueError(
                                f"Files larger than {file_size_limit_mib} MiB "
                                f"cannot be uploaded for this account."
                            )

                    await queue.put((idx, chunk))

                    async with uploaded_lock:
                        uploaded_size = uploaded_size + len(chunk)

                    idx += 1

                for _ in range(workers):
                    await queue.put(None)
            except Exception:
                for _ in range(workers):
                    await queue.put(None)
                raise

        async def worker(session: pyrogram.session.session.Session):
            max_retries = 8
            while True:
                item = await queue.get()
                if item is None:
                    queue.task_done()
                    break
                part_idx, data = item
                retries = 0
                while True:
                    try:
                        rpc = pyrogram.raw.functions.upload.SaveBigFilePart(
                            file_id=file_id,
                            file_part=part_idx,
                            file_total_parts=file_total_parts,
                            bytes=data,
                        )
                        await session.invoke(rpc)
                        break
                    except FloodWait as e:
                        await asyncio.sleep(int(getattr(e, "value", 1)))
                        retries += 1
                        if retries > max_retries:
                            queue.task_done()
                            raise
                    except Exception:
                        queue.task_done()
                        raise
                queue.task_done()

        prod_task = asyncio.create_task(producer())
        work_tasks = [asyncio.create_task(worker(s)) for s in sessions]

        await prod_task
        await queue.join()
        await asyncio.gather(*work_tasks)

        return pyrogram.raw.types.input_file_big.InputFileBig(
            id=file_id,
            parts=file_total_parts,
            name=chunk_gen.name,
        )

    finally:
        for s in sessions:
            try:
                await s.stop()
            except Exception:
                pass


def guess_stdin_name(fallback: str = "stdin.bin") -> str:
    try:
        path = os.readlink("/proc/self/fd/0")
        base = os.path.basename(path)
        return base or fallback
    except Exception:
        return fallback


def install_signal_handlers(cancel_event: asyncio.Event):
    loop = asyncio.get_running_loop()

    def _set_cancel():
        cancel_event.set()

    try:
        loop.add_signal_handler(signal.SIGINT, _set_cancel)
    except NotImplementedError:
        pass
    try:
        loop.add_signal_handler(signal.SIGTERM, _set_cancel)
    except (AttributeError, NotImplementedError):
        pass


def build_proxy_from_url(url: str, rdns: bool = False) -> dict:
    u = urlparse(url)
    if not u.scheme or not u.hostname or not u.port:
        raise ValueError("Invalid --proxy URL. Expected: http[s]://[user:pass@]host:port")
    proxy = {
        "scheme": u.scheme.lower(),
        "hostname": u.hostname,
        "port": u.port,
    }
    if u.username:
        proxy["username"] = unquote(u.username)
    if u.password:
        proxy["password"] = unquote(u.password)
    if rdns:
        proxy["rdns"] = True
    return proxy


def proxy_to_logline(p: dict | None) -> str:
    if not p:
        return "direct (no proxy)"
    scheme = p.get("scheme", "?")
    host = p.get("hostname", "?")
    port = p.get("port", "?")
    auth = " +auth" if p.get("username") or p.get("password") else ""
    rdns = " rdns" if p.get("rdns") else ""
    return f"{scheme}://{host}:{port}{auth}{rdns}"


async def run():
    parser = argparse.ArgumentParser(
        description="Send file from stdin to Telegram with pretty progress and safe cancel."
    )
    parser.add_argument("--token", required=True, help="Bot token")
    parser.add_argument("--chat_id", required=True, help="Target chat id or @username")
    parser.add_argument("--name", default=None, help="Override file name")
    parser.add_argument("--no-transient", action="store_true",
                        help="Do not clear progress after completion")
    parser.add_argument("--size", type=int, default=None,
                        help="Exact size of input stream in bytes (for pipes)")
    parser.add_argument("--caption", default=None, help="Caption text to send with the document")
    parser.add_argument("--proxy", default=None,
                        help="Proxy URL, e.g. http://user:pass@host:port (leave unset to disable)")
    parser.add_argument("--proxy-rdns", action="store_true",
                        help="Resolve DNS via proxy (sets rdns=True)")
    parser.add_argument("--workers", type=int, default=DEFAULT_WORKERS,
                        help="Number of parallel upload workers")
    parser.add_argument("--buffer-parts", type=int, default=DEFAULT_BUFFER_PARTS,
                        help="Max queued parts to buffer in memory")
    parser.add_argument("--log-progress-interval", type=float, default=1.0, help="Seconds between progress lines when not a TTY")

    args = parser.parse_args()

    if args.size is not None:
        total_size = int(args.size)
    else:
        try:
            total_size = os.fstat(0).st_size
        except Exception:
            total_size = -1

    if total_size <= 0:
        raise RuntimeError(
            "stdin is empty or its size is unknown. Use a redirection like:  script.py < file.bin"
        )

    file_name = args.name or guess_stdin_name()
    cancel_event = asyncio.Event()
    install_signal_handlers(cancel_event)

    proxy = None
    if args.proxy:
        proxy = build_proxy_from_url(args.proxy, rdns=args.proxy_rdns)

    client_kwargs = {
        "name": "bot",
        "api_id": "1",
        "api_hash": "b6b154c3707471f5339bd661645ed3d6",
        "bot_token": args.token,
    }
    if proxy:
        client_kwargs["proxy"] = proxy

    app = pyrogram.Client(**client_kwargs)
    app.save_file = types.MethodType(save_file_custom_wrapper, app)

    await app.start()

    chat = args.chat_id
    if chat.startswith("-") or chat.lstrip("-").isdigit():
        try:
            chat = int(chat)
        except ValueError:
            pass

    if sys.stdout.isatty():
        progress_cm = make_progress(transient=not args.no_transient)
    else:
        progress_cm = LineProgress(console, interval=args.log_progress_interval)

    console.rule("[bold blue]Telegram Upload")
    console.log(f":satellite: Network: [bold]{proxy_to_logline(proxy)}[/]")
    console.log(f":rocket: Sending [bold]{file_name}[/] → [bold]{chat}[/]  ({total_size} bytes)")
    console.log(f":gear: Workers: [bold]{args.workers}[/], buffer parts: [bold]{args.buffer_parts}[/]")
    typing_task = await gen_typing(app, chat, True)

    try:
        if total_size >= MIN_BIG_FILE_SIZE:
            if progress_cm is not None:
                with progress_cm as progress:
                    task_id = progress.add_task(f"{file_name}", total=total_size)
                    gen = StdinChunkGen(
                        total_size=total_size,
                        name=file_name,
                        progress=progress,
                        task_id=task_id,
                        cancel_event=cancel_event,
                    )
                    gen.workers = int(args.workers)
                    gen.buffer_parts = int(args.buffer_parts)

                    await asyncio.shield(
                        app.send_document(chat_id=chat, document=gen, caption=args.caption)
                    )
            else:
                gen = StdinChunkGen(
                    total_size=total_size,
                    name=file_name,
                    cancel_event=cancel_event,
                )
                gen.workers = int(args.workers)
                gen.buffer_parts = int(args.buffer_parts)

                await asyncio.shield(
                    app.send_document(chat_id=chat, document=gen, caption=args.caption)
                )
        else:
            data = sys.stdin.buffer.read()
            if not data:
                raise RuntimeError("stdin is empty — nothing to send.")
            bio = io.BytesIO(data)
            bio.name = file_name
            await asyncio.shield(
                app.send_document(chat_id=chat, document=bio, caption=args.caption)
            )
        console.log(":white_check_mark: Done — file sent.")
    except (asyncio.CancelledError, KeyboardInterrupt):
        console.log(":warning: Cancel: interrupting download and closing connections...")
        raise
    finally:
        try:
            await app.stop()
            await gen_typing(app, chat, typing_task)
        except Exception:
            pass


def main():
    try:
        asyncio.run(run())
    except KeyboardInterrupt:
        console.print("[yellow]Cancelled by user (KeyboardInterrupt).[/]")
        sys.exit(130)
    except asyncio.CancelledError:
        console.print("[yellow]Cancelled by user (Cancelled).[/]")
        sys.exit(130)
    except Exception as e:
        console.print(f"[red]Error:[/] {e}")
        sys.exit(1)


if __name__ == "__main__":
    main()