#include "vmhook.hpp"

namespace vm
{
    namespace handler
    {
        table_t::table_t( u64 module_base, u64 image_base, u32 table_rva, vm::handler::edit_entry_t edit_entry,
                          vm::decrypt_handler_t decrypt_handler, vm::encrypt_handler_t encrypt_handler )
            : module_base( module_base ), image_base( image_base ), table_rva( table_rva ), edit_entry( edit_entry ),
              decrypt_handler( decrypt_handler ), encrypt_handler( encrypt_handler )
        {
            table_addr = reinterpret_cast< u64 * >( module_base + table_rva );
            for ( auto idx = 0u; idx < 256; ++idx )
            {
                entry_t entry;
                entry.decrypted = decrypt( idx );
                entry.encrypted = get_entry( idx );
                entry.virt = ( entry.decrypted - image_base ) + module_base;
                entry.callback = nullptr;
                handlers[ idx ] = entry;
            }
        }

        u64 table_t::get_entry( u8 idx ) const
        {
            return table_addr[ idx ];
        }

        entry_t table_t::get_meta_data( u8 idx ) const
        {
            return handlers[ idx ];
        }

        void table_t::set_entry( u8 idx, u64 entry )
        {
            edit_entry( table_addr + idx, entry );
        }

        void table_t::set_meta_data( u8 idx, const entry_t &entry )
        {
            handlers[ idx ] = entry;
        }

        void table_t::set_callback( u8 idx, entry_callback_t callback )
        {
            handlers[ idx ].callback = callback;
        }

        u64 table_t::decrypt( u8 idx )
        {
            return decrypt_handler( get_entry( idx ) );
        }

        u64 table_t::encrypt( u64 val )
        {
            return encrypt_handler( val );
        }
    } // namespace handler

    hook_t::hook_t( void ) : table_count( 0 )
    {
        for ( auto idx = 0u; idx < 10; ++idx )
            handler_tables[ idx ] = nullptr;
    }

    void hook_t::add_table( vm::handler::table_t *table )
    {
        handler_tables[ table_count ] = table;
        ++table_count;
    }

    void hook_t::start( void )
    {
        for ( auto idx = 0u; idx < table_count; ++idx )
        {
            auto enc_trap_hndlr = handler_tables[ idx ]->encrypt(
                ( reinterpret_cast< std::uintptr_t >( &__vtrap ) - handler_tables[ idx ]->module_base ) +
                handler_tables[ idx ]->image_base );

            for ( auto table_idx = 0u; table_idx < 256; ++table_idx )
                handler_tables[ idx ]->set_entry( table_idx, enc_trap_hndlr );
        }
    }

    void hook_t::stop( void )
    {
        for ( auto idx = 0u; idx < table_count; ++idx )
            for ( auto table_idx = 0u; table_idx < 256; ++table_idx )
                handler_tables[ idx ]->set_entry( table_idx, // restore vm handler table encrypted values...
                                                  handler_tables[ idx ]->get_meta_data( table_idx ).encrypted );
    }

} // namespace vm

void vtrap_wrapper( vm::registers *regs, u8 handler_idx )
{
    // r12: vm handler linear virtual address
    // r13: module base linear virtual address
    auto table_rva = regs->r12 - ( regs->r13 + vm::g_vmctx->handler_tables[ 0 ]->image_base );

    // can only be a max of 10 vms... so idx < 10...
    for ( auto idx = 0u; idx < 10; ++idx )
    {
        if ( vm::g_vmctx->handler_tables[ idx ] && vm::g_vmctx->handler_tables[ idx ]->table_rva == table_rva )
        {
            regs->vm_handler = vm::g_vmctx->handler_tables[ idx ]->get_meta_data( handler_idx ).virt;
            const auto callback = vm::g_vmctx->handler_tables[ idx ]->get_meta_data( handler_idx ).callback;

            // per-virtual instruction callbacks...
            if ( callback )
                callback( regs, handler_idx );

            return;
        }
    }
}