#!/usr/bin/env python

from hashlib import sha512
from time import sleep
from uuid import UUID
import binascii
import argparse
import struct
import sys

DUMMY_BYTE = 255
MEMORY_SIZE = 256 * 256 * 256


class FakeSPI(object):
    def xfer(self, data):
        print("--> [%s]" % ", ".join(("%x" % c for c in data)))


def null_output(*args):
    pass

def std_output(o):
    print(o)

def task_read_rsa_key(spi, ca_file, uuid_only=False):
    from fluxmonitor.security import _security

    output = null_output if uuid_only else std_output

    output("=" * 32)
    output(">> Read Serial from SPI OPT Section")
    opt_section = read_OPT(spi)
    sn = opt_section[16:].strip("\x00")
    uuid = UUID(bytes=opt_section[:16])
    model = struct.unpack(">H", opt_section[32:34])
    output("S/N: %s" % sn)
    output("UUID: %s" % uuid)
    output("Model: %s" % repr(model))
    output("-" * 32)


    output(">> Read RSA Key from SPI (Page 31)")
    payload = read(spi, 31)
    l = struct.unpack("<H", payload[:2])[0]

    try:
        rsaobj = _security.RSAObject(der=payload[2:2 + l])
        output(">> Key vaild")
    except:
        output(">> Can not load key:")
        output("length=%i" % l)
        output(repr(payload[2:2 + l]))
        if uuid_only:
            return
        else:
            raise

    output(">> Read Signature from SPI (Page 30)")
    payload = read(spi, 30)
    l = struct.unpack("<H", payload[:2])[0]
    output(">> %s" % binascii.b2a_hex(payload[2:2 + l]))

    if uuid_only:
        print(uuid.hex)
    return 0


def task_write_rsa_key(spi, sn_file):
    if sn_file == "-":
        print("Input S/N document from console...")
        data = ""
        while True:
            buf = sys.stdin.read()
            if buf:
                data += buf
            else:
                break
    else:
        with open(sn_file, "r") as f:
            data = f.read()

    doc = data.split("\n")

    model_id = binascii.a2b_hex(doc[0])
    sn = doc[1]
    uuid = UUID(hex=doc[2])
    pkey_der = binascii.a2b_hex(doc[3])
    sign = binascii.a2b_hex(doc[4])

    print("=" * 32)
    print("S/N: %s" % sn)
    print("UUID: %s" % uuid)
    print("Model ID: %s" % struct.unpack(">H", model_id)[0])
    print("-" * 32)

    opt_section = struct.pack("<16s16s2s30s",
        uuid.bytes, sn.encode(), model_id, "")
    print(">> Write Srrial to SPI OPT Section")
    write_OPT(spi, opt_section)

    payload = struct.pack("<H", len(pkey_der)) + pkey_der
    print(">> Write RSA Key to SPI (Page 31)")
    write(spi, 31, payload)

    payload = struct.pack("<H", len(sign)) + sign
    print(">> Write Signature to SPI (Page 30)")
    write(spi, 30, payload)

    print("=" * 32)

def read_OPT(spi):
    spi.xfer([0xB1])
    buf = _read(spi, 0, 64)
    spi.xfer([0xC1])
    return buf


def write_OPT(spi, data):
    if len(data) != 64:
        raise RuntimeError("OPE len error", len(data))

    # Enter OTP
    spi.xfer([0xB1])

    spi.xfer([0x06])
    payload = [0x02, 0x00, 0x00, 0x00] + [ord(c) for c in data]
    spi.xfer(payload)
    spi.xfer([0x04])
    sleep(0.2)

    # Quit OTP
    spi.xfer([0xC1])

    spi.xfer([0x06])
    spi.xfer([0x2F, 0x03])
    spi.xfer([0x04])


def write(spi, sector, data):
    if not isinstance(sector, int) or sector < 0 or sector > 511:
        raise RuntimeError("sector must be int and < 512 or >= 0")

    if len(data) > 4096:
        raise RuntimeError("Data length must <= 4096")

    data = data + (b"\x00" * (4096 - len(data)))
    position = 4096 * sector
    _earse(spi, position)

    offset = 0
    while offset < 4096:
        _write(spi, position + offset, data[offset:offset + 256])
        sleep(0.1)
        offset += 256


def read(spi, sector):
    if not isinstance(sector, int) or sector < 0 or sector > 511:
        raise RuntimeError("sector must be int and < 512 or >= 0")

    position = 4096 * sector
    buf1 = _read(spi, position, 2048)
    buf2 = _read(spi, position + 2048, 2048)
    return buf1 + buf2


def _earse(spi, position):
    spi.xfer([0x06])
    spi.xfer([
        0x20,
        position // 65536,
        position //256 % 256,
        position % 256
    ])
    spi.xfer([0x04])
    sleep(1.0)


def _read(spi, position, length):
    payload = [
        0x03,
        position // 65536,
        position //256 % 256,
        position % 256
    ] + [DUMMY_BYTE] * length

    output = spi.xfer(payload)[4:]
    buf =  bytearray(length)

    for i in xrange(length):
        buf[i] = output[i]
    return bytes(buf)


def _write(spi, position, data):
    if position + len(data) > MEMORY_SIZE:
        raise RuntimeError("Write data size error")

    offset = position % 256
    if 256 - offset != len(data):
        raise RuntimeError("Data length must == %i" % (256 - offset))

    spi.xfer([0x06])

    payload = [
        0x02,
        position // 65536,
        position //256 % 256,
        position % 256
    ] + [ord(c) for c in data]

    spi.xfer(payload)
    spi.xfer([0x04])


def get_spi():
    import spidev

    # Select device
    spi = spidev.SpiDev(0, 1)

    spi.xfer([0xB9])
    sleep(0.02)
    spi.xfer([0xAB])
    sleep(0.02)
    spi.xfer([0x06])

    # Get information
    ret = spi.xfer([0x90, DUMMY_BYTE, DUMMY_BYTE, 0x00, DUMMY_BYTE, DUMMY_BYTE])
    manufacturer_id = ret[4]
    device_id = ret[5]

    assert manufacturer_id == 0xc2, manufacturer_id
    return spi


def main():
    parser = argparse.ArgumentParser(description='Manage SPI contents')
    parser.add_argument('--dryrun', dest='dryrun', action='store_const',
                        const=True, default=False, help='Dry run')
    mutex_group = parser.add_mutually_exclusive_group()
    mutex_group.add_argument('-w', dest='write', type=str, default=None,
                             help='Write S/N, give S/N file')
    mutex_group.add_argument('-r', dest='read', action='store_const',
                             const=True, default=False, help='Read S/N')
    mutex_group.add_argument('-uuid', dest='read_uuid', action='store_const',
                             const=True, default=False, help='Read S/N')
    mutex_group.add_argument('--ca', dest='ca_file', action='store_const',
                             const=True, default=None, help='CA file')

    options = parser.parse_args()

    if options.dryrun:
        spi = FakeSPI()
    else:
        spi = get_spi()

    if options.write:
        task_write_rsa_key(spi, options.write)
    elif options.read_uuid:
        task_read_rsa_key(spi, options.ca_file, uuid_only=True)
    elif options.read:
        task_read_rsa_key(spi, options.ca_file)
    else:
        print("Select a task")


if __name__ == "__main__":
    main()
