#!/usr/bin/python
#

import argparse
import json
import os
import re

# Some Data
PREFIX = "gluon-ffho"
VENDORS = {
    '8devices' : '8 Devices',
    'a5' : 'A5',
    'alfa-network' : 'Alfa Network',
    'allnet' : 'Allnet',
    'avm' : 'AVM',
    'buffalo' : 'Buffalo',
    'd-link' : 'D-Link',
    'gl' : 'GL',
    'gl-inet' : 'GL.Inet',
    'lemarker' : 'Lemarker',
    'linksys' : 'Linksys',
    'netgear' : 'Netgear',
    'ocedo' : 'Ocedo',
    'onion' : 'Onion',
    'openmesh' : 'Openmesh',
    'raspberri-pi' : 'Raspberri Pi',
    'tp-link' : 'TP-Link',
    'ubiquiti' : 'Ubiquiti',
    'ubnt' : 'UBNT',
    'vocore' : 'Vocore',
    'wd-my-net' : 'WD My Net',
    'x86' : 'x86',
    'zbt' : 'ZBT',
    'zyxel' : 'Zyxel'
}

BLACKLIST_MODELS = [
    'cpe510-v1.0',
    'cpe510-v1.1',
    'cpe520-v1.1',
    'wbs510-v1.20',
    'bullet-m5',
    'nanostation-m5',
    'nanostation-m5-xw',
    'nanostation-loco-m5',
    'nanostation-loco-m5-xw',
    'rocket-m5',
    'rocket-m5-ti',
    'rocket-m5-xw',
    'unifi-ac-mesh',
    'unifi-ac-mesh-pro',
    'unifiap-outdoor',
    'unifiap-outdoor+'
]

def longestPrefixKeyMatch(string, dictionary):
    maxlen = 0
    retval = None
    for key in dictionary:
        if string.startswith(key):
            if len(key) > maxlen:
                maxlen = len(key)
                retval = key
    return retval

def vendorPP(vendor):
    if vendor in VENDORS:
	return VENDORS[vendor]
    else:
        return vendor

def findFiles(directory):
    retval = {}
    for imtype in ['sysupgrade', 'factory', 'bootloader']:
        searchpath = directory + "/"+imtype+"/"
        if imtype == 'sysupgrade':
            pattern = PREFIX + "-(([0-9]+\.?)+)-(.+)-sysupgrade\.((img|bin|img.gz|ubi|chk|vdi|vmdk|tar))"
        elif imtype == 'bootloader':
            pattern = PREFIX + "-(([0-9]+\.?)+)-(.+)-bootloader\.((img|bin|img.gz|ubi|chk|vdi|vmdk|tar))"
        else:  
            pattern = PREFIX + "-(([0-9]+\.?)+)-(.+)\.((img|bin|img.gz|ubi|chk|vdi|vmdk|tar))"

        for f in os.listdir(searchpath):
            m = re.search(pattern, f)
            if m:
                # extract info from filename
		file = m.group(0)
                firmware = m.group(3)
                imageversion = m.group(1)
                imagetype = m.group(4)

                vendor = 'Unknown'
                model = 'Unknown'
                vendor_key = longestPrefixKeyMatch(firmware, VENDORS)
                if vendor_key:
                    vendor = VENDORS[vendor_key]
                    n = re.search(vendor_key + "-(.+)$", firmware)
                    if n: 
                        model = n.group(1)

                if model in BLACKLIST_MODELS:
                    continue

		vendor = vendorPP(vendor)
                entry = {'vendor' : vendor,
                    'model' : model,
                    'file' : file
                    }
                if vendor not in retval:
                    retval.update({vendor : {}})
                if model not in retval[vendor]:
                    retval[vendor].update({model : {}})
                retval[vendor][model].update({'modelname' : model , imtype : entry})

    return retval

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Build models.json file.')
    parser.add_argument('directory', help='Directory where the firmware files are')
    parser.add_argument('outputfile', help='Write output to file', nargs='?')

    args = parser.parse_args()
    data = findFiles(args.directory)

    if args.outputfile:
        with open(args.outputfile, "w") as fd:
            json.dump(data, fd, sort_keys=True)
    else:
        print json.dumps(data, sort_keys=True)