Angelo Gladding
lahacker.net

dlv5vbq7lzlthol5 4b942a3185b37d00

angelo@lahacker.net

South Pasadena, California, United States currently feels like 68.15°F

Home CodecanopyFiles

canopy_torrents.py

Download raw file

"""
BitTorrent support for the Canopy

"""

# TODO async_add_torrent/add_torrent_alert

import pathlib
import signal
import sqlite3
import sys
import time

import libtorrent


save_path = pathlib.Path("./")
port = 6881
dht_bootstraps = ["router.bittorrent.com:6881"]
max_download, max_upload = -1, -1  # -1 = no limit; 4000 = 4kb/s
statistics = ("state", "download_rate", "upload_rate",
              "total_download", "total_upload")


def main():
    """

    """
    dbconn = sqlite3.connect("torrents.db")
    db = dbconn.cursor()

    session = libtorrent.session()
    try:
        with (save_path / ".bt-session").open("rb") as fp:
            session.load_state(libtorrent.bdecode(fp.read()))
    except FileNotFoundError:
        pass
    settings = libtorrent.session_settings()
    settings.user_agent = "bt/{}".format(libtorrent.version)
    session.set_settings(settings)
    session.listen_on(port, port + 10)
    for dht_server in dht_bootstraps:
        dht_host, dht_port = dht_server.split(":")
        session.add_dht_router(dht_host, int(dht_port))
    session.start_dht()
    session.set_download_rate_limit(max_download)
    session.set_upload_rate_limit(max_upload)
    session.set_severity_level(libtorrent.alert.severity_levels.info)
    session.add_extension(libtorrent.create_ut_pex_plugin)
    session.add_extension(libtorrent.create_ut_metadata_plugin)
    session.add_extension(libtorrent.create_metadata_plugin)
    # XXX session.add_extension(lambda x: PythonExtension(alerts))

    def add_torrent(**details):
        storage_mode = libtorrent.storage_mode_t.storage_mode_sparse
        details.update(storage_mode=storage_mode, auto_managed=True,
                       duplicate_is_error=True, save_path=str(save_path))
        handle = session.add_torrent(details)
        handle.set_max_connections(60)
        handle.set_max_uploads(-1)
        handle.set_ratio(0)
        return handle

    for torrent in db.execute("SELECT * FROM torrents"):
        add_torrent(url=torrent[1], resume_data=torrent[3])

    def handle_exit(*args):
        for handle in session.get_torrents():
            # if not handle.is_valid() or not handle.has_metadata():
            #     continue
            handle.pause()
            db.execute("""UPDATE torrents SET fast_resume = ?
                          WHERE info_hash = ?""",
                       (libtorrent.bencode(handle.write_resume_data()),
                        str(handle.info_hash())))
            dbconn.commit()
        with (save_path / ".bt-session").open("wb") as fp:
            fp.write(libtorrent.bencode(session.save_state()))
        dbconn.close()
        sys.exit(0)

    signal.signal(signal.SIGINT, handle_exit)
    signal.signal(signal.SIGTERM, handle_exit)
    try:
        while True:
            for handle in session.get_torrents():
                status = {k: str(getattr(handle.status(), k))
                          for k in statistics}
                db.execute("""UPDATE torrents
                              SET state = ?, download_rate = ?,
                                upload_rate = ?, total_download = ?,
                                total_upload = ?
                              WHERE info_hash = ?""",
                           (status["state"], status["download_rate"],
                            status["upload_rate"],
                            status["total_download"], status["total_upload"],
                            str(handle.info_hash())))
                dbconn.commit()
            while True:
                alert = session.pop_alert()
                if not alert:
                    break
                try:
                    alert_message = alert.message()
                except AttributeError:
                    alert_message = alert
                if alert_message.startswith(":"):
                    continue
                print("!", alert_message)
            for magnet, in db.execute("SELECT magnet FROM incoming"):
                try:
                    handle = add_torrent(url=magnet)
                except RuntimeError:  # torrent already exists in session
                    continue
                # TODO pause until torrent info has been fetched from DHT?
                # while not handle.has_metadata():
                #     time.sleep(.1)
                print("added torrent:", handle.name())
                db.execute("""INSERT INTO torrents (info_hash, magnet, name)
                              VALUES (?, ?, ?)""",
                           (str(handle.info_hash()), magnet,
                            str(handle.name())))
                db.execute("DELETE FROM incoming WHERE magnet = ?", (magnet,))
                dbconn.commit()
            time.sleep(3)
    except KeyboardInterrupt:
        handle_exit()


if __name__ == "__main__":
    main()