#!/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

KeyType='RSA2048'
MagicNum=b'SIMLOCK_MAGIC_S'
KeyPath = "SimlocKeys/"
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":
            if kf == "":
                kf = i
            j = j + 1
        if j > 1:
            kf = ""
    return j, kf

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

def InsertKey(kpath,kid,img):
    try:
        f = open(kpath, 'rb')
    except IOError:
        print("Can\'t find key: \""+kpath+"\"")
    else:
        k = RSA.importKey(f.read())
        f.close
    Mlen = (k.publickey().size()+7)//8
    Nlen = len(long_to_bytes(k.publickey().n))
    Elen = len(long_to_bytes(k.publickey().e))
    img.write(struct.pack("<I", kid))
    img.write(struct.pack("<I", Mlen))
    img.write(struct.pack("7s",bytes(KeyType.encode('utf-8'))))
    img.seek(1,1)
    img.write(struct.pack("<I", Nlen))
    for x in array.array("B", long_to_bytes(k.publickey().n)):
        data_n = struct.pack("<B",x)
        img.write(data_n)
    img.seek(512-Nlen, 1)
    img.write(struct.pack("<I", Elen))
    for x in array.array("B", long_to_bytes(k.publickey().e)):
        data_e = struct.pack("<B",x)
        img.write(data_e)
    img.seek(512-Elen, 1)
    print("Insert key["+hex(kid)+"]:\""+os.path.basename(kpath)+"\" success")
    return 1

def main():
    filename = sys.argv[1]
    FullKeyPath = os.path.join(sys.path[0], KeyPath)
    print("Find the path to the PubKey and os.path.exists is:",FullKeyPath)
    if not os.path.exists(FullKeyPath):
        print("Warning: not find simlock-key dir,FullKeyPath is:",FullKeyPath)
        return 0
    else:
        print("simlockkey_dir: ",FullKeyPath)
    if os.path.exists(filename):
        f1 = open(filename,'rb+')
        bindata = f1.read()
        bin_len = len(bindata)
        offset = get_head_offset(bindata,bin_len)
        if offset == 0:
            print("Can\'t find head magic num")
        else:
            print("Magic Offset is " + str(offset))
            f1.seek(offset+16,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("Warning: not find "+FullKeyPath+KeysDirs[n])
        f1.close
    else:
        print("File is not accessible!!!")

if __name__ == "__main__":
        main()