from filetypes.base import *
import malcat
import struct
from struct import unpack_from



def decompile_script(data):
    """
    This is an almost direct adaptation of AutoIt-Ripper (https://github.com/nazywam/AutoIt-Ripper)

    License:
    MIT License
    Copyright (c) 2020 Michał Praszmo

    Permission is hereby granted, free of charge, to any person obtaining a copy
    of this software and associated documentation files (the "Software"), to deal
    in the Software without restriction, including without limitation the rights
    to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
    copies of the Software, and to permit persons to whom the Software is
    furnished to do so, subject to the following conditions:

    The above copyright notice and this permission notice shall be included in all
    copies or substantial portions of the Software.

    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
    SOFTWARE.
    """

    KEYWORDS = (
        "<Dummy>",
        "And",
        "Or",
        "Not",
        "If",
        "Then",
        "Else",
        "ElseIf",
        "EndIf",
        "While",
        "WEnd",
        "Do",
        "Until",
        "For",
        "Next",
        "To",
        "Step",
        "In",
        "ExitLoop",
        "ContinueLoop",
        "Select",
        "Case",
        "EndSelect",
        "Switch",
        "EndSwitch",
        "ContinueCase",
        "Dim",
        "ReDim",
        "Local",
        "Global",
        "Const",
        "Static",
        "Func",
        "EndFunc",
        "Return",
        "Exit",
        "ByRef",
        "With",
        "EndWith",
        "True",
        "False",
        "Default",
        "Null",
        "Volatile",
        "Enum",
    )

    FUNCTIONS = (
        "Abs",
        "ACos",
        "AdlibRegister",
        "AdlibUnRegister",
        "Asc",
        "AscW",
        "ASin",
        "Assign",
        "ATan",
        "AutoItSetOption",
        "AutoItWinGetTitle",
        "AutoItWinSetTitle",
        "Beep",
        "Binary",
        "BinaryLen",
        "BinaryMid",
        "BinaryToString",
        "BitAND",
        "BitNOT",
        "BitOR",
        "BitRotate",
        "BitShift",
        "BitXOR",
        "BlockInput",
        "Break",
        "Call",
        "CDTray",
        "Ceiling",
        "Chr",
        "ChrW",
        "ClipGet",
        "ClipPut",
        "ConsoleRead",
        "ConsoleWrite",
        "ConsoleWriteError",
        "ControlClick",
        "ControlCommand",
        "ControlDisable",
        "ControlEnable",
        "ControlFocus",
        "ControlGetFocus",
        "ControlGetHandle",
        "ControlGetPos",
        "ControlGetText",
        "ControlHide",
        "ControlListView",
        "ControlMove",
        "ControlSend",
        "ControlSetText",
        "ControlShow",
        "ControlTreeView",
        "Cos",
        "Dec",
        "DirCopy",
        "DirCreate",
        "DirGetSize",
        "DirMove",
        "DirRemove",
        "DllCall",
        "DllCallAddress",
        "DllCallbackFree",
        "DllCallbackGetPtr",
        "DllCallbackRegister",
        "DllClose",
        "DllOpen",
        "DllStructCreate",
        "DllStructGetData",
        "DllStructGetPtr",
        "DllStructGetSize",
        "DllStructSetData",
        "DriveGetDrive",
        "DriveGetFileSystem",
        "DriveGetLabel",
        "DriveGetSerial",
        "DriveGetType",
        "DriveMapAdd",
        "DriveMapDel",
        "DriveMapGet",
        "DriveSetLabel",
        "DriveSpaceFree",
        "DriveSpaceTotal",
        "DriveStatus",
        "DUMMYSPEEDTEST",
        "EnvGet",
        "EnvSet",
        "EnvUpdate",
        "Eval",
        "Execute",
        "Exp",
        "FileChangeDir",
        "FileClose",
        "FileCopy",
        "FileCreateNTFSLink",
        "FileCreateShortcut",
        "FileDelete",
        "FileExists",
        "FileFindFirstFile",
        "FileFindNextFile",
        "FileFlush",
        "FileGetAttrib",
        "FileGetEncoding",
        "FileGetLongName",
        "FileGetPos",
        "FileGetShortcut",
        "FileGetShortName",
        "FileGetSize",
        "FileGetTime",
        "FileGetVersion",
        "FileInstall",
        "FileMove",
        "FileOpen",
        "FileOpenDialog",
        "FileRead",
        "FileReadLine",
        "FileReadToArray",
        "FileRecycle",
        "FileRecycleEmpty",
        "FileSaveDialog",
        "FileSelectFolder",
        "FileSetAttrib",
        "FileSetEnd",
        "FileSetPos",
        "FileSetTime",
        "FileWrite",
        "FileWriteLine",
        "Floor",
        "FtpSetProxy",
        "FuncName",
        "GUICreate",
        "GUICtrlCreateAvi",
        "GUICtrlCreateButton",
        "GUICtrlCreateCheckbox",
        "GUICtrlCreateCombo",
        "GUICtrlCreateContextMenu",
        "GUICtrlCreateDate",
        "GUICtrlCreateDummy",
        "GUICtrlCreateEdit",
        "GUICtrlCreateGraphic",
        "GUICtrlCreateGroup",
        "GUICtrlCreateIcon",
        "GUICtrlCreateInput",
        "GUICtrlCreateLabel",
        "GUICtrlCreateList",
        "GUICtrlCreateListView",
        "GUICtrlCreateListViewItem",
        "GUICtrlCreateMenu",
        "GUICtrlCreateMenuItem",
        "GUICtrlCreateMonthCal",
        "GUICtrlCreateObj",
        "GUICtrlCreatePic",
        "GUICtrlCreateProgress",
        "GUICtrlCreateRadio",
        "GUICtrlCreateSlider",
        "GUICtrlCreateTab",
        "GUICtrlCreateTabItem",
        "GUICtrlCreateTreeView",
        "GUICtrlCreateTreeViewItem",
        "GUICtrlCreateUpdown",
        "GUICtrlDelete",
        "GUICtrlGetHandle",
        "GUICtrlGetState",
        "GUICtrlRead",
        "GUICtrlRecvMsg",
        "GUICtrlRegisterListViewSort",
        "GUICtrlSendMsg",
        "GUICtrlSendToDummy",
        "GUICtrlSetBkColor",
        "GUICtrlSetColor",
        "GUICtrlSetCursor",
        "GUICtrlSetData",
        "GUICtrlSetDefBkColor",
        "GUICtrlSetDefColor",
        "GUICtrlSetFont",
        "GUICtrlSetGraphic",
        "GUICtrlSetImage",
        "GUICtrlSetLimit",
        "GUICtrlSetOnEvent",
        "GUICtrlSetPos",
        "GUICtrlSetResizing",
        "GUICtrlSetState",
        "GUICtrlSetStyle",
        "GUICtrlSetTip",
        "GUIDelete",
        "GUIGetCursorInfo",
        "GUIGetMsg",
        "GUIGetStyle",
        "GUIRegisterMsg",
        "GUISetAccelerators",
        "GUISetBkColor",
        "GUISetCoord",
        "GUISetCursor",
        "GUISetFont",
        "GUISetHelp",
        "GUISetIcon",
        "GUISetOnEvent",
        "GUISetState",
        "GUISetStyle",
        "GUIStartGroup",
        "GUISwitch",
        "Hex",
        "HotKeySet",
        "HttpSetProxy",
        "HttpSetUserAgent",
        "HWnd",
        "InetClose",
        "InetGet",
        "InetGetInfo",
        "InetGetSize",
        "InetRead",
        "IniDelete",
        "IniRead",
        "IniReadSection",
        "IniReadSectionNames",
        "IniRenameSection",
        "IniWrite",
        "IniWriteSection",
        "InputBox",
        "Int",
        "IsAdmin",
        "IsArray",
        "IsBinary",
        "IsBool",
        "IsDeclared",
        "IsDllStruct",
        "IsFloat",
        "IsFunc",
        "IsHWnd",
        "IsInt",
        "IsKeyword",
        "IsMap",
        "IsNumber",
        "IsObj",
        "IsPtr",
        "IsString",
        "Log",
        "MapAppend",
        "MapExists",
        "MapKeys",
        "MapRemove",
        "MemGetStats",
        "Mod",
        "MouseClick",
        "MouseClickDrag",
        "MouseDown",
        "MouseGetCursor",
        "MouseGetPos",
        "MouseMove",
        "MouseUp",
        "MouseWheel",
        "MsgBox",
        "Number",
        "ObjCreate",
        "ObjCreateInterface",
        "ObjEvent",
        "ObjGet",
        "ObjName",
        "OnAutoItExitRegister",
        "OnAutoItExitUnRegister",
        "Opt",
        "Ping",
        "PixelChecksum",
        "PixelGetColor",
        "PixelSearch",
        "ProcessClose",
        "ProcessExists",
        "ProcessGetStats",
        "ProcessList",
        "ProcessSetPriority",
        "ProcessWait",
        "ProcessWaitClose",
        "ProgressOff",
        "ProgressOn",
        "ProgressSet",
        "Ptr",
        "Random",
        "RegDelete",
        "RegEnumKey",
        "RegEnumVal",
        "RegRead",
        "RegWrite",
        "Round",
        "Run",
        "RunAs",
        "RunAsWait",
        "RunWait",
        "Send",
        "SendKeepActive",
        "SetError",
        "SetExtended",
        "ShellExecute",
        "ShellExecuteWait",
        "Shutdown",
        "Sin",
        "Sleep",
        "SoundPlay",
        "SoundSetWaveVolume",
        "SplashImageOn",
        "SplashOff",
        "SplashTextOn",
        "Sqrt",
        "SRandom",
        "StatusbarGetText",
        "StderrRead",
        "StdinWrite",
        "StdioClose",
        "StdoutRead",
        "String",
        "StringAddCR",
        "StringCompare",
        "StringFormat",
        "StringFromASCIIArray",
        "StringInStr",
        "StringIsAlNum",
        "StringIsAlpha",
        "StringIsASCII",
        "StringIsDigit",
        "StringIsFloat",
        "StringIsInt",
        "StringIsLower",
        "StringIsSpace",
        "StringIsUpper",
        "StringIsXDigit",
        "StringLeft",
        "StringLen",
        "StringLower",
        "StringMid",
        "StringRegExp",
        "StringRegExpReplace",
        "StringReplace",
        "StringReverse",
        "StringRight",
        "StringSplit",
        "StringStripCR",
        "StringStripWS",
        "StringToASCIIArray",
        "StringToBinary",
        "StringTrimLeft",
        "StringTrimRight",
        "StringUpper",
        "Tan",
        "TCPAccept",
        "TCPCloseSocket",
        "TCPConnect",
        "TCPListen",
        "TCPNameToIP",
        "TCPRecv",
        "TCPSend",
        "TCPShutdown",
        "TCPStartup",
        "TimerDiff",
        "TimerInit",
        "ToolTip",
        "TrayCreateItem",
        "TrayCreateMenu",
        "TrayGetMsg",
        "TrayItemDelete",
        "TrayItemGetHandle",
        "TrayItemGetState",
        "TrayItemGetText",
        "TrayItemSetOnEvent",
        "TrayItemSetState",
        "TrayItemSetText",
        "TraySetClick",
        "TraySetIcon",
        "TraySetOnEvent",
        "TraySetPauseIcon",
        "TraySetState",
        "TraySetToolTip",
        "TrayTip",
        "UBound",
        "UDPBind",
        "UDPCloseSocket",
        "UDPOpen",
        "UDPRecv",
        "UDPSend",
        "UDPShutdown",
        "UDPStartup",
        "VarGetType",
        "WinActivate",
        "WinActive",
        "WinClose",
        "WinExists",
        "WinFlash",
        "WinGetCaretPos",
        "WinGetClassList",
        "WinGetClientSize",
        "WinGetHandle",
        "WinGetPos",
        "WinGetProcess",
        "WinGetState",
        "WinGetText",
        "WinGetTitle",
        "WinKill",
        "WinList",
        "WinMenuSelectItem",
        "WinMinimizeAll",
        "WinMinimizeAllUndo",
        "WinMove",
        "WinSetOnTop",
        "WinSetState",
        "WinSetTitle",
        "WinSetTrans",
        "WinWait",
        "WinWaitActive",
        "WinWaitClose",
        "WinWaitNotActive",
    )

    MACROS = (
        "extended",
        "MSEC",
        "SEC",
        "MIN",
        "error",
        "HOUR",
        "MDAY",
        "MON",
        "YEAR",
        "WDAY",
        "YDAY",
        "ProgramFilesDir",
        "CommonFilesDir",
        "MyDocumentsDir",
        "AppDataCommonDir",
        "DesktopCommonDir",
        "DocumentsCommonDir",
        "FavoritesCommonDir",
        "ProgramsCommonDir",
        "StartMenuCommonDir",
        "StartupCommonDir",
        "LocalAppDataDir",
        "AppDataDir",
        "DesktopDir",
        "FavoritesDir",
        "ProgramsDir",
        "StartMenuDir",
        "StartupDir",
        "ComputerName",
        "WindowsDir",
        "SystemDir",
        "SW_HIDE",
        "SW_MINIMIZE",
        "SW_MAXIMIZE",
        "SW_RESTORE",
        "SW_SHOW",
        "SW_SHOWDEFAULT",
        "SW_ENABLE",
        "SW_DISABLE",
        "SW_SHOWMAXIMIZED",
        "SW_SHOWMINIMIZED",
        "SW_SHOWMINNOACTIVE",
        "SW_SHOWNA",
        "SW_SHOWNOACTIVATE",
        "SW_SHOWNORMAL",
        "SW_LOCK",
        "SW_UNLOCK",
        "TrayIconVisible",
        "TrayIconFlashing",
        "ScriptFullPath",
        "ScriptName",
        "ScriptDir",
        "ScriptLineNumber",
        "WorkingDir",
        "OSType",
        "OSVersion",
        "OSBuild",
        "OSServicePack",
        "OSLang",
        "ProcessorArch",
        "OSArch",
        "CPUArch",
        "KBLayout",
        "AutoItVersion",
        "AutoItExe",
        "IPAddress1",
        "IPAddress2",
        "IPAddress3",
        "IPAddress4",
        "CR",
        "LF",
        "CRLF",
        "DesktopWidth",
        "DesktopHeight",
        "DesktopDepth",
        "DesktopRefresh",
        "Compiled",
        "ComSpec",
        "TAB",
        "UserName",
        "TempDir",
        "UserProfileDir",
        "HomeDrive",
        "HomePath",
        "HomeShare",
        "LogonServer",
        "LogonDomain",
        "LogonDNSDomain",
        "InetGetBytesRead ",
        "InetGetActive ",
        "NumParams",
        "HotKeyPressed",
        "AutoItPID",
        "AutoItUnicode",
        "AutoItX64",
        "Unicode ",
        "MUILang",
        "COM_EventObj",
        "GUI_CtrlId",
        "GUI_CtrlHandle",
        "GUI_DragId",
        "GUI_DragFile",
        "GUI_DropId",
        "GUI_WinHandle",
    )

    KEYWORDS_INVERT_CASE = {i.upper(): i for i in KEYWORDS}
    FUNCTIONS_INVERT_CASE = {i.upper(): i for i in FUNCTIONS}
    MACROS_INVERT_CASE = {i.upper(): i for i in MACROS}






    class ByteStream:
        def __init__(self, data: bytes) -> None:
            self._data = data
            self._offset = 0

        def skip_bytes(self, num) -> None:
            self._offset += num

        def get_bytes(self, num) -> bytes:
            if num is None:
                num = len(self._data) - self._offset

            data = self._data[self._offset: self._offset + num]
            self._offset += num
            return data

        def _int(self, len: int, signed: bool) -> int:
            return int.from_bytes(
                bytes=self.get_bytes(len), byteorder="little", signed=signed
            )

        def _int_be(self, len: int, signed: bool) -> int:
            return int.from_bytes(bytes=self.get_bytes(len), byteorder="big", signed=signed)

        def f64(self) -> float:
            return unpack_from("<d", self.get_bytes(8))[0]

        def u8(self) -> int:
            return self._int(1, False)

        def i8(self) -> int:
            return self._int(1, True)

        def u16(self) -> int:
            return self._int(2, False)

        def i16(self) -> int:
            return self._int(2, True)

        def u32(self) -> int:
            return self._int(4, False)

        def u32be(self) -> int:
            return self._int_be(4, False)

        def i32(self) -> int:
            return self._int(4, True)

        def u64(self) -> int:
            return self._int(8, False)

        def i64(self) -> int:
            return self._int(8, True)



    class TokenStream(ByteStream):
        def __init__(self, data: bytes) -> None:
            super().__init__(data)
            self.indent = 0
            self.next_indent = 0
            self.suffix = ""
            self.prefix = ""

        def get_xored_string(self) -> str:
            key = self.u32()
            if key > len(self._data):
                raise Exception("Read xor string out of bounds")

            ret = bytearray(key * 2)
            for i in range(key):
                c = self.u16() ^ key
                ret[i * 2 + 0] = c & 0xFF
                ret[i * 2 + 1] = (c >> 8) & 0xFF
            return ret.decode("utf-16")

        def peek_next_opcode(self) -> int:
            return self._data[self._offset]
        


    @staticmethod
    def xor(data: bytes, key: bytes) -> bytes:
        return bytes(a ^ b for a, b in zip(data, cycle(key)))
 

    def read_keyword_id(stream) -> str:
        keyword_no = stream.i32()
        if keyword_no > len(KEYWORDS):
            raise Exception("Token not found")

        keyword = KEYWORDS[keyword_no]
        apply_keyword_indent(stream, keyword)
        return keyword


    def read_keyword(stream) -> str:
        keyword = KEYWORDS_INVERT_CASE[stream.get_xored_string()]
        apply_keyword_indent(stream, keyword)
        return keyword


    def capitalize_function(data) -> str:
        if data in FUNCTIONS_INVERT_CASE:
            return FUNCTIONS_INVERT_CASE[data]
        return data


    def capitalize_macro(data) -> str:
        if data in MACROS_INVERT_CASE:
            return MACROS_INVERT_CASE[data]
        return data


    def escape_string(string: str) -> str:
        # escape double qutoes
        string = string.replace('"', '""')
        return f'"{string}"'


    def apply_keyword_indent(stream: TokenStream, keyword: str) -> None:
        if keyword in ("While", "Do", "For", "Select", "Switch", "Func", "If"):
            stream.next_indent += 1

        if keyword in ("Case", "Else", "ElseIf"):
            stream.indent -= 1

        if keyword in (
            "WEnd",
            "Until",
            "Next",
            "EndSelect",
            "EndSwitch",
            "EndFunc",
            "EndIf",
        ):
            stream.next_indent -= 1
            stream.indent -= 1

        if keyword in ("Then",):
            if stream.peek_next_opcode() != 0x7F:
                stream.next_indent -= 1

        if keyword in ("Func",):
            stream.prefix = "\n"

        if keyword in ("EndFunc",):
            stream.next_indent = 0
            stream.suffix += "\n"




    OPCODES = {
            # Keyword
            0x00: read_keyword_id,
            # Function
            0x01: lambda x: FUNCTIONS[x.i32()],
            # Numbers
            0x05: lambda x: str(x.u32()),
            0x10: lambda x: str(x.u64()),
            0x20: lambda x: str(x.f64()),
            # Statements
            0x30: read_keyword,
            0x31: lambda x: capitalize_function(x.get_xored_string()),
            0x32: lambda x: "@" + capitalize_macro(x.get_xored_string()),
            0x33: lambda x: "$" + x.get_xored_string(),
            0x34: lambda x: x.get_xored_string(),
            0x35: lambda x: "." + x.get_xored_string(),
            0x36: lambda x: escape_string(x.get_xored_string()),
            0x37: lambda x: x.get_xored_string(),
            # Operators
            0x40: lambda x: ",",
            0x41: lambda x: "=",
            0x42: lambda x: ">",
            0x43: lambda x: "<",
            0x44: lambda x: "<>",
            0x45: lambda x: ">=",
            0x46: lambda x: "<=",
            0x47: lambda x: "(",
            0x48: lambda x: ")",
            0x49: lambda x: "+",
            0x4A: lambda x: "-",
            0x4B: lambda x: "/",
            0x4C: lambda x: "*",
            0x4D: lambda x: "&",
            0x4E: lambda x: "[",
            0x4F: lambda x: "]",
            0x50: lambda x: "==",
            0x51: lambda x: "^",
            0x52: lambda x: "+=",
            0x53: lambda x: "-=",
            0x54: lambda x: "/=",
            0x55: lambda x: "*=",
            0x56: lambda x: "&=",
            0x57: lambda x: "?",
            0x58: lambda x: ":",
        }

    INDENT_STR = "\t"
    res = []
    line_items = []
    curline = 0
    try:
        stream = TokenStream(data)
        numlines = stream.u32()
        while curline < numlines or numlines == 0:
            opcode = stream.u8()
            if opcode in OPCODES:
                line_items.append(OPCODES[opcode](stream))
            elif opcode == 0x7F:
                curline += 1
                final_line = stream.prefix + INDENT_STR * stream.indent + " ".join(line_items) + "\n" + stream.suffix
                stream.suffix = ""
                stream.prefix = ""
                line_items = []
                stream.indent = stream.next_indent
                res.append(final_line)
            else:
                raise Exception(f"Unsupported opcode: {hex(opcode)}")
    except Exception as e:
        res.append(f"\n!AU3 decompilation aborted: {e}")
    return "".join(res)

####################################################################

class BitStream:
    def __init__(self, data: bytes) -> None:
        self.data = "".join(bin(x)[2:].zfill(8) for x in data)
        self.ptr = 0

    def get_bits(self, num: int) -> int:
        if self.ptr + num > len(self.data):
            raise KeyError
        out = int(self.data[self.ptr:self.ptr+num], 2)
        self.ptr += num
        return out


def au3_decompress(data):
    _, uncomp_size = struct.unpack(">II", data[:8])
    if uncomp_size > 100000000:
        raise ValueError("script too big to be decompressed")
    bin_data = BitStream(data[8:])
    out_data = bytearray(uncomp_size)   # small performance hack
    cur_output = 0
    while cur_output < uncomp_size:
        addme = 0
        # version changes...
        if bin_data.get_bits(1):
            out_data[cur_output] = bin_data.get_bits(8)
            cur_output += 1
        else:
            bb = bin_data.get_bits(15)
            bs = bin_data.get_bits(2)
            if bs == 3:
                addme = 3
                bs = bin_data.get_bits(3)
                if bs == 7:
                    addme = 10
                    bs = bin_data.get_bits(5)
                    if bs == 31:
                        addme = 41
                        bs = bin_data.get_bits(8)
                        if bs == 255:
                            addme = 296
                            while True:
                                bs = bin_data.get_bits(8)
                                if bs != 255:
                                    break
                                addme += 255
            bs += 3 + addme
            i = cur_output - bb
            while bs > 0 and cur_output < uncomp_size:
                out_data[cur_output] = out_data[i]
                cur_output += 1
                i += 1
                bs -= 1
    return out_data



# based on https://github.com/dzzie/myaut_contrib
class RanRot:

    def __init__(self, seed):
        self.randbuffer = []
        for i in range(17):
            seed = (seed * 2891336453 + 1) & 0xFFFFFFFF
            self.randbuffer.append(seed)
        self.p1 = 0
        self.p2 = 10
        for i in range(9):
            self.random()

    def random(self):
        def lrotl(n, d):
            return ((n << d) & 0xFFFFFFFF) | (n >> (32 - d))
        x = (lrotl(self.randbuffer[self.p2], 13) + lrotl(self.randbuffer[self.p1], 9)) & 0xFFFFFFFF
        self.randbuffer[self.p1] = x
        if self.p1:
            self.p1 = self.p1 - 1
        else:
            self.p1 = 16
        if self.p2:
            self.p2 = self.p2 - 1
        else:
            self.p2 = 16
        x = ((x << 20 ) & 0xFFFFFFFF) | ((x >> 12) | 0x3FF00000) << 32
        x = struct.unpack("<d", struct.pack("<Q", x))[0] - 1.0
        return x

    def random8(self):
        self.random()
        return min(255, int(self.random() * 256))


def au3_decrypt(data, rr_seed, fmt=None):
    rr = RanRot(rr_seed)
    if fmt is not None:
        data = struct.pack(fmt, data)
    res = bytearray(data)
    for i in range(len(res)):
        r8 = rr.random8()
        res[i] = res[i] ^ r8
    if fmt is not None:
        res, = struct.unpack(fmt, res)
    return res


class Au3FileHeader(Struct):

    def parse(self):
        yield Bytes(16, name="Password", comment="password md5 (encrypted)")
        restype_crypted = yield Bytes(4, name="ResType", comment="resource type (encrypted)")
        restype = au3_decrypt(restype_crypted, 0x18ee)
        if restype != b"FILE":
            raise FatalError("Invalid resource type: {}".format(restype))
        prolog_len = yield UInt32Xor(0xADBC, name="PrologLen", comment="len of prolog string")
        prolog = yield Bytes(prolog_len*2, name="Prolog", comment="script prolog (encrypted)")
        prolog = au3_decrypt(prolog, 0xb33f + prolog_len).decode("utf-16-le", errors="ignore")
        path_len = yield UInt32Xor(0xF820, name="PathLen", comment="len of compiled path string")
        if path_len:
            path = yield Bytes(path_len*2, name="Path", comment="script path (encrypted)")
            path = au3_decrypt(path, 0xf479 + path_len).decode("utf-16-le", errors="ignore")
        compressed = yield UInt8(name="IsCompressed", comment="is the script compressed ?")
        script_size_compressed = yield UInt32Xor(0x87BC, name="CompressedSize", comment="len of compressed script")
        script_size = yield UInt32Xor(0x87BC, name="UncompressedSize", comment="len of uncompressed script")
        yield UInt32Xor(0xa685, name="Crc32", comment="crc32 of plain-text script")
        if script_size_compressed:
            yield UInt32(name="CreationTimeHigh")
            yield UInt32(name="CreationTimeLow")
            yield UInt32(name="ModificationTimeHigh")
            yield UInt32(name="ModificationTimeLow")




class AutoItAnalyzer(FileTypeAnalyzer):
    category = malcat.FileType.ARCHIVE
    name = "AU3"
    regexp = r"\xA3\x48\x4B\xBE\x98\x6C\x4A\xA9\x99\x4C\x53\x0A\x86\xD6\x48\x7DAU3!"

    def __init__(self, *args, **kwargs):
        FileTypeAnalyzer.__init__(self, *args, **kwargs)
        self.filesystem = {}

    def unpack(self, vfile, password=None):
        if vfile.size == 0:
            return b"EMPTY FILE"
        return self.unpack_file(self.filesystem[vfile.path])
       

    def unpack_file(self, hdr):
        data = self.read(hdr.offset + hdr.size, hdr["CompressedSize"])
        decrypted = au3_decrypt(data, 0x2477)
        if decrypted[:4] != b"EA06":
            raise ValueError("Not a valid AU3.3+ script")
        if hdr["IsCompressed"] == 0:
            return decrypted
        return au3_decompress(decrypted)


    def parse(self, hint):
        yield Bytes(24, name="Magic", category=Type.HEADER)
        while True:
            data = self.read(self.tell(), 24)[16:]
            if data == b"\x41\x55\x33\x21\x45\x41\x30\x36":
                self.jump(self.tell() + 16)
                yield Bytes(8, name="MagicEnd", category=Type.HEADER)
                break
            else:
                hdr = yield Au3FileHeader(category=Type.HEADER)
                self.confirm()
                sz = hdr["CompressedSize"] - 16 # pwd len
                if sz > 0:
                    self.add_section(f"file{len(self.filesystem)}", hdr.offset, sz + hdr.size, r=True)
                    name = au3_decrypt(hdr["Path"], 0xf479 + hdr["PathLen"]).decode("utf-16-le", errors="ignore")
                    self.filesystem[name] = hdr
                    ftype = ""
                    if name.endswith(".tok") or name.endswith(".tmp"):
                        ftype = "AU3.Tokens"
                    self.add_file(name, hdr["UncompressedSize"], "unpack", ftype)
                    if hdr["CompressedSize"]:
                        yield Bytes(sz, name="File", category=Type.DATA)




class TokensHeader(Struct):

    def parse(self):
        yield UInt32(name="NumberOfLines")





class AutoItTokensAnalyzer(FileTypeAnalyzer):
    category = malcat.FileType.PROGRAM
    name = "AU3.Tokens"

    def __init__(self, *args, **kwargs):
        FileTypeAnalyzer.__init__(self, *args, **kwargs)
        self.decompiled = ""

    def decompile(self):
        if not self.decompiled:
            data = self.read(0, self.size())
            self.decompiled = decompile_script(data)
        return self.decompiled

    def parse(self, hint):
        hdr = yield TokensHeader(name="Header")
        self.set_architecture(malcat.Architecture.AU3)
        
                

