#!/usr/bin/env python3
import binascii
import sys
import traceback

import Pyro4
import serpent
from trezorlib import cosi, tools

PORT = 5001
indexmap = {"bootloader": 0, "vendorheader": 1, "firmware": 2}


def get_trezor():
    from trezorlib.client import TrezorClient
    from trezorlib.transport import get_transport
    from trezorlib.ui import ClickUI

    return TrezorClient(get_transport(), ui=ClickUI())


def get_path(index):
    return "10018'/%d'" % indexmap[index]


@Pyro4.expose
class KeyctlProxy(object):
    def get_commit(self, index, digest):
        digest = serpent.tobytes(digest)
        path = get_path(index)
        commit = None
        while commit is None:
            try:
                t = get_trezor()
                print(
                    "\n\n\nCommiting to hash %s with path %s:"
                    % (binascii.hexlify(digest).decode(), path)
                )
                commit = cosi.commit(t, tools.parse_path(path), digest)
            except Exception as e:
                print(e)
                traceback.print_exc()
                print("Trying again ...")
        pk = commit.pubkey
        R = commit.commitment
        print("Commitment sent!")
        return (pk, R)

    def get_signature(self, index, digest, global_R, global_pk):
        digest, global_R, global_pk = (
            serpent.tobytes(digest),
            serpent.tobytes(global_R),
            serpent.tobytes(global_pk),
        )
        path = get_path(index)
        signature = None
        while signature is None:
            try:
                t = get_trezor()
                print(
                    "\n\n\nSigning hash %s with path %s:"
                    % (binascii.hexlify(digest).decode(), path)
                )
                signature = cosi.sign(
                    t, tools.parse_path(path), digest, global_R, global_pk
                )
            except Exception as e:
                print(e)
                traceback.print_exc()
                print("Trying again ...")
        sig = signature.signature
        print("Signature sent!")
        return sig


if __name__ == "__main__":
    if len(sys.argv) > 1:
        ipaddr = sys.argv[1]
    else:
        print("Usage: keyctl-proxy ipaddress")
        sys.exit(1)
    daemon = Pyro4.Daemon(host=ipaddr, port=PORT)
    proxy = KeyctlProxy()
    uri = daemon.register(proxy, "keyctl")
    print('keyctl-proxy running at URI: "%s"' % uri)
    daemon.requestLoop()