#!/usr/bin/env python3
import binascii
import struct
import click
import pyblake2
from trezorlib import ed25519raw, ed25519cosi

indexmap = {
    'bootloader': 0,
    'vendorheader': 1,
    'firmware': 2,
}


def header_digest(index, filename):
    data = open(filename, 'rb').read()
    z = bytes(65 * [0x00])
    if index == 'bootloader':
        header = data[:0x03BF] + z
    elif index == 'vendorheader':
        header = data[:-65] + z
    elif index == 'firmware':
        vhdrlen = struct.unpack('<I', data[4:8])[0]
        header = data[vhdrlen:vhdrlen + 0x03BF] + z
    else:
        raise ValueError('Unknown index "%s"' % index)
    return pyblake2.blake2s(header).digest()


@click.group()
def cli():
    pass


@cli.command(help='')
@click.argument('index', type=click.Choice(indexmap.keys()))
@click.argument('filename')
@click.argument('seckeys', nargs=-1)
def sign(index, filename, seckeys):
    # compute header digest
    digest = header_digest(index, filename)
    # collect commits
    pks, Rs = [], []
    for ctr, seckey in enumerate(seckeys):
        sk = binascii.unhexlify(seckey)
        pk = ed25519raw.publickey(sk)
        _, R = ed25519cosi.get_nonce(sk, digest, ctr)
        pks.append(pk)
        Rs.append(R)
    # compute global commit
    global_pk = ed25519cosi.combine_keys(pks)
    global_R = ed25519cosi.combine_keys(Rs)
    # collect signatures
    sigs = []
    for ctr, seckey in enumerate(seckeys):
        sk = binascii.unhexlify(seckey)
        r, _ = ed25519cosi.get_nonce(sk, digest, ctr)
        h = ed25519raw.H(sk)
        b = ed25519raw.b
        a = 2 ** (b - 2) + sum(2 ** i * ed25519raw.bit(h, i) for i in range(3, b - 2))
        S = (r + ed25519raw.Hint(global_R + global_pk + digest) * a) % ed25519raw.l
        sig = ed25519raw.encodeint(S)
        sigs.append(sig)
    # compute global signature
    sig = ed25519cosi.combine_sig(global_R, sigs)
    ed25519raw.checkvalid(sig, digest, global_pk)
    print(binascii.hexlify(sig).decode())


if __name__ == '__main__':
    cli()