playground/torrentify.py

#!/usr/bin/env -S uv run --script
# /// script
# dependencies = ["click", "httpx", "torf"]
# ///

import base64
import json
import logging
import sys
import time
from pathlib import Path

import click
import httpx
import torf

ANNOUNCE_URLS = ("http://tracker.empornium.sx:2710/tegqucis10uanp672qh6nhn393xlkncs/announce",)
AUTHOR = "zzzp"




def format_size(size: int) -> str:
    if size < 1000:
        return f"{size:,} bytes"
    scaled_size = float(size)
    for unit in ("KB", "MB", "GB", "TB", "PB", "EB"):
        scaled_size /= 1000
        if scaled_size < 1000 or unit == "EB":
            return f"{scaled_size:.2f} {unit} ({size:,} bytes)"
    raise ValueError(f"Cannot format size: {size}")


def create_torrent(
    paths: tuple[Path, ...],
    output: Path | None = None,
    announce_urls: tuple[str, ...] = ANNOUNCE_URLS,
    author: str = AUTHOR,
    private: bool = True,
    name: str | None = None,
    overwrite: bool = False,
) -> Path:
    sources = list(dict.fromkeys(path.expanduser().absolute() for path in paths))
    if not sources:
        raise ValueError("Provide at least one file or directory")
    for source in sources:
        if not source.is_file() and not source.is_dir():
            raise ValueError(f"Not a file or directory: {source}")
    if len({source.parent for source in sources}) > 1:
        raise ValueError("All inputs must be in the same directory")
    torrent = torf.Torrent(
        trackers=announce_urls,
        created_by=author,
        private=private,
        creation_date=int(time.time()),
    )
    if len(sources) == 1:
        torrent.path = sources[0]
    else:
        torrent.filepaths = sources
    if name is not None:
        torrent.name = name
    if torrent.size <= 0:
        raise ValueError("The selected paths contain no nonempty files")
    torrent_path = output or Path.cwd() / f"{torrent.name}.torrent"
    if torrent_path.exists() and not overwrite:
        raise FileExistsError(f"Output already exists: {torrent_path}")
    if torrent_path.absolute() in {Path(path).absolute() for path in torrent.filepaths}:
        raise ValueError("Output must not overwrite a source file")
    with click.progressbar(length=torrent.pieces, label="Hashing", show_percent=True, file=sys.stderr) as progress:

        def report_progress(
            torrent: torf.Torrent,
            filepath: str,
            pieces_done: int,
            pieces_total: int,
        ) -> None:
            progress.update(n_steps=pieces_done - progress.pos)

        if not torrent.generate(callback=report_progress, interval=0.2):
            raise ValueError("Torrent hashing did not complete")
    torrent.write(filepath=torrent_path, overwrite=overwrite)
    return torrent_path


def modify_torrent(
    torrent_path: Path,
    output: Path | None = None,
    announce_urls: tuple[str, ...] = (),
    author: str | None = None,
    private: bool | None = None,
    name: str | None = None,
) -> Path:
    torrent = torf.Torrent.read(filepath=torrent_path)
    if announce_urls:
        torrent.trackers = announce_urls
    if author is not None:
        torrent.created_by = author
    if private is not None:
        torrent.private = private
    if name is not None:
        torrent.name = name
    destination = output or torrent_path
    torrent.write(filepath=destination, overwrite=destination.absolute() == torrent_path.absolute())
    return destination


def enqueue_torrent(
    transmission_url: str,
    credentials: str | tuple[str, str] | None,
    torrent_path: Path,
    download_path: str | Path,
) -> str:
    if isinstance(credentials, str):
        username, separator, password = credentials.partition(":")
        if not separator:
            raise ValueError("Credentials must be username:password")
        credentials = (username, password)
    auth = httpx.BasicAuth(*credentials) if credentials is not None else None
    rpc_url = transmission_url.rstrip("/")
    if not httpx.URL(url=rpc_url).path.endswith("/rpc"):
        rpc_url = f"{rpc_url}/transmission/rpc"
    payload = {
        "method": "torrent-add",
        "arguments": {
            "download-dir": str(download_path),
            "metainfo": base64.b64encode(torrent_path.read_bytes()).decode("ascii"),
            "paused": False,
        },
    }
    with httpx.Client(auth=auth, timeout=20) as client:
        for _ in range(3):
            response = client.post(url=rpc_url, json=payload)
            if response.status_code == 409:
                client.headers["X-Transmission-Session-Id"] = response.headers["X-Transmission-Session-Id"]
                continue
            response.raise_for_status()
            result = response.json()
            if result["result"] not in ("success", "duplicate-torrent"):
                raise ValueError(f"Transmission rejected torrent: {result['result']}")
            added = result["arguments"].get("torrent-added") or result["arguments"].get("torrent-duplicate")
            if not added:
                raise ValueError("Transmission returned no torrent details")
            return added["hashString"]
    raise ValueError("Transmission session negotiation failed")


@click.group()
def main() -> None:
    """Create, modify, inspect, and enqueue torrent files."""
    logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")


@main.command(name="create")
@click.argument("paths", nargs=-1, required=True, type=click.Path(exists=True, path_type=Path))
@click.option("--output", type=click.Path(dir_okay=False, path_type=Path), help="Destination; defaults to NAME.torrent in the current directory.")
@click.option("--announce-url", "announce_urls", multiple=True, help="Tracker URL; repeat for multiple trackers. Defaults to the hardcoded URLs.")
@click.option("--author", default=AUTHOR, show_default=True)
@click.option("--private/--no-private", default=True, show_default=True)
@click.option("--name", help="Torrent name; defaults to the file or common directory name.")
@click.option("--overwrite", is_flag=True, help="Replace an existing output torrent.")
def create(
    paths: tuple[Path, ...],
    output: Path | None,
    announce_urls: tuple[str, ...],
    author: str,
    private: bool,
    name: str | None,
    overwrite: bool,
) -> None:
    """Create one torrent from FILES or DIRECTORIES, retaining their relative paths."""
    try:
        destination = create_torrent(
            paths=paths,
            output=output,
            announce_urls=announce_urls or ANNOUNCE_URLS,
            author=author,
            private=private,
            name=name,
            overwrite=overwrite,
        )
    except (OSError, ValueError, torf.TorfError) as error:
        raise click.ClickException(message=str(error)) from error
    logging.info(f"Saved torrent at {destination}")


@main.command(name="modify")
@click.argument("torrent_path", type=click.Path(exists=True, dir_okay=False, path_type=Path))
@click.option("--output", type=click.Path(dir_okay=False, path_type=Path), help="Destination; defaults to editing the input torrent in place.")
@click.option("--announce-url", "announce_urls", multiple=True, help="Replace trackers; repeat for multiple URLs.")
@click.option("--author", help="Replace the creator.")
@click.option("--private/--no-private", default=None, help="Set or clear the private flag.")
@click.option("--name", help="Replace the torrent name.")
def modify(
    torrent_path: Path,
    output: Path | None,
    announce_urls: tuple[str, ...],
    author: str | None,
    private: bool | None,
    name: str | None,
) -> None:
    """Change metadata without rehashing the content; omitted fields stay unchanged."""
    try:
        destination = modify_torrent(
            torrent_path=torrent_path,
            output=output,
            announce_urls=announce_urls,
            author=author,
            private=private,
            name=name,
        )
    except (OSError, ValueError, torf.TorfError) as error:
        raise click.ClickException(message=str(error)) from error
    logging.info(f"Saved torrent at {destination}")


@main.command(name="enqueue")
@click.argument("torrent_path", type=click.Path(exists=True, dir_okay=False, path_type=Path))
@click.option("--transmission-url", "--host", required=True, envvar="TRANSMISSION_HOST", help="Transmission host or RPC URL; also reads TRANSMISSION_HOST.")
@click.option("--credentials", "--auth", envvar="TRANSMISSION_AUTH", help="Username:password; also reads TRANSMISSION_AUTH.")
@click.option("--download-path", required=True, type=click.Path(path_type=Path), help="Download directory on the remote Transmission host.")
def enqueue(
    torrent_path: Path,
    transmission_url: str,
    credentials: str | None,
    download_path: Path,
) -> None:
    """Add a torrent to Transmission at the given remote download directory."""
    try:
        torrent_hash = enqueue_torrent(
            transmission_url=transmission_url,
            credentials=credentials,
            torrent_path=torrent_path,
            download_path=download_path,
        )
    except (OSError, ValueError, httpx.HTTPError) as error:
        raise click.ClickException(message=str(error)) from error
    logging.info(f"Enqueued torrent {torrent_hash} at {download_path}")




@main.command(name="inspect")
@click.argument("torrent_path", type=click.Path(exists=True, dir_okay=False, path_type=Path))
@click.option("--json", "json_output", is_flag=True, help="Output JSON instead of plain text.")
def inspect(torrent_path: Path, json_output: bool) -> None:
    """Show torrent properties and content paths with readable sizes."""
    try:
        torrent = torf.Torrent.read(filepath=torrent_path)
    except (OSError, ValueError, torf.TorfError) as error:
        raise click.ClickException(message=str(error)) from error
    properties = {
        "name": torrent.name,
        "author": torrent.created_by,
        "private": torrent.private,
        "infohash": torrent.infohash,
        "size": torrent.size,
        "piece_size": torrent.piece_size,
        "pieces": torrent.pieces,
        "creation_date": torrent.creation_date.isoformat() if torrent.creation_date else None,
        "comment": torrent.comment,
        "source": torrent.source,
        "announce_urls": [list(tier) for tier in torrent.trackers],
        "webseeds": list(torrent.webseeds),
        "httpseeds": list(torrent.httpseeds),
    }
    files = [{"path": str(file), "size": file.size} for file in torrent.files]
    if json_output:
        click.echo(message=json.dumps({"properties": properties, "files": files}, indent=2, ensure_ascii=False))
        return
    for key, value in properties.items():
        label = key.replace("_", " ").capitalize()
        if isinstance(value, list):
            click.echo(message=f"{label}:" if value else f"{label}: (none)")
            for index, item in enumerate(value, start=1):
                if isinstance(item, list):
                    click.echo(message=f"  Tier {index}: {', '.join(item)}")
                else:
                    click.echo(message=f"  {item}")
        elif value is None:
            click.echo(message=f"{label}: (unset)")
        elif isinstance(value, bool):
            click.echo(message=f"{label}: {'yes' if value else 'no'}")
        elif key in ("size", "piece_size"):
            click.echo(message=f"{label}: {format_size(size=value)}")
        else:
            click.echo(message=f"{label}: {value}")
    click.echo(message=f"\nFiles ({len(files)}):")
    for file in files:
        click.echo(message=f"  {file['path']}: {format_size(size=file['size'])}")


if __name__ == "__main__":
    main()