#include <ZydisExportConfig.h>
#include <filesystem>
#include <fstream>
#include <iostream>
#include <algorithm>

#include <nt/image.hpp>
#include <cli-parser.hpp>
#include <vmprofiler.hpp>

auto open_binary_file( const std::string &file, std::vector< uint8_t > &data ) -> bool
{
    std::ifstream fstr( file, std::ios::binary );

    if ( !fstr.is_open() )
        return false;

    fstr.unsetf( std::ios::skipws );
    fstr.seekg( 0, std::ios::end );

    const auto file_size = fstr.tellg();

    fstr.seekg( NULL, std::ios::beg );
    data.reserve( static_cast< uint32_t >( file_size ) );
    data.insert( data.begin(), std::istream_iterator< uint8_t >( fstr ), std::istream_iterator< uint8_t >() );
    return true;
}

int __cdecl main( int argc, const char *argv[] )
{
    argparse::argument_parser_t parser( "vmprofiler-cli", "virtual machine information inspector" );
    parser.add_argument().names( { "--bin", "--vmpbin" } ).description( "unpacked binary protected with VMProtect 2" );
    parser.add_argument().names( { "--vmentry", "--entry" } ).description( "rva to push prior to a vm_entry" );
    parser.add_argument().name( "--showhandlers" ).description( "show all vm handlers..." );
    parser.add_argument().name( "--showhandler" ).description( "show a specific vm handler given its index..." );
    parser.add_argument().name( "--vmp2file" ).description( "path to .vmp2 file..." );
    parser.add_argument()
        .name( "--indexes" )
        .description( "displays vm handler table indexes for a given vm handler name such as 'READQ', or 'WRITEQ'..." );
    parser.add_argument()
        .name( "--scanfortables" )
        .description( "scans all executable sections for vm handler tables..." );
    parser.add_argument()
        .name( "--showvirtinstrs" )
        .description( "show the virtual instructions of a specific virtual routine..." );
    parser.add_argument().name( "--showallrtns" ).description( "shows all information for all virtual routines..." );
    parser.add_argument()
        .name( "--rtnaddr" )
        .description( "rva of the virtualized routine to apply another action upon such as --showallblocks..." );

    parser.enable_help();
    auto err = parser.parse( argc, argv );

    if ( err )
    {
        std::cout << err << std::endl;
        return -1;
    }

    if ( parser.exists( "help" ) )
    {
        parser.print_help();
        return 0;
    }

    if ( parser.exists( "bin" ) && parser.exists( "vmentry" ) )
    {
        if ( !std::filesystem::exists( parser.get< std::string >( "bin" ) ) )
        {
            std::printf( "> path to protected file is invalid... check your cli args...\n" );
            return -1;
        }

        std::vector<std::uint8_t> module_data, tmp;
        if(!open_binary_file(parser.get< std::string >( "bin" ), module_data))
        {
            std::printf("[!] failed to open binary = %s\n", parser.get< std::string >( "bin" ).c_str());
            return -1;
        }

        auto img = reinterpret_cast<win::image_t<>*>(module_data.data());
        auto image_size = img->get_nt_headers()->optional_header.size_image;

        tmp.resize(image_size);
        std::memcpy(tmp.data(), module_data.data(), 0x1000);
        std::for_each(img->get_nt_headers()->get_sections(), 
            img->get_nt_headers()->get_sections() + img->get_nt_headers()->file_header.num_sections, 
            [&](const auto& section_header)
        {
            std::memcpy(tmp.data() + section_header.virtual_address, 
                module_data.data() + section_header.ptr_raw_data,
                section_header.size_raw_data);
        }); 

        const auto module_base = reinterpret_cast<std::uintptr_t>(tmp.data());
        const auto vm_entry_rva = std::strtoull( parser.get< std::string >( "vmentry" ).c_str(), nullptr, 16 );
        const auto image_base = img->get_nt_headers()->optional_header.image_base;

        std::printf( "> module base = 0x%p, image base = 0x%p, image size = 0x%p\n", module_base, image_base,
                     image_size );

        vm::ctx_t vmctx( module_base, image_base, image_size, vm_entry_rva );

        if ( !vmctx.init() )
        {
            std::printf( "[!] failed to init vm::ctx_t... make sure all cli arguments are correct!\n" );
            return -1;
        }

        std::puts( "======================== [vm entry] ========================\n" );
        vm::util::print( vmctx.vm_entry );
        std::puts( "======================== [calc jmp] ========================\n" );
        vm::util::print( vmctx.calc_jmp );
        std::puts( "============================================================\n" );
        std::printf( "> vip advancement = %s\n\n",
                     vmctx.exec_type == vmp2::exec_type_t::forward ? "forward" : "backward" );

        if ( parser.exists( "showhandlers" ) )
        {
            for ( auto idx = 0u; idx < vmctx.vm_handlers.size(); ++idx )
            {
                std::printf( "======================== [%s #%d] ========================\n",
                             vmctx.vm_handlers[ idx ].profile ? vmctx.vm_handlers[ idx ].profile->name : "UNK", idx );

                vm::util::print( vmctx.vm_handlers[ idx ].instrs );

                // if there is no imm then there are no transforms...
                if ( !vmctx.vm_handlers[ idx ].imm_size )
                {
                    std::puts( "\n" );
                    continue;
                }

                std::puts( "======================== [transforms] ========================\n" );
                for ( auto &[ mnemonic, instr ] : vmctx.vm_handlers[ idx ].transforms )
                {
                    if ( instr.mnemonic == ZYDIS_MNEMONIC_INVALID )
                        continue;

                    vm::util::print( instr );
                }
                std::puts( "\n" );
            }
        }
        else if ( parser.exists( "showhandler" ) )
        {
            const auto vm_handler_idx = std::strtoul( parser.get< std::string >( "showhandler" ).c_str(), nullptr, 10 );

            if ( vm_handler_idx > 256 )
            {
                std::printf( "> invalid vm handler index... too large...\n" );
                return -1;
            }

            std::printf( "======================== [%s #%d] ========================\n",
                         vmctx.vm_handlers[ vm_handler_idx ].profile ? vmctx.vm_handlers[ vm_handler_idx ].profile->name
                                                                     : "UNK",
                         vm_handler_idx );

            vm::util::print( vmctx.vm_handlers[ vm_handler_idx ].instrs );

            // if there is no imm then there are no transforms...
            if ( !vmctx.vm_handlers[ vm_handler_idx ].imm_size )
            {
                std::puts( "\n" );
                return {};
            }

            std::puts( "======================== [transforms] ========================\n" );
            for ( auto &[ mnemonic, instr ] : vmctx.vm_handlers[ vm_handler_idx ].transforms )
            {
                if ( instr.mnemonic == ZYDIS_MNEMONIC_INVALID )
                    continue;

                vm::util::print( instr );
            }
            std::puts( "\n" );
        }
        else if ( parser.exists( "indexes" ) )
        {
            const auto handler_name = parser.get< std::string >( "indexes" );
            std::printf( "{\n" );
            for ( auto idx = 0u; idx < vmctx.vm_handlers.size(); ++idx )
                if ( vmctx.vm_handlers[ idx ].profile &&
                     !strcmp( vmctx.vm_handlers[ idx ].profile->name, handler_name.c_str() ) )
                    std::printf( "\t0x%x,\n", idx );
            std::printf( "}\n" );
        }
    }

    if ( !parser.exists( "vmp2file" ) )
        return {};

    std::vector< std::uint8_t > vmp2file;
    if ( !open_binary_file( parser.get< std::string >( "vmp2file" ), vmp2file ) )
    {
        std::printf( "[!] failed to open vmp2 file...\n" );
        return -1;
    }

    const auto file_header = reinterpret_cast< vmp2::v4::file_header * >( vmp2file.data() );

    if ( file_header->version != vmp2::version_t::v4 )
    {
        std::printf( "[!] invalid vmp2 file version... this build uses v4...\n" );
        return -1;
    }

    auto first_rtn = reinterpret_cast< vmp2::v4::rtn_t * >( reinterpret_cast< std::uintptr_t >( file_header ) +
                                                            file_header->rtn_offset );

    if ( parser.exists( "showallrtns" ) )
    {
        for ( auto [ rtn_block, rtn_idx ] = std::pair{ first_rtn, 0ull }; rtn_idx < file_header->rtn_count;
              ++rtn_idx, rtn_block = reinterpret_cast< vmp2::v4::rtn_t * >(
                             reinterpret_cast< std::uintptr_t >( rtn_block ) + rtn_block->size ) )
        {
            std::printf( "[rtn_0x%p (vm enter rva = 0x%x)] number of code blocks = %d\n",
                         rtn_block->code_blocks[ 0 ].vip_begin, rtn_block->vm_enter_offset,
                         rtn_block->code_block_count );

            for ( auto [ code_block, block_idx ] = std::pair{ &rtn_block->code_blocks[ 0 ], 0ull };
                  block_idx < rtn_block->code_block_count;
                  ++block_idx, code_block = reinterpret_cast< vmp2::v4::code_block_t * >(
                                   reinterpret_cast< std::uintptr_t >( code_block ) + code_block->next_block_offset ) )
            {
                std::printf( "    [blk_0x%p] number of virtual instructions = %d\n", code_block->vip_begin,
                             code_block->vinstr_count );
            }
        }
    }

    if ( parser.exists( "showvirtinstrs" ) && parser.exists( "rtnaddr" ) )
    {
        auto rtn_addr = std::stoull( parser.get< std::string >( "rtnaddr" ).c_str(), nullptr, 16 );
        for ( auto [ rtn_block, rtn_idx ] = std::pair{ first_rtn, 0ull }; rtn_idx < file_header->rtn_count;
              ++rtn_idx, rtn_block = reinterpret_cast< vmp2::v4::rtn_t * >(
                             reinterpret_cast< std::uintptr_t >( rtn_block ) + rtn_block->size ) )
        {
            if ( rtn_block->code_blocks[ 0 ].vip_begin == rtn_addr )
            {
                for ( auto [ code_block, block_idx ] = std::pair{ &rtn_block->code_blocks[ 0 ], 0ull };
                      block_idx < rtn_block->code_block_count;
                      ++block_idx, code_block = reinterpret_cast< vmp2::v4::code_block_t * >(
                                       reinterpret_cast< std::uintptr_t >( code_block ) +
                                       code_block->next_block_offset ) )
                {
                    std::printf( "[blk_0x%p] number of virtual instructions = %d\n", code_block->vip_begin,
                                 code_block->vinstr_count );

                    std::printf( "> -----------------------------------------------------------------------\n" );
                    std::printf( "> opcode | virtual instructions             | virtual instruction pointer\n" );
                    std::printf( "> -----------------------------------------------------------------------\n" );

                    auto block_vinstrs = reinterpret_cast< vm::instrs::virt_instr_t * >(
                        reinterpret_cast< std::uintptr_t >( code_block ) + sizeof(vmp2::v4::code_block_t) +
                        ( code_block->num_block_addrs * 8 ) );

                    for ( auto idx = 0u; idx < code_block->vinstr_count; ++idx )
                    {
                        const auto vinstr = &block_vinstrs[ idx ];
                        const auto vm_profile = vm::handler::get_profile( vinstr->mnemonic_t );
                        if ( vinstr->operand.has_imm )
                        {
                            std::printf( "> %-6x | %-15s %-15p | 0x%p\n", vinstr->opcode,
                                         vm_profile ? vm_profile->name : "UNK", vinstr->operand.imm.u,
                                         ABS_TO_IMG( vinstr->trace_data.vip, file_header->module_base,
                                                     file_header->image_base ) );
                        }
                        else
                        {
                            std::printf( "> %-6x | %-32s | 0x%p\n", vinstr->opcode,
                                         vm_profile ? vm_profile->name : "UNK",
                                         ABS_TO_IMG( vinstr->trace_data.vip, file_header->module_base,
                                                     file_header->image_base ) );
                        }
                    }
                }
            }
        }
    }
}