Converting rtorrent_fast_resume.pl to Python 3

rtorrent_fast_resume.py

#!/usr/bin/env python3

"""
og perl source:
https://github.com/rakshasa/rtorrent/blob/79765768ff3d93a511ceb90df24a8dbb0a62e4f5/doc/rtorrent_fast_resume.pl
"""

"""
Usage:
    rtorrent_fast_resume.py [base-directory] < plain.torrent > with_fast_resume.torrent
    -OR-
    rtorrent_fast_resume.py [base-directory] plain.torrent [with_fast_resume.torrent]
"""

"""
TODO: $/ and $| equivalences (unbuffered input/output basically)
"""

import collections
import math
import os
import sys
import time

DEBUG = False

def eprint(m):
    print(m, file=sys.stderr)

def die(m):
    eprint(m)
    sys.exit(1)

try:
    import bencode
except ImportError:
    die(f"Requires bencode.py: pip install --user bencode.py")
    
# use .buffer for byte input/output
DEFAULT_IN  = sys.stdin.buffer
DEFAULT_OUT = sys.stdout.buffer
DEFAULT_STREAMS = (DEFAULT_IN, DEFAULT_OUT)

def rel2abs(p):
    return os.path.join(os.getcwd(), p)

def maybe_open(maybe_stream, mode=""):
    if maybe_stream in DEFAULT_STREAMS:
        return maybe_stream
    return open(maybe_stream, mode)

def parse_args(args):
    params = {"basedir": "./", "input": DEFAULT_IN, "output" : DEFAULT_OUT}
    if not args:
        return params
    output_input_or_basedir = args.pop()
    # last param was basedir
    if os.path.isdir(output_input_or_basedir):
        basedir = output_input_or_basedir
        params["basedir"] = basedir if basedir.startswith("/") else f"./{basedir}"
        # args _should_ now be empty
        if len(args):
            die(f"last param must be either input or output torrent")
    # last param was input file
    elif os.path.isfile(output_input_or_basedir):
        _input = output_input_or_basedir
        params["input"] = _input
        # basedir remains
        if len(args):
            basedir = args.pop()
            if not os.path.exists(basedir):
                die("basedir must exist")
            params["basedir"] = basedir
            if len(args):
                die("too many args left")
    # last param was output file
    else:
        output = output_input_or_basedir
        params["output"] = output
        if not len(args):
            die("not enough args")
        _input = args.pop()
        if not os.path.exists(_input):
            die(f"input file {_input} must exist")
        if os.path.isdir(_input):
            die(f"input {_input} must be file not directory")
        params["input"] = _input
        # first param was basedir
        if len(args):
            basedir = args.pop()
            basedir = basedir if basedir.startswith("/") else f"./{basedir}"
            if not os.path.exists(basedir):
                die(f"basedir {basedir} must exist")
            params["basedir"] = basedir
    return params

def get_files_from_torrent(torrent):
    if "files" in torrent["info"] and len(torrent["info"]["files"]) > 1:
        eprint(f"Multi file torrent: {torrent['info']['name']}")
        torrent_files = torrent["info"]["files"]
        t_files = [
            os.path.join(torrent["info"]["name"], f["path"])
            for f in torrent_files
        ]
        t_size = sum(f["length"] for f in torrent_files)
    elif "files" in torrent["info"] and len(torrent["info"]["files"]) == 1:
        eprint(f"Single file torrent: {torrent['info']['files'][0]['path']}")
        t_files = torrent["info"]["files"]
        t_size = torrent["info"]["files"][0]["length"]
    else:
        eprint(f"Single file torrent: {torrent['info']['name']}")
        t_files = [torrent["info"]["name"]]
        t_size = torrent["info"]["length"]
    chunks = num_chunks(t_size, torrent["info"]["piece length"])
    return t_files, t_size, chunks

def consistent_piece_information(chunks, pieces):
    return (chunks*20 == len(pieces))

def num_chunks(t_size, piece_size):
    return int((t_size + piece_size - 1) / piece_size)


def generate_resume_data(torrent, basedir, out):
    t_files, t_size, chunks = get_files_from_torrent(torrent)

    eprint(f"Total size: {t_size / 1024**2} MB; {chunks} chunks; {t_files} files.\n")

    consistent_piece_information(chunks, torrent["info"]["pieces"]) or die("Inconsistent piece information!")

    libtorrent_resume = collections.OrderedDict()
    libtorrent_resume["bitfield"] = chunks
    piece_size = torrent["info"]["piece length"]
    pmod = 0
    for i, f in enumerate(t_files):
        fpath = f["path"][0]
        full_fpath = os.path.join(basedir, fpath)
        os.path.exists(full_fpath) or die(f"{full_fpath} not found")
        mtime = int(os.path.getmtime(full_fpath))

        # compute number of chunks per file
        fsize = torrent["info"]["files"][i]["length"] if "files" in torrent["info"] else 1
        fchunks = 1 if pmod == 1 else 0
        if pmod >= fsize:
            fsize, pmod = 0, (pmod - fsize)
        else:
            pmod, fsize = 0, (fsize - pmod)
        fchunks += int(math.ceil(fsize / piece_size))
        pmod = pmod or (piece_size - (fsize % piece_size))

        fs = collections.OrderedDict()
        fs["priority"] = 0
        fs["mtime"] = mtime
        fs["completed"] = fchunks
        if "files" not in libtorrent_resume:
            libtorrent_resume["files"] = []
        libtorrent_resume["files"].append(fs)
    # Some extra information to re-enforce the fact that this is a finished torrent
    rt_basedir = basedir
    if "files" in torrent["info"]:
        rt_basedir = os.path.join(rt_basedir, torrent["info"]["name"])
    rt = generate_rtorrent_section(rt_basedir, chunks, out)
    return libtorrent_resume, rt

def generate_rtorrent_section(basedir, chunks, out_file):
    rt = collections.OrderedDict()
    rt["state"] = 1 # started
    rt["state_changed"] = int(time.time())
    rt["state_counter"] = 1
    rt["chunks_wanted"] = 0
    rt["chunks_done"] = chunks
    rt["complete"] = 1
    rt["hashing"] = 0 # Not hashing
    rt["directory"] = basedir if os.path.isabs(basedir) else rel2abs(basedir)
    rt["tied_to_file"] = "" if out_file is DEFAULT_OUT else (out_file if os.path.isabs(out_file) else rel2abs(out_file))
    rt["timestamp.finished"] = 0
    rt["timestamp.started"] = int(time.time())
    return rt

if __name__ == "__main__":
    args = sys.argv[1:] if len(sys.argv) > 1 else []
    params = parse_args(args)
    if params["input"] is not DEFAULT_IN:
        input_size = os.path.getsize(params["input"]) / 1024
        eprint(f"Total input torrent size: {input_size} KB")
        eprint(f"Decoding {params['input']}...")
    else:
        eprint(f"Decoding torrent from standard input")
    if DEBUG: eprint(params)

    with maybe_open(params["input"], "rb") as f:
        torrent = bencode.bread(f)

    "info" in torrent or die("No info key.")
    "piece length" in torrent["info"] or die("No piece length key.")
    eprint("done")

    eprint("Adding fast resume information...")

    torrent["libtorrent_resume"], torrent["rtorrent"] = generate_resume_data(torrent, params["basedir"], params["output"])

    eprint("done")

    if torrent["rtorrent"]["tied_to_file"]:
        eprint(f"Encoding {params['output']}...")
    else:
        eprint("Encoding torrent from standard output...")

    with maybe_open(params["output"], "wb") as f:
        bencode.bwrite(torrent, f)

    eprint("done")