disseminator.py

import argparse
import asyncio
import logging
from pathlib import Path
from types import SimpleNamespace

import libtorrent
from pony.orm import db_session

from tribler_core.modules.libtorrent.torrentdef import TorrentDef
from tribler_core.modules.metadata_store.community.gigachannel_community import GigaChannelCommunity
from tribler_core.modules.metadata_store.orm_bindings.channel_node import NEW
from tribler_core.utilities.tiny_tribler_service import TinyTriblerService

_logger = logging.getLogger('Disseminator')


def parse_args():
    parser = argparse.ArgumentParser(description='Disseminate data by using the Tribler network')

    parser.add_argument('-s', '--source', type=str, help='path to data folder', default='.')
    parser.add_argument('-t', '--tribler_dir', type=str, help='path to data folder', default='$HOME/seedbox/.Tribler')
    parser.add_argument('-v', '--verbosity', help='increase output verbosity', action='store_true')

    return parser.parse_args()


def setup_logger(verbosity):
    logging_level = logging.DEBUG if verbosity else logging.INFO
    logging.basicConfig(level=logging_level)


class ChannelHelper:
    def __init__(self, community):
        self.community = community
        self.directories = SimpleNamespace(tree={}, directory=None)

    @db_session
    def create_root_channel(self, name, description=''):
        _logger.info(f'Creating channel: {name}')
        channels = self.community.mds.ChannelMetadata

        if len(channels.get_channels_by_title(name)) >= 1:
            _logger.warning(f'Channel with name {name} already exists')
            return False

        self.directories.directory = channels.create_channel(name, description)
        self.flush()

        return True

    @db_session
    def add_torrent(self, file, relative_path):
        _logger.info(f'Add torrent: {file}')

        directory = self.get_directory(relative_path)
        decoded_torrent = libtorrent.bdecode(file.read_bytes())
        directory.add_torrent_to_channel(TorrentDef(metainfo=decoded_torrent), None)

    @db_session
    def get_directory(self, path):
        current = self.directories

        for part in path.parts[:-1]:
            next_directory = current.tree.get(part, None)
            if next_directory is not None:
                current = next_directory
                continue

            next_directory = SimpleNamespace(
                tree={},
                directory=self.community.mds.CollectionNode(title=part, origin_id=current.directory.id_) # , status=NEW
            )

            current.tree[part] = next_directory
            current = next_directory
            self.flush()

            _logger.info(f'Directory created: {part}')

        return current.directory

    @db_session
    def commit(self):
        self.community.mds.CollectionNode.commit_all_channels()

    @db_session
    def flush(self):
        self.community.mds._db.flush()


class Service(TinyTriblerService):
    def __init__(self, source_dir, working_dir, config_path):
        super().__init__(Service.create_config(working_dir, config_path),
                         working_dir=working_dir,
                         config_path=config_path)
        self.source_dir = Path(source_dir)

    @staticmethod
    def create_config(working_dir, config_path):
        config = TinyTriblerService.create_default_config(working_dir, config_path)

        config.set_libtorrent_enabled(True)
        config.set_ipv8_enabled(True)
        config.set_chant_enabled(True)

        return config

    def get_torrents_from_source(self):
        return [(file, file.relative_to(self.source_dir)) for file in self.source_dir.rglob('*.torrent')]

    async def create_channel(self, community):
        channel_helper = ChannelHelper(community)
        channel_name = self.source_dir.name

        if not channel_helper.create_root_channel(channel_name):
            return

        torrents = self.get_torrents_from_source()

        for file, relative_path in torrents:
            channel_helper.add_torrent(file, relative_path)

        _logger.info(f'{len(torrents)} torrents where added')

    async def on_tribler_started(self):
        await super().on_tribler_started()
        await self.create_channel(self.session.ipv8.get_overlay(GigaChannelCommunity))


def run_tribler(arguments):
    service = Service(
        source_dir=Path(arguments.source),
        working_dir=Path(arguments.tribler_dir),
        config_path=Path('./tribler.conf')
    )

    loop = asyncio.get_event_loop()
    loop.create_task(service.start_tribler())
    try:
        loop.run_forever()
    finally:
        loop.run_until_complete(loop.shutdown_asyncgens())
        loop.close()


if __name__ == "__main__":
    _arguments = parse_args()
    print(f"Arguments: {_arguments}")

    setup_logger(_arguments.verbosity)
    run_tribler(_arguments)