#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