lbry-sdk/scripts/encrypt_blob.py

"""Encrypt a single file using the given key and iv"""
import argparse
import logging
import sys

from twisted.internet import defer
from twisted.internet import reactor

from lbrynet import conf
from lbrynet.core import log_support
from lbrynet.core.HashAnnouncer import DummyHashAnnouncer
from lbrynet.core.BlobManager import DiskBlobManager
from lbrynet.cryptstream.CryptStreamCreator import CryptStreamCreator


log = logging.getLogger('decrypt_blob')


def main():
    conf.initialize_settings()
    parser = argparse.ArgumentParser()
    parser.add_argument('filename')
    parser.add_argument('hex_key')
    parser.add_argument('hex_iv')
    args = parser.parse_args()
    log_support.configure_console(level='DEBUG')

    run(args)
    reactor.run()


@defer.inlineCallbacks
def run(args):
    try:
        yield encrypt_blob(args.filename, args.hex_key, args.hex_iv)
    except Exception:
        log.exception('Failed to encrypt blob')
    finally:
        reactor.callLater(0, reactor.stop)


@defer.inlineCallbacks
def encrypt_blob(filename, key, iv):
    dummy_announcer = DummyHashAnnouncer()
    manager = DiskBlobManager(dummy_announcer, '.', '.')
    yield manager.setup()
    creator = CryptStreamCreator(manager, filename, key, iv_generator(iv))
    with open(filename, 'r') as infile:
        data = infile.read(2**14)
        while data:
            yield creator.write(data)
            data = infile.read(2**14)
    yield creator.stop()


def iv_generator(iv):
    iv = int(iv, 16)
    while 1:
        iv += 1
        yield ("%016d" % iv)[-16:]


if __name__ == '__main__':
    sys.exit(main())