persistent storage work

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

lib/msf/base.rb

@@ -30,3 +30,6 @@ require 'msf/base/sessions/command_shell'
 
 # Serialization
 require 'msf/base/serializer/readable_text'
+
+# Persistent Storage
+require 'msf/base/persistent_storage'

lib/msf/base/persistent_storage.rb

@@ -0,0 +1,69 @@
+module Msf
+
+###
+#
+# This class provides a generalized interface to persisting information,
+# either in whole or in part, about the state of the framework.  This can
+# be used to store data that can later be reinitialized in a new instance
+# of the framework or to provide a simple mechanism for generating reports
+# of some form.
+#
+###
+class PersistentStorage
+
+	@@storage_classes = {}
+
+	#
+	# Creates an instance of the storage class with the supplied name.  The
+	# array supplied as an argument is passed to the constructor of the
+	# associated class as a means of generic initialization.
+	#
+	def self.create(name, *params)
+		if (klass = @@storage_classes[name])
+			klass.new(*params)
+		else
+			nil
+		end
+	end
+
+	#
+	# Stub initialization routine that takes the params passed to create.
+	#
+	def initialize(*params)
+	end
+
+	#
+	# This methods stores all or part of the current state of the supplied
+	# framework instance to whatever medium the derived class implements.
+	# If the derived class does not implement this method, the
+	# NotImplementedError is raised.
+	#
+	def store(framework)
+		raise NotImplementedError
+	end
+
+	#
+	# This method initializes the supplied framework instance with the state
+	# that is stored in the persisted backing that the derived class
+	# implements.  If the derived class does not implement this method, the
+	# NotImplementedError is raised.
+	#
+	def fetch(framework)
+		raise NotImplementedError
+	end
+
+	#
+	# This method adds a new storage class to the hash of storage classes that
+	# can be created through create
+	#
+	def self.add_storage_class(name, klass)
+		@@storage_classes[name] = klass
+	end
+
+protected
+
+end
+
+end
+	
+require 'msf/base/persistent_storage/flatfile'

lib/msf/base/persistent_storage/flatfile.rb

@@ -0,0 +1,102 @@
+module Msf
+class PersistentStorage
+
+###
+#
+# This class persists the state of the framework to a flatfile in a human
+# readable format.  At the moment, the level of information it conveys is
+# rather basic and ugly, but this is just a prototype, so it will be improved.
+# Oh yes, it will be improved.
+#
+###
+class Flatfile < PersistentStorage
+
+	#
+	# Initializes the flatfile for storage based on the parameters specified.
+	# The hash must contain a FilePath attribute.
+	#
+	def initialize(*params)
+		raise ArgumentError, "You must specify a file path" if (params.length == 0)
+
+		self.path = params[0]
+	end
+
+	#
+	# This method stores the current state of the framework in human readable
+	# form to a flatfile.  This can be used as a reporting mechanism.
+	#
+	def store(framework)
+		# Open the supplied file path for writing.
+		self.fd = File.new(self.path, "w")
+
+		begin
+			store_general(framework)
+			store_recon(framework)
+		ensure
+			self.fd.close
+		end
+	end
+
+protected
+
+	attr_accessor :fd, :path # :nodoc:
+
+	#
+	# This method stores general information about the current state of the
+	# framework instance.
+	#
+	def store_general(framework)
+		fd.print(
+			"\n" +
+			"Metasploit Framework Report\n" +
+			"===========================\n\n" +
+			"Generated: #{Time.now}\n\n")
+	end
+
+	#
+	# This method stores the recon information that has been collected by this
+	# framework instance.
+	#
+	def store_recon(framework)
+		fd.print(
+			"Reconnaissance Information\n" +
+			"==========================\n\n")
+
+		framework.reconmgr.each_host { |address, host|
+			store_recon_host(framework, host)
+		}
+	end
+
+	#
+	# This method stores information about a specific host and its services.
+	#
+	def store_recon_host(framework, host)
+		fd.print(
+			"Host: #{host.address}\n")
+
+		store_recon_host_services(framework, host)
+
+		fd.print("\n")
+	end
+
+	#
+	# This method stores information about the services running on a host, if
+	# any.
+	#
+	def store_recon_host_services(framework, host)
+		host.services.entities.each { |name, proto|
+			if (proto.kind_of?(Msf::Recon::Entity::Group) == true)
+				proto.entities.each { |name, serv|
+					fd.print(
+						"\tService: #{serv.port} (#{serv.proto})\n")
+				}
+			end
+		}
+	end
+
+	Msf::PersistentStorage.add_storage_class('flatfile', self)
+
+end
+
+end
+end

lib/msf/core/recon/entity/container.rb

@@ -24,7 +24,11 @@ module Container
 	# This routine adds a sub-container of entities to this entity container.
 	#
 	def add_entity_subcontainer(name, container = Group.new)
-		add_entity(name, container)
+		if (self._entity_hash[name] == nil)
+			add_entity(name, container)
+		end
+
+		self._entity_hash[name]
 	end
 
 	#

lib/msf/core/recon_manager.rb

@@ -152,6 +152,16 @@ class ReconManager
 		host_hash[host_hash_key(address)]
 	end
 
+	#
+	# This method iterates the host hash calling the supplied block with the
+	# address and host instances of each entry in the host hash.
+	#
+	def each_host(&block)
+		host_hash.each_pair { |k, host|
+			block.call(host.address, host)
+		}
+	end
+
 protected
 
 	#

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

@@ -24,6 +24,11 @@ class Core
 		"-h" => [ false, "Help banner."                                   ],
 		"-k" => [ true,  "Terminate the specified job name."              ],
 		"-l" => [ false, "List all running jobs."                         ])
+	
+	@@persist_opts = Rex::Parser::Arguments.new(
+		"-s" => [ true,  "Storage medium to be used (ex: flatfile)."      ],
+		"-r" => [ false, "Restore framework state."                       ],
+		"-h" => [ false, "Help banner."                                   ])
 
 	# Returns the list of commands supported by this command dispatcher
 	def commands
@@ -35,6 +40,7 @@ class Core
 			"help"    => "Help menu",
 			"info"    => "Displays information about one or more module",
 			"jobs"    => "Displays and manages jobs",
+			"persist" => "Persist or restore framework state information",
 			"quit"    => "Exit the console",
 			"route"   => "Route traffic through a session",
 			"save"    => "Saves the active datastores",
@@ -170,6 +176,50 @@ class Core
 		}
 	end
 
+	#
+	# This method persists or restores framework state from a persistent
+	# storage medium, such as a flatfile.
+	#
+	def cmd_persist(*args)
+		if (args.length == 0)
+			args.unshift("-h")
+		end
+
+		arg_idx = 0
+		restore = false
+		storage = 'flatfile'
+
+		@@persist_opts.parse(args) { |opt, idx, val|
+			case opt
+				when "-s"
+					storage = val
+					arg_idx = idx + 2
+				when "-r"
+					restore = true
+					arg_idx = idx + 1
+				when "-h"
+					print(
+						"Usage: persist [-r] -s storage arg1 arg2 arg3 ...\n\n" +
+						"Persist or restore framework state information.\n" +
+						@@persist_opts.usage())
+					return false
+			end
+		}
+
+		# Chop off all the non-arguments
+		args = args[arg_idx..-1]
+
+		begin
+			if (inst = Msf::PersistentStorage.create(storage, *args))
+				inst.store(framework)
+			else
+				print_error("Failed to find storage medium named '#{storage}'")
+			end
+		rescue
+			log_error("Failed to persist to #{storage}: #{$!}")
+		end
+	end
+
 	#
 	# This method handles the route command which allows a user to specify
 	# which session a given subnet should route through.

lib/rex/logging/sinks/flatfile.rb

@@ -33,6 +33,8 @@ class Flatfile
 					code = 'e'
 				when LOG_INFO
 					code = 'i'
+				when LOG_WARN
+					code = 'w'
 			end
 			fd.write("[#{get_current_timestamp}] [#{code}(#{level})] #{src}: #{msg}\n")
 		end