import ecdsa
from helper import *
from ecdsa import SECP256k1
from ecdsa.keys import SigningKey, VerifyingKey
import secrets

CURVE_ORDER = SECP256k1.order
KEY_ID_LENGTH = 32
HALF_KEY_ID_LENGTH = int(KEY_ID_LENGTH/2)
CRED_ID_LENGTH = KEY_ID_LENGTH * 2


class HDKey(object):
    ''' extended key '''

    def __init__(self, keyid, prikey, ccode, pubkey, is_prikey, parentId=None, depth=0):
        self.depth = depth
        self.is_prikey = is_prikey
        self.keyid = keyid
        self.credid = (parentId + keyid) if parentId else keyid
        ccode_int = int.from_bytes(ccode, 'big')

        if not ccode or ccode_int > CURVE_ORDER:
            raise Exception('ccode must less than {}'.format(CURVE_ORDER))

        self.ccode = ccode[:]

        if is_prikey:
            if not isinstance(prikey, SigningKey):
                raise Exception('need prikey')
            self.prikey = prikey
            self.pubkey = prikey.get_verifying_key()
        else:
            self.pubkey = VerifyingKey.from_string(
                pubkey.to_string(), curve=pubkey.curve)

    def _child_key(self, keyid, include_prikey=False):
        '''generate child key'''
        if include_prikey:
            if not self.is_prikey:
                raise Exception('this key doesn\'t include prikey')
            return self._child_key_from_prikey(keyid)
        else:
            pubkey = self.pubkey
            ccode = self.ccode
            deltakey, child_ccode = deltakey_and_ccode(
                keyid, pubkey.to_string(), ccode)
            deltakey_point = deltakey.get_verifying_key().pubkey.point
            point = pubkey.pubkey.point + deltakey_point
            child_key = ecdsa.VerifyingKey.from_public_point(
                point, curve=SECP256k1)
            return HDKey(keyid=keyid, prikey=None, ccode=child_ccode, is_prikey=False, pubkey=child_key, parentId=self.keyid, depth=self.depth+1)

    def _child_key_from_prikey(self, keyid):
        ''' generate childkey from prikey and chain code'''
        prikey = self.prikey
        ccode = self.ccode
        pubkey = prikey.get_verifying_key().to_string()

        delta_key, child_ccode = deltakey_and_ccode(keyid, pubkey, ccode)

        child_key_str = add_secret_keys(
            prikey.to_string(), delta_key.to_string(), order=SECP256k1.order)
        child_key = ecdsa.SigningKey.from_string(
            child_key_str, curve=SECP256k1)
        return HDKey(keyid=keyid, prikey=child_key, ccode=child_ccode, pubkey=None, parentId=self.keyid, is_prikey=True, depth=self.depth+1)

    def _checksum(self, source, appid_hash=None):
        if appid_hash:
            s = source + appid_hash
        else:
            s = source
        return hmac512(self.ccode, s)[:HALF_KEY_ID_LENGTH]

    def _generateRandomKeyId(self, appid_hash=None):
        keyid_L = secrets.token_bytes(HALF_KEY_ID_LENGTH)

        return keyid_L + self._checksum(keyid_L, appid_hash)

    def _child_key_from_id(self, keyid, appid_hash=None):
        if self.is_child_key_id(keyid, appid_hash):
            return self._child_key(keyid, include_prikey=self.is_prikey)
        else:
            raise Exception('invalid keyid {}'.format(keyid.hex()))

    def sign(self, source):
        return self.prikey.sign(source)

    def app_prikey(self, credid, appid_hash):
        if not self.is_prikey:
            raise Exception('this key doesn\'t prikey')

        if len(credid) == CRED_ID_LENGTH:
            childkey = self._child_key_from_id(credid[:KEY_ID_LENGTH])
            prikey = childkey._child_key_from_id(
                credid[KEY_ID_LENGTH:], appid_hash)
            return prikey
        else:
            return None

    def pubkey_seed(self):
        child_keyid = self._generateRandomKeyId()
        return self._child_key(child_keyid, include_prikey=False)

    def app_pubkey(self, appid_hash):
        if not self.depth == 1:
            raise Exception('app pubkey should be generated by child key')
        elif not appid_hash:
            raise Exception('required appid_hash to generate app pubkey')
        else:
            child_keyid = self._generateRandomKeyId(appid_hash=appid_hash)
            return self._child_key(child_keyid, include_prikey=False)

    def is_child_key_id(self, keyid, appid_hash=None):
        keyid_L = keyid[:HALF_KEY_ID_LENGTH]
        keyid_R = keyid[HALF_KEY_ID_LENGTH:]

        return keyid_R == self._checksum(keyid_L, appid_hash=appid_hash)

    def print_debug(self):
        print(str(self))

    def __str__(self):
        s = '''is_prikey: {}
depth    : {}
keyid    : {}
prikey   : {}
pubkey   : {}
credid   : {}
ccode    : {}
'''
        return s.format(self.is_prikey, self.depth, self.keyid.hex(), self.prikey.to_string().hex() if self.is_prikey else None, self.pubkey.to_string().hex(), self.credid.hex(), self.ccode.hex())


m_key, m_ccode = prikey_and_ccode('webauthn', 'seed')

master_key = HDKey(keyid=b'0', prikey=m_key, ccode=m_ccode,
                   pubkey=None, is_prikey=True)

print("======== master_key ==========")

master_key.print_debug()


pubkey_seed = master_key.pubkey_seed()


print("======== pubkey_seed ==========")

pubkey_seed.print_debug()

print("======== pubkey ==========")

pubkey = pubkey_seed.app_pubkey(b'aaaaa')

pubkey.print_debug()

print("======== private key ==========")

prikey = master_key.app_prikey(pubkey.credid, b'aaaaa')

keyid = pubkey_seed.keyid + pubkey.keyid

prikey.print_debug()
