#include "dsprot/encryptor.h"

#include <nitro/os/cache.h>

#include "dsprot/bss.h"
#include "dsprot/encoding_constants.h"
#include "dsprot/rc4.h"

#define ROTL(x, a) ((a) == 0 ? (x) : (((x) << (a)) | ((x) >> (32 - (a)))))

static void clearDataAndInstructionCache(void *start_addr, u32 num_bytes);

static void clearDataAndInstructionCache(void *start_addr, u32 num_bytes) {
    DC_FlushRange(start_addr, num_bytes);
    IC_InvalidateRange(start_addr, num_bytes);
}

u32 Encryptor_CategorizeInstruction(u32 instruction) {
    u8 opcode;

    opcode = instruction >> INS_OPCODE_SHIFT;

    // Branch instruction
    if ((opcode & 0x0E) == 0x0A) {
        // BLX immediate type
        if ((opcode & 0xF0) == 0xF0) {
            return INS_TYPE_BLXIMM;
        }

        // Link bit
        if (opcode & INS_OPCODE_LINKBIT) {
            return INS_TYPE_BL;
        } else {
            return INS_TYPE_B;
        }
    }

    return INS_TYPE_OTHER;
}

void Encryptor_DecodeFunctionTable(FuncInfo *functions) {
    u32 *addr;
    u32 size;
    u32 *end_addr;

    if (functions == NULL) {
        return;
    }

    for (; functions->obfs_addr != 0; functions++) {
        addr = (u32 *) functions->obfs_addr;
        size = functions->obfs_size - (u32) &DSProt_BSS - ENC_VAL_1;

        if (addr == NULL) {
            break;
        }

        addr     = (u32 *) ((u32) addr - ENC_VAL_1);
        end_addr = addr + (size / 4);

        for (; addr < end_addr; addr++) {
            switch (Encryptor_CategorizeInstruction(*addr)) {
                case INS_TYPE_BLXIMM:
                case INS_TYPE_BL: {
                    u32 opcode   = (*addr & INS_OPCODE_MASK) ^ (INS_OPCODE_LINKBIT << INS_OPCODE_SHIFT);
                    u32 operands = ((*addr & INS_OPERANDS_MASK) - ENC_VAL_1) & INS_OPERANDS_MASK;

                    *addr = opcode | operands;
                } break;

                case INS_TYPE_B: {
                    u32 opcode   = (*addr & INS_OPCODE_MASK) ^ (INS_OPCODE_LINKBIT << INS_OPCODE_SHIFT);
                    u32 operands = ((*addr & INS_OPERANDS_MASK) - ENC_VAL_2) & INS_OPERANDS_MASK;

                    *addr = opcode | operands;
                } break;

                default: {
                    u8 *addr_bytes = (u8 *) addr;
                    *addr          = (addr_bytes[0] ^ ENC_BYTE_A) | ((addr_bytes[1] ^ ENC_BYTE_B) << 8) |
                            ((addr_bytes[2] ^ ENC_BYTE_C) << 16) | ((addr_bytes[3] ^ ENC_BYTE_D) << 24);
                } break;
            }
        }

        clearDataAndInstructionCache((void *) (functions->obfs_addr - ENC_VAL_1), size);
    }
}

static inline void expandRC4Key(u32 seed_key, u32 size, u32 *expanded_key) {
    expanded_key[0] = ROTL(seed_key, 0) ^ size;
    expanded_key[1] = ROTL(seed_key, 8) ^ size;
    expanded_key[2] = ROTL(seed_key, 16) ^ size;
    expanded_key[3] = ROTL(seed_key, 24) ^ size;
}

void *Encryptor_DecryptFunction(u32 key, u32 func_addr, u32 size) {
    u32 expanded_key[4];
    void *func_ptr;

    // Deobfuscate arguments
    func_ptr = (void *) func_addr;
    func_ptr -= ENC_VAL_1;

    size -= (u32) &DSProt_BSS + ENC_VAL_1;
    key -= (u32) &DSProt_BSS + ENC_VAL_1;

    expandRC4Key(key, size, &expanded_key[0]);
    RC4_InitAndDecryptInstructions(&expanded_key[0], func_ptr, func_ptr, size);
    clearDataAndInstructionCache(func_ptr, size);

    return func_ptr;
}

u32 Encryptor_EncryptFunction(u32 key, u32 func_addr, u32 size) {
    u32 expanded_key[4];
    void *func_ptr;

    // Deobfuscate arguments
    func_ptr = (void *) func_addr;
    func_ptr -= ENC_VAL_1;

    size -= (u32) &DSProt_BSS + ENC_VAL_1;
    key -= (u32) &DSProt_BSS + ENC_VAL_1;

    // Change key
    key += func_addr >> 20;

    expandRC4Key(key, size, &expanded_key[0]);
    RC4_InitAndEncryptInstructions(&expanded_key[0], func_ptr, func_ptr, size);
    clearDataAndInstructionCache(func_ptr, size);

    // Re-obfuscate key
    key += (u32) &DSProt_BSS + ENC_VAL_1;

    return key;
}