Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add TransmissionBt client implementation #26

Merged
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
140 changes: 140 additions & 0 deletions src/clients/transmission.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,140 @@
import base64
import json
from enum import Enum
from http import HTTPStatus

import requests

from requests.auth import HTTPBasicAuth
from requests.structures import CaseInsensitiveDict

from ..filesystem import sane_join
from ..parser import get_bencoded_data, calculate_infohash
from ..errors import TorrentClientError, TorrentClientAuthenticationError, TorrentExistsInClientError
from .torrent_client import TorrentClient


class StatusEnum(Enum):
STOPPED = 0
QUEUED_VERIFY = 1
VERIFYING = 2
QUEUE_DOWNLOAD = 3
DOWNLOADING = 4
QUEUED_SEED = 5
SEEDING = 6


class TransmissionBt(TorrentClient):
X_TRANSMISSION_SESSION_ID = "X-Transmission-Session-Id"

def __init__(self, rpc_url):
super().__init__()
transmission_url_parts = self._extract_credentials_from_url(rpc_url, "transmission/rpc")
self._base_url = transmission_url_parts[0]
self._basic_auth = HTTPBasicAuth(transmission_url_parts[1], transmission_url_parts[2])
self._transmission_session_id = None

def setup(self):
self.__authenticate()
return self

def get_torrent_info(self, infohash):
response = self.__wrap_request(
"torrent-get",
arguments={"fields": ["labels", "downloadDir", "percentDone", "status", "doneDate", "name"], "ids": [infohash]},
)

if response:
try:
parsed_response = json.loads(response)
except json.JSONDecodeError as json_parse_error:
raise TorrentClientError("Client returned malformed json response") from json_parse_error

if not parsed_response.get("arguments", {}).get("torrents", []):
raise TorrentClientError(f"Torrent not found in client ({infohash})")

torrent = parsed_response["arguments"]["torrents"][0]
torrent_completed = (torrent["percentDone"] == 1.0 or torrent["doneDate"] > 0) and torrent["status"] in [
StatusEnum.SEEDING.value,
StatusEnum.QUEUED_SEED.value,
]

return {
"complete": torrent_completed,
"label": torrent["labels"],
"save_path": torrent["downloadDir"],
"content_path": sane_join(torrent["downloadDir"], torrent["name"]),
}
else:
raise TorrentClientError("Client returned unexpected response")

def inject_torrent(self, source_torrent_infohash, new_torrent_filepath, save_path_override=None):
source_torrent_info = self.get_torrent_info(source_torrent_infohash)

if not source_torrent_info["complete"]:
raise TorrentClientError("Cannot inject a torrent that is not complete")

new_torrent_infohash = calculate_infohash(get_bencoded_data(new_torrent_filepath)).lower()
new_torrent_already_exists = self.__does_torrent_exist_in_client(new_torrent_infohash)
if new_torrent_already_exists:
raise TorrentExistsInClientError(f"New torrent already exists in client ({new_torrent_infohash})")

self.__wrap_request(
"torrent-add",
arguments={
"download-dir": save_path_override if save_path_override else source_torrent_info["save_path"],
"metainfo": base64.b64encode(open(new_torrent_filepath, "rb").read()).decode("utf-8"),
"labels": source_torrent_info["label"],
},
)

return new_torrent_infohash

def __authenticate(self):
try:
# This method specifically does not use the __wrap_request method
# because we want to avoid an infinite loop of re-authenticating
response = requests.post(self._base_url, auth=self._basic_auth)
# TransmissionBt returns a 409 status code if the session id is invalid
# (which it is on your first request) and includes a new session id in the response headers.
if response.status_code == HTTPStatus.CONFLICT:
self._transmission_session_id = response.headers.get(self.X_TRANSMISSION_SESSION_ID)
else:
response.raise_for_status()
except requests.RequestException as e:
raise TorrentClientAuthenticationError(f"TransmissionBt login failed: {e}")

if not self._transmission_session_id:
raise TorrentClientAuthenticationError("TransmissionBt login failed: Invalid username or password")

def __wrap_request(self, method, arguments, files=None):
try:
return self.__request(method, arguments, files)
except TorrentClientAuthenticationError:
self.__authenticate()
return self.__request(method, arguments, files)

def __request(self, method, arguments=None, files=None):
try:
response = requests.post(
self._base_url,
auth=self._basic_auth,
headers=CaseInsensitiveDict({self.X_TRANSMISSION_SESSION_ID: self._transmission_session_id}),
json={"method": method, "arguments": arguments},
files=files,
)

response.raise_for_status()

return response.text
except requests.RequestException as e:
if e.response.status_code == HTTPStatus.CONFLICT:
raise TorrentClientAuthenticationError("Failed to authenticate with TransmissionBt")

raise TorrentClientError(f"TransmissionBt request to '{self._base_url}' for method '{method}' failed: {e}")

def __does_torrent_exist_in_client(self, infohash):
try:
return bool(self.get_torrent_info(infohash))
zivkovic marked this conversation as resolved.
Show resolved Hide resolved
except TorrentClientError:
return False
5 changes: 5 additions & 0 deletions src/config.py
Original file line number Diff line number Diff line change
Expand Up @@ -23,6 +23,7 @@ def build_config_dict(cls, config_filepath: str, env_vars: dict):
"port": env_vars.get("PORT"),
"inject_torrents": True if env_vars.get("INJECT_TORRENTS", "").lower().strip() == "true" else False,
"deluge_rpc_url": env_vars.get("DELUGE_RPC_URL"),
"transmission_rpc_url": env_vars.get("TRANSMISSION_RPC_URL"),
"qbittorrent_url": env_vars.get("QBITTORRENT_URL"),
"injection_link_directory": env_vars.get("INJECTION_LINK_DIRECTORY"),
}.items()
Expand Down Expand Up @@ -50,6 +51,10 @@ def server_port(self) -> str:
def deluge_rpc_url(self) -> ParseResult | None:
return self._config.get("deluge_rpc_url")

@property
def transmission_rpc_url(self) -> ParseResult | None:
return self._config.get("transmission_rpc_url")

@property
def qbittorrent_url(self) -> ParseResult | None:
return self._config.get("qbittorrent_url")
Expand Down
14 changes: 13 additions & 1 deletion src/config_validator.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@

class ConfigValidator:
REQUIRED_KEYS = ["red_key", "ops_key"]
TORRENT_CLIENT_KEYS = ["deluge_rpc_url", "qbittorrent_url"]
TORRENT_CLIENT_KEYS = ["deluge_rpc_url", "transmission_rpc_url", "qbittorrent_url"]

def __init__(self, config_dict):
self.config_dict = config_dict
Expand All @@ -16,6 +16,7 @@ def __init__(self, config_dict):
"ops_key": self.__is_valid_ops_key,
"port": self.__is_valid_port,
"deluge_rpc_url": self.__is_valid_deluge_url,
"transmission_rpc_url": self.__is_valid_transmission_rpc_url,
"qbittorrent_url": self.__is_valid_qbit_url,
"inject_torrents": self.__is_boolean,
"injection_link_directory": assert_path_exists,
Expand Down Expand Up @@ -97,6 +98,17 @@ def __is_valid_deluge_url(url):
return parsed_url.geturl() # return the parsed URL
raise ValueError(f'Invalid "deluge_rpc_url" provided: {url}')

@staticmethod
def __is_valid_transmission_rpc_url(url):
parsed_url = urlparse(url)
if parsed_url.scheme and parsed_url.netloc:
if not parsed_url.password:
raise Exception(
"You need to define a password in the TransmissionBt RPC URL. (e.g. http://:<PASSWORD>@localhost:51413/transmission/rpc)"
)
return parsed_url.geturl() # return the parsed URL
raise ValueError(f'Invalid "transmission_rpc_url" provided: {url}')

@staticmethod
def __is_boolean(value):
coerced = value.lower().strip()
Expand Down
5 changes: 4 additions & 1 deletion src/injection.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@

from .clients.deluge import Deluge
from .clients.qbittorrent import Qbittorrent
from .clients.transmission import TransmissionBt
from .config import Config
from .errors import TorrentInjectionError
from .parser import calculate_infohash, get_bencoded_data
Expand Down Expand Up @@ -39,7 +40,7 @@ def __validate_config(config: Config):
if not config.injection_link_directory:
raise TorrentInjectionError("No injection link directory specified in the config file.")

if (not config.deluge_rpc_url) and (not config.qbittorrent_url):
if (not config.deluge_rpc_url) and (not config.transmission_rpc_url) and (not config.qbittorrent_url):
raise TorrentInjectionError("No torrent client configuration specified in the config file.")

return config
Expand All @@ -48,6 +49,8 @@ def __validate_config(config: Config):
def __determine_torrent_client(config: Config):
if config.deluge_rpc_url:
return Deluge(config.deluge_rpc_url)
elif config.transmission_rpc_url:
return TransmissionBt(config.transmission_rpc_url)
elif config.qbittorrent_url:
return Qbittorrent(config.qbittorrent_url)

Expand Down
176 changes: 176 additions & 0 deletions tests/clients/test_transmission.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,176 @@
import re
import pytest
import requests_mock

from tests.helpers import SetupTeardown, get_torrent_path

from src.errors import TorrentClientError, TorrentClientAuthenticationError, TorrentExistsInClientError
from src.clients.transmission import TransmissionBt


@pytest.fixture
def transmission_client():
return TransmissionBt("http://admin:supersecret@localhost:51314")


@pytest.fixture
def torrent_info_response():
return {
"arguments": {
"torrents": [
{
"name": "foo.torrent",
"percentDone": 1.0,
"doneDate": 0,
"status": 6,
"labels": ["bar"],
"downloadDir": "/tmp/baz",
}
]
}
}


class TestInit(SetupTeardown):
def test_initializes_with_url_parts(self):
transmission_client = TransmissionBt("http://admin:supersecret@localhost:51314")

assert transmission_client._base_url == "http://localhost:51314/transmission/rpc"
assert transmission_client._basic_auth.username == "admin"
assert transmission_client._basic_auth.password == "supersecret"


class TestSetup(SetupTeardown):
def test_sets_session_id(self, transmission_client):
assert transmission_client._transmission_session_id is None

with requests_mock.Mocker() as m:
m.post(re.compile("transmission/rpc"), headers={"X-Transmission-Session-Id": "1234"}, status_code=409)

transmission_client.setup()

assert transmission_client._transmission_session_id == "1234"

def test_raises_exception_on_failed_auth(self, transmission_client):
with requests_mock.Mocker() as m:
m.post(re.compile("transmission/rpc"), status_code=403)

with pytest.raises(TorrentClientAuthenticationError):
transmission_client.setup()

def test_raises_exception_if_no_session_id(self, transmission_client):
with requests_mock.Mocker() as m:
m.post(re.compile("transmission/rpc"), status_code=409)

with pytest.raises(TorrentClientAuthenticationError):
transmission_client.setup()


class TestGetTorrentInfo(SetupTeardown):
def test_returns_torrent_info(self, transmission_client, torrent_info_response):
with requests_mock.Mocker() as m:
m.post(re.compile("transmission/rpc"), json=torrent_info_response)

response = transmission_client.get_torrent_info("infohash")

assert response == {
"complete": True,
"label": ["bar"],
"save_path": "/tmp/baz",
"content_path": "/tmp/baz/foo.torrent",
}

def test_passes_headers_to_request(self, transmission_client, torrent_info_response):
with requests_mock.Mocker() as m:
m.post(re.compile("transmission/rpc"), headers={"X-Transmission-Session-Id": "1234"}, status_code=409)
transmission_client.setup()

with requests_mock.Mocker() as m:
m.post(re.compile("transmission/rpc"), json=torrent_info_response)

transmission_client.get_torrent_info("infohash")

assert m.last_request.headers["X-Transmission-Session-Id"] == transmission_client._transmission_session_id

def test_passes_json_body_to_request(self, transmission_client, torrent_info_response):
with requests_mock.Mocker() as m:
m.post(re.compile("transmission/rpc"), json=torrent_info_response)

transmission_client.get_torrent_info("infohash")

assert m.last_request.json() == {
"method": "torrent-get",
"arguments": {
"ids": ["infohash"],
"fields": ["labels", "downloadDir", "percentDone", "status", "doneDate", "name"],
},
}

def test_raises_if_json_error(self, transmission_client):
with requests_mock.Mocker() as m:
m.post(re.compile("transmission/rpc"), text="not json")

with pytest.raises(TorrentClientError, match="Client returned malformed json response"):
transmission_client.get_torrent_info("infohash")

def test_raises_if_no_torrents_found(self, transmission_client):
with requests_mock.Mocker() as m:
m.post(re.compile("transmission/rpc"), json={"arguments": {"torrents": []}})

with pytest.raises(TorrentClientError, match="Torrent not found in client"):
transmission_client.get_torrent_info("infohash")

def test_raises_on_unexpected_response(self, transmission_client):
with requests_mock.Mocker() as m:
m.post(re.compile("transmission/rpc"), text="")

with pytest.raises(TorrentClientError, match="Client returned unexpected response"):
transmission_client.get_torrent_info("infohash")


class TestInjectTorrent(SetupTeardown):
def test_injects_torrent(self, transmission_client, torrent_info_response):
torrent_path = get_torrent_path("red_source")

with requests_mock.Mocker() as m:
m.post(
re.compile("transmission/rpc"), [{"json": torrent_info_response}, {"json": {"arguments": {"torrents": []}}}]
)

transmission_client.inject_torrent("foo", torrent_path)

assert b'"method": "torrent-add"' in m.request_history[-1].body
assert b'"download-dir": "/tmp/baz"' in m.request_history[-1].body
assert b'"labels": ["bar"]' in m.request_history[-1].body
assert b'"metainfo"' in m.request_history[-1].body

def test_uses_save_path_override_if_present(self, transmission_client, torrent_info_response):
torrent_path = get_torrent_path("red_source")

with requests_mock.Mocker() as m:
m.post(
re.compile("transmission/rpc"), [{"json": torrent_info_response}, {"json": {"arguments": {"torrents": []}}}]
)

transmission_client.inject_torrent("foo", torrent_path, "/tmp/override/")

assert b'"download-dir": "/tmp/override/"' in m.request_history[-1].body

def test_raises_if_source_torrent_isnt_found_in_client(self, transmission_client):
with requests_mock.Mocker() as m:
m.post(
re.compile("transmission/rpc"),
[{"json": {"arguments": {"torrents": []}}}, {"json": {"arguments": {"torrents": []}}}],
)

with pytest.raises(TorrentClientError, match="Torrent not found in client"):
transmission_client.inject_torrent("foo", "bar.torrent")

def test_raises_if_destination_torrent_is_found_in_client(self, transmission_client, torrent_info_response):
torrent_path = get_torrent_path("red_source")

with requests_mock.Mocker() as m:
m.post(re.compile("transmission/rpc"), [{"json": torrent_info_response}, {"json": torrent_info_response}])

with pytest.raises(TorrentExistsInClientError, match="New torrent already exists in client"):
transmission_client.inject_torrent("foo", torrent_path)
Loading
Loading