// stallker trace
class InstHelper {

    inited
    supportInst

    // 基本都是涉及pc的，需要不断完善
    static testSupport(inst) {
        if (!InstHelper.inited) {
            InstHelper.inited = true
            InstHelper.supportInst = new Set()
            InstHelper.supportInst.add('add')
            InstHelper.supportInst.add('blx')
            InstHelper.supportInst.add('bl')
        }
        if (!InstHelper.is_pc_relative(inst)) {
            return true
        }
        if (InstHelper.supportInst.has(inst.mnemonic)) {
            return true
        } else {
            return false
        }
    }

   
    static is_pc_relative(inst) {

        if (inst.regsRead.includes('pc')) {
            return true
        }
        if (inst.regsWritten.includes('pc')) {
            return true
        }
        if (inst.opStr.includes('pc')) {
            return true
        }
        if (inst.groups.includes('call')) {
            return true
        }
        if (inst.groups.includes('branch_relative')) {
            return true
        }


        return false

    }

    // thumb or arm or arm64
    static getInstType(inst) {
        if (Process.arch == 'arm64') {
            return 'arm64'
        }
        for (let type of inst.groups) {
            if (type.toLowerCase() == 'thumb') {
                return 'thumb'
            } else if (type.toLowerCase() == 'arm') {
                return 'arm'
            } else {
                /*  */
            }
        }

    }

    // 获得内存操作数的地址,detail为异常中的包含寄存的上下文，或者任何包含了寄存器的上下文都行
    static getMemAddr(memOprand, detail) {
        // console.log(JSON.stringify(memOprand))
        // console.log(JSON.stringify(detail))
        // console.log(JSON.stringify(memOprand),JSON.stringify(detail))

        var base = detail[memOprand['base']]
        var numDisp = memOprand.disp
        // console.log(base, disp)
        // memAddr = BinSupport.hexAdd(base, disp,false,true)

        base = BigInt(base)
        var disp = BigInt(numDisp)
        var memAddr = base + disp

        // memAddr = BigInt(base.toString())+numDisp
        // console.log('memAddr', memAddr)
        return '0x' + memAddr.toString(16)
    }

    // 模拟执行，执行完之后会把结果修正到detail中
    static simulateInst(inst, detail) {
        console.log('simulateInst enter')
        switch (inst.mnemonic) {
            case 'add':
                console.log('simulateInst', inst.mnemonic)
                var tempRes = 0
                tempRes = tempRes.toString(16)
                for (let i = 1; i < inst.operands.length; i++) {
                    console.log(JSON.stringify(inst.operands[i]))
                    // 寄存器操作数
                    if (build_re('reg').test(inst.operands[i].type)) {
                        tempRes = BinSupport.hexAdd(tempRes, detail[inst.operands[i].value])
                    } else if (build_re('imm').test(inst.operands[i].type)) {
                        tempRes = BinSupport.hexAdd(tempRes, inst.operands[i].value.toString(16))
                    } else if (build_re('mem').test(inst.operands[i].type)) {
                        var addr = InstHelper.getMemAddr(inst.operands[i].value)
                        tempRes = BinSupport.hexAdd(tempRes, addr.readPointer())
                    } else {
                        console.log('simulateInst error', JSON.stringify(inst))
                    }
                }
                detail.context[inst.operands[0].value] = tempRes
                break
            case 'blx':
                console.log( inst.mnemonic)
                var uNextPc = ExceptionHook.getNextPc(parseInt(detail.address, 16))
                detail.context.pc = uNextPc
             
                console.log('detail.context.pc', detail.context.pc)
                break
            case 'bl':
                console.log( inst.mnemonic)
                var uNextPc = ExceptionHook.getNextPc(parseInt(detail.address, 16))
                console.log('uNextPc.toString(16)', uNextPc.toString(16))
                detail.context.pc = uNextPc
                // detail.context.pc = 0
                detail.context.lr = Number(detail.address)+4
                console.log('detail.context.pc', detail.context.pc)
                break
            default:
                console.log('暂未支持', inst.mnemonic)
                break
        }
    }

    static test() {
        var addr = Memory.alloc(4).writeU32(0xaf03)
        var inst = Instruction.parse(addr.add(1))
        var detail = { 'r7': '0xa23', 'sp': '0xa23' }
        console.log(JSON.stringify(inst))
        console.log(JSON.stringify(detail))
        InstHelper.simulateInst(inst, detail)
        console.log(JSON.stringify(detail))
    }


}


class Trace {
    constructor() {
        this.moduleBase = null;
        this.isFirstIn = true;
        this.pre_context_jstr = null;
        this.infoMap = new Map();
        this.detailInsMap = new Map();
        this.lastOffset = undefined;
        this.currentIndex = 0;
        this.byteToHex = [];
        this.lastAddr = undefined;
        this.currentIndex = 0;
        this.TAG = 'giao '
        this.excludeRegs = new Set()
        this.excludeRegs.add('pc')
        this.excludeRegs.add('nzcv')
        for (let n = 0; n <= 0xff; ++n) {
            const hexOctet = n.toString(16).padStart(2, "0");
            this.byteToHex.push(hexOctet);
        }

    }

    parserNextAddr(ins) {
        let s = JSON.stringify(ins);
        let address = ins.address;
        let offset = address - this.moduleBase;
        let s1 = (offset).toString(16);
        let entity = {};
        entity.address = offset;
        return s1;
    }


    // 把寄存器当做指针，读字符串或者u64
    tryReadMemory(addr) {
        if (!addr) {
            return
        }
        var changeString = ''
        var logInfo = addr.toString(16) + '  '
        // 读字符串
        try {
            let nativePointer = new NativePointer(addr);
            changeString = nativePointer.readUtf8String();

        }
        catch (e) {
            changeString = "";
        }
        if (changeString !== "") {
            logInfo = logInfo + "(" + changeString + ")";
            return logInfo
        }

        // 读u64
        try {
            let nativePointer = new NativePointer(addr);
            changeString = nativePointer.readU64().toString(16);
        }
        catch (e) {
            changeString = "";
        }
        if (changeString !== "") {
            logInfo = logInfo + "(" + changeString + ")";
            return logInfo
        }

        return logInfo
    }

    // 寄存器变化
    getRegChange(context_jstr, isLastBranch) {

        var logInfo = ''
        var context = JSON.parse(context_jstr)
        var pre_context = JSON.parse(this.pre_context_jstr)
        // console.log('key',Object.values(context))
        for (let key in context) {
            // console.log('key',key)
            if (this.excludeRegs.has(key)) {
                continue
            }
            if (key.includes('q') || key.includes('d')) {
                var pre = JSON.stringify(pre_context[key])
                var cur = JSON.stringify(context[key])
                if (pre != cur) {
                    logInfo = logInfo + key + '=' + pre + ' --> ' + cur + '\t'
                    // console.log()
                }

            } else {

                if (pre_context[key] != context[key]) {
                    logInfo = logInfo + key + '=' + pre_context[key] + ' --> ' + this.tryReadMemory(context[key]) + '\t'
                }
            }

        }



        return logInfo
    }

    context_change_by_inst(context, ins) {

        let parse = JSON.parse(ins);
        // 上一条指令
        let lastParse = Instruction.parse(ptr(parse.address).sub(4))
        let isLastBranch = lastParse.mnemonic[0] == 'b'

        let logInfo = "";
        var context_jstr = JSON.stringify(context)
        // 如果上一条指令是跳转指令，则把x0-x7打印出来，即便他们没发生变化
        logInfo = this.getRegChange(context_jstr, isLastBranch);

        // 内存的变化

        let mnemonic = parse.mnemonic;

        if (mnemonic === "str" || mnemonic === "stp" || mnemonic === "strb" || mnemonic === "strh" || mnemonic === "stur") {
            let strParams = this.getMemoryChange(parse, context_jstr);
            logInfo = logInfo + strParams;
        }
        else if (mnemonic === "cmp") {
            let cmpParams = this.getCmpParams(parse, context_jstr);
            logInfo = logInfo + cmpParams;
        }
        else if (mnemonic === "b.gt" || mnemonic === "tbnz" || mnemonic === "b.le" || mnemonic === "b.eq" || mnemonic === "b.ne" || mnemonic === "b") {
            let bgtAddr = this.getbgtAddr(parse, context_jstr);
            logInfo = logInfo + bgtAddr;
        } else if (mnemonic === "ret") {
            let retOffset = this.getRetOffset(parse, context_jstr);
            logInfo = logInfo + retOffset;
        }
        let entity = {};
        entity.info = logInfo;
        this.pre_context_jstr = context_jstr;
        return entity;
    }

    ZY_trace_Stalker_begin(soname, addr, size) {
        let module = Process.findModuleByName(soname);
        this.moduleBase = module.base;
        console.log(this.TAG + JSON.stringify(module));
        console.log(this.TAG + "addr = " + addr);
        console.log(this.TAG + "size = " + size);
        var thiz = this
        Interceptor.attach(module.base.add(addr), {
            onEnter: function (args) {
                this.pid = Process.getCurrentThreadId();
                Stalker.follow(this.pid, {
                    events: {
                        call: false,
                        ret: false,
                        exec: false,

                        block: false,
                        compile: false
                    },
                    onReceive: function (events) {
                    },

                    transform: function (iterator) {

                        const instruction = iterator.next();
                        let startAddress = instruction.address;
                        if (size === 0) {
                            size = module.size;
                            addr = 0;
                        }
                        const isModuleCode = startAddress.compare(module.base.add(addr)) >= 0 &&
                            startAddress.compare(module.base.add(addr).add(size)) < 0;
                        // console.log('isModuleCode', isModuleCode)
                        do {
                            if (isModuleCode) {

                                // 就是拿指令的偏移
                                let s = thiz.parserNextAddr(instruction);
                                let address = instruction.address;
                                let offset = address - thiz.moduleBase;

                                // console.log(offset.toString(16),instruction)

                                let logInfo = s.toString(16) + "  " + instruction
                                thiz.detailInsMap.set(offset, JSON.stringify(instruction));
                                thiz.infoMap.set(offset, logInfo);

                                iterator.putCallout(function (context) {
                                    // context的pc指向的指令，是准备执行的，但是还没有执行;所以这个context保存的是执行完上条指令的状态,因此对应的是上条指令

                                    // console.log(JSON.stringify(context))

                                    // console.log(regs,'\n')
                                    if (thiz.isFirstIn) {

                                        thiz.isFirstIn = false;
                                        thiz.pre_context_jstr = JSON.stringify(context)
                                        thiz.lastOffset = Number(context.pc) - thiz.moduleBase

                                    }
                                    else {
                                        let offset = thiz.lastOffset
                                        let logInfo = thiz.infoMap.get(offset);
                                        if (!logInfo) {
                                            return
                                        }
                                        // 从第二条开始，打印本条指令运行后的影响,实际上用的是上条指令，本条context
                                        let detailIns = thiz.detailInsMap.get(offset);
                                        let entity = thiz.context_change_by_inst(context, detailIns);
                                        thiz.lastOffset = Number(context.pc) - thiz.moduleBase
                                        console.log(thiz.TAG + logInfo + " ; " + entity.info);
                                    }
                                });
                            }
                            iterator.keep();
                        } while (iterator.next() != null);
                    },
                });
            },
            onLeave: function (ret) {
                Stalker.unfollow(this.pid);
                console.log(thiz.TAG + "ret:" + ret);
            }
        });
    }


    getColor() {
        const colors = {
            colorize: (str, cc) => `\x1b${cc}${str}\x1b[0m`,
            red: str => colors.colorize(str, '[31m'),
            green: str => colors.colorize(str, '[32m'),
            yellow: str => colors.colorize(str, '[33m'),
            blue: str => colors.colorize(str, '[34m'),
            cyan: str => colors.colorize(str, '[36m'),
            white: str => colors.colorize(str, '[37m'),
        };
        if (this.currentIndex > 1) {
            this.currentIndex = 0;
        }
        if (this.currentIndex === 0) {
            return "C35";  // logger_1.LogColor.C35;
        }
        else if (this.currentIndex === 1) {
            return "C97";  // logger_1.LogColor.C97;
        }
        else if (this.currentIndex === 2) {
            return "C97";  // logger_1.LogColor.C97;
        }
    }


    getbgtAddr(parser, context) {

        let bgtAddr = "";
        let operands = parser.operands;
        for (let i = 0; i < operands.length; i++) {
            let operand = operands[i];
            if (operand.type === "imm") {

                let value = operand.value;
                let number = value - this.moduleBase;
                if (parser.mnemonic.includes('tbnz')) {
                    if (i > 0) {
                        // console.log(value, number)
                        // console.log(JSON.stringify(parser))
                        bgtAddr = "\t block addr:" + number.toString(16);
                        break;
                    }

                } else {
                    bgtAddr = "\t block addr:" + number.toString(16);
                    break;
                }


            }
        }
        return bgtAddr;
    }
    getRetOffset(parser, context_jstr) {
        let context = JSON.parse(context_jstr)
        let numOffset = Number(context.lr) - this.moduleBase
        return numOffset.toString(16)

    }

    getMemoryChange(parser, context_jstr) {

        var context = JSON.parse(context_jstr)
        var pre_context = JSON.parse(this.pre_context_jstr)
        // console.log('getMemoryChange',JSON.stringify(regs))
        let operands = parser.operands;
        let loginfo = '\t memory changed '
        var mem_addr = ''
        for (let op of operands) {
            if (op.type == 'mem') {
                mem_addr = InstHelper.getMemAddr(op.value, context)
            }
        }

        var changeString = this.tryReadMemory(mem_addr)
        loginfo = loginfo + changeString + '\t'

        if (parser.mnemonic == 'stp') {
            var ptr2 = ptr(mem_addr).add(8)
            var changeString = this.tryReadMemory(ptr2)
            loginfo = loginfo + changeString + '\t'
        }
        return loginfo
    }


    getCmpParams(parser, context_jstr) {
        var context = JSON.parse(context_jstr)
        var pre_context = JSON.parse(this.pre_context_jstr)
        let operands = parser.operands;
        let cmpInfo = "";
        for (let i = 0; i < operands.length; i++) {
            let operand = operands[i];
            if (operand.type === "reg") {
                let value = operand.value;
                let replace = value.replace("w", "x");

                let reg_value = context[replace];
                let changeString = "";
                try {
                    let nativePointer = new NativePointer(reg_value);
                    changeString = nativePointer.readUtf8String();
                }
                catch (e) {
                    changeString = "";
                }
                if (changeString !== "") {
                    reg_value = reg_value + "(" + changeString + ")";
                }
                cmpInfo = cmpInfo + value + " = " + reg_value + "\t";
            }
        }
        return cmpInfo;
    }


}


var trace = new Trace()
trace.ZY_trace_Stalker_begin('libnative-lib.so', 0xF5D4, 0);