rucio 37.0.0rc1__py3-none-any.whl

rucio/transfertool/__init__.py

@@ -0,0 +1,13 @@
+# Copyright European Organization for Nuclear Research (CERN) since 2012
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

rucio/transfertool/bittorrent.py

@@ -0,0 +1,200 @@
+# Copyright European Organization for Nuclear Research (CERN) since 2012
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import base64
+import logging
+from os import path
+from typing import TYPE_CHECKING, Any, Optional
+
+from rucio.common.bittorrent import construct_torrent
+from rucio.common.config import config_get
+from rucio.common.constants import RseAttr
+from rucio.common.extra import import_extras
+from rucio.core.did_meta_plugins import get_metadata
+from rucio.transfertool.transfertool import TransferStatusReport, Transfertool, TransferToolBuilder
+
+if TYPE_CHECKING:
+    from collections.abc import Mapping, Sequence
+
+    from rucio.common import types
+    from rucio.core.request import DirectTransfer
+    from rucio.core.rse import RseData
+
+    from .bittorrent_driver import BittorrentDriver
+
+DRIVER_NAME_RSE_ATTRIBUTE = 'bittorrent_driver'
+DRIVER_CLASSES_BY_NAME: dict[str, type['BittorrentDriver']] = {}
+
+EXTRA_MODULES = import_extras(['qbittorrentapi'])
+
+if EXTRA_MODULES['qbittorrentapi']:
+    from .bittorrent_driver_qbittorrent import QBittorrentDriver
+    DRIVER_CLASSES_BY_NAME[QBittorrentDriver.external_name] = QBittorrentDriver
+
+
+class BittorrentTransfertool(Transfertool):
+    """
+    Use bittorrent to perform the peer-to-peer transfer.
+    """
+    external_name = 'bittorrent'
+    supported_schemes = {'magnet'}
+
+    required_rse_attrs = (DRIVER_NAME_RSE_ATTRIBUTE, )
+
+    def __init__(self, external_host: str, logger: 'types.LoggerFunction' = logging.log) -> None:
+        super().__init__(external_host=external_host, logger=logger)
+
+        self._drivers_by_rse_id = {}
+        self.ca_cert, self.ca_key = None, None
+
+        self.tracker = config_get('transfers', 'bittorrent_tracker_addr', raise_exception=False, default=None)
+
+    @classmethod
+    def _pick_management_api_driver_cls(cls: "type[BittorrentTransfertool]", rse: "RseData") -> Optional[type['BittorrentDriver']]:
+        driver_cls = DRIVER_CLASSES_BY_NAME.get(rse.attributes.get(DRIVER_NAME_RSE_ATTRIBUTE, ''))
+        if driver_cls is None:
+            return None
+        if not all(rse.attributes.get(attribute) is not None for attribute in driver_cls.required_rse_attrs):
+            return None
+        return driver_cls
+
+    def _driver_for_rse(self, rse: "RseData") -> Optional['BittorrentDriver']:
+        driver = self._drivers_by_rse_id.get(rse.id)
+        if driver:
+            return driver
+
+        driver_cls = self._pick_management_api_driver_cls(rse)
+        if not driver_cls:
+            return None
+
+        driver = driver_cls.make_driver(rse)
+        self._drivers_by_rse_id[rse.id] = driver
+        return driver
+
+    @staticmethod
+    def _get_torrent_meta(scope: "types.InternalScope", name: str) -> tuple[bytes, bytes, int]:
+        meta = get_metadata(scope=scope, name=name, plugin='all')
+        pieces_root = base64.b64decode(meta.get('bittorrent_pieces_root', ''))
+        pieces_layers = base64.b64decode(meta.get('bittorrent_pieces_layers', ''))
+        piece_length = meta.get('bittorrent_piece_length', 0)
+        return pieces_root, pieces_layers, piece_length
+
+    @classmethod
+    def submission_builder_for_path(
+            cls: "type[BittorrentTransfertool]",
+            transfer_path: "list[DirectTransfer]",
+            logger: 'types.LoggerFunction' = logging.log
+    ) -> "tuple[list[DirectTransfer], Optional[TransferToolBuilder]]":
+        hop = transfer_path[0]
+        if hop.rws.byte_count == 0:
+            logger(logging.INFO, f"Bittorrent cannot transfer fully empty torrents. Skipping {hop}")
+            return [], None
+
+        if not cls.can_perform_transfer(hop.src.rse, hop.dst.rse):
+            logger(logging.INFO, f"The required RSE attributes are not set. Skipping {hop}")
+            return [], None
+
+        for rse in [hop.src.rse, hop.dst.rse]:
+            driver_cls = cls._pick_management_api_driver_cls(rse)
+            if not driver_cls:
+                logger(logging.INFO, f"The rse '{rse}' is not configured correctly for bittorrent")
+                return [], None
+
+        pieces_root, _pieces_layers, piece_length = cls._get_torrent_meta(hop.rws.scope, hop.rws.name)
+        if not pieces_root or not piece_length:
+            logger(logging.INFO, "The required bittorrent metadata not set on the DID")
+            return [], None
+
+        return [hop], TransferToolBuilder(cls, external_host='Bittorrent Transfertool')
+
+    def group_into_submit_jobs(self, transfer_paths: "Sequence[list[DirectTransfer]]") -> list[dict[str, Any]]:
+        return [{'transfers': transfer_path, 'job_params': {}} for transfer_path in transfer_paths]
+
+    @staticmethod
+    def _connect_directly(torrent_id: str, peers_drivers: 'Sequence[BittorrentDriver]') -> None:
+        peer_addr = []
+        for i, driver in enumerate(peers_drivers):
+            peer_addr.append(driver.listen_addr())
+
+        for driver in peers_drivers:
+            driver.add_peers(torrent_id=torrent_id, peers=peer_addr)
+
+    def submit(self, transfers: "Sequence[DirectTransfer]", job_params: dict[str, str], timeout: Optional[int] = None) -> str:
+        [transfer] = transfers
+        rws = transfer.rws
+
+        tracker = transfer.dst.rse.attributes.get(RseAttr.BITTORRENT_TRACKER_ADDR, self.tracker)
+
+        src_drivers = {}
+        for source in transfer.sources:
+            driver = self._driver_for_rse(source.rse)
+            if driver:
+                src_drivers[source] = driver
+
+        dst_driver = self._driver_for_rse(transfer.dst.rse)
+
+        if not dst_driver or not src_drivers:
+            raise Exception('Cannot initialize bittorrent drivers to submit transfers')
+
+        pieces_root, pieces_layers, piece_length = self._get_torrent_meta(rws.scope, rws.name)
+        torrent_id, torrent = construct_torrent(
+            scope=str(rws.scope),
+            name=rws.name,
+            length=rws.byte_count,
+            piece_length=piece_length,
+            pieces_root=pieces_root,
+            pieces_layers=pieces_layers,
+            trackers=[tracker] if tracker else None,
+        )
+
+        for source, driver in src_drivers.items():
+            source_protocol = transfer.source_protocol(source)
+            [lfn] = source_protocol.parse_pfns([transfer.source_url(source)]).values()
+            driver.add_torrent(
+                file_name=rws.name,
+                file_content=torrent,
+                download_location=lfn['prefix'] + path.dirname(lfn['path']),
+                seed_mode=True,
+            )
+
+        dest_protocol = transfer.dest_protocol()
+        [lfn] = dest_protocol.parse_pfns([transfer.dest_url]).values()
+        dst_driver.add_torrent(
+            file_name=rws.name,
+            file_content=torrent,
+            download_location=lfn['prefix'] + lfn['path'],
+        )
+
+        self._connect_directly(torrent_id, [dst_driver] + list(src_drivers.values()))
+        return torrent_id
+
+    def bulk_query(self, requests_by_eid, timeout: Optional[int] = None) -> 'Mapping[str, Mapping[str, TransferStatusReport]]':
+        response = {}
+        for transfer_id, requests in requests_by_eid.items():
+            for request_id, request in requests.items():
+                driver = self._driver_for_rse(request['dst_rse'])
+                if not driver:
+                    self.logger(f'Cannot instantiate BitTorrent driver for {request["dest_rse"]}')
+                    continue
+                response.setdefault(transfer_id, {})[request_id] = driver.get_status(request_id=request_id, torrent_id=transfer_id)
+        return response
+
+    def query(self, transfer_ids: 'Sequence[str]', details: bool = False, timeout: Optional[int] = None) -> None:
+        pass
+
+    def cancel(self, transfer_ids: 'Sequence[str]', timeout: Optional[int] = None) -> None:
+        pass
+
+    def update_priority(self, transfer_id: str, priority: int, timeout: Optional[int] = None) -> None:
+        pass

rucio/transfertool/bittorrent_driver.py

@@ -0,0 +1,50 @@
+# Copyright European Organization for Nuclear Research (CERN) since 2012
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from abc import ABCMeta, abstractmethod
+from typing import TYPE_CHECKING, Optional
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from rucio.common import types
+    from rucio.core.rse import RseData
+    from rucio.transfertool.transfertool import TransferStatusReport
+
+
+class BittorrentDriver(metaclass=ABCMeta):
+    external_name = ''
+    required_rse_attrs = tuple()
+
+    @classmethod
+    @abstractmethod
+    def make_driver(cls: "type[BittorrentDriver]", rse: "RseData", logger: 'types.LoggerFunction' = logging.log) -> "Optional[BittorrentDriver]":
+        pass
+
+    @abstractmethod
+    def listen_addr(self) -> tuple[str, int]:
+        pass
+
+    @abstractmethod
+    def add_torrent(self, file_name: str, file_content: bytes, download_location: str, seed_mode: bool = False) -> None:
+        pass
+
+    @abstractmethod
+    def add_peers(self, torrent_id: str, peers: 'Sequence[tuple[str, int]]') -> None:
+        pass
+
+    @abstractmethod
+    def get_status(self, request_id: str, torrent_id: str) -> "TransferStatusReport":
+        pass

rucio/transfertool/bittorrent_driver_qbittorrent.py

@@ -0,0 +1,134 @@
+# Copyright European Organization for Nuclear Research (CERN) since 2012
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import logging
+from typing import TYPE_CHECKING, Optional, cast
+from urllib.parse import urlparse
+
+import qbittorrentapi
+
+from rucio.common.config import get_rse_credentials
+from rucio.common.constants import RseAttr
+from rucio.common.utils import resolve_ip
+from rucio.core.oidc import request_token
+from rucio.db.sqla.constants import RequestState
+from rucio.transfertool.transfertool import TransferStatusReport
+
+from .bittorrent_driver import BittorrentDriver
+
+if TYPE_CHECKING:
+    from collections.abc import Sequence
+
+    from sqlalchemy.orm import Session
+
+    from rucio.common import types
+    from rucio.core.rse import RseData
+
+
+class QBittorrentTransferStatusReport(TransferStatusReport):
+
+    supported_db_fields = [
+        'state',
+        'external_id',
+    ]
+
+    def __init__(self, request_id: str, external_id: str, qbittorrent_response: Optional[qbittorrentapi.TorrentDictionary]) -> None:
+        super().__init__(request_id)
+
+        if qbittorrent_response and qbittorrent_response.state_enum.is_complete == 1:
+            new_state = RequestState.DONE
+        else:
+            new_state = RequestState.SUBMITTED
+
+        self.state = new_state
+        self.external_id = None
+        if new_state in [RequestState.FAILED, RequestState.DONE]:
+            self.external_id = external_id
+
+    def initialize(self, session: "Session", logger: 'types.LoggerFunction' = logging.log) -> None:
+        pass
+
+    def get_monitor_msg_fields(self, session: "Session", logger: 'types.LoggerFunction' = logging.log) -> dict[str, str]:
+        return {'protocol': 'qbittorrent'}
+
+
+class QBittorrentDriver(BittorrentDriver):
+
+    external_name = 'qbittorrent'
+    required_rse_attrs = (RseAttr.QBITTORRENT_MANAGEMENT_ADDRESS, )
+
+    @classmethod
+    def make_driver(cls: "type[QBittorrentDriver]", rse: "RseData", logger: 'types.LoggerFunction' = logging.log) -> "Optional[BittorrentDriver]":
+
+        address = rse.attributes.get(RseAttr.QBITTORRENT_MANAGEMENT_ADDRESS)
+        if not address:
+            return None
+
+        url = urlparse(address)
+        token = None
+        if url.scheme.lower() == 'https':
+            token = request_token(audience=url.hostname, scope='qbittorrent_admin')
+        else:
+            logging.debug(f'{cls.external_name} will not try token authentication. Requires HTTPS.')
+
+        rse_cred = get_rse_credentials().get(rse.id, {})
+        username = rse_cred.get('qbittorrent_username')
+        password = rse_cred.get('qbittorrent_password')
+
+        if not (token or (username and password)):
+            return None
+
+        return cls(
+            address=address,
+            username=username,
+            password=password,
+            token=token,
+            logger=logger,
+        )
+
+    def __init__(self, address: str, username: str, password: str, token: Optional[str] = None, logger: 'types.LoggerFunction' = logging.log) -> None:
+        extra_headers = None
+        if token:
+            extra_headers = {'Authorization': 'Bearer ' + token}
+
+        self.client = qbittorrentapi.Client(
+            host=address,
+            username=username,
+            password=password,
+            EXTRA_HEADERS=extra_headers,
+            FORCE_SCHEME_FROM_HOST=True,
+        )
+        self.logger = logger
+
+    def listen_addr(self) -> tuple[str, int]:
+        preferences = self.client.app_preferences()
+        port = cast('int', preferences['listen_port'])
+        ip = resolve_ip(urlparse(self.client.host).hostname or self.client.host)
+        return ip, port
+
+    def add_torrent(self, file_name: str, file_content: bytes, download_location: str, seed_mode: bool = False) -> None:
+        self.client.torrents_add(
+            rename=file_name,
+            torrent_files=file_content,
+            save_path=download_location,
+            is_skip_checking=seed_mode,
+            is_sequential_download=True,
+        )
+
+    def add_peers(self, torrent_id: str, peers: 'Sequence[tuple[str, int]]') -> None:
+        self.client.torrents_add_peers(torrent_hashes=[torrent_id], peers=[f'{ip}:{port}' for ip, port in peers])
+
+    def get_status(self, request_id: str, torrent_id: str) -> TransferStatusReport:
+        info = self.client.torrents_info(torrent_hashes=[torrent_id])
+        return QBittorrentTransferStatusReport(request_id, external_id=torrent_id, qbittorrent_response=info[0] if info else None)