package com.ollvm;

import capstone.Capstone;
import com.github.unidbg.AndroidEmulator;
import com.github.unidbg.LibraryResolver;
import com.github.unidbg.arm.ARM;
import com.github.unidbg.arm.backend.Backend;
import com.github.unidbg.arm.backend.BlockHook;
import com.github.unidbg.linux.android.AndroidEmulatorBuilder;
import com.github.unidbg.linux.android.AndroidResolver;
import com.github.unidbg.linux.android.dvm.DalvikModule;
import com.github.unidbg.linux.android.dvm.DvmClass;
import com.github.unidbg.linux.android.dvm.StringObject;
import com.github.unidbg.linux.android.dvm.VM;
import com.github.unidbg.memory.Memory;
import javafx.util.Pair;
import keystone.Keystone;
import keystone.KeystoneArchitecture;
import keystone.KeystoneEncoded;
import keystone.KeystoneMode;

import java.io.*;
import java.util.ArrayList;
import java.util.Arrays;

public class ollvmTesttest4 {
    private final AndroidEmulator emulator;
    private final VM vm;
    private final DvmClass mainActivityDvm;

    public static long startaddresstag= 0x10DBC;
    public static long endaddresstag= 0x1167C;
    public static int whileaddresstag = 0x11678;

    public static String inputmodulename = "unidbg-android/src/test/resources/example_binaries/ollvm_fla/libmyapplication4.so";
    public static String outputmodulename = "unidbg-android/src/test/resources/example_binaries/ollvm_fla/libmyapplication4out.so";
    public static void main(String[] args) {
        ollvmTesttest4 bcfTest = new ollvmTesttest4();
        bcfTest.call_calckey();
    }


    private ollvmTesttest4(){
        emulator = AndroidEmulatorBuilder
                .for64Bit()
                .build();
        Memory memory = emulator.getMemory();
        LibraryResolver resolver = new AndroidResolver(23);
        memory.setLibraryResolver(resolver);
        vm = emulator.createDalvikVM(null);
        vm.setVerbose(false);

        mainActivityDvm = vm.resolveClass("com/example/myapplication/MainActivity");
        DalvikModule dm = vm.loadLibrary(new File(inputmodulename), false);
        dm.callJNI_OnLoad(emulator);

    }

    //判断真实块
    private boolean OpstrContains(String opstr){
        ArrayList<String> flags= new ArrayList(Arrays.asList("ldur","ldr","str","b","movz","movk","cmp","b.eq"));
        if(flags.contains(opstr)){
            return true;
        }
        return false;
    }

    public static String bytesToHexString(byte[] src){
        StringBuilder stringBuilder = new StringBuilder("");
        if (src == null || src.length <= 0) {
            return null;
        }
        for (int i = 0; i < src.length; i++) {
            int v = src[i] & 0xFF;
            String hv = Integer.toHexString(v);
            if (hv.length() < 2) {
                stringBuilder.append(0);
            }
            stringBuilder.append(hv);
        }
        return stringBuilder.toString();
    }

    public static byte[] readFile(String strFile){
        try{
            InputStream is = new FileInputStream(strFile);
            int iAvail = is.available();
            byte[] bytes = new byte[iAvail];
            is.read(bytes);
            is.close();
            return bytes;
        }catch(Exception e){
            e.printStackTrace();
        }
        return null ;
    }

    public static void writeFile(byte[] data,String savefile){
        try {
            FileOutputStream fos=new FileOutputStream(savefile);
            BufferedOutputStream bos=new BufferedOutputStream(fos);
            bos.write(data,0,data.length);
            bos.flush();
            bos.close();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
    //真实块列表
    ArrayList<Pair<Long, Capstone.CsInsn[]>> readlyBlock=new ArrayList<Pair<Long, Capstone.CsInsn[]>>();
    //真实块的分支执行块

    //根据执行流程得到的所有块，过滤出真实块，并且将块合并成ida中的那样。在这里会用bl和b分割，所以bl分割开来的块我们要合并起来
    private void LoadReadlyAddress(){
        readlyBlock.add(blocks.get(0));
        for(int i=1;i<blocks.size();i++){
            Pair<Long, Capstone.CsInsn[]> pdata=blocks.get(i);
            Capstone.CsInsn[] insns=pdata.getValue();
            Long address=pdata.getKey();
            boolean isReadly=false;
            for(Capstone.CsInsn ins :insns){
                if(readlyBlock.contains(address)){
                    continue;
                }
                if(!OpstrContains(ins.mnemonic)){
                    isReadly=true;
                    String opstr= ARM.assembleDetail(emulator,ins,address,false,false);
                    System.out.println(String.format("block readly address:0x%x opstr:%s",ins.address,opstr) );
                    break;
                }
            }
            if(isReadly){
                //先把当前块添加为真实块，并且检查下一个块是否需要合并
                ArrayList<Capstone.CsInsn> arrayInsns=new ArrayList(Arrays.asList(insns));
                //如果这个块的最后一个指令是bl。则把下一个块也给合进来。直到非bl的结束块
                Pair<Long, Capstone.CsInsn[]> curBlock=blocks.get(i);
                while(true){
                    if(i>blocks.size()){
                        break;
                    }
                    Capstone.CsInsn[] curInsns=curBlock.getValue();
                    String mnemonic=curInsns[curInsns.length-1].mnemonic;
                    // 这里有个问题，如果是bl，合并了下一个快，不应该把bl移除吗
                    if(mnemonic.equals("bl")){
                        Pair<Long, Capstone.CsInsn[]> nextBlock=blocks.get(++i);
                        arrayInsns.addAll(Arrays.asList(nextBlock.getValue()));
                        curBlock=nextBlock;
                    }else{
                        break;
                    }
                }
                readlyBlock.add(new Pair<Long,Capstone.CsInsn[]>(address,(Capstone.CsInsn[])arrayInsns.toArray(new Capstone.CsInsn[arrayInsns.size()])));
            }
            if(i>blocks.size()){
                break;
            }
        }

    }

    //获取block的最后一个指令地址
    private Long GetEndAddress(Capstone.CsInsn[] insns){
        Capstone.CsInsn ins=insns[insns.length-1];
        return ins.address-0x40000000;
    }
    //获取block的第一个指令地址
    private Long GetStartAddress(Capstone.CsInsn[] insns){
        Capstone.CsInsn ins=insns[0];
        return ins.address-0x40000000;
    }


    //用来保存所有执行过的block
    ArrayList<Pair<Long, Capstone.CsInsn[]>> blocks=new ArrayList<Pair<Long, Capstone.CsInsn[]>>();
    //保存分支的所有block
    //主动调用目标函数
    private void call_calckey(){
        //这里BlockHook就是按照一个block的触发
        emulator.getBackend().hook_add_new(new BlockHook() {
            @Override
            public void hook(Backend backend, long address, int size, Object user) {
                //这里的insns是整个block。
                Capstone.CsInsn[] insns = emulator.disassemble(address, size,0);
//                System.out.println(String.format("address:0x%x size:0x%x",address,insns.length));
                //如果当前块中有分支
                boolean isexit = false;
                for (Pair<Long, Capstone.CsInsn[]> tempins : blocks) {
                    if (tempins.getKey() == address) {
                        isexit = true;
                        break;
                    }
                }
                if (!isexit && insns.length > 1) {
                    blocks.add(new Pair<Long, Capstone.CsInsn[]>(address, insns));
                }
            }
        },0x40000000+startaddresstag,0x40000000+endaddresstag,null);

        //调用一个返回值为object的静态的jni函数
        StringObject res = mainActivityDvm.callStaticJniMethodObject(emulator, "stringFromJNI()Ljava/lang/String;");
//        System.out.println(res.toString());
//        System.out.println(res2.toString());
        //筛选出真实块
        LoadReadlyAddress();
        String modulePath=inputmodulename;
        byte[] sodata=readFile(modulePath);
        //遍历真实块。然后直接修改成跳转真实块
        for(int i=0;i<readlyBlock.size();i++){
            if(i<readlyBlock.size()-1){
                //获取当前真实块
                Pair<Long, Capstone.CsInsn[]>block=readlyBlock.get(i);
                System.out.println(String.format("curBlock address:0x%x",block.getKey()));
                //获取下一个真实块
                Pair<Long, Capstone.CsInsn[]>nextBlock=readlyBlock.get(i+1);

                //取出当前真实块最后一个指令的地址
                int end_address= GetEndAddress(block.getValue()).intValue();
                if(end_address<=0){
                    continue;
                }
                //获取下一个真实块第一个指令的地址
                int start_address=GetStartAddress(nextBlock.getValue()).intValue();

                //准备转换汇编代码进行替换
                try (Keystone keystone = new Keystone(KeystoneArchitecture.Arm64, KeystoneMode.LittleEndian)) {
                    int subAddress=start_address-end_address;
                    //用来patch修改的asm指令。这里是要计算出当前地址的相对地址跳转。所以上面要减一下。
                    String asmStr=String.format("b #0x%x",subAddress);
                    //这个是我们显示日志看结果的。看看和我们之前手动分析的是不是差不多
                    KeystoneEncoded encoded = keystone.assemble(asmStr);
                    byte[] patch = encoded.getMachineCode();
                    if (patch.length <=0) {
                        System.out.println("转换汇编失败");
                        return;
                    }
//                    System.out.println(bytesToHexString(patch));
                    //替换原来的字节数据
                    for(int y =0;y<patch.length;y++){
                        sodata[end_address+y]=patch[y];
                    }
                }
            }
        }
        //循环的地方给nop掉
        byte[] nop_byte=new byte[]{0x1F,0x20,0x03,(byte)0xD5};
        int nop_address=whileaddresstag;
        for(int y =0;y<nop_byte.length;y++){
            sodata[nop_address+y]=nop_byte[y];
        }
        String savepath=outputmodulename;
        writeFile(sodata,savepath);
        System.out.println("处理完成");
    }
}