#include <vmprofiler.hpp>

namespace vm
{
    namespace util
    {
        namespace reg
        {
            zydis_register_t to64( zydis_register_t reg )
            {
                return ZydisRegisterGetLargestEnclosing( ZYDIS_MACHINE_MODE_LONG_64, reg );
            }

            bool compare( zydis_register_t a, zydis_register_t b )
            {
                return to64( a ) == to64( b );
            }
        } // namespace reg

        bool get_fetch_operand( const zydis_routine_t &routine, zydis_instr_t &fetch_instr )
        {
            const auto result =
                std::find_if( routine.begin(), routine.end(), []( const zydis_instr_t &instr_data ) -> bool {
                    // mov/movsx/movzx rax/eax/ax/al, [rsi]
                    return instr_data.instr.operand_count > 1 &&
                           ( instr_data.instr.mnemonic == ZYDIS_MNEMONIC_MOV ||
                             instr_data.instr.mnemonic == ZYDIS_MNEMONIC_MOVSX ||
                             instr_data.instr.mnemonic == ZYDIS_MNEMONIC_MOVZX ) &&
                           instr_data.instr.operands[ 0 ].type == ZYDIS_OPERAND_TYPE_REGISTER &&
                           util::reg::to64( instr_data.instr.operands[ 0 ].reg.value ) == ZYDIS_REGISTER_RAX &&
                           instr_data.instr.operands[ 1 ].type == ZYDIS_OPERAND_TYPE_MEMORY &&
                           instr_data.instr.operands[ 1 ].mem.base == ZYDIS_REGISTER_RSI;
                } );

            if ( result == routine.end() )
                return false;

            fetch_instr = *result;
            return true;
        }

        std::optional< zydis_routine_t::iterator > get_fetch_operand( zydis_routine_t &routine )
        {
            auto result = std::find_if( routine.begin(), routine.end(), []( const zydis_instr_t &instr_data ) -> bool {
                // mov/movsx/movzx rax/eax/ax/al, [rsi]
                return instr_data.instr.operand_count > 1 &&
                       ( instr_data.instr.mnemonic == ZYDIS_MNEMONIC_MOV ||
                         instr_data.instr.mnemonic == ZYDIS_MNEMONIC_MOVSX ||
                         instr_data.instr.mnemonic == ZYDIS_MNEMONIC_MOVZX ) &&
                       instr_data.instr.operands[ 0 ].type == ZYDIS_OPERAND_TYPE_REGISTER &&
                       util::reg::to64( instr_data.instr.operands[ 0 ].reg.value ) == ZYDIS_REGISTER_RAX &&
                       instr_data.instr.operands[ 1 ].type == ZYDIS_OPERAND_TYPE_MEMORY &&
                       instr_data.instr.operands[ 1 ].mem.base == ZYDIS_REGISTER_RSI;
            } );

            if ( result == routine.end() )
                return {};

            return result;
        }

        void print( const zydis_decoded_instr_t &instr )
        {
            char buffer[ 256 ];
            ZydisFormatter formatter;
            ZydisFormatterInit( &formatter, ZYDIS_FORMATTER_STYLE_INTEL );
            ZydisFormatterFormatInstruction( &formatter, &instr, buffer, sizeof( buffer ), 0u );
            std::puts( buffer );
        }

        void print( zydis_routine_t &routine )
        {
            char buffer[ 256 ];
            ZydisFormatter formatter;
            ZydisFormatterInit( &formatter, ZYDIS_FORMATTER_STYLE_INTEL );

            for ( auto [ instr, raw, addr ] : routine )
            {
                ZydisFormatterFormatInstruction( &formatter, &instr, buffer, sizeof( buffer ), addr );
                std::printf( "> 0x%p %s", addr, buffer );
            }
        }

        bool is_jmp( const zydis_decoded_instr_t &instr )
        {
            switch ( instr.mnemonic )
            {
            case ZYDIS_MNEMONIC_JB:
            case ZYDIS_MNEMONIC_JBE:
            case ZYDIS_MNEMONIC_JCXZ:
            case ZYDIS_MNEMONIC_JECXZ:
            case ZYDIS_MNEMONIC_JKNZD:
            case ZYDIS_MNEMONIC_JKZD:
            case ZYDIS_MNEMONIC_JL:
            case ZYDIS_MNEMONIC_JLE:
            case ZYDIS_MNEMONIC_JMP:
            case ZYDIS_MNEMONIC_JNB:
            case ZYDIS_MNEMONIC_JNBE:
            case ZYDIS_MNEMONIC_JNL:
            case ZYDIS_MNEMONIC_JNLE:
            case ZYDIS_MNEMONIC_JNO:
            case ZYDIS_MNEMONIC_JNP:
            case ZYDIS_MNEMONIC_JNS:
            case ZYDIS_MNEMONIC_JNZ:
            case ZYDIS_MNEMONIC_JO:
            case ZYDIS_MNEMONIC_JP:
            case ZYDIS_MNEMONIC_JRCXZ:
            case ZYDIS_MNEMONIC_JS:
            case ZYDIS_MNEMONIC_JZ:
                return true;
            default:
                break;
            }
            return false;
        }

        bool flatten( zydis_routine_t &routine, std::uintptr_t routine_addr, bool keep_jmps )
        {
            ZydisDecoder decoder;
            zydis_decoded_instr_t instr;
            ZydisDecoderInit( &decoder, ZYDIS_MACHINE_MODE_LONG_64, ZYDIS_ADDRESS_WIDTH_64 );

            while ( ZYAN_SUCCESS(
                ZydisDecoderDecodeBuffer( &decoder, reinterpret_cast< void * >( routine_addr ), 0x1000, &instr ) ) )
            {
                // detect if we have already been at this instruction... if so that means there is a loop and we are
                // going to just return...
                if ( std::find_if( routine.begin(), routine.end(), [ & ]( const zydis_instr_t &zydis_instr ) -> bool {
                         return zydis_instr.addr == routine_addr;
                     } ) != routine.end() )
                    return true;

                std::vector< u8 > raw_instr;
                raw_instr.insert( raw_instr.begin(), ( u8 * )routine_addr, ( u8 * )routine_addr + instr.length );

                if ( is_jmp( instr ) )
                {
                    if ( instr.operands[ 0 ].type == ZYDIS_OPERAND_TYPE_REGISTER )
                    {
                        routine.push_back( { instr, raw_instr, routine_addr } );
                        return true;
                    }

                    if ( keep_jmps )
                        routine.push_back( { instr, raw_instr, routine_addr } );

                    ZydisCalcAbsoluteAddress( &instr, &instr.operands[ 0 ], routine_addr, &routine_addr );
                }
                else if ( instr.mnemonic == ZYDIS_MNEMONIC_RET )
                {
                    routine.push_back( { instr, raw_instr, routine_addr } );
                    return true;
                }
                else
                {
                    routine.push_back( { instr, raw_instr, routine_addr } );
                    routine_addr += instr.length;
                }
            }
            return false;
        }

        void deobfuscate( zydis_routine_t &routine )
        {
            static const auto _uses = []( ZydisDecodedOperand &op, zydis_register_t reg ) -> bool {
                switch ( op.type )
                {
                case ZYDIS_OPERAND_TYPE_MEMORY:
                {
                    return reg::compare( op.mem.base, reg ) || reg::compare( op.mem.index, reg );
                }
                case ZYDIS_OPERAND_TYPE_REGISTER:
                {
                    return reg::compare( op.reg.value, reg );
                }
                default:
                    break;
                }
                return false;
            };

            static const auto _writes = []( zydis_decoded_instr_t &inst ) -> bool {
                for ( auto idx = 0; idx < inst.operand_count; ++idx )
                    if ( inst.operands[ idx ].actions & ZYDIS_OPERAND_ACTION_MASK_WRITE )
                        return true;

                return false;
            };

            static const auto _remove = []( zydis_routine_t &routine, zydis_routine_t::iterator itr,
                                            zydis_register_t reg, u32 opcode_size ) -> void {
                for ( ; itr >= routine.begin(); --itr )
                {
                    const auto instruction = &itr->instr;
                    auto stop = false;

                    if ( instruction->mnemonic == ZYDIS_MNEMONIC_JMP )
                        continue;

                    for ( auto op_idx = 0u; op_idx < instruction->operand_count; ++op_idx )
                    {
                        const auto op = &instruction->operands[ op_idx ];

                        if ( !_uses( *op, reg ) )
                            continue;

                        if ( op->type == ZYDIS_OPERAND_TYPE_MEMORY )
                        {
                            stop = true;
                            continue;
                        }

                        if ( opcode_size < 32 && op->size > opcode_size )
                            continue;

                        op->actions &ZYDIS_OPERAND_ACTION_MASK_WRITE ? op->actions &= ~ZYDIS_OPERAND_ACTION_MASK_WRITE
                                                                     : stop = true;
                    }

                    if ( !_writes( *instruction ) )
                        routine.erase( itr );

                    else if ( stop )
                        break;
                }
            };

            for ( const auto &instr_data : routine )
            {
                if ( routine.empty() || routine.size() == 1 || instr_data.instr.mnemonic == ZYDIS_MNEMONIC_JMP )
                    continue;

                for ( auto itr = routine.begin() + 1; itr != routine.end(); itr++ )
                {
                    if ( itr->instr.mnemonic == ZYDIS_MNEMONIC_JMP || itr->instr.mnemonic == ZYDIS_MNEMONIC_RET )
                        break;

                    // find the write operations that happen...
                    for ( auto idx = 0u; idx < itr->instr.operand_count; ++idx )
                    {
                        const auto op = &itr->instr.operands[ idx ];
                        // if its a read, continue to next opcode...
                        if ( op->actions & ZYDIS_OPERAND_ACTION_MASK_READ )
                            continue;

                        // if its not a write then continue to next opcode...
                        if ( !( op->actions & ZYDIS_OPERAND_ACTION_MASK_WRITE ) )
                            continue;

                        // if this operand is not a register then we continue...
                        if ( op->type != ZYDIS_OPERAND_TYPE_REGISTER )
                            continue;

                        // else we see if we can remove dead writes to this register...
                        _remove( routine, itr - 1, op->reg.value, op->size );
                    }
                }
            }
        }
    } // namespace util
} // namespace vm