# -*- coding: binary -*-

require 'rex/post/meterpreter/extensions/stdapi/constants'
require 'rex/post/meterpreter/extensions/stdapi/sys/registry'

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

###
#
# Class wrapper around a logical registry key on the remote side
#
###
class RegistryKey


  #
  # Initializes an instance of a registry key using the supplied properties
  # and HKEY handle from the server.
  #
  def initialize(client, root_key, base_key, perm, hkey)
    self.client   = client
    self.root_key = root_key
    self.base_key = base_key
    self.perm     = perm
    self.hkey     = hkey

    # Ensure the remote object is closed when all references are removed
    ObjectSpace.define_finalizer(self, self.class.finalize(client, hkey))
  end

  def self.finalize(client,hkey)
    proc do
      deferred_close_proc = proc do
        begin
          self.close(client,hkey)
        rescue => e
          elog("finalize method for RegistryKey failed", error: e)
        end
      end

      # Schedule the finalizing logic out-of-band; as this logic might be called in the context of a Signal.trap, which can't synchronize mutexes
      client.framework.sessions.schedule(deferred_close_proc)
    end
  end

  ##
  #
  # Enumerators
  #
  ##

  #
  # Enumerates all of the child keys within this registry key.
  #
  def each_key(&block)
    return enum_key.each(&block)
  end

  #
  # Enumerates all of the child values within this registry key.
  #
  def each_value(&block)
    return enum_value.each(&block)
  end

  #
  # Retrieves all of the registry keys that are direct descendents of
  # the class' registry key.
  #
  def enum_key()
    return self.client.sys.registry.enum_key(self.hkey)
  end

  #
  # Retrieves all of the registry values that exist within the opened
  # registry key.
  #
  def enum_value()
    return self.client.sys.registry.enum_value(self.hkey)
  end


  ##
  #
  # Registry key interaction
  #
  ##

  #
  # Opens a registry key that is relative to this registry key.
  #
  def open_key(base_key, perm = KEY_READ)
    return self.client.sys.registry.open_key(self.hkey, base_key, perm)
  end

  #
  # Creates a registry key that is relative to this registry key.
  #
  def create_key(base_key, perm = KEY_READ)
    return self.client.sys.registry.create_key(self.hkey, base_key, perm)
  end

  #
  # Deletes a registry key that is relative to this registry key.
  #
  def delete_key(base_key, recursive = true)
    return self.client.sys.registry.delete_key(self.hkey, base_key, recursive)
  end

  #
  # Closes the open key.  This must be called if the registry
  # key was opened.
  #
  def self.close(client, hkey)
    if hkey != nil
      return client.sys.registry.close_key(hkey)
    end

    return false
  end

  # Instance method for the same
  def close()
    unless self.hkey.nil?
      ObjectSpace.undefine_finalizer(self)
      self.class.close(self.client, self.hkey)
      self.hkey = nil
    end
  end

  ##
  #
  # Registry value interaction
  #
  ##

  #
  # Sets a value relative to the opened registry key.
  #
  def set_value(name, type, data)
    return self.client.sys.registry.set_value(self.hkey, name, type, data)
  end

  #
  # Queries the attributes of the supplied registry value relative to
  # the opened registry key.
  #
  def query_value(name)
    return self.client.sys.registry.query_value(self.hkey, name)
  end

  #
  # Queries the class of the specified key
  #
  def query_class
    return self.client.sys.registry.query_class(self.hkey)
  end

  #
  # Delete the supplied registry value.
  #
  def delete_value(name)
    return self.client.sys.registry.delete_value(self.hkey, name)
  end

  ##
  #
  # Serializers
  #
  ##

  #
  # Returns the path to the key.
  #
  def to_s
    if self.base_key.nil?
      self.root_key.to_s + "\\"
    else
      self.root_key.to_s + "\\" + self.base_key
    end
  end

  #
  # The open handle to the key on the server.
  #
  attr_reader   :hkey
  #
  # The root key name, such as HKEY_LOCAL_MACHINE.
  #
  attr_reader   :root_key
  #
  # The base key name, such as Software\Foo.
  #
  attr_reader   :base_key
  #
  # The permissions that the key was opened with.
  #
  attr_reader   :perm

protected

  attr_accessor :client # :nodoc:
  attr_writer   :hkey, :root_key, :base_key, :perm # :nodoc:
end

end; end; end; end; end; end; end

