#pragma once
#include <optional>
#include <transform.hpp>
#include <vmp2.hpp>

namespace vm
{
    namespace handler
    {
        using instr_callback_t = bool ( * )( const zydis_decoded_instr_t &instr );

        enum mnemonic_t
        {
            INVALID,
            LRFLAGS,
            PUSHVSP,
            MULQ,
            DIVQ,
            CALL,
            JMP,
            VMEXIT,

            SREGQ,
            SREGDW,
            SREGW,

            LREGQ,
            LREGDW,

            LCONSTQ,
            LCONSTBZXW,
            LCONSTBSXQ,
            LCONSTBSXDW,
            LCONSTDWSXQ,
            LCONSTWSXQ,
            LCONSTWSXDW,
            LCONSTDW,
            LCONSTW,

            READQ,
            READDW,
            READW,

            WRITEQ,
            WRITEDW,
            WRITEW,
            WRITEB,

            ADDQ,
            ADDDW,
            ADDW,

            SHLQ,
            SHLDW,

            SHRQ,
            SHRW,

            NANDQ,
            NANDDW,
            NANDW
        };

        enum extention_t
        {
            none,
            sign_extend,
            zero_extend
        };

        struct profile_t
        {
            const char *name;
            mnemonic_t mnemonic;
            u8 imm_size;
            std::vector< instr_callback_t > signature;
            extention_t extention;
        };

        struct handler_t
        {
            u8 imm_size; // size in bits...
            vm::transform::map_t transforms;
            vm::handler::profile_t *profile;
            zydis_routine_t instrs;
            std::uintptr_t address;
        };

        bool has_imm( const zydis_routine_t &vm_handler );
        std::uint8_t imm_size( const zydis_routine_t &vm_handler );
        bool get( zydis_routine_t &vm_entry, zydis_routine_t &vm_handler, std::uintptr_t handler_addr );

        // may throw an exception...
        bool get_all( std::uintptr_t module_base, std::uintptr_t image_base, zydis_routine_t &vm_entry,
                      std::uintptr_t *vm_handler_table, std::vector< handler_t > &vm_handlers );

        // can be used on calc_jmp...
        bool get_operand_transforms( const zydis_routine_t &vm_handler, transform::map_t &transforms );
        vm::handler::profile_t *get_profile( handler_t &vm_handler );
        vm::handler::profile_t *get_profile( vm::handler::mnemonic_t mnemonic );

        namespace table
        {
            std::uintptr_t *get( const zydis_routine_t &vm_entry );
            bool get_transform( const zydis_routine_t &vm_entry, zydis_decoded_instr_t *transform_instr );

            std::uint64_t encrypt( zydis_decoded_instr_t &transform_instr, std::uint64_t val );
            std::uint64_t decrypt( zydis_decoded_instr_t &transform_instr, std::uint64_t val );
        } // namespace table

        namespace profile
        {
            extern vm::handler::profile_t sregq;
            extern vm::handler::profile_t sregdw;
            extern vm::handler::profile_t sregw;

            extern vm::handler::profile_t lregq;
            extern vm::handler::profile_t lregdw;

            extern vm::handler::profile_t lconstq;
            extern vm::handler::profile_t lconstdw;
            extern vm::handler::profile_t lconstw;

            extern vm::handler::profile_t lconstbzxw;
            extern vm::handler::profile_t lconstbsxdw;
            extern vm::handler::profile_t lconstbsxq;
            extern vm::handler::profile_t lconstdwsxq;
            extern vm::handler::profile_t lconstwsxq;
            extern vm::handler::profile_t lconstwsxdw;

            extern vm::handler::profile_t addq;
            extern vm::handler::profile_t adddw;
            extern vm::handler::profile_t addw;

            extern vm::handler::profile_t shlq;
            extern vm::handler::profile_t shldw;

            extern vm::handler::profile_t nandq;
            extern vm::handler::profile_t nanddw;
            extern vm::handler::profile_t nandw;

            extern vm::handler::profile_t writeq;
            extern vm::handler::profile_t writedw;
            extern vm::handler::profile_t writeb;

            extern vm::handler::profile_t readq;
            extern vm::handler::profile_t readdw;

            extern vm::handler::profile_t shrq;
            extern vm::handler::profile_t shrw;

            extern vm::handler::profile_t lrflags;
            extern vm::handler::profile_t call;
            extern vm::handler::profile_t pushvsp;
            extern vm::handler::profile_t mulq;
            extern vm::handler::profile_t divq;
            extern vm::handler::profile_t jmp;
            extern vm::handler::profile_t vmexit;

            inline std::vector< vm::handler::profile_t * > all = {
                &sregq,      &sregdw,      &sregw,      &lregq,       &lregdw,   &lconstq, &lconstbzxw, &lconstbsxdw,
                &lconstbsxq, &lconstdwsxq, &lconstwsxq, &lconstwsxdw, &lconstdw, &lconstw, &addq,       &adddw,
                &addw,

                &shlq,       &shldw,       &writeq,     &writedw,     &writeb,   &nandq,   &nanddw,     &nandw,

                &shrq,       &shrw,        &readq,      &readdw,      &mulq,     &pushvsp, &divq,       &jmp,
                &lrflags,    &vmexit,      &call };
        } // namespace profile
    }     // namespace handler

    class ctx_t
    {
      public:
        explicit ctx_t( std::uintptr_t module_base, std::uintptr_t image_base, std::uintptr_t image_size,
                        std::uintptr_t vm_entry_rva );

        bool init();
        const std::uintptr_t module_base, image_base, vm_entry_rva, image_size;
        vmp2::exec_type_t exec_type;
        zydis_routine_t vm_entry, calc_jmp;
        std::vector< vm::handler::handler_t > vm_handlers;
    };
} // namespace vm

namespace vm
{
    namespace instrs
    {
        struct virt_instr_t
        {
            vm::handler::mnemonic_t mnemonic_t;
            std::uint8_t opcode; // aka vm handler idx...

            // can be used to look at values on the stack...
            vmp2::v2::entry_t trace_data;

            struct
            {
                bool has_imm;
                struct
                {
                    std::uint8_t imm_size; // size in bits...
                    union
                    {
                        std::int64_t s;
                        std::uint64_t u;
                    };
                } imm;
            } operand;
        };

        enum class jcc_type
        {
            none,
            branching,
            absolute
        };

        struct jcc_data
        {
            bool has_jcc;
            jcc_type type;
            std::uintptr_t block_addr[ 2 ];
        };

        struct code_block_t
        {
            std::uintptr_t vip_begin;
            jcc_data jcc;
            std::vector< virt_instr_t > vinstrs;
        };
    } // namespace instrs
} // namespace vm

namespace vmp2
{
    namespace v3
    {
        struct file_header
        {
            u32 magic; // VMP2
            u64 epoch_time;
            version_t version;

            u64 module_base;
            u64 image_base;
            u64 vm_entry_rva;

            u32 module_offset;
            u32 module_size;

            u32 code_block_offset;
            u32 code_block_count;
        };

        struct code_block_t
        {
            std::uintptr_t vip_begin;
            std::uintptr_t next_block_offset;
            vm::instrs::jcc_data jcc;

            // serialized from std::vector<virt_instr_t>...
            std::uint32_t vinstr_count;
            vm::instrs::virt_instr_t vinstr[];
        };
    } // namespace v3
} // namespace vmp2

namespace vm
{
    namespace instrs
    {
        // decrypt transformations for encrypted virtual instruction rva...
        bool get_rva_decrypt( const zydis_routine_t &vm_entry, std::vector< zydis_decoded_instr_t > &transform_instrs );

        std::pair< std::uint64_t, std::uint64_t > decrypt_operand( transform::map_t &transforms, std::uint64_t operand,
                                                                   std::uint64_t 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 );

        /// <summary>
        /// get virt_instr_t filled in with data given a vmp2 trace entry and vm context...
        /// </summary>
        /// <param name="ctx">current vm context</param>
        /// <param name="entry">vmp2 trace entry containing all of the native/virtual register/stack values...</param>
        /// <returns>returns a filled in virt_instr_t on success...</returns>
        std::optional< virt_instr_t > get( vm::ctx_t &ctx, vmp2::v2::entry_t &entry );

        /// <summary>
        /// gets the second operand (imm) given vip and vm::ctx_t...
        /// </summary>
        /// <param name="ctx">vm context</param>
        /// <param name="imm_size">immediate value size in bits...</param>
        /// <param name="vip">virtual instruction pointer, linear virtual address...</param>
        /// <returns>returns immediate value if imm_size is not 0...</returns>
        std::optional< std::uint64_t > get_imm( vm::ctx_t &ctx, std::uint8_t imm_size, std::uintptr_t vip );

        /// <summary>
        /// get jcc data out of a code block... this function will loop over the code block
        /// and look for the last two NANDW in the virtual instructions, then it will look
        /// for the last LCONSTW which is the xor key...
        ///
        /// it will then loop and look for all PUSHVSP's, checking each to see if the stack
        /// contains two encrypted rva's to each branch.. if there is not two encrypted rva's
        /// then the virtual jmp instruction only has one dest...
        /// </summary>
        /// <param name="ctx">vm context</param>
        /// <param name="code_block">code block that does not have its jcc_data yet</param>
        /// <returns>if last lconstdw is found, return filled in jcc_data structure...</returns>
        std::optional< jcc_data > get_jcc_data( vm::ctx_t &ctx, code_block_t &code_block );

        /// <summary>
        /// the top of the stack will contain the lower 32bits of the RVA to the virtual instructions
        /// that will be jumping too... the RVA is image based (not module based, but optional header image
        /// based)... this means the value ontop of the stack could be "40007fd8" with image base being
        /// 0x140000000... as you can see the 0x100000000 is missing... the below statement deals with this...
        /// </summary>
        /// <param name="ctx">vm context</param>
        /// <param name="entry">current trace entry for virtual JMP instruction</param>
        /// <returns>returns linear virtual address of the next code block...</returns>
        std::uintptr_t code_block_addr( const vm::ctx_t &ctx, const vmp2::v2::entry_t &entry );

        /// <summary>
        /// same routine as above except lower_32bits is passed directly and not extracted from the stack...
        /// </summary>
        /// <param name="ctx">vm context</param>
        /// <param name="lower_32bits">lower 32bits of the relative virtual address...</param>
        /// <returns>returns full linear virtual address of code block...</returns>
        std::uintptr_t code_block_addr( const vm::ctx_t &ctx, const std::uint32_t lower_32bits );
    } // namespace instrs

    namespace calc_jmp
    {
        bool get( const zydis_routine_t &vm_entry, zydis_routine_t &calc_jmp );

        std::optional< vmp2::exec_type_t > get_advancement( const zydis_routine_t &calc_jmp );
    } // namespace calc_jmp
} // namespace vm