lbry-sdk/lbrynet/pointtraderclient/pointtraderclient.py

from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization, hashes
from cryptography.hazmat.primitives.asymmetric import rsa, padding, utils

from lbrynet import conf

from twisted.web.client import Agent, FileBodyProducer, Headers, ResponseDone
from twisted.internet import threads, defer, protocol
from hashlib import sha1
from StringIO import StringIO
import time
import json
import binascii


def gen_rsa_key(bits):
    PUBLIC_EXPOENT = 65537  # http://www.daemonology.net/blog/2009-06-11-cryptographic-right-answers.html
    return rsa.generate_private_key(public_exponent=PUBLIC_EXPOENT,
                                    key_size=4096, backend=default_backend())


def sign(private_key, recipient_public_key=None, amount=None):
    encoded_public_key = private_key.public_key().public_bytes(serialization.Encoding.PEM,
                                                               serialization.PublicFormat.PKCS1)
    timestamp = time.time()
    h = sha1()
    h.update(encoded_public_key)
    if amount and recipient_public_key:
        h.update(recipient_public_key)
        h.update(str(amount))
    h.update(str(timestamp))
    signature = private_key.sign(h.digest(), padding.PSS(mgf=padding.MGF1(hashes.SHA1()),
                                                         salt_length=padding.PSS.MAX_LENGTH),
                                 utils.Prehashed(hashes.SHA1()))
    return encoded_public_key, timestamp, binascii.hexlify(signature)


class BeginningPrinter(protocol.Protocol):
    def __init__(self, finished):
        self.finished = finished
        self.data = ""

    def dataReceived(self, bytes):
        self.data = self.data + bytes

    def connectionLost(self, reason):
        if reason.check(ResponseDone) is not None:
            self.finished.callback(str(self.data))
        else:
            self.finished.errback(reason)


def read_body(response):
    d = defer.Deferred()
    response.deliverBody(BeginningPrinter(d))
    return d


def get_body(response):
    if response.code != 200:
        print "\n\n\n\nbad error code\n\n\n\n"
        raise ValueError(response.phrase)
    else:
        return read_body(response)


def get_body_from_request(path, data):

    from twisted.internet import reactor

    jsondata = FileBodyProducer(StringIO(json.dumps(data)))
    agent = Agent(reactor)
    d = agent.request(
        'POST', conf.settings['pointtrader_server'] + path,
        Headers({'Content-Type': ['application/json']}), jsondata)
    d.addCallback(get_body)
    return d


def print_response(response):
    pass


def print_error(err):
    print err.getTraceback()
    return err


def register_new_account(private_key):
    data = {}
    encoded_public_key = private_key.public_key().public_bytes(serialization.Encoding.PEM,
                                                               serialization.PublicFormat.PKCS1)
    data['pub_key'] = encoded_public_key

    def get_success_from_body(body):
        r = json.loads(body)
        if not 'success' in r or r['success'] is False:
            return False
        return True

    d = get_body_from_request('/register/', data)

    d.addCallback(get_success_from_body)
    return d


def send_points(private_key, recipient_public_key, amount):
    encoded_public_key, timestamp, signature = sign(private_key, recipient_public_key, amount)

    data = {}
    data['sender_pub_key'] = encoded_public_key
    data['recipient_pub_key'] = recipient_public_key
    data['amount'] = amount
    data['timestamp'] = timestamp
    data['signature'] = signature

    def get_success_from_body(body):
        r = json.loads(body)
        if not 'success' in r or r['success'] is False:
            return False
        return True

    d = get_body_from_request('/send-points/', data)

    d.addCallback(get_success_from_body)

    return d


def get_recent_transactions(private_key):
    encoded_public_key, timestamp, signature = sign(private_key)

    data = {}
    data['pub_key'] = encoded_public_key
    data['timestamp'] = timestamp
    data['signature'] = signature
    data['end_time'] = 0
    data['start_time'] = 120

    def get_transactions_from_body(body):
        r = json.loads(body)
        if "transactions" not in r:
            raise ValueError("Invalid response: no 'transactions' field")
        else:
            return r['transactions']

    d = get_body_from_request('/get-transactions/', data)

    d.addCallback(get_transactions_from_body)

    return d


def get_balance(private_key):
    encoded_public_key, timestamp, signature = sign(private_key)

    data = {}
    data['pub_key'] = encoded_public_key
    data['timestamp'] = timestamp
    data['signature'] = signature

    def get_balance_from_body(body):
        r = json.loads(body)
        if not 'balance' in r:
            raise ValueError("Invalid response: no 'balance' field")
        else:
            return float(r['balance'])

    d = get_body_from_request('/get-balance/', data)

    d.addCallback(get_balance_from_body)

    return d


def run_full_test():

    keys = []

    def save_key(private_key):
        keys.append(private_key)
        return private_key

    def check_balances_and_transactions(unused, bal1, bal2, num_transactions):

        def assert_balance_is(actual, expected):
            assert abs(actual - expected) < .05
            print "correct balance. actual:", str(actual), "expected:", str(expected)
            return True

        def assert_transaction_length_is(transactions, expected_length):
            assert len(transactions) == expected_length
            print "correct transaction length"
            return True

        d1 = get_balance(keys[0])
        d1.addCallback(assert_balance_is, bal1)

        d2 = get_balance(keys[1])
        d2.addCallback(assert_balance_is, bal2)

        d3 = get_recent_transactions(keys[0])
        d3.addCallback(assert_transaction_length_is, num_transactions)

        d4 = get_recent_transactions(keys[1])
        d4.addCallback(assert_transaction_length_is, num_transactions)

        dl = defer.DeferredList([d1, d2, d3, d4])
        return dl

    def do_transfer(unused, amount):
        encoded_public_key = keys[1].public_key().public_bytes(serialization.Encoding.PEM,
                                                               serialization.PublicFormat.PKCS1)
        d = send_points(keys[0], encoded_public_key, amount)
        return d

    d1 = threads.deferToThread(gen_rsa_key, 4096)
    d1.addCallback(save_key)
    d1.addCallback(register_new_account)
    d2 = threads.deferToThread(gen_rsa_key, 4096)
    d2.addCallback(save_key)
    d2.addCallback(register_new_account)
    dlist = defer.DeferredList([d1, d2])
    dlist.addCallback(check_balances_and_transactions, 1000, 1000, 0)
    dlist.addCallback(do_transfer, 50)
    dlist.addCallback(check_balances_and_transactions, 950, 1050, 1)
    dlist.addCallback(do_transfer, 75)
    dlist.addCallback(check_balances_and_transactions, 875, 1125, 2)
    dlist.addErrback(print_error)


if __name__ == "__main__":
    conf.initialize_settings()

    from twisted.internet import reactor

    reactor.callLater(1, run_full_test)
    reactor.callLater(25, reactor.stop)
    reactor.run()