libtorrent/bindings/python/client.py

#!/usr/bin/env python3
from __future__ import annotations

# Copyright Daniel Wallin 2006. Use, modification and distribution is
# subject to the Boost Software License, Version 1.0. (See accompanying
# file LICENSE_1_0.txt or copy at http://www.boost.org/LICENSE_1_0.txt)

import atexit
import os.path
import sys
import time
from abc import ABC, abstractmethod
from collections.abc import Sequence, Mapping
from optparse import Values
from typing import Any, cast

import libtorrent as lt


class ConsoleABC(ABC):
    @abstractmethod
    def clear(self) -> None:
        pass

    @abstractmethod
    def write(self, line: str) -> None:
        pass

    @abstractmethod
    def sleep_and_input(self, seconds: float) -> bytes | str | None:
        pass


class WindowsConsole(ConsoleABC):
    def __init__(self) -> None:
        import Console

        self.console = Console.getconsole()

    def clear(self) -> None:
        self.console.page()

    def write(self, line: str) -> None:
        self.console.write(line)

    def sleep_and_input(self, seconds: float) -> bytes | None:
        import msvcrt

        time.sleep(seconds)
        if msvcrt.kbhit():
            return msvcrt.getch()
        return None


class UnixConsole(ConsoleABC):
    def __init__(self) -> None:
        import termios

        self.fd = sys.stdin
        self.old = termios.tcgetattr(self.fd.fileno())
        new = termios.tcgetattr(self.fd.fileno())
        new[3] = new[3] & ~termios.ICANON
        new[6][termios.VTIME] = 0
        new[6][termios.VMIN] = 1
        termios.tcsetattr(self.fd.fileno(), termios.TCSADRAIN, new)

        atexit.register(self._onexit)

    def _onexit(self) -> None:
        import termios

        termios.tcsetattr(self.fd.fileno(), termios.TCSADRAIN, self.old)

    def clear(self) -> None:
        sys.stdout.write('\033[2J\033[0;0H')
        sys.stdout.flush()

    def write(self, line: str) -> None:
        sys.stdout.write(line)
        sys.stdout.flush()

    def sleep_and_input(self, seconds: float) -> str | None:
        import select

        read, __, __ = select.select(
            [self.fd.fileno()], [], [], seconds)
        if len(read) > 0:
            return self.fd.read(1)
        return None


def write_line(console: ConsoleABC, line: str) -> None:
    console.write(line)


def add_suffix(val: float) -> str:
    prefix = ['B', 'kB', 'MB', 'GB', 'TB']
    for i in range(len(prefix)):
        if abs(val) < 1000:
            if i == 0:
                return '%5.3g%s' % (val, prefix[i])
            else:
                return '%4.3g%s' % (val, prefix[i])
        val /= 1000

    return '%6.3gPB' % val


def progress_bar(progress: float, width: int) -> str:
    assert(progress <= 1)
    progress_chars = int(progress * width + 0.5)
    return progress_chars * '#' + (width - progress_chars) * '-'


def print_peer_info(console: ConsoleABC, peers: list[lt.peer_info]) -> None:

    out = (' down    (total )   up     (total )'
           '  q  r flags  block progress  client\n')

    for p in peers:

        out += '%s/s ' % add_suffix(p.down_speed)
        out += '(%s) ' % add_suffix(p.total_download)
        out += '%s/s ' % add_suffix(p.up_speed)
        out += '(%s) ' % add_suffix(p.total_upload)
        out += '%2d ' % p.download_queue_length
        out += '%2d ' % p.upload_queue_length

        out += 'I' if p.flags & lt.peer_info.interesting else '.'
        out += 'C' if p.flags & lt.peer_info.choked else '.'
        out += 'i' if p.flags & lt.peer_info.remote_interested else '.'
        out += 'c' if p.flags & lt.peer_info.remote_choked else '.'
        out += 'e' if p.flags & lt.peer_info.supports_extensions else '.'
        out += 'l' if p.flags & lt.peer_info.local_connection else 'r'
        out += ' '

        if p.downloading_piece_index >= 0:
            assert(p.downloading_progress <= p.downloading_total)
            out += progress_bar(float(p.downloading_progress) /
                                p.downloading_total, 15)
        else:
            out += progress_bar(0, 15)
        out += ' '

        if p.flags & lt.peer_info.handshake:
            id = 'waiting for handshake'
        elif p.flags & lt.peer_info.connecting:
            id = 'connecting to peer'
        else:
            try:
                id = p.client.decode()
            except ValueError:
                id = p.client.hex()

        out += '%s\n' % id[:10]

    write_line(console, out)


def print_download_queue(console: ConsoleABC, download_queue: Sequence[Mapping[str, Any]]) -> None:
    out = ""

    for e in download_queue:
        out += '%4d: [' % e['piece_index']
        for b in e['blocks']:
            s = b['state']
            if s == 3:
                out += '#'
            elif s == 2:
                out += '='
            elif s == 1:
                out += '-'
            else:
                out += ' '
        out += ']\n'

    write_line(console, out)


def add_torrent(ses: lt.session, filename: str, options: Values) -> None:
    atp = lt.add_torrent_params()
    save_path = cast(str, options.save_path)
    if filename.startswith('magnet:'):
        atp = lt.parse_magnet_uri(filename)
    else:
        ti = lt.torrent_info(filename)
        resume_file = os.path.join(save_path, ti.name() + '.fastresume')
        try:
            atp = lt.read_resume_data(open(resume_file, 'rb').read())
        except Exception as e:
            print('failed to open resume file "%s": %s' % (resume_file, e))
        atp.ti = ti

    atp.save_path = save_path
    atp.storage_mode = lt.storage_mode_t.storage_mode_sparse
    atp.flags |= lt.torrent_flags.duplicate_is_error \
        | lt.torrent_flags.auto_managed \
        | lt.torrent_flags.duplicate_is_error
    ses.async_add_torrent(atp)


def main() -> None:
    from optparse import OptionParser

    parser = OptionParser()

    parser.add_option('-p', '--port', type='int', help='set listening port')

    parser.add_option(
        '-i', '--listen-interface', type='string',
        help='set interface for incoming connections', )

    parser.add_option(
        '-o', '--outgoing-interface', type='string',
        help='set interface for outgoing connections')

    parser.add_option(
        '-d', '--max-download-rate', type='float',
        help='the maximum download rate given in kB/s. 0 means infinite.')

    parser.add_option(
        '-u', '--max-upload-rate', type='float',
        help='the maximum upload rate given in kB/s. 0 means infinite.')

    parser.add_option(
        '-s', '--save-path', type='string',
        help='the path where the downloaded file/folder should be placed.')

    parser.add_option(
        '-r', '--proxy-host', type='string',
        help='sets HTTP proxy host and port (separated by \':\')')

    parser.set_defaults(
        port=6881,
        listen_interface='0.0.0.0',
        outgoing_interface='',
        max_download_rate=0,
        max_upload_rate=0,
        save_path='.',
        proxy_host=''
    )

    (options, args) = parser.parse_args()

    if options.port < 0 or options.port > 65525:
        options.port = 6881

    options.max_upload_rate *= 1000
    options.max_download_rate *= 1000

    if options.max_upload_rate <= 0:
        options.max_upload_rate = -1
    if options.max_download_rate <= 0:
        options.max_download_rate = -1

    params = lt.session_params()
    params.settings.update({
        "user_agent": 'python_client/' + lt.__version__,
        "listen_interfaces": '%s:%d' % (options.listen_interface, options.port),
        "download_rate_limit": int(options.max_download_rate),
        "upload_rate_limit": int(options.max_upload_rate),
        "alert_mask": lt.alert.category_t.all_categories,
        "outgoing_interfaces": options.outgoing_interface,
    })

    if options.proxy_host != '':
        params.settings["proxy_hostname"] = options.proxy_host.split(':')[0]
        params.settings["proxy_type"] = lt.proxy_type_t.http
        params.settings["proxy_port"] = options.proxy_host.split(':')[1]

    ses = lt.session(params)

    # map torrent_handle to torrent_status
    torrents: dict[lt.torrent_handle, lt.torrent_status] = {}
    alerts_log = []

    for f in args:
        add_torrent(ses, f, options)

    console: ConsoleABC
    if os.name == 'nt':
        console = WindowsConsole()
    else:
        console = UnixConsole()

    alive = True
    while alive:
        console.clear()

        out = ''

        for h, t in torrents.items():
            out += 'name: %-40s\n' % t.name[:40]

            if t.state != lt.torrent_status.seeding:
                state_str = ['queued', 'checking', 'downloading metadata',
                             'downloading', 'finished', 'seeding',
                             '', 'checking fastresume']
                out += state_str[t.state] + ' '

                out += '%5.4f%% ' % (t.progress * 100)
                out += progress_bar(t.progress, 49)
                out += '\n'

                out += 'total downloaded: %d Bytes\n' % t.total_done
                out += 'peers: %d seeds: %d distributed copies: %d\n' % \
                    (t.num_peers, t.num_seeds, t.distributed_copies)
                out += '\n'

            out += 'download: %s/s (%s) ' \
                % (add_suffix(t.download_rate), add_suffix(t.total_download))
            out += 'upload: %s/s (%s) ' \
                % (add_suffix(t.upload_rate), add_suffix(t.total_upload))

            if t.state != lt.torrent_status.seeding:
                out += 'info-hash: %s\n' % t.info_hashes
                out += 'next announce: %s\n' % t.next_announce
                out += 'tracker: %s\n' % t.current_tracker

            write_line(console, out)

            print_peer_info(console, t.handle.get_peer_info())
            print_download_queue(console, t.handle.get_download_queue())

            if t.state != lt.torrent_status.seeding:
                try:
                    out = '\n'
                    fp = h.file_progress()
                    ti = cast(lt.torrent_info, t.torrent_file)
                    for idx, p in enumerate(fp):
                        out += progress_bar(p / float(ti.files().file_size(idx)), 20)
                        out += ' ' + ti.files().file_path(idx) + '\n'
                    write_line(console, out)
                except Exception:
                    pass

        write_line(console, 76 * '-' + '\n')
        write_line(console, '(q)uit), (p)ause), (u)npause), (r)eannounce\n')
        write_line(console, 76 * '-' + '\n')

        alerts = ses.pop_alerts()
        for a in alerts:
            alerts_log.append(a.message())

            # add new torrents to our list of torrent_status
            if isinstance(a, lt.add_torrent_alert):
                h = a.handle
                h.set_max_connections(60)
                h.set_max_uploads(-1)
                torrents[h] = h.status()

            # update our torrent_status array for torrents that have
            # changed some of their state
            if isinstance(a, lt.state_update_alert):
                for s in a.status:
                    torrents[s.handle] = s

        if len(alerts_log) > 20:
            alerts_log = alerts_log[-20:]

        for a_log in alerts_log:
            write_line(console, a_log + '\n')

        c = console.sleep_and_input(0.5)

        ses.post_torrent_updates()
        if not c:
            continue

        if c == 'r':
            for h in torrents:
                h.force_reannounce()
        elif c == 'q':
            alive = False
        elif c == 'p':
            for h in torrents:
                h.pause()
        elif c == 'u':
            for h in torrents:
                h.resume()

    ses.pause()
    for h, t in torrents.items():
        if not h.is_valid() or not t.has_metadata:
            continue
        h.save_resume_data()

    while len(torrents) > 0:
        alerts = ses.pop_alerts()
        for a in alerts:
            if isinstance(a, lt.save_resume_data_alert):
                print(a)
                data = lt.write_resume_data_buf(a.params)
                h = a.handle
                if h in torrents:
                    open(os.path.join(options.save_path, torrents[h].name + '.fastresume'), 'wb').write(data)
                    del torrents[h]

            if isinstance(a, lt.save_resume_data_failed_alert):
                h = a.handle
                if h in torrents:
                    print('failed to save resume data for ', torrents[h].name)
                    del torrents[h]
        time.sleep(0.5)


main()