#!/usr/bin/env python3

import os
import sys
import textwrap
import uuid
import hashlib
import binascii
import sqlite3
import time

import nacl.signing
import nacl.encoding

from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend


PRESIG_VERSION = '000.001'

def format_hex(data, indent=4, wrap=True):
    indent = ' '*indent
    par = ', '.join(f'0x{b:02x}' for b in data)
    par = textwrap.fill(par, width=120,
        initial_indent=indent, subsequent_indent=indent,
        replace_whitespace=False, drop_whitespace=False)
    if wrap:
        return f'{{\n{par}\n}}'
    return par

def domain_string(domain_name, value, serial):
    return f'smart reset domain string v{PRESIG_VERSION}: domain:{domain_name}={value}@{serial}'

if __name__ == '__main__':
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument('keyfile', help='Key file to use')
    parser.add_argument('presig_db', nargs='?', help='sqlite3 dbfile for generated presig authorization keys')
    parser.add_argument('-g', '--generate', action='store_true', help='Generate signing keypair')
    parser.add_argument('-v', '--vendor', type=str, default='Darthenschmidt Cyberei und Verschleierungstechnik GmbH', help='Vendor name for vendor domain')
    parser.add_argument('-s', '--series', type=str, default='Frobnicator v0.23.7', help='Series identifier for series domain')
    parser.add_argument('-r', '--region', type=str, default='Neuland', help='Region name for region domain')
    parser.add_argument('-c', '--country', type=str, default='Germany', help='Country name for country domain')
    parser.add_argument('-p', '--start-serial', type=int, default=0, help='First presig serial number to use')
    parser.add_argument('-n', '--presig-count', type=int, default=3, help='Number of presigs to generate')
    parser.add_argument('-i', '--iv', type=str, default='safety reset oob presig iv', help='IV for presig generation')
    args = parser.parse_args()

    if args.generate:
        if os.path.exists(args.keyfile):
            print("Error: keyfile already exists. We won't overwrite it. Instead please remove it manually.",
                    file=sys.stderr)
            sys.exit(1)
        signing_key = nacl.signing.SigningKey.generate()
        with open(args.keyfile, 'wb') as f:
            f.write(signing_key.encode(encoder=nacl.encoding.Base64Encoder))
            f.write(b'\n')
        sys.exit(0)

    with open(args.keyfile, 'r') as f:
        signing_key = nacl.signing.SigningKey(f.read().strip(), encoder=nacl.encoding.Base64Encoder)
        pubkey_bytes = signing_key.verify_key.encode(encoder=nacl.encoding.RawEncoder)
        pubkey_hash = hashlib.sha512(pubkey_bytes).digest()[:16]

    if not args.presig_db:
        print('The presig_db parameter is required.', file=sys.stderr)
        sys.exit(1)

    db = sqlite3.connect(args.presig_db)
    db.execute('CREATE TABLE IF NOT EXISTS presig_authkey (timestamp, pubkey_hash, bundle_id, presig_ver, domain, value, serial, authkey)')

    bundle_id = uuid.uuid4().bytes

    print('#include <stdint.h>')
    print('#include <assert.h>')
    print()
    print('#include "crypto.h"')
    print()

    print(f'/* bundle id {binascii.hexlify(bundle_id).decode()} */')
    print(f'uint8_t presig_bundle_id[16] = {format_hex(bundle_id)};')
    print(f'int presig_first_serial = {args.start_serial};')
    print()

    print(f'uint8_t oob_trigger_pubkey[crypto_sign_PUBLICKEYBYTES] = {format_hex(pubkey_bytes)};')
    print()

    print('uint8_t presig_messages[_TRIGGER_DOMAIN_COUNT][PRESIG_STORE_SIZE][PRESIG_MSG_LEN] = {')
    device_domains = {
            'all': 'all',
            'country': args.country,
            'region': args.region,
            'vendor': args.vendor,
            'series': args.series
            }
    presigs = { dom: [] for dom in device_domains }
    for dom, val in device_domains.items():
        print('    {')
        for i in range(args.presig_count):
            serial = args.start_serial + i
            ds = domain_string(dom, val, serial)
            ds_hash = hashlib.sha512(ds.encode()).digest()[:16]
            presigs[dom].append((ds_hash, val, serial))
            print(f'        {{ /* "{ds}" */')
            print(format_hex(ds_hash, indent=8, wrap=False))
            print(f'        }},')
        print('    },')
    print('};')
    print()

    presig_iv = hashlib.sha512(args.iv.encode()).digest()[:16]
    print(f'uint8_t oob_presig_iv[16] = {{ /* sha512("{args.iv}")[:16] */')
    print(format_hex(presig_iv, wrap=False))
    print(f'}};')
    print()


    print('uint8_t presig_store[_TRIGGER_DOMAIN_COUNT][PRESIG_STORE_SIZE][crypto_sign_BYTES] = {')
    for dom, hashes in presigs.items():
        print(f'    {{ /* domain {dom} */')

        for ds_hash, val, serial in hashes:
            authkey = os.urandom(16)
            cipher = Cipher(algorithms.AES(authkey), modes.CTR(presig_iv), backend=default_backend())
            enc = cipher.encryptor()
            ciphertext = enc.update(ds_hash)
            assert len(enc.finalize()) == 0

            with db:
                db.execute('INSERT INTO presig_authkey VALUES (?, ?, ?, ?, ?, ?, ?, ?)',
                        (int(time.time()*1000), pubkey_hash, binascii.hexlify(bundle_id).decode(), PRESIG_VERSION, dom,
            print(format_hex(ciphertext, indent=8, wrap=False))
            print(f'        }},')

        print(f'    }},')
    print(f'}};')

    print()
    print('static inline void __hack_asserts_only(void) {')
    print(f'    static_assert(_TRIGGER_DOMAIN_COUNT == {len(presigs)});')
    print(f'    static_assert(PRESIG_STORE_SIZE == {args.presig_count});')
    print('}')
    print()