# -*- coding: binary -*-

require 'rex/post/process'
require 'rex/post/meterpreter/packet'
require 'rex/post/meterpreter/client'
require 'rex/post/meterpreter/channels/pools/stream_pool'
require 'rex/post/meterpreter/extensions/stdapi/stdapi'

require 'rex/post/meterpreter/extensions/stdapi/sys/process_subsystem/image'
require 'rex/post/meterpreter/extensions/stdapi/sys/process_subsystem/io'
require 'rex/post/meterpreter/extensions/stdapi/sys/process_subsystem/memory'
require 'rex/post/meterpreter/extensions/stdapi/sys/process_subsystem/thread'

module Rex
module Post
module Meterpreter
module Extensions
module Stdapi
module Sys

##
#
# This class implements the Rex::Post::Process interface.
#
##
class Process < Rex::Post::Process

  include Rex::Post::Meterpreter::ObjectAliasesContainer

  ##
  #
  # Class methods
  #
  ##

  class << self
    attr_accessor :client
  end

  #
  # Returns the process identifier of the process supplied in key if it's
  # valid.
  #
  def Process.[](key)
    each_process { |p|
      if (p['name'].downcase == key.downcase)
        return p['pid']
      end
    }

    return nil
  end

  #
  # Attachs to the supplied process with a given set of permissions.
  #
  def Process.open(pid = nil, perms = nil)
    real_perms = 0

    if (perms == nil)
      perms = PROCESS_ALL
    end

    if (perms & PROCESS_READ)
      real_perms |= PROCESS_VM_OPERATION | PROCESS_VM_READ | PROCESS_QUERY_INFORMATION
    end

    if (perms & PROCESS_WRITE)
      real_perms |= PROCESS_SET_SESSIONID | PROCESS_VM_WRITE | PROCESS_DUP_HANDLE | PROCESS_SET_QUOTA | PROCESS_SET_INFORMATION
    end

    if (perms & PROCESS_EXECUTE)
      real_perms |= PROCESS_TERMINATE | PROCESS_CREATE_THREAD | PROCESS_CREATE_PROCESS | PROCESS_SUSPEND_RESUME
    end

    return _open(pid, real_perms)
  end

  #
  # Low-level process open.
  #
  def Process._open(pid, perms, inherit = false)
    request = Packet.create_request('stdapi_sys_process_attach')

    if (pid == nil)
      pid = 0
    end

    # Populate the request
    request.add_tlv(TLV_TYPE_PID, pid)
    request.add_tlv(TLV_TYPE_PROCESS_PERMS, perms)
    request.add_tlv(TLV_TYPE_INHERIT, inherit)

    # Transmit the request
    response = self.client.send_request(request)
    handle   = response.get_tlv_value(TLV_TYPE_HANDLE)

    # If the handle is valid, allocate a process instance and return it
    if (handle != nil)
      return self.new(pid, handle)
    end

    return nil
  end

  #
  # Executes an application using the arguments provided
  #
  # Hash arguments supported:
  #
  #   Hidden      => true/false
  #   Channelized => true/false
  #   Suspended   => true/false
  #   InMemory    => true/false
  #
  def Process.execute(path, arguments = nil, opts = nil)
    request = Packet.create_request('stdapi_sys_process_execute')
    flags   = 0

    # If we were supplied optional arguments...
    if (opts != nil)
      if (opts['Hidden'])
        flags |= PROCESS_EXECUTE_FLAG_HIDDEN
      end
      if (opts['Channelized'])
        flags |= PROCESS_EXECUTE_FLAG_CHANNELIZED
      end
      if (opts['Suspended'])
        flags |= PROCESS_EXECUTE_FLAG_SUSPENDED
      end
      if (opts['UseThreadToken'])
        flags |= PROCESS_EXECUTE_FLAG_USE_THREAD_TOKEN
      end
      if (opts['Desktop'])
        flags |= PROCESS_EXECUTE_FLAG_DESKTOP
      end
      if (opts['Session'])
        flags |= PROCESS_EXECUTE_FLAG_SESSION
        request.add_tlv( TLV_TYPE_PROCESS_SESSION, opts['Session'] )
      end
      inmem = opts['InMemory']
      if inmem

        # add the file contents into the tlv
        f = ::File.new(path, 'rb')
        request.add_tlv(TLV_TYPE_VALUE_DATA, f.read(f.stat.size))
        f.close

        # replace the path with the "dummy"
        path = inmem.kind_of?(String) ? inmem : 'cmd'
      end
    end

    request.add_tlv(TLV_TYPE_PROCESS_PATH, client.unicode_filter_decode( path ));

    # If process arguments were supplied
    if (arguments != nil)
      request.add_tlv(TLV_TYPE_PROCESS_ARGUMENTS, arguments);
    end

    request.add_tlv(TLV_TYPE_PROCESS_FLAGS, flags);

    response = client.send_request(request)

    # Get the response parameters
    pid        = response.get_tlv_value(TLV_TYPE_PID)
    handle     = response.get_tlv_value(TLV_TYPE_PROCESS_HANDLE)
    channel_id = response.get_tlv_value(TLV_TYPE_CHANNEL_ID)
    channel    = nil

    # If we were creating a channel out of this
    if (channel_id != nil)
      channel = Rex::Post::Meterpreter::Channels::Pools::StreamPool.new(client,
          channel_id, "stdapi_process", CHANNEL_FLAG_SYNCHRONOUS)
    end

    # Return a process instance
    return self.new(pid, handle, channel)
  end

  #
  # Kills one or more processes.
  #
  def Process.kill(*args)
    request = Packet.create_request('stdapi_sys_process_kill')

    args.each { |id|
      request.add_tlv(TLV_TYPE_PID, id)
    }

    client.send_request(request)

    return true
  end

  #
  # Gets the process id that the remote side is executing under.
  #
  def Process.getpid
    request = Packet.create_request('stdapi_sys_process_getpid')

    response = client.send_request(request)

    return response.get_tlv_value(TLV_TYPE_PID)
  end

  #
  # Enumerates all of the elements in the array returned by get_processes.
  #
  def Process.each_process(&block)
    self.get_processes.each(&block)
  end

  #
  # Returns a ProcessList of processes as Hash objects with keys for 'pid',
  # 'ppid', 'name', 'path', 'user', 'session' and 'arch'.
  #
  def Process.get_processes
    request   = Packet.create_request('stdapi_sys_process_get_processes')
    processes = ProcessList.new

    response = client.send_request(request)

    response.each(TLV_TYPE_PROCESS_GROUP) { |p|
    arch = ""

    pa = p.get_tlv_value( TLV_TYPE_PROCESS_ARCH )
    if( pa != nil )
      if pa == 1 # PROCESS_ARCH_X86
        arch = ARCH_X86
      elsif pa == 2 # PROCESS_ARCH_X64
        arch = ARCH_X86_64
      end
    end

    processes <<
        {
          'pid'      => p.get_tlv_value(TLV_TYPE_PID),
          'ppid'     => p.get_tlv_value(TLV_TYPE_PARENT_PID),
          'name'     => client.unicode_filter_encode( p.get_tlv_value(TLV_TYPE_PROCESS_NAME) ),
          'path'     => client.unicode_filter_encode( p.get_tlv_value(TLV_TYPE_PROCESS_PATH) ),
          'session'  => p.get_tlv_value(TLV_TYPE_PROCESS_SESSION),
          'user'     => client.unicode_filter_encode( p.get_tlv_value(TLV_TYPE_USER_NAME) ),
          'arch'     => arch
        }
    }

    return processes
  end

  #
  # An alias for get_processes.
  #
  def Process.processes
    self.get_processes
  end

  ##
  #
  # Instance methods
  #
  ##

  #
  # Initializes the process instance and its aliases.
  #
  def initialize(pid, handle, channel = nil)
    self.client  = self.class.client
    self.handle  = handle
    self.channel = channel

    # If the process identifier is zero, then we must lookup the current
    # process identifier
    if (pid == 0)
      self.pid = client.sys.process.getpid
    else
      self.pid = pid
    end

    initialize_aliases(
      {
        'image'  => Rex::Post::Meterpreter::Extensions::Stdapi::Sys::ProcessSubsystem::Image.new(self),
        'io'     => Rex::Post::Meterpreter::Extensions::Stdapi::Sys::ProcessSubsystem::IO.new(self),
        'memory' => Rex::Post::Meterpreter::Extensions::Stdapi::Sys::ProcessSubsystem::Memory.new(self),
        'thread' => Rex::Post::Meterpreter::Extensions::Stdapi::Sys::ProcessSubsystem::Thread.new(self),
      })

    ObjectSpace.define_finalizer( self, self.class.finalize(self.client, self.handle) )
  end

  def self.finalize(client,handle)
    proc { self.close(client,handle) }
  end

  #
  # Returns the executable name of the process.
  #
  def name
    return get_info()['name']
  end

  #
  # Returns the path to the process' executable.
  #
  def path
    return get_info()['path']
  end

  #
  # Closes the handle to the process that was opened.
  #
  def self.close(client, handle)
    request = Packet.create_request('stdapi_sys_process_close')
    request.add_tlv(TLV_TYPE_HANDLE, handle)
    response = client.send_request(request, nil)
    handle = nil;
    return true
  end

  #
  # Instance method
  #
  def close(handle=self.handle)
    self.class.close(self.client, handle)
  end

  #
  # Block until this process terminates on the remote side.
  # By default we choose not to allow a packet responce timeout to
  # occur as we may be waiting indefinatly for the process to terminate.
  #
  def wait( timeout = -1 )
    request = Packet.create_request('stdapi_sys_process_wait')

    request.add_tlv(TLV_TYPE_HANDLE, self.handle)

    response = self.client.send_request(request, timeout)

    self.handle = nil

    return true
  end

  attr_reader   :client, :handle, :channel, :pid # :nodoc:
protected
  attr_writer   :client, :handle, :channel, :pid # :nodoc:

  #
  # Gathers information about the process and returns a hash.
  #
  def get_info
    request = Packet.create_request('stdapi_sys_process_get_info')
    info    = {}

    request.add_tlv(TLV_TYPE_HANDLE, handle)

    # Send the request
    response = client.send_request(request)

    # Populate the hash
    info['name'] = client.unicode_filter_encode( response.get_tlv_value(TLV_TYPE_PROCESS_NAME) )
    info['path'] = client.unicode_filter_encode( response.get_tlv_value(TLV_TYPE_PROCESS_PATH) )

    return info
  end

end

#
# Simple wrapper class for storing processes
#
class ProcessList < Array

  #
  # Create a Rex::Ui::Text::Table out of the processes stored in this list
  #
  # +opts+ is passed on to Rex::Ui::Text::Table.new, mostly unmolested
  #
  # Note that this output is affected by Rex::Post::Meterpreter::Client#unicode_filter_encode
  #
  def to_table(opts={})
    if empty?
      return Rex::Ui::Text::Table.new(opts)
    end

    cols = [ "PID", "PPID", "Name", "Arch", "Session", "User", "Path" ]
    # Arch and Session are specific to native Windows, PHP and Java can't do
    # ppid.  Cut columns from the list if they aren't there.  It is conceivable
    # that processes might have different columns, but for now assume that the
    # first one is representative.
    cols.delete_if { |c| !( first.has_key?(c.downcase) ) or first[c.downcase].nil? }

    opts = {
      "Header"  => "Process List",
      "Columns" => cols
    }.merge(opts)

    tbl = Rex::Ui::Text::Table.new(opts)
    each { |process|
      tbl << cols.map {|c| process[c.downcase] }.compact
    }

    tbl
  end
end

end; end; end; end; end; end

