##
## $Id$
##

$:.unshift(File.join(File.expand_path(File.dirname(__FILE__)), '..', 'lib', 'lab'))

require 'yaml'
require 'vm_controller'

module Msf

class Plugin::Lab < Msf::Plugin
	class LabCommandDispatcher
		include Msf::Ui::Console::CommandDispatcher

		attr_accessor :controller
		
		def initialize(driver)
			super(driver)
			@controller = nil
		end

		#
		# Returns the hash of commands supported by this dispatcher.
		#
		def commands
		{
			"lab_help" => "lab_help <lab command> - Show that command's description.",
			"lab_show" => "lab_show - show all vms in the lab.",
			"lab_show_running" => "lab_show_running - show running vms.",
			"lab_load" => "lab_load [file] - load a lab definition from disk.", 			
			"lab_save" => "lab_save [filename] - persist a lab definition in a file.",
			"lab_load_running" => "lab_load_running - use the running vms to create a lab.", 
			"lab_add_running" => "lab_add_running - add the running vms to the current lab.", 
			"lab_load_dir" => "lab_load_dir [directory] - create a lab from a specified directory.",
			"lab_add_dir" => "lab_add_dir [directory] - add vms in a specified directory.",
			"lab_clear" => "lab_clear - clear the running lab.",	
			"lab_start" => "lab_start [vmid+|all] start the specified vm.",
			"lab_reset" => "lab_reset [vmid+|all] reset the specified vm.",
			"lab_suspend" => "lab_suspend [vmid+|all] suspend the specified vm.",
			"lab_stop" => "lab_stop [vmid+|all] stop the specified vm.",
			"lab_revert" => "lab_revert [vmid+|all] [snapshot] revert the specified vm.",
			"lab_snapshot" => "lab_snapshot [vmid+|all] [snapshot] snapshot all targets for this exploit."
			#"lab_run_command" => "lab_run_command [vmid+|all] [command] run a command on all targets.",
			#"lab_browse_to" => "lab_browse_to [vmid+|all] [uri] use the default browser to browse to a uri."
		}
		end

		def name
			"Lab"
		end

		##
		## Commands for help
		##
		
		def longest_cmd_size
			commands.keys.map {|x| x.size}.sort.last
		end

		# No extended help yet, but this is where more detailed documentation
		# on particular commands would live. Key is command, (not cmd_command),
		# value is the documentation.
		def extended_help
			{
				"lab_fake_cmd" =>              "This is a fake command. It's got its own special docs." +
					(" " * longest_cmd_size) + "It might be long so so deal with formatting somehow."
			}
		end

		# Map for usages
		def lab_usage
			caller[0][/`cmd_(.*)'/]
			cmd = $1
			if extended_help[cmd] || commands[cmd]
				cmd_lab_help cmd
			else # Should never really get here...
				print_error "Unknown command. Try 'help'"
			end
		end

		def cmd_lab_help(*args)
			if args.empty?
				commands.each_pair {|k,v| print_line "%-#{longest_cmd_size}s - %s" % [k,v] }
			else
				args.each do |c|
					if extended_help[c] || commands[c]
						print_line "%-#{longest_cmd_size}s - %s" % [c,extended_help[c] || commands[c]]
					else
						print_error "Unknown command '#{c}'"
					end
				end
			end
		end
		
		##
		## Regular Lab Commands
		## 

		def cmd_lab_load(*args)
			if args[0]
				filename = args[0]

				## TODO - check for existence	
				@controller.from_file(filename)
			else
				lab_usage
			end
		end

		def cmd_lab_load_running(*args)
			@controller.build_from_running_workstation(true)
		end

		def cmd_lab_add_running(*args)
			@controller.build_from_running_workstation
		end

		def cmd_lab_load_dir(*args)			
			if args[0]
				@controller.build_from_dir_workstation(args[0],true)
			else
				lab_usage
			end
		end

		def cmd_lab_add_dir(*args)			
			if args[0]
				@controller.build_from_dir_workstation(args[0])
			else
				lab_usage
			end
		end

		def cmd_lab_clear(*args)
			@controller.clear!
	        end

		def cmd_lab_save(*args)		
			return lab_usage if args.empty?
			@controller.to_file(args[0])
		end
		

		## 
		## Commands for dealing with a currently-loaded lab
		## 

		def cmd_lab_show(*args)
			hlp_print_lab
	        end

		def cmd_lab_show_running(*args)
			hlp_print_lab_running
	        end
	        
		def cmd_lab_start(*args)
			return lab_usage if args.empty?
		
			if args[0] == "all"
				@controller.each do |vm| 
					print_line "Starting lab vm #{vm.vmid}."	
					if !vm.running?
						vm.start
					else
						print_line "Lab vm #{vm.vmid} already running."	
					end
				end
			else
				args.each do |arg|
					if @controller.includes_vmid? arg
						vm = @controller.find_by_vmid(arg)	
						if !vm.running?
							print_line "Starting lab vm #{vm.vmid}."	
							vm.start
						else
							print_line "Lab vm #{vm.vmid} already running."	
						end
					end	
				end
			end
		end
	     
		def cmd_lab_stop(*args)
			return lab_usage if args.empty?
		
			if args[0] == "all"
				@controller.each do |vm| 
					print_line "Stopping lab vm #{vm.vmid}."	
					if vm.running?
						vm.stop
					else
						print_line "Lab vm #{vm.vmid} not running."	
					end
				end
			else
				args.each do |arg|
					if @controller.includes_vmid? arg
						vm = @controller.find_by_vmid(arg)	
						if vm.running?
							print_line "Stopping lab vm #{vm.vmid}."	
							vm.stop
						else
							print_line "Lab vm #{vm.vmid} not running."	
						end
					end	
				end
			end
	        end

		def cmd_lab_suspend(*args)
			return lab_usage if args.empty?
					
			if args[0] == "all"
				@controller.each{ |vm| vm.start }
			else
				args.each do |arg|
					if @controller.includes_vmid? arg
						print_line "Suspending lab vm #{arg}."
						@controller.find_by_vmid(arg).suspend	
					end	
				end
			end
	        end

		def cmd_lab_reset(*args)
			return lab_usage if args.empty?
		
			if args[0] == "all"
				print_line "Resetting all lab vms."
				@controller.each{ |vm| vm.start }
			else
				args.each do |arg|
					if @controller.includes_vmid? arg
						print_line "Resetting lab vm #{arg}."
						@controller.find_by_vmid(arg).reset	
					end	
				end
			end
	        end


		def cmd_lab_snapshot(*args)
			return lab_usage if args.count < 2
			snapshot = args[args.count-1] 	
		
			if args[0] == "all"
				print_line "Snapshotting all lab vms to snapshot: #{snapshot}."
				@controller.each{ |vm| vm.snapshot(snapshot) }
			else
				args.each_index do |index|
					if !index == args.count-1 ##skip the last argument
						if @controller.includes_vmid? args[index]
							print_line "Snapshotting #{args[index]} to snapshot: #{snapshot}."
							@controller.find_by_vmid(args[index]).snapshot(snapshot)	
						end	
					end
				end
			end
	        end


		def cmd_lab_revert(*args)
			return lab_usage if args.count < 2
			snapshot = args[args.count-1] 		

			if args[0] == "all"
				print_line "Reverting all lab vms to snapshot: #{snapshot}."
				@controller.each{ |vm| vm.revert(snapshot) }
			else
				args.each_index do |index|
					if !index == args.count-1 ##skip the last argument
						if @controller.includes_vmid? args[index]
							print_line "Reverting #{args[index]} to snapshot: #{snapshot}."
							@controller.find_by_vmid(args[index]).revert(snapshot)	
						end
					end	
				end
			end
	        end


		def cmd_lab_run_command(*args)
			return lab_usage if args.empty?
			command = args[args.count-1] ## gimmie the loot
			print_line "not implemented"
	        end

		def cmd_lab_browse_to(*args)
			return lab_usage if args.empty?
			uri = args[args.count-1] ## gimmie the loot
			print_line "not implemented"
		end
	
		private
			def hlp_print_lab
				indent = '    '

				tbl = Rex::Ui::Text::Table.new(
					'Header'  => 'Available Lab VMs',
					'Indent'  => indent.length,
					'Columns' => [ 'Vmid', 'Location', "Powered On" ]
				)

				@controller.each do |vm| 
					tbl << [ 	vm.vmid, 
							vm.location,
							vm.running?]
				end
			
				print_line tbl.to_s
			end
				
			def hlp_print_lab_running
				indent = '    ' 

				tbl = Rex::Ui::Text::Table.new(
					'Header'  => 'Running Lab VMs',
					'Indent'  => indent.length,
					'Columns' => [ 'vmid', 'file', 'powered on' ]
				)

				@controller.each do |vm|
					if vm.running? 
						tbl << [ 	vm.vmid, 
								vm.location,
								vm.running?]
					end	
				end
				print_line tbl.to_s
			end


	end
	
	#
	# The constructor is called when an instance of the plugin is created.  The
	# framework instance that the plugin is being associated with is passed in
	# the framework parameter.  Plugins should call the parent constructor when
	# inheriting from Msf::Plugin to ensure that the framework attribute on
	# their instance gets set.
	#
	def initialize(framework, opts)
		super

		## Register the commands above
		console_dispatcher = add_console_dispatcher(LabCommandDispatcher)

		@controller = VmController.new

		## Share the vms
		console_dispatcher.controller = @controller
	end


	#
	# The cleanup routine for plugins gives them a chance to undo any actions
	# they may have done to the framework.  For instance, if a console
	# dispatcher was added, then it should be removed in the cleanup routine.
	#
	def cleanup
		# If we had previously registered a console dispatcher with the console,
		# deregister it now.
		remove_console_dispatcher('Lab')
	end

	#
	# This method returns a short, friendly name for the plugin.
	#
	def name
		"lab"
	end

	#
	# This method returns a brief description of the plugin.  It should be no
	# more than 60 characters, but there are no hard limits.
	#
	def desc
		"Adds the ability to manage VMs"
	end
	
end ## End Class
end ## End Module