#!/usr/bin/env python

from __future__ import print_function
import os, sys, json, argparse, struct
from Crypto.PublicKey import RSA
from smbus import SMBus
from subprocess import getoutput
from time import sleep

print_error = lambda msg: print(msg, file=sys.stderr)

# load hardware descriptor
with open(os.getenv('MOD_HARDWARE_DESCRIPTOR',
                    '/etc/mod-hardware-descriptor.json'), 'r') as fh:
    hw_desc_data = json.load(fh)

i2c_main_bus = hw_desc_data.get('i2c_main_bus', None)
if i2c_main_bus is None:
    os.environ['SKIP_HW_DESC_INIT'] = "1"
    i2c_main_bus = int(getoutput("source /etc/mod-hardware-descriptor.env && echo ${I2C_MAIN_BUS}"))

class AT24CS08(object):
    def __init__(self, bus):
        self.bus = bus
        self.dev = SMBus(bus)

        # uid mapping
        self.uid_address = 0x58
        self.uid_offset = 0x40
        self.uid_size = 16

        # serial tag mapping
        self.tag_address = 0x50
        self.tag_offset = 0x00

        # openssl key mapping
        self.key_address = 0x50
        self.key_offset = 0x20

    def __read_data(self, address, offset, size):
        n = 16
        data = []
        page_size = 256
        while size > 0:
            if offset + n >= page_size:
                to_read = (page_size - offset)
                data += self.dev.read_i2c_block_data(address, offset, to_read)
                sleep(0.01)
                offset = 0
                address += 1
            else:
                to_read = size if size < n else n
                data += self.dev.read_i2c_block_data(address, offset, to_read)
                sleep(0.01)
                offset += to_read

            size -= to_read

        return data

    def read_uid(self):
        data = self.__read_data(self.uid_address, self.uid_offset, self.uid_size)
        uid = '%02X:'*self.uid_size % tuple(data)
        return uid[:-1]

    def read_tag(self):
        serial_format = getoutput("source /etc/mod-hardware-descriptor.env && echo ${SERIAL_FORMAT}")
        size = self.__read_data(self.tag_address, self.tag_offset, 1)
        tag_size = struct.unpack('<B', bytes(size))[0]
        data = self.__read_data(self.tag_address, self.tag_offset + 1, tag_size)
        if serial_format == 'simplified':
            data = struct.unpack('<2scBcBI', bytes(data))
            tag = '%s%c%02i%c%02i-%05i' % tuple([i.decode('utf-8') if isinstance(i, bytes) else i for i in data])
        else:
            data = struct.unpack('<2scIBBHI', bytes(data))
            tag = '%s%s-%s-%s-%02i-%03i-%04i' % tuple([i.decode('utf-8') if isinstance(i, bytes) else i for i in data])
        return tag

    def read_key(self):
        size = self.__read_data(self.key_address, self.key_offset, 2)
        key_size = struct.unpack('<H', bytes(size))[0]
        data = self.__read_data(self.key_address, self.key_offset + 2, key_size)
        key = RSA.importKey(data).exportKey().decode('utf-8')
        return key

def get_device_path(info):
    # set device information to make the searching
    if hw_desc_data['gpio_platform'] == 'sunxi':
        if 'gpio' in info:
            dev_match = info.lower() + '_'
            dev_search_path = '/sys/devices/platform/gpio-sunxi/gpio'
            dev_path = '/sys/class/gpio/%s'
        elif 'uart' in info:
            dev_match = 'tty'
            dev_search_path = '/sys/devices/platform/sunxi-' + info + '/tty'
            dev_path = '/dev/%s'
        elif 'i2c' in info:
            dev_match = 'i2c'
            dev_search_path = '/sys/devices/platform/sunxi-' + info
            dev_path = '/dev/%s'
        else:
            dev_match = None
            dev_search_path = None
            dev_path = None

        # try to get the device file descriptor
        if dev_search_path:
            try:
                for dev in os.listdir(dev_search_path):
                    if dev_match in dev:
                        return dev_path % dev
            except OSError:
                print_error("%s was configured in mod-hardware-descriptor but it " \
                            "does not exist in the file system" % info)
                exit(1)

    # fallback
    else:
        dev_path = '/dev/i2c-{}'.format(i2c_main_bus)
        return dev_path

    return info

def init():
    try:
        mem = AT24CS08(i2c_main_bus)

        # cache device info (tag and uid first, then key)
        if not os.path.exists('/var/cache/mod'):
            os.makedirs('/var/cache/mod')
        with open('/var/cache/mod/uid', 'w') as fh:
            fh.write(mem.read_uid())
        with open('/var/cache/mod/tag', 'w') as fh:
            fh.write(mem.read_tag())
        with open('/var/cache/mod/key', 'w') as fh:
            fh.write(mem.read_key())

    except:
        print_error('error reading memory')

def main():
    # quick init for fast boot, don't parse arguments
    if len(sys.argv) == 2 and sys.argv[1] == "init":
        return init()

    # get user arguments
    parser = argparse.ArgumentParser(description='query information about MOD hardware')
    parser.add_argument('hw_resource')
    parser.add_argument('resource_info', nargs='?')
    args = parser.parse_args()

    hw_resource = args.hw_resource.lower()
    resource_info = args.resource_info

    # check if it's memory related query
    if hw_resource in ['uid', 'tag', 'key']:
        mem = AT24CS08(i2c_main_bus)

        value = None
        if hw_resource == 'uid': value = mem.read_uid()
        if hw_resource == 'tag': value = mem.read_tag()
        if hw_resource == 'key': value = mem.read_key()

        if value:
            print(value)

        return

    # check if resource is valid
    if not hw_resource in hw_desc_data.keys():
        print_error('%s not found as hardware resource' % args.hw_resource)
        print_error('use one of the following: %s' % list(hw_desc_data.keys()))
        exit(1)

    # check if requested information is valid
    if (isinstance(hw_desc_data[hw_resource], dict) and not resource_info) or \
       (resource_info and not resource_info in hw_desc_data[hw_resource].keys()):
        print_error('%s not found as resource information' % args.resource_info)
        print_error('use one of the following: %s' %
              list(hw_desc_data[hw_resource].keys()))
        exit(1)

    if resource_info:
        info = hw_desc_data[hw_resource][resource_info]
    else:
        info = hw_desc_data[hw_resource]

    # default value to return
    ret = get_device_path(info)

    print(ret)

if __name__ == "__main__":
    main()
