from filetypes.base import *
import re
import struct
from filetypes.PE import Resource, PEAnalyzer

def convert_signature_re(r):
    last_num=None
    reg=b""
    for c in r.upper():
        if c in "0123456789ABCDEF":
            if not last_num:
                last_num=c
            else:
                c="{}{}".format(last_num,c)
                reg+=re.escape(bytes.fromhex(c))
                last_num = None
        else:
            reg+=c.encode("ascii")
    return reg

def unpack_script_v1(self, vfile, password=None):
    from transforms.stream import RC4
    from transforms.arithmetic import Neg
    if not vfile.path in self.bat2exe_scripts:
        raise KeyError("Unknown script")
    rsrc, rc4key = self.bat2exe_scripts[vfile.path]
    rsrc_data = self.read(self.rva2off(rsrc.rva), rsrc.size)
    rsrc_data = Neg().run(rsrc_data)
    rsrc_data = RC4().run(rsrc_data, key=rc4key)
    rsrc_data = unpack_zlb1(rsrc_data)
    return rsrc_data


def unpack_script_v2(self, vfile, password=None):
    from transforms.stream import RC4
    if not vfile.path in self.bat2exe_scripts:
        raise KeyError("Unknown script")
    rsrc, rc4key = self.bat2exe_scripts[vfile.path]
    rsrc_data = self.read(self.rva2off(rsrc.rva), rsrc.size)
    rsrc_data = RC4().run(rsrc_data, key=rc4key)
    return rsrc_data


def unpack_zlb1(data):
    from transforms.compress import Zlb1Decompress
    if len(data) < 4:
        return data
    size_unpack, = struct.unpack("<I", data[:4])
    if size_unpack:
        data = Zlb1Decompress().run(data[4:])
        try:
            data = data.decode("latin1").encode("utf8")
        except: pass
        return data
    else:
        return data[4:]
    


def parse_bat2exe(pe):

    num_rcdata = 0
    for rsrc_name in pe.resources.keys():
        if rsrc_name.startswith("RCDATA/"):
            num_rcdata += 1
    if num_rcdata <= 1:
        return

    import hashlib
    import datetime
    from transforms.stream import RC4
    from transforms.arithmetic import Neg

    BAT2EXE_VERSION2_X86 = re.compile(convert_signature_re("680000000068001000006800000000E8....A3....E8....B8(....)A3"), re.DOTALL)
    BAT2EXE_VERSION2_X64 = re.compile(convert_signature_re("488905....4D31C048C7C2001000004831C9E8....488905....48B8(........)488905...."), re.DOTALL)
                                                            
    epcode = pe.read(pe.rva2off(pe["OptionalHeader"]["AddressOfEntryPoint"]), 1024)
    if pe.is64:
        matched_ep = BAT2EXE_VERSION2_X64.search(epcode)
    else:
        matched_ep = BAT2EXE_VERSION2_X86.search(epcode)
    if not matched_ep:
        return

    pe.bat2exe_scripts = {}
    # TODO: better API
    from types import MethodType
    pe.bat2exe_unpack_script_v1 = MethodType(unpack_script_v1, pe)
    pe.bat2exe_unpack_script_v2 = MethodType(unpack_script_v2, pe)

    if len(matched_ep.group(1)) == 4:
        offset, = struct.unpack("<I", matched_ep.group(1))
    else:
        offset, = struct.unpack("<Q", matched_ep.group(1))

    subkey_v1 = hashlib.md5(pe.read(pe.va2off(offset), 20)).hexdigest().lower().encode("ascii")
    subkey_v2 = hashlib.md5(pe.read(pe.va2off(offset), 8)).hexdigest().upper().encode("ascii")

    rsrc_name_script_v1 = "RCDATA/" + hashlib.md5(subkey_v1+b"B").hexdigest().upper() + "/unk"
    rsrc_name_filename_v1 = "RCDATA/" + hashlib.md5(subkey_v1+b"N").hexdigest().upper() + "/unk"
    rsrc_name_options_v1 = "RCDATA/" + hashlib.md5(subkey_v1+b"O").hexdigest().upper() + "/unk"

    rsrc_name_script_v2 = "RCDATA/" + hashlib.md5(subkey_v2).hexdigest().upper() + "/unk"
    rsrc_name_filename_v2 = "RCDATA/" + hashlib.sha1(subkey_v2).hexdigest().upper() + "/unk"
    rsrc_name_options_v2 = "RCDATA/" + hashlib.md5(hashlib.md5(subkey_v2).hexdigest().upper().encode("ascii")).hexdigest().upper() + "/unk"


    if rsrc_name_script_v1 in pe.resources:
        # V1
        if rsrc_name_filename_v1 in pe.resources:
            rsrc = pe.resources[rsrc_name_filename_v1]
            filename_data = pe.read(pe.rva2off(rsrc.rva), rsrc.size)
            filename_data = Neg().run(filename_data)
            filename_data = RC4().run(filename_data, key=subkey_v1)
            filename_data = unpack_zlb1(filename_data)
            zero = filename_data.find(b"\x00")
            if zero >= 0:
                filename_data = filename_data[:zero]
            filename = filename_data.decode("utf8", errors="replace")
        else:
            filename = "<no name>"
        path = f"#BAT2EXE/{filename}"
        pe.bat2exe_scripts[path] = (pe.resources[rsrc_name_script_v1], subkey_v1)
        pe.add_file(path, pe.resources[rsrc_name_script_v1].size, "bat2exe_unpack_script_v1")

    elif rsrc_name_script_v2 in pe.resources:
        # V2
        if pe["PE"]["TimeDateStamp"] <= datetime.datetime(year=2018, month=1, day=20, tzinfo=datetime.timezone.utc):
            rc4key = subkey_v2[0]  # buggy RC4 implementation
        else:
            rc4key = subkey_v2
        if rsrc_name_filename_v2 in pe.resources:
            rsrc = pe.resources[rsrc_name_filename_v2]
            filename_data = pe.read(pe.rva2off(rsrc.rva), rsrc.size)
            filename_data = RC4().run(filename_data, key=rc4key)
            zero = filename_data.find(b"\x00")
            if zero >= 0:
                filename_data = filename_data[:zero]
            filename = filename_data.decode("utf8", errors="replace")
        else:
            filename = "<no name>"
        path = f"#BAT2EXE/{filename}"
        pe.bat2exe_scripts[path] = (pe.resources[rsrc_name_script_v2], rc4key)
        pe.add_file(path, pe.resources[rsrc_name_script_v2].size, "bat2exe_unpack_script_v2")

