#!/usr/bin/env python3
import sys
import struct
import binascii
import hashlib
import ed25519


# loader/firmware headers specification: https://github.com/trezor/trezor-core/blob/master/docs/bootloader.md


# converts 8-bit bitmap to tuple of values
def bitmap_to_tuple(b):
    return tuple([ i + 1 for i in range(8) if b & (1 << i) ])


# converts tuple of values to 8-bit bitmap
def tuple_to_bitmap(t):
    b = 0
    for i in t:
        if i >= 1 and i <= 8:
            b |= (1 << (i - 1))
    return b


def get_sig(data):
    print('Enter privkey: ', end='')
    seckey = binascii.unhexlify(input())
    signkey = ed25519.SigningKey(seckey)
    digest = hashlib.sha256(data).digest()
    sigidx = (1, )
    sig = signkey.sign(digest)
    return sigidx, sig


class LoaderImage:

    def __init__(self, data):
        header = struct.unpack('<4sIIIBBBB171sB64s', data[:256])
        self.magic, self.hdrlen, self.expiry, self.codelen, \
            self.vmajor, self.vminor, self.vpatch, self.vbuild, \
            self.reserved, self.sigidx, self.sig = header
        assert self.magic == b'TRZL'
        assert self.hdrlen == 256
        assert self.codelen + self.hdrlen >= 4 * 1024
        assert self.codelen + self.hdrlen <= 64 * 1024 + 7 * 128 * 1024
        assert (self.codelen + self.hdrlen) % 512 == 0
        assert self.reserved == 171 * b'\x00'
        self.sigidx = bitmap_to_tuple(self.sigidx)
        self.code = data[self.hdrlen:]
        assert len(self.code) == self.codelen

    def print(self):
        print('TREZOR Loader Image')
        print('  * magic   :', self.magic.decode('ascii'))
        print('  * hdrlen  :', self.hdrlen)
        print('  * expiry  :', self.expiry)
        print('  * codelen :', self.codelen)
        print('  * version : %d.%d.%d.%d' % (self.vmajor, self.vminor, self.vpatch, self.vbuild))
        print('  * sigidx  :', self.sigidx)
        print('  * sig     :', binascii.hexlify(self.sig).decode('ascii'))

    def header(self, sig=True):
        header = struct.pack('<4sIIIBBBB171s', \
            self.magic, self.hdrlen, self.expiry, self.codelen, \
            self.vmajor, self.vminor, self.vpatch, self.vbuild, \
            self.reserved)
        if sig:
            sigidx = tuple_to_bitmap(self.sigidx)
            header += struct.pack('<B64s', sigidx, self.sig)
        else:
            header += 65 * b'\x00'
        assert len(header) == self.hdrlen
        return header

    def sign(self):
        header = self.header(sig=False)
        data = header + self.code
        assert len(data) == self.hdrlen + self.codelen
        self.sigidx, self.sig = get_sig(data)

    def write(self, filename):
        with open(filename, 'wb') as f:
            f.write(self.header())
            f.write(self.code)


class VendorHeader:

    def __init__(self, data):
        header = struct.unpack('<4sIIBBBB', data[:16])
        self.magic, self.hdrlen, self.expiry, \
            self.vmajor, self.vminor, \
            self.vsig_m, self.vsig_n = header
        assert self.magic == b'TRZF'
        assert self.vsig_m > 0 and self.vsig_m <= self.vsig_n
        assert self.vsig_n > 0 and self.vsig_n <= 8
        p = 16
        self.vpub = []
        for _ in range(self.vsig_n):
            self.vpub.append(data[p:p + 32])
            p += 32
        self.vstr_len = data[p]
        p += 1
        self.vstr = data[p:p + self.vstr_len]
        p += self.vstr_len
        self.vimg_len, _ = struct.unpack('<H', data[p: p + 2])
        p += 2
        self.vimg = data[p:p + self.vimg_len]
        p += self.vimg_len
        self.sigidx = bitmap_to_tuple(data[p])
        p += 1
        self.sig = data[p:p+64]
        assert len(data) == 4 + 4 + 4 + 1 + 1 + 1 + 1 + \
                            32 * len(self.vpub) + \
                            1 + self.vstr_len + \
                            2 + self.vimg_len + \
                            1 + 64

    def print(self):
        print('TREZOR Vendor Header')
        print('  * magic   :', self.magic.decode('ascii'))
        print('  * hdrlen  :', self.hdrlen)
        print('  * expiry  :', self.expiry)
        print('  * version : %d.%d' % (self.vmajor, self.vminor))
        print('  * scheme  : %d out of %d' % (self.vsig_m, self.vsig_n))
        for i in range(self.vsig_n):
            print('  * vpub #%d :' % (i + 1), binascii.hexlify(self.vpub[i]).decode('ascii'))
        print('  * vstr     :', self.vstr.decode('ascii'))
        print('  * vimg     : (%d bytes)', len(self.vimg))
        print('  * sigidx  :', self.sigidx)
        print('  * sig     :', binascii.hexlify(self.sig).decode('ascii'))

    def header(self, sig=True):
        header = struct.pack('<4sIIBBBB', \
            self.magic, self.hdrlen, self.expiry, \
            self.vmajor, self.vminor, \
            self.vsig_m, self.vsig_n)
        for i in range(self.vsig_n):
            header += self.vpub[i]
        header += struct.pack('<B', self.vstr_len) + self.vstr
        header += struct.pack('<H', self.vimg_len) + self.vimg
        if sig:
            sigidx = tuple_to_bitmap(self.sigidx)
            header += struct.pack('<B64s', sigidx, self.sig)
        else:
            header += 65 * b'\x00'
        assert len(header) == self.hdrlen
        return header

    def sign(self):
        header = self.header(sig=False)
        self.sigidx, self.sig = get_sig(header)

    def write(self, filename):
        with open(filename, 'wb') as f:
            f.write(self.header())


class FirmwareImage:

    def __init__(self, data):
        header = struct.unpack('<4sIIIBBBB171sB64s', data[:256])
        self.magic, self.hdrlen, self.expiry, self.codelen, \
            self.vmajor, self.vminor, self.vpatch, self.vbuild, \
            self.reserved, self.sigidx, self.sig = header
        assert self.magic == b'TRZF'
        assert self.hdrlen == 256
        assert self.codelen % 4 == 0
        assert self.reserved == 171 * b'\x00'
        self.sigidx = bitmap_to_tuple(self.sigidx)
        self.code = data[self.hdrlen:]
        assert len(self.code) == self.codelen

    def print(self):
        print('TREZOR Firmware Image')
        print('  * magic   :', self.magic.decode('ascii'))
        print('  * hdrlen  :', self.hdrlen)
        print('  * expiry  :', self.expiry)
        print('  * codelen :', self.codelen)
        print('  * version : %d.%d.%d.%d' % (self.vmajor, self.vminor, self.vpatch, self.vbuild))
        print('  * sigidx  :', self.sigidx)
        print('  * sig     :', binascii.hexlify(self.sig).decode('ascii'))

    def header(self, sig=True):
        header = struct.pack('<4sIIIBBBB171s', \
            self.magic, self.hdrlen, self.expiry, self.codelen, \
            self.vmajor, self.vminor, self.vpatch, self.vbuild, \
            self.reserved)
        if sig:
            sigidx = tuple_to_bitmap(self.sigidx)
            header += struct.pack('<B64s', sigidx, self.sig)
        else:
            header += 65 * b'\x00'
        assert len(header) == self.hdrlen
        return header

    def sign(self):
        header = self.header(sig=False)
        data = header + self.code
        assert len(data) == self.hdrlen + self.codelen
        self.sigidx, self.sig = get_sig(data)

    def write(self, filename):
        with open(filename, 'wb') as f:
            f.write(self.header())
            f.write(self.code)


def binopen(filename):
    print()
    print('Opening file', filename)
    print()
    data = open(filename, 'rb').read()
    magic = data[:4]
    if magic == b'TRZL':
        return LoaderImage(data)
    if magic == b'TRZV':
        return VendorHeader(data)
    if magic == b'TRZF':
        return FirmwareImage(data)
    raise Exception('Unknown file format')


def main():
    if len(sys.argv) < 2:
        print('Usage: firmwarectl file.bin [-s]')
        return 1
    fn = sys.argv[1]
    sign = len(sys.argv) > 2 and sys.argv[2] == '-s'
    b = binopen(fn)
    b.print()
    if sign:
        print()
        b.sign()
        print()
        b.print()
        b.write(fn)


if __name__ == '__main__':
    main()