from Crypto.Cipher import AES
import hashlib
import os
import json
import urllib

#微信小程序文件格式
#文件头
#1字节 一定是190
#4字节 一定是0
#4字节 索引段长度
#4字节 数据段长度
#1字节 一定是237
#4字节 文件总个数
#索引段
#4字节 文件名长度
#N字节 文件名
#4字节 文件在数据段中的位置（相对于header的0偏移，而不是如下数据段的0偏移）
#4字节 文件长度
#数据段

root = "output"
iv = "the iv: 16 bytes"
salt = "saltiest"

def decrypt(wx_package, wxid):
    fsize = os.path.getsize(wx_package)

    f = open(wx_package, "rb")
    f.seek(6)
    wx_header = f.read(1024)
    wx_others = f.read(fsize - 6 - 1024)
    f.close()

    aes_key = hashlib.pbkdf2_hmac('sha1', wxid.encode(), salt.encode(), 1000, 32)
    cipher = AES.new(aes_key, AES.MODE_CBC, iv.encode())
    decrypted_wx_header = cipher.decrypt(wx_header)
    n = decrypted_wx_header[-1]
    if n > 0:
        decrypted_wx_header = decrypted_wx_header[:-n]

    xor_key = ord(str(wxid[-2]))
    decrypted_wx_others = []
    for b in wx_others:
        decrypted_wx_others.append(b ^ xor_key)

    return decrypted_wx_header + bytes(decrypted_wx_others)

def write_file(fname, buf):
    items= fname.split("/")
    path = root
    for i in range(1,len(items)-1):
        path += "/"+items[i]
        md(path)


    f=open(path+"/"+items[-1],"wb")
    f.write(buf)
    f.close()

def process(wx):
    index = 0
    print("magic number is " + str(ord(wx[index:1])))
    index += 1
    print("always o is {0}".format(int.from_bytes(wx[index:index + 4], "big")))
    index += 4
    index_seg_length = int.from_bytes(wx[index:index + 4], "big")
    print("index segments length is {0}".format(index_seg_length))
    index += 4
    body_seg_length = int.from_bytes(wx[index:index + 4], "big")
    print("body segments length is {0}".format(body_seg_length))
    index += 4
    print("last mask must be 237 ---> {0}".format(int.from_bytes(wx[index:index + 1], "big")))
    index += 1
    file_count = int.from_bytes(wx[index:index + 4], "big")
    print("file count is {0}".format(file_count))
    index += 4

    index_length = 0
    for fcount in range(0, file_count):
        if index_length + 4 >= index_seg_length:  # 如果用while true
            break
        filename_length = int.from_bytes(wx[index:index + 4], "big")
        index += 4
        fname = wx[index:index + filename_length].decode()
        index += filename_length
        offset_of_file_in_segment = int.from_bytes(wx[index:index + 4], "big")
        index += 4
        file_size = int.from_bytes(wx[index:index + 4], "big")
        index += 4

        print("File name length ={0}, file name = {1}, offset in segment = {2}, file size = {3}".format(filename_length,
                                                                                                        fname,
                                                                                                        offset_of_file_in_segment,
                                                                                                        file_size))

        write_file(fname,wx[offset_of_file_in_segment:offset_of_file_in_segment+file_size])
        index_length += 4 * 3 + filename_length

def md(dir):
    if os.path.exists(dir) is False:
        os.mkdir(dir)

def process_app_config_json():
    f=open(root+"/app-config.json","r")
    lines = f.readlines()[0]
    f.close()
    all_json = json.loads(lines)

    app_json = json.loads("{}")
    app_json["pages"] = all_json["pages"]
    app_json["window"] = all_json["global"]["window"]
    app_json["tabBar"] = all_json["tabBar"]
    app_json["networkTimeout"] = all_json["networkTimeout"]
    if all_json.get("subPackages") is not None:
        app_json["subPackages"] = all_json["subPackages"]
        for sub_root in app_json["subPackages"]:
            i = 0
            while i<len(app_json["pages"]):
                if app_json["pages"][i].startswith(sub_root["root"]):
                    app_json["pages"].remove(app_json["pages"][i])
                    i = 0
                else:
                    i += 1

    if all_json.get("navigateToMiniProgramAppIdList") is not None:
        app_json["navigateToMiniProgramAppIdList"] = all_json["navigateToMiniProgramAppIdList"]

#if(e.extAppid)
#        wu.save(path.resolve(dir,'ext.json'),JSON.stringify({extEnable:true,extAppid:e.extAppid,ext:e.ext},null,4));
    if all_json.get("debug") is not None:
        app_json["debug"] = all_json["debug"]

    entry_page_path = all_json["entryPagePath"]

    for page in all_json["page"].keys():
        json_file = root+"/"+page.replace(".html",".json")
        with open(json_file,"w") as f:
            f.writelines(urllib.parse.unquote(json.dumps(all_json["page"][page]["window"],indent=4)))

    with open(root+"/app.json","w") as f:
        f.writelines(urllib.parse.unquote(json.dumps(app_json,indent=4)))

    a = lines

def get_package_content(wx_package):
    fsize = os.path.getsize(wx_package)

    f = open(wx_package, "rb")
    buf = f.read(fsize)
    f.close()

    return buf

def main():
    md(root)

    wxid = "wx56d1401d645efccb"
    wxapkg = ["1.wxapkg"]

    for apkg in wxapkg:
        buf = get_package_content(apkg)
        if int(buf[0])!=190 and buf[0:6].decode() == "V1MMWX":
            buf = decrypt(apkg,wxid)
        if int(buf[0])!=190:
            print("Error file format.")
            return

        process(buf)

    process_app_config_json()

if __name__ == "__main__":
    main()