# -*- coding: utf-8 -*-
# IDAPython extraction script for IDA Pro 6.6
# Equivalent to extract.idc — same JSON output structure
import sys
import traceback
import os
import json

import idc
import idaapi
import idautils

OUTPUT = "output.json"
DEBUG_LOG = "extract_debug.log"

# Cross-reference type constants (IDA SDK values)
FL_CF = 16  # Call Far
FL_CN = 17  # Call Near
FL_JF = 18  # Jump Far
FL_JN = 19  # Jump Near

# Data xref type constants
try:
    DR_O = idaapi.dr_O
    DR_W = idaapi.dr_W
    DR_R = idaapi.dr_R
    DR_T = idaapi.dr_T
except AttributeError:
    DR_O = 1
    DR_W = 2
    DR_R = 3
    DR_T = 4

_DATA_XREF_TYPE_NAMES = {
    DR_O: "OFFSET",
    DR_W: "WRITE",
    DR_R: "READ",
    DR_T: "TEXT",
}

# Member type names (upper nibble of member flag)
_MEMBER_TYPE_NAMES = {
    0x00000000: "BYTE",
    0x10000000: "WORD",
    0x20000000: "DWORD",
    0x30000000: "QWORD",
    0x40000000: "TBYTE",
    0x50000000: "ASCII",
    0x60000000: "STRUCT",
    0x70000000: "OWORD",
    0x80000000: "FLOAT",
    0x90000000: "DOUBLE",
    0xB0000000: "ALIGN",
    0xC0000000: "CUSTOM",
}

# String type constants (may not be in idc namespace in all IDA versions)
try:
    _ASCSTR_C = idc.ASCSTR_C
    _ASCSTR_UNICODE = idc.ASCSTR_UNICODE
    _ASCSTR_ULEN2 = idc.ASCSTR_ULEN2
    _ASCSTR_ULEN4 = idc.ASCSTR_ULEN4
    _ASCSTR_PASCAL = idc.ASCSTR_PASCAL
    _ASCSTR_LEN2 = idc.ASCSTR_LEN2
    _ASCSTR_LEN4 = idc.ASCSTR_LEN4
except AttributeError:
    _ASCSTR_C = 0
    _ASCSTR_UNICODE = 4
    _ASCSTR_ULEN2 = 5
    _ASCSTR_ULEN4 = 6
    _ASCSTR_PASCAL = 1
    _ASCSTR_LEN2 = 2
    _ASCSTR_LEN4 = 3

_STRING_TYPE_NAMES = {
    _ASCSTR_C: "ASCSTR_C",
    _ASCSTR_UNICODE: "ASCSTR_UNICODE",
    _ASCSTR_ULEN2: "ASCSTR_ULEN2",
    _ASCSTR_ULEN4: "ASCSTR_ULEN4",
    _ASCSTR_PASCAL: "ASCSTR_PASCAL",
    _ASCSTR_LEN2: "ASCSTR_LEN2",
    _ASCSTR_LEN4: "ASCSTR_LEN4",
}


def msg(text):
    idc.Message(text)
    sys.stderr.write(text)
    sys.stderr.flush()


def dbg(text):
    try:
        with open(DEBUG_LOG, "a") as df:
            df.write(text)
            df.flush()
    except:
        pass


def _json_string(s):
    """Escape a string for JSON — returns the quoted representation.
    Handles raw bytes in disassembly that aren't valid UTF-8."""
    if isinstance(s, unicode):
        return json.dumps(s)
    try:
        return json.dumps(s)
    except (UnicodeDecodeError, UnicodeEncodeError):
        return json.dumps(s.decode('latin-1'))


def _json_hex(addr):
    """Format an address as 0x-prefixed hex string (uppercase, 8 digits)."""
    return "0x%08X" % addr


def extract_func(f, ea):
    """Extract a single function. Returns True on success, False on failure."""
    func_name = idc.GetFunctionName(ea)
    func_end = idc.FindFuncEnd(ea)
    func_flags = idc.GetFunctionFlags(ea)

    f.write("    {\n")
    f.write("      \"function\": %s,\n" % _json_string(func_name))
    f.write("      \"address\": \"%s\",\n" % _json_hex(ea))
    f.write("      \"size\": %d,\n" % (func_end - ea))

    # ---- flags ----
    flags_parts = []
    if func_flags & idc.FUNC_LIB:
        flags_parts.append("LIBRARY")
    if func_flags & idc.FUNC_THUNK:
        flags_parts.append("THUNK")
    if func_flags & idc.FUNC_FAR:
        flags_parts.append("FAR")
    if func_flags & idc.FUNC_FRAME:
        flags_parts.append("FRAME")
    f.write("      \"flags\": %s,\n" % _json_string("|".join(flags_parts)))

    # ---- xrefs (who calls/jumps to this function) ----
    xref_parts = []
    for xref in idautils.XrefsTo(ea):
        if xref.type in (FL_CF, FL_CN, FL_JF, FL_JN):
            caller_name = idc.GetFunctionName(xref.frm)
            if caller_name != "":
                if xref.type == FL_CN or xref.type == FL_CF:
                    type_str = "CALL"
                elif xref.type == FL_JN or xref.type == FL_JF:
                    type_str = "JUMP"
                else:
                    type_str = "UNKNOWN"
                xref_parts.append("%s:%s" % (caller_name, type_str))
    f.write("      \"xrefs\": %s,\n" % _json_string(" ".join(xref_parts)))
    f.flush()

    # ---- assembly ----
    asm_lines = []
    curr_inst = ea
    last_inst = 0
    while curr_inst < func_end:
        if curr_inst == last_inst:
            break
        last_inst = curr_inst
        disasm = idc.GetDisasm(curr_inst)
        if disasm:
            # Strip all inline annotations appended by IDA
            # e.g. "; block" (label), "; '@'" (char), "; UINT" (type)
            while True:
                idx = disasm.rfind("; ")
                if idx < 0:
                    break
                after = disasm[idx + 2:].strip()
                if after and (
                    after.startswith("'") or not after.startswith("0x")
                ):
                    disasm = disasm[:idx].rstrip()
                else:
                    break
            asm_lines.append("%s  %s" % (_json_hex(curr_inst), disasm))
        curr_inst = curr_inst + idc.ItemSize(curr_inst)
    f.write("      \"assembly\": %s,\n" % _json_string("\n".join(asm_lines)))
    f.flush()

    # ---- labels (named data/code references from this function) ----
    label_parts = []
    curr_inst = ea
    last_inst = 0
    while curr_inst < func_end:
        if curr_inst == last_inst:
            break
        last_inst = curr_inst
        for xref in idautils.XrefsFrom(curr_inst):
            if xref.type in (FL_CF, FL_CN, FL_JF, FL_JN):
                label_name = idc.Name(xref.to)
                if label_name != "":
                    if xref.type == FL_CN or xref.type == FL_CF:
                        type_str = "CALL"
                    elif xref.type == FL_JN or xref.type == FL_JF:
                        type_str = "JUMP"
                    else:
                        type_str = "UNKNOWN"
                    label_parts.append(
                        "%s:%s:%s" % (label_name, _json_hex(xref.to), type_str)
                    )
        curr_inst = curr_inst + idc.ItemSize(curr_inst)
    f.write("      \"labels\": %s,\n" % _json_string(" ".join(label_parts)))

    # ---- stack frame ----
    frame_members = []
    try:
        frame_id = idc.GetFrame(ea)
        if frame_id and frame_id != idc.BADADDR:
            fsize = idc.GetStrucSize(frame_id)
            if fsize > 4:
                offset = idc.GetFirstMember(frame_id)
                if offset != idc.BADADDR and offset != -1:
                    current = offset
                    safety = 0
                    while current < fsize and safety < 100:
                        safety += 1
                        mname = idc.GetMemberName(frame_id, current)
                        msize = idc.GetMemberSize(frame_id, current)
                        if msize <= 0:
                            msize = 4
                        frame_members.append({
                            "offset": current,
                            "size": msize,
                            "name": mname if mname else ""
                        })
                        current += msize
    except:
        pass
    f.write("      \"stackFrame\": %s,\n" % _json_string(json.dumps(frame_members)))

    # ---- comments ----
    comment_entries = []
    curr_inst = ea
    last_inst = 0
    while curr_inst < func_end:
        if curr_inst == last_inst:
            break
        last_inst = curr_inst
        cmt = idc.GetCommentEx(curr_inst, 0)
        if cmt and cmt.strip():
            comment_entries.append({
                "address": _json_hex(curr_inst),
                "text": cmt,
                "repeatable": False
            })
        rpt = idc.GetCommentEx(curr_inst, 1)
        if rpt and rpt.strip():
            comment_entries.append({
                "address": _json_hex(curr_inst),
                "text": rpt,
                "repeatable": True
            })
        curr_inst = curr_inst + idc.ItemSize(curr_inst)
    f.write("      \"comments\": %s,\n" % _json_string(json.dumps(comment_entries)))

    # ---- cfg blocks ----
    cfg_blocks = []
    try:
        f_obj = idaapi.get_func(ea)
        if f_obj:
            fc = idaapi.FlowChart(f_obj, flags=idaapi.FC_PREDS)
            for block in fc:
                cfg_blocks.append({
                    "id": block.id,
                    "start": _json_hex(block.startEA),
                    "end": _json_hex(block.endEA),
                    "type": block.type,
                    "succs": [s.id for s in block.succs()],
                    "preds": [p.id for p in block.preds()]
                })
    except:
        pass
    f.write("      \"cfgBlocks\": %s,\n" % _json_string(json.dumps(cfg_blocks)))

    # ---- signature ----
    sig = idc.GetType(ea)
    if not sig:
        sig = idc.GuessType(ea)
    f.write("      \"signature\": %s\n" % _json_string(sig if sig else ""))
    f.write("    }")
    f.flush()
    return True


def extract(f):
    entry_ea = idc.MinEA()
    entry_name = idc.Name(entry_ea)
    if entry_name == "":
        entry_name = "entry_point"
    f.write(
        "{\n  \"entry_point\": \"%s:%s\",\n" % (_json_hex(entry_ea), entry_name)
    )

    # ---- segments ----
    f.write("  \"segments\": [\n")
    segs = []
    seg = idc.FirstSeg()
    while seg != idc.BADADDR:
        segs.append(
            "    {\"name\": %s, \"start\": \"%s\", \"end\": \"%s\"}"
            % (
                _json_string(idc.SegName(seg)),
                _json_hex(idc.SegStart(seg)),
                _json_hex(idc.SegEnd(seg)),
            )
        )
        seg = idc.NextSeg(seg)
    f.write(",\n".join(segs))
    f.write("\n  ],\n")
    f.flush()

    # ---- strings ----
    f.write("  \"strings\": [\n")
    f.flush()
    try:
        strings = list(idautils.Strings())
    except Exception as e:
        dbg("idautils.Strings() failed: %s\n" % str(e))
        strings = []
    string_count = 0
    for i, s in enumerate(strings):
        try:
            str_value = str(s)
            if len(str_value) < 2:
                continue
            str_type_id = idc.GetStringType(s.ea)
            if str_type_id is None:
                str_type_id = -1
            encoding = _STRING_TYPE_NAMES.get(str_type_id, "UNKNOWN")
            ref_parts = []
            for xref in idautils.XrefsTo(s.ea):
                func_name = idc.GetFunctionName(xref.frm)
                if func_name:
                    ref_parts.append(func_name)
            referenced_by = " ".join(ref_parts)
            if string_count > 0:
                f.write(",\n")
            f.write(
                "    {\"address\": \"%s\", \"value\": %s, \"encoding\": \"%s\", \"referencedBy\": %s}"
                % (_json_hex(s.ea), _json_string(str_value), encoding, _json_string(referenced_by))
            )
            string_count += 1
        except Exception as e:
            dbg("String at %s failed: %s\n" % (_json_hex(s.ea), str(e)))
    f.write("\n  ],\n")
    f.flush()
    msg("[extract.py] Extraidas %d strings.\n" % string_count)

    # ---- imports ----
    f.write("  \"imports\": [\n")
    f.flush()
    import_entries = []
    try:
        n_modules = idaapi.get_import_module_qty()
        for i in range(n_modules):
            dll = idaapi.get_import_module_name(i)
            if not dll:
                continue
            by_name = {}
            def collect(ea, name, ordinal):
                key = "%s:%s" % (dll, name)
                if key not in by_name or ordinal > by_name[key][2]:
                    by_name[key] = (ea, name, ordinal)
                return True
            idaapi.enum_import_names(i, collect)
            for ea, name, ordinal in by_name.values():
                import_entries.append({
                    "dll": dll,
                    "name": name,
                    "address": _json_hex(ea),
                    "ordinal": ordinal
                })
    except Exception as e:
        dbg("Imports extraction failed: %s\n" % str(e))
    for j, entry in enumerate(import_entries):
        if j > 0:
            f.write(",\n")
        f.write(
            "    {\"dll\": %s, \"name\": %s, \"address\": \"%s\", \"ordinal\": %d}"
            % (_json_string(entry["dll"]), _json_string(entry["name"]), entry["address"], entry["ordinal"])
        )
    f.write("\n  ],\n")
    f.flush()
    msg("[extract.py] Extraidos %d imports.\n" % len(import_entries))

    # ---- global variables ----
    f.write("  \"globalVars\": [\n")
    f.flush()
    global_var_entries = []
    try:
        for ea, name in idautils.Names():
            try:
                flags = idc.GetFlags(ea)
                if not idc.isData(flags):
                    continue
                seg_name = idc.SegName(ea)
                size = idc.ItemSize(ea)
                dxrefs = []
                for xref in idautils.XrefsTo(ea):
                    if xref.type in (DR_O, DR_W, DR_R, DR_T):
                        func_name = idc.GetFunctionName(xref.frm)
                        dxrefs.append({
                            "function": func_name if func_name else "",
                            "address": _json_hex(xref.frm),
                            "type": _DATA_XREF_TYPE_NAMES.get(xref.type, "UNKNOWN")
                        })
                global_var_entries.append({
                    "name": name,
                    "address": _json_hex(ea),
                    "segment": seg_name,
                    "size": size,
                    "dataXrefs": json.dumps(dxrefs)
                })
            except:
                pass
    except Exception as e:
        dbg("Global vars extraction failed: %s\n" % str(e))
    for j, entry in enumerate(global_var_entries):
        if j > 0:
            f.write(",\n")
        f.write(
            "    {\"name\": %s, \"address\": \"%s\", \"segment\": %s, \"size\": %d, \"dataXrefs\": %s}"
            % (_json_string(entry["name"]), entry["address"], _json_string(entry["segment"]),
               entry["size"], _json_string(entry["dataXrefs"]))
        )
    f.write("\n  ],\n")
    f.flush()
    msg("[extract.py] Extraidas %d global vars.\n" % len(global_var_entries))

    # ---- structures ----
    f.write("  \"structures\": [\n")
    f.flush()
    struct_entries = []
    try:
        idx = idc.GetFirstStrucIdx()
        while idx != idc.BADADDR:
            try:
                sid = idc.GetStrucId(idx)
                name = idc.GetStrucName(sid)
                size = idc.GetStrucSize(sid)
                is_union = idc.IsUnion(sid)
                members = []
                off = idc.GetFirstMember(sid)
                while off != idc.BADADDR and off != -1:
                    mname = idc.GetMemberName(sid, off)
                    msize = idc.GetMemberSize(sid, off)
                    if msize is None or msize <= 0:
                        off = idc.GetStrucNextOff(sid, off)
                        continue
                    flag = idc.GetMemberFlag(sid, off)
                    if (flag & 0xF0000000) == 0x60000000:
                        nested_sid = idc.GetMemberStrId(sid, off)
                        mtype = idc.GetStrucName(nested_sid) if nested_sid != idc.BADADDR else "STRUCT"
                    else:
                        mtype = _MEMBER_TYPE_NAMES.get(flag & 0xF0000000, "UNKNOWN")
                    members.append({
                        "offset": off,
                        "name": mname if mname else "",
                        "size": msize,
                        "type": mtype
                    })
                    off = idc.GetStrucNextOff(sid, off)
                struct_entries.append({
                    "name": name,
                    "size": size,
                    "isUnion": is_union,
                    "members": json.dumps(members)
                })
            except:
                pass
            idx = idc.GetNextStrucIdx(idx)
    except Exception as e:
        dbg("Structures extraction failed: %s\n" % str(e))
    for j, entry in enumerate(struct_entries):
        if j > 0:
            f.write(",\n")
        f.write(
            "    {\"name\": %s, \"size\": %d, \"isUnion\": %s, \"members\": %s}"
            % (_json_string(entry["name"]), entry["size"],
               "true" if entry["isUnion"] else "false", _json_string(entry["members"]))
        )
    f.write("\n  ],\n")
    f.flush()
    msg("[extract.py] Extraidas %d structures.\n" % len(struct_entries))

    # ---- enums ----
    f.write("  \"enums\": [\n")
    f.flush()
    enum_entries = []
    try:
        n_enums = idc.GetEnumQty()
        for i in range(n_enums):
            try:
                eid = idc.GetnEnum(i)
                name = idc.GetEnumName(eid)
                width = idc.GetEnumWidth(eid)
                constants = []
                val = idc.GetFirstConst(eid, -1)
                while val != idc.BADADDR:
                    cid = idc.GetConstEx(eid, val, 0, -1)
                    cname = idc.GetConstName(cid)
                    constants.append({
                        "name": cname if cname else "",
                        "value": val
                    })
                    val = idc.GetNextConst(eid, val, -1)
                enum_entries.append({
                    "name": name,
                    "width": width,
                    "constants": json.dumps(constants)
                })
            except:
                pass
    except Exception as e:
        dbg("Enums extraction failed: %s\n" % str(e))
    for j, entry in enumerate(enum_entries):
        if j > 0:
            f.write(",\n")
        f.write(
            "    {\"name\": %s, \"width\": %d, \"constants\": %s}"
            % (_json_string(entry["name"]), entry["width"], _json_string(entry["constants"]))
        )
    f.write("\n  ],\n")
    f.flush()
    msg("[extract.py] Extraidos %d enums.\n" % len(enum_entries))

    # ---- functions ----
    f.write("  \"functions\": [\n")
    f.flush()

    ea = idc.NextFunction(0)
    func_count = 0
    error_count = 0

    while ea != idc.BADADDR:
        if func_count > 0:
            f.write(",\n")
            f.flush()

        try:
            if extract_func(f, ea):
                func_count += 1
            else:
                error_count += 1
                msg("[extract.py] Aviso: falha ao extrair funcao em %s\n"
                    % _json_hex(ea))
        except Exception as e:
            error_count += 1
            tb_str = traceback.format_exc().strip().replace("\n", "\\n")
            msg("[extract.py] Erro na funcao %s: %s\n"
                % (_json_hex(ea), str(e)))
            # Write a placeholder with the error info so JSON remains valid
            f.write(
                "    {\"function\": \"__ERROR__\","
                " \"address\": \"%s\","
                " \"error\": %s}"
                % (_json_hex(ea), _json_string(tb_str))
            )
            f.flush()

        ea = idc.NextFunction(ea)

    msg("[extract.py] Extraidas %d funcoes, %d erros.\n"
        % (func_count, error_count))

    f.write("\n  ]\n}\n")
    f.flush()


def main():
    msg("[extract.py] Script iniciado. CWD=%s\n" % os.getcwd())

    idaapi.autoWait()
    msg("[extract.py] Analise concluida. Extraindo...\n")

    try:
        # Line-buffered — flushes after every newline, prevents data loss
        f = open(OUTPUT, "w", 1)
    except Exception as e:
        msg("[extract.py] ERRO ao criar %s: %s\n" % (OUTPUT, str(e)))
        idc.Exit(1)
        return

    try:
        extract(f)
    except Exception as e:
        msg("[extract.py] ERRO na extracao: %s\n" % str(e))
        msg(traceback.format_exc() + "\n")
        try:
            dbg("FATAL in extract(): %s\n%s\n" % (str(e), traceback.format_exc()))
        except:
            pass
    finally:
        f.close()

    if os.path.exists(OUTPUT):
        msg("[extract.py] %s gerado com sucesso (%d bytes)!\n" %
            (OUTPUT, os.path.getsize(OUTPUT)))
    else:
        msg("[extract.py] FALHA: %s nao foi criado.\n" % OUTPUT)

    idc.Exit(0)


if __name__ == "__main__":
    main()
