#!/usr/bin/env python
from __future__ import print_function

from keepkeylib.client import KeepKeyClient
from keepkeylib.transport_hid import HidTransport

import sys
import binascii
import hashlib
import urllib2
import json
import base64

def xor_strings(xs, ys):
    return "".join(chr(ord(x) ^ ord(y)) for x, y in zip(xs, ys))

def getApiToken(client):
    apiTokenNode = client.get_public_node(client.expand_path("3579891166/3592008662"))
    publicKey = apiTokenNode.node.public_key
    sha256 = hashlib.sha256()
    sha256.update(publicKey)
    hash = sha256.digest()

    token = binascii.hexlify(xor_strings(hash[:16], hash[16:]))

    # print('api token:', token)
    return token

def decryptNodePath(client, encryptedNodePath):
    return client.decrypt_keyvalue(client.expand_path(''), 'node-location', encryptedNodePath,
                                           ask_on_decrypt=False, ask_on_encrypt=False)

def main():
    message = sys.argv[1]

    devices = HidTransport.enumerate()
    if len(devices) == 0:
        print('No KeepKey found')
        return
    transport = HidTransport(devices[0])
    client = KeepKeyClient(transport)

    token = getApiToken(client)
    walletIdsRaw = urllib2.urlopen('https://api.blockcypher.com/v1/btc/main/wallets?token=%s' % token).read()
    walletIds = json.loads(walletIdsRaw)['wallet_names']

    for walletId in walletIds:
        walletRaw = urllib2\
            .urlopen('https://api.blockcypher.com/v1/btc/main/addrs/%s/meta?token=%s&private=true' % (walletId, token))\
            .read()

        encryptedNodePath = binascii.unhexlify(json.loads(walletRaw)['encrypted_node_path'])
        nodePath = binascii.hexlify(decryptNodePath(client, encryptedNodePath))

        npLength = int(nodePath[0:2], 16)

        path = ''
        for depth in range(npLength):
            s = 2 + depth * 8
            e = s + 8
            path = "%s/%s" % (path, int(nodePath[s:e], 16))
        path = path[1:]

        walletBalanceDataRaw = urllib2\
            .urlopen('https://api.blockcypher.com/v1/btc/main/addrs/%s?token=%s&unspentOnly=true' % (walletId, token))\
            .read()

        txrefs = json.loads(walletBalanceDataRaw).get('txrefs')

        if txrefs != None:
            walletInfoRaw = urllib2\
                .urlopen('https://api.blockcypher.com/v1/btc/main/wallets/hd/%s?token=%s' % (walletId, token))\
                .read()

            walletAddrs = [
                addr
                for chain in json.loads(walletInfoRaw).get('chains')
                if 'chain_addresses' in chain
                for addr in chain['chain_addresses']
            ]

            addressesWithBalance = [
                txref['address']
                for txref in txrefs
                if 'address' in txref
            ]

            for addr in addressesWithBalance:
                nodePath = path + '/' + [a['path'] for a in walletAddrs if addr == a['address']][0][2:]
                msg = client.sign_message('Bitcoin', client.expand_path(nodePath), message)
                assert(msg.address == addr)
                print(msg.address, base64.b64encode(msg.signature))

    client.close()

if __name__ == '__main__':
    main()

