some more madness

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

lib/msf/base/sessions/command_shell.rb

@@ -25,6 +25,14 @@ class CommandShell
 	#
 	include Msf::Session::Provider::SingleCommandShell
 
+	def desc
+		"Command Shell"
+	end
+
+	def type
+		"shell"
+	end
+
 	#
 	# The shell will have been initialized by default
 	#

lib/msf/core/framework.rb

@@ -30,10 +30,12 @@ class Framework
 	end
 
 	require 'msf/core/module_manager'
+	require 'msf/core/session_manager'
 
 	def initialize()
 		self.events    = EventDispatcher.new
 		self.modules   = ModuleManager.new(self)
+		self.sessions  = SessionManager.new(self)
 		self.datastore = DataStore.new
 	end
 
@@ -74,12 +76,14 @@ class Framework
 
 	attr_reader   :events
 	attr_reader   :modules
+	attr_reader   :sessions
 	attr_reader   :datastore
 
 protected
 
 	attr_writer   :events
 	attr_writer   :modules
+	attr_writer   :sessions
 	attr_writer   :datastore
 
 end

lib/msf/core/handler.rb

@@ -102,6 +102,9 @@ module Handler
 		if (self.session)
 			s = self.session.new(conn)
 
+			# Pass along the framework context
+			s.framework = framework
+
 			# If the session is valid, register it with the framework and
 			# notify any waiters we may have.
 			if (s)
@@ -140,9 +143,11 @@ protected
 	# new session.
 	#
 	def register_session(session)
-		session_waiter_event.notify(session)
+		# Register the session with the framework
+		framework.sessions.register(session)
 
-		# TODO: register with the framework
+		# Notify waiters that they should be ready to rock
+		session_waiter_event.notify(session)
 	end
 
 	attr_accessor :session_waiter_event

lib/msf/core/handler/reverse_tcp.rb

@@ -80,7 +80,11 @@ module ReverseTcp
 			# as the input and output pipe.  Client's are expected
 			# to implement the Stream interface.
 			conn_threads << Thread.new {
-				handle_connection(client)
+				begin
+					handle_connection(client)
+				rescue
+					elog("Exception raised from handle_connection: #{$!}")
+				end
 			}
 		}
 	end

lib/msf/core/session.rb

@@ -52,14 +52,58 @@ module Session
 	require 'msf/core/session/provider/multi_command_execution'
 	require 'msf/core/session/provider/single_command_shell'
 	require 'msf/core/session/provider/multi_command_shell'
+	#
+	# Returns the session's name if it's been assigned one, otherwise
+	# the sid is returned.
+	#
+	def name
+		return sname || sid
+	end
 
 	#
-	# By default, sessions are not interactive.
+	# Sets the session's name
 	#
-	def interactive?
-		false
+	def name=(name)
+		self.sname = name
 	end
 
+	#
+	# Returns the description of the session
+	#
+	def desc
+	end
+
+	#
+	# Returns the type of session in use
+	#
+	def type
+	end
+
+	#
+	# Returns the local side of the tunnel
+	#
+	def tunnel_local
+	end
+
+	#
+	# Returns the peer side of the tunnel
+	#
+	def tunnel_peer
+	end
+
+	#
+	# Returns a pretty representation of the tunnel
+	#
+	def tunnel_to_s
+		"#{(tunnel_local || '??').to_s} -> #{(tunnel_peer || '??').to_s}"
+	end
+
+	##
+	#
+	# Core interface
+	#
+	##
+
 	#
 	# Perform session-specific cleanup
 	#
@@ -67,21 +111,13 @@ module Session
 	end
 
 	#
-	# Returns the session's name if it's been assigned one, otherwise
-	# the sid is returned.
+	# By default, sessions are not interactive.
 	#
-	def sname
-		return name || sid
+	def interactive?
+		false
 	end
 
-	#
-	# Sets the session's name
-	#
-	def sname=(name)
-		self.name = name
-	end
-
-	attr_accessor :framework, :sid, :name
+	attr_accessor :framework, :sid, :sname
 
 protected
 

lib/msf/core/session/basic.rb

@@ -30,6 +30,34 @@ module Basic
 		true
 	end
 
+	#
+	# Description of the session
+	#
+	def desc
+		"Basic Session"
+	end
+
+	#
+	# Basic session
+	#
+	def type
+		"basic"
+	end
+	
+	#
+	# Returns the local information
+	#
+	def tunnel_local
+		rstream.localinfo
+	end
+
+	#
+	# Returns the remote peer information
+	#
+	def tunnel_peer
+		rstream.peerinfo
+	end
+	
 	#
 	# Closes rstream.
 	#
@@ -44,6 +72,8 @@ module Basic
 	# rstream to loutput.
 	#
 	def interact
+		eof = false
+
 		callcc { |ctx|
 			while true
 				begin
@@ -52,18 +82,23 @@ module Basic
 				# abort the interaction.  If they do, then we return out of
 				# the interact function and call it a day.
 				rescue Interrupt
-					loutput.print("\nStop interacting with session #{sname}? [y/N]  ")
+					loutput.print("\nStop interacting with session #{name}? [y/N]  ")
 
 					r = linput.gets
 
 					# Break out of the continuation
 					ctx.call if (r =~ /^y/i)
 				rescue EOFError
-					loutput.print_line("Session #{sname} terminating...")
+					loutput.print_line("Session #{name} terminating...")
+					eof = true
 					ctx.call
 				end
 			end
 		}
+
+		# If we hit end-of-file, then that means we should finish off this
+		# session and call it a day.
+		framework.sessions.deregister(self) if (eof == true)
 	end
 
 	#

lib/msf/core/session_manager.rb

@@ -16,37 +16,15 @@ module Msf
 # with.
 #
 ###
-class SessionManager
+class SessionManager < Hash
 
-	include Enumerable
+	include Framework::Offspring
 
 	def initialize(framework)
 		self.framework = framework
-		self.sessions  = {}
 		self.sid_pool  = 0
 	end
 
-	#
-	# Returns the session object that is associated with the supplied sid
-	#
-	def [](sid)
-		return get(sid)
-	end
-
-	#
-	# Register the supplied session
-	#
-	def <<(session)
-		return register(session)
-	end
-
-	#
-	# Implement for Enumerable
-	#
-	def each(&block)
-		sessions.each(&block)
-	end
-
 	#
 	# Registers the supplied session object with the framework and returns
 	# a unique session identifier to the caller.
@@ -60,7 +38,7 @@ class SessionManager
 		next_sid = (self.sid_pool += 1)
 
 		# Insert the session into the session hash table
-		sessions[next_sid] = session
+		self[next_sid] = session
 
 		# Initialize the session's sid and framework instance pointer
 		session.sid       = next_sid
@@ -80,7 +58,7 @@ class SessionManager
 		framework.events.on_session_close(session)
 
 		# Remove it from the hash
-		sessions.delete(session.sid)
+		self.delete(session.sid)
 
 		# Close it down
 		session.cleanup
@@ -90,12 +68,12 @@ class SessionManager
 	# Returns the session associated with the supplied sid, if any
 	#
 	def get(sid)
-		return sessions[sid]
+		return self[sid]
 	end
 
 protected
 	
-	attr_accessor :sid_pool, :sessions, :framework
+	attr_accessor :sid_pool, :sessions
 
 end
 

lib/msf/ui/console/driver.rb

@@ -1,6 +1,7 @@
 require 'msf/core'
 require 'msf/base'
 require 'msf/ui'
+require 'msf/ui/console/framework_event_manager'
 require 'msf/ui/console/command_dispatcher'
 require 'msf/ui/console/shell'
 require 'msf/ui/console/table'
@@ -23,6 +24,14 @@ class Driver < Msf::Ui::Driver
 
 	ConfigGroup = "framework/ui/console"
 
+	#
+	# The console driver processes various framework notified events.
+	#
+	include FrameworkEventManager
+
+	#
+	# The console driver is a command shell.
+	#
 	include Msf::Ui::Console::Shell
 
 	def initialize(prompt = "%umsf", prompt_char = ">%c")
@@ -39,6 +48,9 @@ class Driver < Msf::Ui::Driver
 
 		# Initialize the super
 		super
+
+		# Register event handlers
+		register_event_handlers
 		
 		# Process things before we actually display the prompt and get rocking
 		on_startup

lib/msf/ui/console/framework_event_manager.rb

@@ -0,0 +1,49 @@
+module Msf
+module Ui
+module Console
+
+###
+#
+# FrameworkEventManager
+# ---------------------
+#
+# Handles events of various types that are sent from the framework.
+#
+###
+module FrameworkEventManager
+
+	include Msf::SessionEvents
+
+	#
+	# Subscribes to the framework as a subscriber of various events.
+	#
+	def register_event_handlers
+		framework.events.add_session_subscriber(self)
+	end
+
+	#
+	# Unsubscribes from the framework.
+	#
+	def deregister_event_handlers
+		framework.events.remove_session_subscriber(self)
+	end
+
+	#
+	# Called when a session is registered with the framework.
+	#
+	def on_session_open(session)
+		output.print_status("#{session.desc} session opened (#{session.tunnel_to_s})")
+	end
+
+	#
+	# Called when a session is closed and removed from the framework.
+	#
+	def on_session_close(session)
+		output.print_status("#{session.desc} session closed (#{session.tunnel_to_s})")
+	end
+
+end
+
+end
+end
+end

lib/rex/io/stream.rb

@@ -73,6 +73,18 @@ module Stream
 		poll_fd
 	end
 
+	#
+	# Returns peer information, such as host and port.
+	#
+	def peerinfo
+	end
+
+	#
+	# Returns local information, such as host and port.
+	#
+	def localinfo
+	end
+
 	##
 	#
 	# Common methods

lib/rex/socket.rb

@@ -114,6 +114,22 @@ class Socket
 		return self.sock
 	end
 
+	#
+	# Returns local connection information.
+	#
+	def getlocalname
+		return Socket.from_sockaddr(self.sock.getsockname)
+	end
+
+	alias getsockname getlocalname
+
+	#
+	# Return peer connection information.
+	#
+	def getpeername
+		return Socket.from_sockaddr(self.sock.getpeername)
+	end
+
 	attr_reader :sock
 	attr_reader :peerhost, :peerport, :localhost, :localport
 

lib/rex/socket/tcp.rb

@@ -39,12 +39,16 @@ class Rex::Socket::Tcp < Rex::Socket
 	#
 	##
 
+	#
+	# Writes to the TCP connection.
+	#
 	def write(buf, opts = {})
 		return sock.syswrite(buf)
 	end
 
 	#
-	# Raises EOFError if it reaches end-of-file
+	# Reads from the TCP connection and raises EOFError if there is no data
+	# left.
 	#
 	def read(length = nil, opts = {})
 		length = 16384 unless length
@@ -52,22 +56,52 @@ class Rex::Socket::Tcp < Rex::Socket
 		return sock.sysread(length)
 	end
 
+	#
+	# Calls shutdown on the TCP connection.
+	#
 	def shutdown(how = SHUT_RDWR)
 		return (sock.shutdown(how) == 0)
 	end
 
-	def has_read_data?(timeout = 0)
+	#
+	# Checks to see if the connection has read data.
+	#
+	def has_read_data?(timeout = nil)
 		timeout = timeout.to_i if (timeout)
 	
 		return (select([ poll_fd ], nil, nil, timeout) != nil)
 	end
 
+	#
+	# Closes the connection.
+	#
 	def close
 		self.sock.close if (self.sock)
 	end
 
+	#
+	# Returns the file descriptor to use with calls to select.
+	#
 	def poll_fd
 		return self.sock
 	end
 
+	#
+	# Returns peer information (host + port) in host:port format.
+	#
+	def peerinfo
+		if (pi = getpeername)
+			return pi[1] + ':' + pi[2].to_s
+		end
+	end
+
+	#
+	# Returns local information (host + port) in host:port format.
+	#
+	def localinfo
+		if (pi = getlocalname)
+			return pi[1] + ':' + pi[2].to_s
+		end
+	end
+
 end