session shit

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

lib/msf/core.rb

@@ -22,6 +22,7 @@ require 'Msf/Core/OptionContainer'
 
 # Framework context and core classes
 require 'Msf/Core/ModuleManager'
+require 'Msf/Core/SessionManager'
 require 'Msf/Core/Session'
 require 'Msf/Core/Framework'
 

lib/msf/core.rb.ts.rb

@@ -3,12 +3,14 @@
 require 'test/unit'
 require 'Msf/Core'
 require 'Msf/Core/OptionContainer.rb.ut'
+require 'Msf/Core/SessionManager.rb.ut'
 
 class Msf::TestSuite
 	def self.suite
 		suite = Test::Unit::TestSuite.new
 
 		suite << Msf::OptionContainer::UnitTest.suite
+		suite << Msf::SessionManager::UnitTest.suite
 
 		return suite;
 	end

lib/msf/core/session.rb

@@ -41,37 +41,19 @@ end
 ###
 class Session
 
-	def initialize(conn = nil)
-		self.conn = conn
+	def initialize()
 	end
 
 	#
-	# Read length supplied bytes from the conn
+	# Perform session-specific cleanup
 	#
-	def read(length = nil)
-		return conn.read(length)
+	def cleanup
 	end
 
-	#
-	# Write the supplied buffer to the conn
-	#
-	def write(buf)
-		return conn.write(buf)
-	end
-
-	#
-	# Close the session's conn and perform cleanup as necessary
-	#
-	def close
-		return conn.close
-	end
-
-	attr_reader   :conn
 	attr_accessor :framework, :sid
 
 protected
 
-	attr_writer   :conn
 end
 
 end

lib/msf/core/session_manager.rb

@@ -18,6 +18,8 @@ module Msf
 ###
 class SessionManager
 
+	include Enumerable
+
 	def initialize(framework)
 		self.framework = framework
 		self.sessions  = {}
@@ -31,6 +33,20 @@ class SessionManager
 		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.
@@ -41,7 +57,7 @@ class SessionManager
 			return nil
 		end
 
-		next_sid = (sid_pool += 1)
+		next_sid = (self.sid_pool += 1)
 
 		# Insert the session into the session hash table
 		sessions[next_sid] = session
@@ -67,7 +83,7 @@ class SessionManager
 		sessions.delete(session.sid)
 
 		# Close it down
-		session.close
+		session.cleanup
 	end
 
 	#
@@ -77,9 +93,9 @@ class SessionManager
 		return sessions[sid]
 	end
 
-protected:
+protected
 	
-	attr_accessor :sid_pool, :sessions
+	attr_accessor :sid_pool, :sessions, :framework
 
 end
 

lib/msf/core/session_manager.rb.ut.rb

@@ -0,0 +1,72 @@
+#!/usr/bin/ruby
+
+$:.unshift(File.join('..', '..', File.dirname(__FILE__)))
+
+require 'test/unit'
+require 'Msf/Core/SessionManager'
+
+module Msf
+
+class SessionManager::UnitTest < Test::Unit::TestCase
+
+class UtSessionEvents
+	include SessionEvents
+
+	def on_session_open(session)
+		self.open_session = session
+	end
+	def on_session_close(session)
+		self.close_session = session
+	end
+
+	attr_accessor :open_session, :close_session
+end
+
+	#
+	# Tests to make sure session registration functions properly
+	#
+	def test_registration
+		framework = Framework.new
+		manager   = SessionManager.new(framework)
+
+		session = Session.new
+
+		assert_equal(1, manager.register(session), "Did not get expected sid")
+		assert_equal(1, session.sid, 
+			"Session sid was not initialized properly")
+		assert_equal(framework, session.framework, 
+			"Session framework was not initialized properly")
+		assert_equal(1, manager.entries.length, 
+			"Number of session manager entries is not one")
+		assert_equal(session, manager[session.sid], 
+			"Index of sid did not return session")
+
+		manager.deregister(session)
+		
+		assert_equal(0, manager.entries.length, 
+			"Number of session manager entries is not zero")
+	end
+
+	#
+	# Test session notification events
+	#
+	def test_notify
+		framework = Framework.new
+		manager   = SessionManager.new(framework)
+		handler   = UtSessionEvents.new
+		session   = Session.new
+
+		framework.events.add_session_subscriber(handler)
+
+		manager.register(session)
+		assert_equal(handler.open_session, session, 
+			"Open session handler not called")
+
+		manager.deregister(session)
+		assert_equal(handler.close_session, session, 
+			"Close session handler not called")
+	end
+
+end
+
+end