route integration into the ui

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

lib/msf/core/session_manager.rb

@@ -61,6 +61,12 @@ class SessionManager < Hash
 		# Tell the framework that we have a parting session
 		framework.events.on_session_close(session)
 
+		# If this session implements the comm interface, remove any routes
+		# that have been created for it.
+		if (session.kind_of?(Msf::Session::Comm))
+			Rex::Socket::SwitchBoard.remove_by_comm(session)
+		end
+
 		# Remove it from the hash
 		self.delete(session.sid.to_i)
 

lib/msf/ui/console/command_dispatcher/core.rb

@@ -36,6 +36,7 @@ class Core
 			"info"    => "Displays information about one or more module",
 			"jobs"    => "Displays and manages jobs",
 			"quit"    => "Exit the console",
+			"route"   => "Route traffic through a session",
 			"save"    => "Saves the active datastores",
 			"search"  => "Adds one or more module search paths",
 			"session" => "Dump session listings and display information about sessions",
@@ -169,6 +170,99 @@ class Core
 		}
 	end
 
+	#
+	# This method handles the route command which allows a user to specify
+	# which session a given subnet should route through.
+	#
+	def cmd_route(*args)
+		if (args.length == 0)
+			print(
+				"Usage: route [add/remove/flush/print] subnet netmask [local/sid]\n\n" +
+				"Route traffic destined to a given subnet through a supplied session.\n")
+			return false
+		end
+
+		case args.shift
+			when "add"
+				if (args.length < 3)
+					print_error("Missing arguments to route add.")
+					return false
+				end
+
+				gw = nil
+
+				# Determine the gateway to use
+				if (args[2] =~ /local/i)
+					gw = Rex::Socket::Comm::Local
+				elsif ((session = framework.sessions.get(args[2])) and
+				       (session.kind_of?(Msf::Session::Comm)))
+					gw = session
+				else
+					print_error("Invalid gateway specified.")
+					return false
+				end
+
+				Rex::Socket::SwitchBoard.add_route(
+					args[0],
+					args[1],
+					gw)
+			when "remove"
+				if (args.length < 3)
+					print_error("Missing arguments to route remove.")
+					return false
+				end
+				
+				gw = nil
+
+				# Determine the gateway to use
+				if (args[2] =~ /local/i)
+					gw = Rex::Socket::Comm::Local
+				elsif ((session = framework.sessions.get(args[2])) and
+				       (session.kind_of?(Msf::Session::Comm)))
+					gw = session
+				else
+					print_error("Invalid gateway specified.")
+					return false
+				end
+
+				Rex::Socket::SwitchBoard.remove_route(
+					args[0],
+					args[1],
+					gw)
+			when "flush"
+				Rex::Socket::SwitchBoard.flush_routes
+			when "print"
+				tbl =	Table.new(
+					Table::Style::Default,
+					'Header'  => "Active Routing Table",
+					'Prefix'  => "\n",
+					'Postfix' => "\n",
+					'Columns' => 
+						[
+							'Subnet',
+							'Netmask',
+							'Gateway',
+						],
+					'ColProps' =>
+						{
+							'Subnet'  => { 'MaxWidth' => 17 },
+							'Netmask' => { 'MaxWidth' => 17 },
+						})
+
+				Rex::Socket::SwitchBoard.each { |route|
+					gw = "Local"
+
+					if (route.comm.kind_of?(Msf::Session))
+						gw = "Session #{route.comm.sid}"
+					end
+
+					tbl << [ route.subnet, route.netmask, gw ]
+				}
+
+				print(tbl.to_s)
+		end
+	end
+
 	#
 	# Saves the active datastore contents to disk for automatic use across
 	# restarts of the console.

lib/rex/socket/switch_board.rb

@@ -89,6 +89,10 @@ class SwitchBoard
 		self.instance.best_comm(addr)
 	end
 
+	def self.remove_by_comm(comm)
+		self.instance.remove_by_comm(comm)
+	end
+
 	##
 	#
 	# Instance methods
@@ -193,6 +197,17 @@ class SwitchBoard
 		comm
 	end
 
+	#
+	# Remove all routes that go through the supplied comm.
+	#
+	def remove_by_comm(comm)
+		mutex.synchronize {
+			routes.delete_if { |route|
+				route.comm == comm
+			}
+		}
+	end
+
 	attr_reader :routes, :mutex
 
 protected

lib/rex/socket/switch_board.rb.ut.rb

@@ -38,4 +38,15 @@ class Rex::Socket::SwitchBoard::UnitTest < Test::Unit::TestCase
 		assert_equal('spec', Klass.best_comm('1.2.3.7'))
 	end
 
+	def test_remove_by_comm
+		Klass.flush_routes
+		Klass.add_route('1.2.3.0', 24, 'foo')
+		Klass.add_route('1.2.4.0', 24, 'dog')
+
+		Klass.remove_by_comm('foo')
+
+		assert_equal('dog', Klass.best_comm('1.2.4.7'))
+		assert_nil(Klass.best_comm('1.2.3.7'))
+	end
+
 end