#include "compiler.h"

namespace vm
{
    compiler_t::compiler_t( vm::ctx_t *vmctx ) : vmctx( vmctx )
    {
        if ( !parse_t::get_instance()->for_each( [ & ]( _vlabel_meta *label_data ) -> bool {
                 std::printf( "> checking label %s for invalid instructions... number of instructions = %d\n",
                              label_data->label_name.c_str(), label_data->vinstrs.size() );

                 const auto result = std::find_if(
                     label_data->vinstrs.begin(), label_data->vinstrs.end(),
                     [ & ]( const _vinstr_meta &vinstr ) -> bool {
                         std::printf( "> vinstr name = %s, has imm = %d, imm = 0x%p\n", vinstr.name.c_str(),
                                      vinstr.has_imm, vinstr.imm );

                         for ( auto &vm_handler : vmctx->vm_handlers )
                             if ( vm_handler.profile && vm_handler.profile->name == vinstr.name )
                                 return false;

                         std::printf( "[!] this vm protected file does not have the vm handler for: %s...\n",
                                      vinstr.name.c_str() );

                         return true;
                     } );

                 return result == label_data->vinstrs.end();
             } ) )
        {
            std::printf( "[!] binary does not have the required vm handlers...\n" );
            exit( -1 );
        }

        if ( !vm::handler::get_operand_transforms( vmctx->calc_jmp, calc_jmp_transforms ) )
        {
            std::printf( "[!] failed to extract calc_jmp transformations...\n" );
            exit( -1 );
        }

        if ( !vm::instrs::get_rva_decrypt( vmctx->vm_entry, encrypt_vinstrs_rva ) )
        {
            std::printf( "[!] failed to extract virtual instruction rva decryption instructions...\n" );
            exit( -1 );
        }

        if ( !vm::transform::inverse_transforms( encrypt_vinstrs_rva ) )
        {
            std::printf( "[!] failed to inverse virtual instruction rva decrypt instructions...\n" );
            exit( -1 );
        }
    }

    std::vector< vlabel_data > *compiler_t::encode()
    {
        parse_t::get_instance()->for_each( [ & ]( _vlabel_meta *label_data ) -> bool {
            virt_labels.push_back( { label_data->label_name } );
            for ( const auto &vinstr : label_data->vinstrs )
            {
                for ( auto itr = vmctx->vm_handlers.begin(); itr != vmctx->vm_handlers.end(); ++itr )
                {
                    if ( itr->profile && itr->profile->name == vinstr.name )
                    {
                        virt_labels.back().vinstrs.push_back( { ( std::uint8_t )( itr - vmctx->vm_handlers.begin() ),
                                                                vinstr.imm, itr->profile->imm_size } );
                        break;
                    }
                }
            }
            return true;
        } );

        return &virt_labels;
    }

    std::vector< compiled_label_data > compiler_t::encrypt()
    {
        std::vector< compiled_label_data > result;
        const auto end_of_module = vmctx->image_size + vmctx->image_base;

        // decryption key starts off as the image
        // base address of the virtual instructions...
        std::uintptr_t decrypt_key = end_of_module, start_addr;
        if ( vmctx->exec_type == vmp2::exec_type_t::backward )
        {
            std::for_each( virt_labels.begin(), virt_labels.end(), [ & ]( const vlabel_data &label ) {
                std::for_each( label.vinstrs.begin(), label.vinstrs.end(), [ & ]( const vinstr_data &vinstr ) {
                    ( ++decrypt_key ) += vinstr.imm_size ? vinstr.imm_size / 8 : 0;
                } );
            } );
        }
        start_addr = decrypt_key;

        // invert the encoded virtual instructions operands if vip advances backward...
        if ( vmctx->exec_type == vmp2::exec_type_t::backward )
            std::reverse( virt_labels.begin(), virt_labels.end() );

        // loop over the instructions and encrypt them...
        for ( auto &label : virt_labels )
        {
            result.push_back( { label.label_name, start_addr } );

            for ( auto &vinstr : label.vinstrs )
            {
                auto vm_handler_idx = vinstr.vm_handler;
                std::tie( vinstr.vm_handler, decrypt_key ) =
                    vm::instrs::encrypt_operand( calc_jmp_transforms, vinstr.vm_handler, decrypt_key );

                if ( !vinstr.imm_size )
                {
                    result.back().vinstrs.push_back( vinstr.vm_handler );
                    continue;
                }

                auto transforms = vmctx->vm_handlers.at( vm_handler_idx ).transforms;

                std::tie( vinstr.operand, decrypt_key ) =
                    vm::instrs::encrypt_operand( transforms, vinstr.operand, decrypt_key );

                // operands must be backwards if VIP advances backward...
                if ( vmctx->exec_type == vmp2::exec_type_t::backward )
                {
                    for ( auto idx = 0u; idx < vinstr.imm_size / 8; ++idx )
                        result.back().vinstrs.push_back( reinterpret_cast< std::uint8_t * >( &vinstr.operand )[ idx ] );

                    result.back().vinstrs.push_back( vinstr.vm_handler );
                }
                else
                {
                    result.back().vinstrs.push_back( vinstr.vm_handler );

                    for ( auto idx = 0u; idx < vinstr.imm_size / 8; ++idx )
                        result.back().vinstrs.push_back( reinterpret_cast< std::uint8_t * >( &vinstr.operand )[ idx ] );
                }
            }

            result.back().enc_alloc_rva = encrypt_rva( start_addr );
            start_addr += result.back().vinstrs.size();
        }

        return result;
    }

    std::uint64_t compiler_t::encrypt_rva( std::uint64_t rva )
    {
        for ( auto &instr : encrypt_vinstrs_rva )
            rva = vm::transform::apply( instr.operands[ 0 ].size, instr.mnemonic, rva,
                                        transform::has_imm( &instr ) ? instr.operands[ 1 ].imm.value.u : 0 );

        return rva;
    }
} // namespace vm