libtorrent/bindings/python/tests/tdummy.py

import hashlib
import random
from typing import List
from typing import Optional
from typing import Type
from typing import TYPE_CHECKING
from typing import TypeVar

from typing_extensions import TypedDict

import libtorrent as lt

if TYPE_CHECKING:
    from libtorrent import TorrentFileDict
    from libtorrent import TorrentFileFileDict
    from libtorrent import TorrentFileInfoDict
else:
    TorrentFileInfoDict, TorrentFileFileDict, TorrentFileDict = dict, dict, dict


class _FParams(TypedDict, total=False):
    length: int
    data: Optional[bytes]
    path: Optional[bytes]
    path_split: Optional[List[bytes]]
    attr: Optional[bytes]


class File:
    def __init__(
        self,
        *,
        length: int,
        start: int,
        stop: int,
        data: Optional[bytes] = None,
        path: Optional[bytes] = None,
        path_split: Optional[List[bytes]] = None,
        attr: Optional[bytes] = None,
    ):
        assert stop - start == length, (start, stop, length)
        if data is not None:
            assert len(data) == length

        if path is None:
            assert path_split is not None
            path = b"/".join(path_split)
        if path_split is None:
            assert path is not None
            path_split = path.split(b"/")

        self._data = data
        self.path = path
        self.path_split = path_split
        self.length = length
        self.attr = attr or b""
        self.start = start
        self.stop = stop

    @property
    def data(self) -> bytes:
        if b"p" in self.attr:
            return b"\x00" * self.length
        if self._data is None:
            self._data = bytes(random.getrandbits(8) for _ in range(self.length))
        return self._data


_T = TypeVar("_T", bound="Torrent")


class Torrent:
    @classmethod
    def single_file(
        cls: Type[_T],
        *,
        length: int,
        piece_length: int = 16384,
        name: Optional[bytes] = None,
        attr: Optional[bytes] = None,
        data: Optional[bytes] = None,
    ) -> _T:
        return cls(
            piece_length=piece_length,
            files=[
                _FParams(length=length, path=name, attr=attr, data=data),
            ],
        )

    def __init__(self, *, files: List[_FParams], piece_length: int = 16384):
        assert piece_length is not None

        self.piece_length = piece_length
        self.files: List[File] = []

        offset = 0
        for file_ in files:
            start = offset
            stop = offset + file_["length"]
            offset = stop
            self.files.append(File(start=start, stop=stop, **file_))
        self.length = sum(f.length for f in self.files)

        self._data: Optional[bytes] = None
        self._pieces: Optional[List[bytes]] = None
        self._info: Optional[TorrentFileInfoDict] = None
        self._dict: Optional[TorrentFileDict] = None
        self._info_hash_bytes: Optional[bytes] = None

    @property
    def data(self) -> bytes:
        if self._data is None:
            self._data = b"".join(f.data for f in self.files)
        return self._data

    @property
    def pieces(self) -> List[bytes]:
        if self._pieces is None:
            self._pieces = [
                self.data[i : i + self.piece_length]
                for i in range(0, self.length, self.piece_length)
            ]
        return self._pieces

    @property
    def info(self) -> TorrentFileInfoDict:
        if self._info is None:
            self._info = TorrentFileInfoDict(
                {
                    b"piece length": self.piece_length,
                    b"length": self.length,
                    b"pieces": b"".join(hashlib.sha1(p).digest() for p in self.pieces),
                }
            )

            if len(self.files) == 1:
                self._info[b"name"] = self.files[0].path
            else:
                assert len({f.path_split[0] for f in self.files}) == 1
                assert all(len(f.path_split) > 1 for f in self.files)
                self._info[b"name"] = self.files[0].path_split[0]
                files_list: list[TorrentFileFileDict] = []
                for file_ in self.files:
                    fdict = TorrentFileFileDict(
                        {
                            b"length": file_.length,
                            b"path": file_.path_split[1:],
                        }
                    )
                    if file_.attr:
                        fdict[b"attr"] = file_.attr
                    files_list.append(fdict)
                self._info[b"files"] = files_list
        return self._info

    @property
    def dict(self) -> TorrentFileDict:
        if self._dict is None:
            self._dict = TorrentFileDict(
                {
                    b"info": self.info,
                }
            )
        return self._dict

    @property
    def info_hash_bytes(self) -> bytes:
        if self._info_hash_bytes is None:
            self._info_hash_bytes = hashlib.sha1(lt.bencode(self.info)).digest()
        return self._info_hash_bytes

    @property
    def info_hash(self) -> str:
        return self.info_hash_bytes.hex()

    @property
    def sha1_hash(self) -> lt.sha1_hash:
        return lt.sha1_hash(self.info_hash_bytes)

    def torrent_info(self) -> lt.torrent_info:
        return lt.torrent_info(self.dict)

    def atp(self) -> lt.add_torrent_params:
        atp = lt.add_torrent_params()
        self.configure_atp(atp)
        return atp

    def configure_atp(self, atp: lt.add_torrent_params) -> None:
        # this is necessary so that
        # atp == read_resume_data(write_resume_data(atp))
        atp.info_hash = self.sha1_hash
        atp.ti = self.torrent_info()
        atp.flags &= ~(lt.torrent_flags.paused | lt.torrent_flags.auto_managed)


def get_default() -> Torrent:
    return Torrent.single_file(
        piece_length=16384, name=b"test.txt", length=16384 * 9 + 1000
    )