#include <vmprofiler.hpp>

namespace vm
{
    namespace instrs
    {
        std::pair< std::uint64_t, std::uint64_t > decrypt_operand( transform::map_t &transforms, std::uint64_t operand,
                                                                   std::uint64_t rolling_key )
        {
            const auto& generic_decrypt_0 = transforms[ transform::type::generic0 ];
            const auto& key_decrypt = transforms[ transform::type::rolling_key ];
            const auto& generic_decrypt_1 = transforms[ transform::type::generic1 ];
            const auto& generic_decrypt_2 = transforms[ transform::type::generic2 ];
            const auto& generic_decrypt_3 = transforms[ transform::type::generic3 ];
            const auto& update_key = transforms[ transform::type::update_key ];

            if ( generic_decrypt_0.mnemonic != ZYDIS_MNEMONIC_INVALID )
            {
                operand = transform::apply(
                    generic_decrypt_0.operands[ 0 ].size, generic_decrypt_0.mnemonic, operand,
                    // check to see if this instruction has an IMM...
                    transform::has_imm( &generic_decrypt_0 ) ? generic_decrypt_0.operands[ 1 ].imm.value.u : 0 );
            }

            // apply transformation with rolling decrypt key...
            operand = transform::apply( key_decrypt.operands[ 0 ].size, key_decrypt.mnemonic, operand, rolling_key );

            // apply three generic transformations...
            {
                operand = transform::apply(
                    generic_decrypt_1.operands[ 0 ].size, generic_decrypt_1.mnemonic, operand,
                    // check to see if this instruction has an IMM...
                    transform::has_imm( &generic_decrypt_1 ) ? generic_decrypt_1.operands[ 1 ].imm.value.u : 0 );

                operand = transform::apply(
                    generic_decrypt_2.operands[ 0 ].size, generic_decrypt_2.mnemonic, operand,
                    // check to see if this instruction has an IMM...
                    transform::has_imm( &generic_decrypt_2 ) ? generic_decrypt_2.operands[ 1 ].imm.value.u : 0 );

                operand = transform::apply(
                    generic_decrypt_3.operands[ 0 ].size, generic_decrypt_3.mnemonic, operand,
                    // check to see if this instruction has an IMM...
                    transform::has_imm( &generic_decrypt_3 ) ? generic_decrypt_3.operands[ 1 ].imm.value.u : 0 );
            }

            // update rolling key...
            auto result =
                transform::apply( update_key.operands[ 0 ].size, update_key.mnemonic, rolling_key, operand );

            // update decryption key correctly...
            switch ( update_key.operands[ 0 ].size )
            {
            case 8:
                rolling_key = ( rolling_key & ~0xFFull ) + result;
                break;
            case 16:
                rolling_key = ( rolling_key & ~0xFFFFull ) + result;
                break;
            default:
                rolling_key = result;
                break;
            }

            return { operand, rolling_key };
        }

        std::pair< std::uint64_t, std::uint64_t > encrypt_operand( transform::map_t &transforms, std::uint64_t operand,
                                                                   std::uint64_t rolling_key )
        {
            transform::map_t inverse;
            inverse_transforms( transforms, inverse );

            const auto& generic_decrypt_0 = inverse[ transform::type::generic0 ];
            const auto& key_decrypt = inverse[ transform::type::rolling_key ];
            const auto& generic_decrypt_1 = inverse[ transform::type::generic1 ];
            const auto& generic_decrypt_2 = inverse[ transform::type::generic2 ];
            const auto& generic_decrypt_3 = inverse[ transform::type::generic3 ];
            const auto& update_key = inverse[ transform::type::update_key ];

            auto result =
                transform::apply( update_key.operands[ 0 ].size, update_key.mnemonic, rolling_key, operand );

            // make sure we update the rolling decryption key correctly...
            switch ( update_key.operands[ 0 ].size )
            {
            case 8:
                rolling_key = ( rolling_key & ~0xFFull ) + result;
                break;
            case 16:
                rolling_key = ( rolling_key & ~0xFFFFull ) + result;
                break;
            default:
                rolling_key = result;
                break;
            }

            {
                operand = transform::apply(
                    generic_decrypt_3.operands[ 0 ].size, generic_decrypt_3.mnemonic, operand,
                    // check to see if this instruction has an IMM...
                    transform::has_imm( &generic_decrypt_3 ) ? generic_decrypt_3.operands[ 1 ].imm.value.u : 0 );

                operand = transform::apply(
                    generic_decrypt_2.operands[ 0 ].size, generic_decrypt_2.mnemonic, operand,
                    // check to see if this instruction has an IMM...
                    transform::has_imm( &generic_decrypt_2 ) ? generic_decrypt_2.operands[ 1 ].imm.value.u : 0 );

                operand = transform::apply(
                    generic_decrypt_1.operands[ 0 ].size, generic_decrypt_1.mnemonic, operand,
                    // check to see if this instruction has an IMM...
                    transform::has_imm( &generic_decrypt_1 ) ? generic_decrypt_1.operands[ 1 ].imm.value.u : 0 );
            }

            operand = transform::apply( key_decrypt.operands[ 0 ].size, key_decrypt.mnemonic, operand, rolling_key );

            if ( generic_decrypt_0.mnemonic != ZYDIS_MNEMONIC_INVALID )
            {
                operand = transform::apply(
                    generic_decrypt_0.operands[ 0 ].size, generic_decrypt_0.mnemonic, operand,
                    // check to see if this instruction has an IMM...
                    transform::has_imm( &generic_decrypt_0 ) ? generic_decrypt_0.operands[ 1 ].imm.value.u : 0 );
            }

            return { operand, rolling_key };
        }

        bool get_rva_decrypt( const zydis_routine_t &vm_entry, std::vector< zydis_decoded_instr_t > &transform_instrs )
        {
            //
            // find mov esi, [rsp+0xA0]
            //

            auto result =
                std::find_if( vm_entry.begin(), vm_entry.end(), []( const zydis_instr_t &instr_data ) -> bool {
                    if ( instr_data.instr.mnemonic == ZYDIS_MNEMONIC_MOV && instr_data.instr.operand_count == 2 &&
                         instr_data.instr.operands[ 0 ].reg.value == ZYDIS_REGISTER_ESI &&
                         instr_data.instr.operands[ 1 ].mem.base == ZYDIS_REGISTER_RSP &&
                         instr_data.instr.operands[ 1 ].mem.disp.has_displacement &&
                         instr_data.instr.operands[ 1 ].mem.disp.value == 0xA0 )
                        return true;
                    return false;
                } );

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

            //
            // find the next three instruction with ESI as the dest...
            //

            for ( auto idx = 0u; idx < 3; ++idx )
            {
                result = std::find_if( ++result, vm_entry.end(), []( const zydis_instr_t &instr_data ) -> bool {
                    return instr_data.instr.operands[ 0 ].reg.value == ZYDIS_REGISTER_ESI;
                } );

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

                transform_instrs.push_back( result->instr );
            }

            return true;
        }
    } // namespace instrs
} // namespace vm