/* Unicorn Emulator Engine */
/* By Nguyen Anh Quynh, 2015 */

/* Sample code to demonstrate how to emulate ARM64 code */

#include <unicorn/unicorn.h>
#include <string.h>

// code to be emulated
#define ARM64_CODE                                                             \
    "\xab\x05\x00\xb8\xaf\x05\x40\x38" // str w11, [x13], #0; ldrb w15, [x13],
                                       // #0
//#define ARM64_CODE_EB "\xb8\x00\x05\xab\x38\x40\x05\xaf" // str w11, [x13];
// ldrb w15, [x13]
#define ARM64_CODE_EB ARM64_CODE

// memory address where emulation starts
#define ADDRESS 0x10000

static void hook_block(uc_engine *uc, uint64_t address, uint32_t size,
                       void *user_data)
{
    printf(">>> Tracing basic block at 0x%" PRIx64 ", block size = 0x%x\n",
           address, size);
}

static void hook_code(uc_engine *uc, uint64_t address, uint32_t size,
                      void *user_data)
{
    printf(">>> Tracing instruction at 0x%" PRIx64
           ", instruction size = 0x%x\n",
           address, size);
}

static void test_arm64_mem_fetch(void)
{
    uc_engine *uc;
    uc_err err;
    uint64_t x1, sp, x0;
    // msr x0, CurrentEL
    unsigned char shellcode0[4] = {64, 66, 56, 213};
    // .text:00000000004002C0                 LDR             X1, [SP,#arg_0]
    unsigned char shellcode[4] = {0xE1, 0x03, 0x40, 0xF9};
    unsigned shellcode_address = 0x4002C0;
    uint64_t data_address = 0x10000000000000;

    printf(">>> Emulate ARM64 fetching stack data from high address %" PRIx64
           "\n",
           data_address);

    // Initialize emulator in ARM mode
    err = uc_open(UC_ARCH_ARM64, UC_MODE_ARM, &uc);
    if (err) {
        printf("Failed on uc_open() with error returned: %u (%s)\n", err,
               uc_strerror(err));
        return;
    }

    uc_mem_map(uc, data_address, 0x30000, UC_PROT_ALL);
    uc_mem_map(uc, 0x400000, 0x1000, UC_PROT_ALL);

    sp = data_address;
    uc_reg_write(uc, UC_ARM64_REG_SP, &sp);
    uc_mem_write(uc, data_address, "\xc8\xc8\xc8\xc8\xc8\xc8\xc8\xc8", 8);
    uc_mem_write(uc, shellcode_address, shellcode0, 4);
    uc_mem_write(uc, shellcode_address + 4, shellcode, 4);

    err = uc_emu_start(uc, shellcode_address, shellcode_address + 4, 0, 0);
    if (err) {
        printf("Failed on uc_emu_start() with error returned: %u\n", err);
    }

    x0 = 0;
    uc_reg_read(uc, UC_ARM64_REG_X0, &x0);
    printf(">>> x0(Exception Level)=%" PRIx64 "\n", x0 >> 2);

    err = uc_emu_start(uc, shellcode_address + 4, shellcode_address + 8, 0, 0);
    if (err) {
        printf("Failed on uc_emu_start() with error returned: %u\n", err);
    }

    uc_reg_read(uc, UC_ARM64_REG_X1, &x1);

    printf(">>> X1 = 0x%" PRIx64 "\n", x1);

    uc_close(uc);
}

static void test_arm64(void)
{
    uc_engine *uc;
    uc_err err;
    uc_hook trace1, trace2;

    int64_t x11 = 0x12345678;    // X11 register
    int64_t x13 = 0x10000 + 0x8; // X13 register
    int64_t x15 = 0x33;          // X15 register

    printf("Emulate ARM64 code\n");

    // Initialize emulator in ARM mode
    err = uc_open(UC_ARCH_ARM64, UC_MODE_ARM, &uc);
    if (err) {
        printf("Failed on uc_open() with error returned: %u (%s)\n", err,
               uc_strerror(err));
        return;
    }

    // map 2MB memory for this emulation
    uc_mem_map(uc, ADDRESS, 2 * 1024 * 1024, UC_PROT_ALL);

    // write machine code to be emulated to memory
    uc_mem_write(uc, ADDRESS, ARM64_CODE, sizeof(ARM64_CODE) - 1);

    // initialize machine registers
    uc_reg_write(uc, UC_ARM64_REG_X11, &x11);
    uc_reg_write(uc, UC_ARM64_REG_X13, &x13);
    uc_reg_write(uc, UC_ARM64_REG_X15, &x15);

    // tracing all basic blocks with customized callback
    uc_hook_add(uc, &trace1, UC_HOOK_BLOCK, hook_block, NULL, 1, 0);

    // tracing one instruction at ADDRESS with customized callback
    uc_hook_add(uc, &trace2, UC_HOOK_CODE, hook_code, NULL, ADDRESS, ADDRESS);

    // emulate machine code in infinite time (last param = 0), or when
    // finishing all the code.
    err = uc_emu_start(uc, ADDRESS, ADDRESS + sizeof(ARM64_CODE) - 1, 0, 0);
    if (err) {
        printf("Failed on uc_emu_start() with error returned: %u\n", err);
    }

    // now print out some registers
    printf(">>> Emulation done. Below is the CPU context\n");
    printf(">>> As little endian, X15 should be 0x78:\n");

    uc_reg_read(uc, UC_ARM64_REG_X15, &x15);
    printf(">>> X15 = 0x%" PRIx64 "\n", x15);

    uc_close(uc);
}

static void test_arm64eb(void)
{
    uc_engine *uc;
    uc_err err;
    uc_hook trace1, trace2;

    int64_t x11 = 0x12345678;    // X11 register
    int64_t x13 = 0x10000 + 0x8; // X13 register
    int64_t x15 = 0x33;          // X15 register

    printf("Emulate ARM64 Big-Endian code\n");

    // Initialize emulator in ARM mode
    err = uc_open(UC_ARCH_ARM64, UC_MODE_ARM + UC_MODE_BIG_ENDIAN, &uc);
    if (err) {
        printf("Failed on uc_open() with error returned: %u (%s)\n", err,
               uc_strerror(err));
        return;
    }

    // map 2MB memory for this emulation
    uc_mem_map(uc, ADDRESS, 2 * 1024 * 1024, UC_PROT_ALL);

    // write machine code to be emulated to memory
    uc_mem_write(uc, ADDRESS, ARM64_CODE_EB, sizeof(ARM64_CODE_EB) - 1);

    // initialize machine registers
    uc_reg_write(uc, UC_ARM64_REG_X11, &x11);
    uc_reg_write(uc, UC_ARM64_REG_X13, &x13);
    uc_reg_write(uc, UC_ARM64_REG_X15, &x15);

    // tracing all basic blocks with customized callback
    uc_hook_add(uc, &trace1, UC_HOOK_BLOCK, hook_block, NULL, 1, 0);

    // tracing one instruction at ADDRESS with customized callback
    uc_hook_add(uc, &trace2, UC_HOOK_CODE, hook_code, NULL, ADDRESS, ADDRESS);

    // emulate machine code in infinite time (last param = 0), or when
    // finishing all the code.
    err = uc_emu_start(uc, ADDRESS, ADDRESS + sizeof(ARM64_CODE_EB) - 1, 0, 0);
    if (err) {
        printf("Failed on uc_emu_start() with error returned: %u\n", err);
    }

    // now print out some registers
    printf(">>> Emulation done. Below is the CPU context\n");
    printf(">>> As big endian, X15 should be 0x78:\n");

    uc_reg_read(uc, UC_ARM64_REG_X15, &x15);
    printf(">>> X15 = 0x%" PRIx64 "\n", x15);

    uc_close(uc);
}

int main(int argc, char **argv, char **envp)
{
    test_arm64_mem_fetch();
    test_arm64();

    printf("-------------------------\n");
    test_arm64eb();

    return 0;
}