client.rb in librex-0.0.66

- old
+ new

@@ -1,6 +1,7 @@
 #!/usr/bin/env ruby
+# -*- coding: binary -*-
 
 require 'socket'
 require 'openssl'
 
 require 'rex/script'
@@ -102,10 +103,11 @@
 		self.ext          = ObjectAliases.new
 		self.ext_aliases  = ObjectAliases.new
 		self.alive        = true
 		self.target_id    = opts[:target_id]
 		self.capabilities = opts[:capabilities] || {}
+		self.commands     = []
 
 
 		self.conn_id      = opts[:conn_id]
 		self.url          = opts[:url]
 		self.ssl          = opts[:ssl]
@@ -279,10 +281,11 @@
 	#
 	# Translates unhandled methods into registered extension aliases
 	# if a matching extension alias exists for the supplied symbol.
 	#
 	def method_missing(symbol, *args)
+		#$stdout.puts("method_missing: #{symbol}")
 		self.ext_aliases.aliases[symbol.to_s]
 	end
 
 	##
 	#
@@ -292,11 +295,13 @@
 
 	#
 	# Loads the client half of the supplied extension and initializes it as a
 	# registered extension that can be reached through client.ext.[extension].
 	#
-	def add_extension(name)
+	def add_extension(name, commands=[])
+		self.commands |= commands
+
 		# Check to see if this extension has already been loaded.
 		if ((klass = self.class.check_ext_hash(name.downcase)) == nil)
 			old = Rex::Post::Meterpreter::Extensions.constants
 			require("rex/post/meterpreter/extensions/#{name.downcase}/#{name.downcase}")
 			new = Rex::Post::Meterpreter::Extensions.constants
@@ -339,10 +344,22 @@
 	# Registers an aliased extension that can be referenced through
 	# client.name.
 	#
 	def register_extension_alias(name, ext)
 		self.ext_aliases.aliases[name] = ext
+		# Whee!  Syntactic sugar, where art thou?
+		#
+		# Create an instance method on this object called +name+ that returns
+		# +ext+.  We have to do it this way instead of simply
+		# self.class.class_eval so that other meterpreter sessions don't get
+		# extension methods when this one does
+		(class << self; self; end).class_eval do
+			define_method(name.to_sym) do
+				ext
+			end
+		end
+		ext
 	end
 
 	#
 	# Registers zero or more aliases that are provided in an array.
 	#
@@ -443,13 +460,18 @@
 	attr_accessor :passive_dispatcher
 	#
 	# Flag indicating whether to hex-encode UTF-8 file names and other strings
 	#
 	attr_accessor :encode_unicode
+	#
+	# A list of the commands
+	#
+	attr_reader :commands
 
 protected
 	attr_accessor :parser, :ext_aliases # :nodoc:
 	attr_writer   :ext, :sock # :nodoc:
+	attr_writer   :commands # :nodoc:
 end
 
 end; end; end
 