packet dispatching/waiting (untested)

git-svn-id: file:///home/svn/incoming/trunk@2347 4d416f70-5f16-0410-b530-b9f4589650da

lib/rex/post/meterpreter/client.rb

@@ -3,6 +3,7 @@
 require 'socket'
 require 'Rex/Post/Meterpreter/Packet'
 require 'Rex/Post/Meterpreter/PacketParser'
+require 'Rex/Post/Meterpreter/PacketDispatcher'
 
 module Rex
 module Post
@@ -19,6 +20,8 @@ module Meterpreter
 ###
 class Client
 
+	include Rex::Post::Meterpreter::PacketDispatcher
+
 	def initialize(sock)
 		self.sock   = sock
 		self.parser = PacketParser.new
@@ -36,47 +39,6 @@ class Client
 	# Packet transmission/reception
 	#
 
-	def dispatch_inbound_packet(packet)
-		printf "Got packet with rid #{packet.rid}, method #{packet.method}\n"
-	end
-
-	def monitor_socket
-
-		# Spawn a new thread that monitors the socket
-		thr = Thread.new {
-			while (true)
-				rv = select([ self.sock ], nil, nil, 2)
-
-				begin
-					packet = receive_packet
-				rescue EOFError
-					puts "EOF reached on socket\n"
-					break
-				end
-
-				if (packet)
-					dispatch_inbound_packet(packet)
-				end
-			end
-		}
-
-	end
-
-	def send_packet(packet)
-		bytes = 0
-		raw   = packet.to_r
-
-		if (raw)
-			bytes = self.sock.write(raw)
-		end	
-
-		return bytes
-	end
-
-	def receive_packet
-		return parser.recv(self.sock)
-	end
-
 	protected
 	attr_accessor :sock, :parser
 end

lib/rex/post/meterpreter/packet.rb

@@ -330,6 +330,19 @@ class Packet < GroupTlv
 		end
 	end
 
+	#
+	# Conditionals
+	#
+
+	def response?
+		return ((self.type == PACKET_TYPE_RESPONSE) ||
+		        (self.type == PACKET_TYPE_PLAIN_RESPONSE))
+	end
+
+	#
+	# Accessors
+	#
+
 	def method?(method)
 		tlv = get_tlv(TLV_TYPE_METHOD)
 

lib/rex/post/meterpreter/packet_dispatcher.rb

@@ -0,0 +1,123 @@
+#!/usr/bin/ruby
+
+require 'Rex/Post/Meterpreter/PacketResponseWaiter'
+
+module Rex
+module Post
+module Meterpreter
+
+###
+#
+# PacketDispatcher
+# ----------------
+#
+# Handles packet transmission, reception, and correlation
+#
+###
+module PacketDispatcher
+	def dispatch_inbound_packet(packet)
+		puts "Inbound packet: rid=#{packet.rid} method=#{packet.method}\n"
+
+		# If the packet is a response, try to notify any potential
+		# waiters
+		if (packet.response?)
+			if (notify_response_waiter(packet))
+				return true
+			end
+		end
+	end
+
+	# Adds a waiter association with the supplied request packet
+	def add_response_waiter(request, completion_routine = nil, completion_param = nil)
+		waiter = PacketResponseWaiter.new(request.rid, completion_routine, completion_param)
+
+		self.waiters << waiter
+	end
+
+	# Notifies a whomever is waiting for a the supplied response,
+	# if anyone
+	def notify_response_waiter(response)
+		self.waiters.each() { |waiter|
+			if (waiter.waiting_for?(response))
+				waiter.notify(response)
+
+				remove_response_waiter(waiter)
+
+				break
+			end
+		}
+	end
+
+	# Removes a waiter from the list of waiters
+	def remove_response_waiter(waiter)
+		self.waiters.delete(waiter)
+	end
+
+	def monitor_socket
+		self.waiters = []
+
+		# Spawn a new thread that monitors the socket
+		thr = Thread.new {
+			while (true)
+				rv = select([ self.sock ], nil, nil, 2)
+
+				begin
+					packet = receive_packet
+				rescue EOFError
+					puts "EOF reached on socket\n"
+					break
+				end
+
+				if (packet)
+					dispatch_inbound_packet(packet)
+				end
+			end
+		}
+	end
+
+	# Sends a packet without waiting for a response
+	def send_packet(packet, completion_routine = nil, completion_param = nil)
+		if (completion_routine)
+			add_response_waiter(packet, completion_routine, completion_param)
+		end
+
+		bytes = 0
+		raw   = packet.to_r
+
+		if (raw)
+			bytes = self.sock.write(raw)
+		end	
+
+		return bytes
+	end
+
+	# Transmits a packet and waits for a response
+	def send_packet_wait_response(packet, t)
+		# First, add the waiter association for the supplied packet
+		waiter = add_response_waiter(packet)
+
+		# Transmit the packet
+		if (send_packet(packet) <= 0)
+			return nil
+		end
+
+		# Wait for the supplied time interval
+		waiter.wait(t)
+
+		# Remove the waiter from the list of waiters in case it wasn't 
+		# removed
+		remove_response_waiter(waiter)
+
+		# Return the response packet, if any
+		return waiter.response
+	end
+
+	def receive_packet
+		return parser.recv(self.sock)
+	end
+
+	protected
+	attr_accessor :waiters
+end
+
+end; end; end

lib/rex/post/meterpreter/packet_response_waiter.rb

@@ -0,0 +1,54 @@
+#!/usr/bin/ruby
+
+module Rex
+module Post
+module Meterpreter
+
+###
+#
+# PacketResponseWaiter
+# --------------------
+#
+# This class handles waiting for a response to a given request
+# and the subsequent response association
+#
+###
+class PacketResponseWaiter
+	def initialize(rid, completion_routine = nil, completion_param = nil)
+		self.rid      = rid
+		self.response = nil
+
+		if (completion_routine)
+			self.completion_routine = completion_routine
+			self.completion_param   = completion_param
+		else
+			self.mutex = Mutex.new
+			self.cond  = ConditionVariable.new
+		end
+	end
+
+	def waiting_for?(packet)
+		return (packet.rid == rid)
+	end
+
+	def notify(response)
+		self.response = response
+
+		if (self.completion_routine)
+			self.completion_routine(self.completion_param, response)
+		else
+			self.cond.signal
+		end
+	end
+
+	def wait(interval)
+		timeout(interval) { self.cond.wait(self.mutex) }
+
+		return self.response
+	end
+
+	attr_accessor :rid, :mutex, :cond, :response
+	attr_accessor :completion_routine, :completion_param
+end
+
+end; end; end