#!/usr/bin/env python
from struct import pack, unpack
import subprocess,os,sys,struct,array
from Crypto.PublicKey import RSA
from Crypto.Util.number import long_to_bytes

KeyMagicNum=b'SIMLOCK_PUBKEY_INFO'
ConfigInfoMagicNum=b'SIMLOCK_CONFIG_INFO'
KeyPath = "SimlocKeysV6/"
ConfigInfoPath = "ConfigInfo/"
KeyIDTable = [0x1011, 0x1012, 0x1013, 0x1014, 0x1015]
KeysDirs = ["authstartkey", "authendkey", "simlockey", "imeikey", "reskey"]

def GetKeyFiles(fdir):
    kf = ""
    fls=os.listdir(fdir)
    j = 0
    for i in fls:
        if os.path.splitext(i)[1] == ".pem":
            #print"Find key: " + i
            if kf == "":
                kf = i
            j = j + 1
        if j > 1:
            kf = ""
    return j, kf

def GetConfigInfoFiles(fdir):
    cf = ""
    fls=os.listdir(fdir)
    j = 0
    for i in fls:
        if os.path.splitext(i)[1] == ".bin":
            #print"Find key: " + i
            if cf == "":
                cf = i
            j = j + 1
        if j > 1:
            cf = ""
    return j, cf

def GetKeyAreaHeadOffset(tos_img,tos_len):
    secpos = tos_img.find(KeyMagicNum)
    if secpos < 0:
        print("there is no head!!!")
        return 0
    else:
        print("Find magic num: " + KeyMagicNum.decode())
        return  secpos

def GetConfigInfoAreaHeadOffset(tos_img,tos_len):
    secpos = tos_img.find(ConfigInfoMagicNum)
    if secpos < 0:
        print("there is no head!!!")
        return 0
    else:
        print("Find magic num: " + ConfigInfoMagicNum.decode())
        return  secpos

def InsertKey(kpath,kid,img):
    ksize = os.path.getsize(kpath)
    print("The length of the file '%s' is %d bytes." % (kpath, ksize))
    try:
        f = open(kpath, 'rb')
    except IOError:
        print("Can\'t find key: \""+kpath+"\"")
    else:
        kpem = f.read()
        f.close
    img.write(struct.pack("<I", kid))
    img.write(struct.pack("<I", ksize))
    img.write(kpem)
    img.seek(3072-ksize, 1)
    print("Insert key["+hex(kid)+"]:\""+os.path.basename(kpath)+"\" success")
    return 1

def InsertConfigInfo(cpath,img):
    csize = os.path.getsize(cpath)
    print("The length of the file '%s' is %d bytes." % (cpath, csize))
    try:
        f = open(cpath, 'rb')
    except IOError:
        print("Can\'t find config info file: \""+cpath+"\"")
    else:
        configfile = f.read()
        f.close
    img.write(struct.pack("<I", csize))
    img.write(configfile)
    img.seek(1024-csize, 1)
    print("Insert simlock config info:\""+os.path.basename(cpath)+"\" success")
    return 1

def main():
    filename = sys.argv[1]
    FullKeyPath = os.path.join(sys.path[0], KeyPath)
    FullConfigInfoPath = os.path.join(sys.path[0], ConfigInfoPath)
    if not os.path.exists(FullKeyPath):
        print("Warning: not find simlock-key dir")
        print("simlockkey_dir: " + FullKeyPath)
        return 0
    else:
        print("simlockkey_dir: " + FullKeyPath)
    if not os.path.exists(FullConfigInfoPath):
        print("Warning: not find Config-info dir")
        return 0
    else:
        print("simlockconfiginfo_dir: " + FullConfigInfoPath)
    if os.path.exists(filename):
        f1 = open(filename,'rb+')
        bindata = f1.read()
        bin_len = len(bindata)
        key_offset = GetKeyAreaHeadOffset(bindata,bin_len)
        config_offet = GetConfigInfoAreaHeadOffset(bindata,bin_len)
        if key_offset == 0 or config_offet == 0:
            print("Can\'t find head magic num")
        else:
            print("Key Magic Offset is " + str(key_offset))
            print("Config Info Magic Offset is " + str(config_offet))
            f1.seek(key_offset+32,0)
            for n in range(0,5):
                if os.path.exists(FullKeyPath+KeysDirs[n]):
                    knum, kfile = GetKeyFiles(FullKeyPath+KeysDirs[n])
                    if knum == 1:
                        InsertKey(FullKeyPath+KeysDirs[n]+"/"+kfile, KeyIDTable[n], f1)
                    else:
                        print("Key file num err, knum: %d" % knum)
                else:
                    print("Warning: not find "+FullKeyPath+KeysDirs[n])
            f1.seek(config_offet+32,0)
            cnum, cfile = GetConfigInfoFiles(FullConfigInfoPath)
            if cnum == 1:
                InsertConfigInfo(FullConfigInfoPath+"/"+cfile, f1)
            else:
                print("ConfigInfo file num err, cnum: %d" % cnum)
        f1.close
    else:
        print("File is not accessible!!!")

if __name__ == "__main__":
        main()