From 73d9135c91fb568b8eccccdc4a8b4b424f1700a1 Mon Sep 17 00:00:00 2001
From: Joshua Drake <github.jdrake@qoop.org>
Date: Tue, 9 Nov 2010 02:31:21 +0000
Subject: [PATCH] merge patch from Larry Wert, fixes #2510

git-svn-id: file:///home/svn/framework3/trunk@10955 4d416f70-5f16-0410-b530-b9f4589650da
---
 .../ui/console/command_dispatcher/exploit.rb  | 66 ++++++++++++++-----
 1 file changed, 49 insertions(+), 17 deletions(-)

diff --git a/lib/msf/ui/console/command_dispatcher/exploit.rb b/lib/msf/ui/console/command_dispatcher/exploit.rb
index 52139378fe..bafc5b284e 100644
--- a/lib/msf/ui/console/command_dispatcher/exploit.rb
+++ b/lib/msf/ui/console/command_dispatcher/exploit.rb
@@ -23,6 +23,10 @@ class Exploit
 		"-t" => [ true,  "The target index to use.  If none is specified, TARGET is used."     ],
 		"-z" => [ false, "Do not interact with the session after successful exploitation."     ])
 
+	@@reload_opts =  Rex::Parser::Arguments.new(
+		'-k' => [ false,  'Stop the current job before reloading.' ],
+		'-h' => [ false,  'Help banner.' ])
+
 	#
 	# Returns the hash of exploit module specific commands.
 	#
@@ -32,6 +36,7 @@ class Exploit
 			"exploit"  => "Launch an exploit attempt",
 			"rcheck"   => "Reloads the module and checks if the target is vulnerable",
 			"rexploit" => "Reloads the module and launches an exploit attempt",
+			"reload"   => "Just reloads the module"
 		}
 	end
 
@@ -140,7 +145,7 @@ class Exploit
 				return
 			end
 		end
-		
+
 		if not payload
 			payload = exploit_choose_payload(mod, target)
 		end
@@ -199,29 +204,27 @@ class Exploit
 	# vulnerable.
 	#
 	def cmd_rcheck(*args)
-		omod = self.mod
-		self.mod = framework.modules.reload_module(mod)
-		if(not self.mod)
-			print_status("Failed to reload module: #{framework.modules.failed[omod.file_path]}")
-			self.mod = omod
-			return
-		end
+		reload()
 
 		self.mod.init_ui(driver.input, driver.output)
 		cmd_check(*args)
 	end
 
 	#
-	# Reloads an exploit module and launches an exploit.
+	# Reload an exploit module, optionally stopping existing job
 	#
-	def cmd_rexploit(*args)
-		if mod.job_id
-			print_status("Stopping existing job...")
+	def reload(should_stop_job=false)
+
+		if should_stop_job and mod.job_id
+			print_status('Stopping existing job...')
 
 			framework.jobs.stop_job(mod.job_id)
 			mod.job_id = nil
 		end
 
+
+		print_status('Reloading module...')
+
 		omod = self.mod
 		self.mod = framework.modules.reload_module(mod)
 
@@ -232,19 +235,48 @@ class Exploit
 		end
 
 		self.mod.init_ui(driver.input, driver.output)
+	end
 
+	#
+	# Handles the command to reload an exploit module.
+	#
+	def cmd_reload(*args)
+		# By default, do not stop the existing job
+		stop_existing = false
+
+		@@reload_opts.parse(args) { |opt, idx, val|
+			case opt
+			when '-k'
+				stop_existing = true
+			when '-h'
+				print_line "Usage: reload [-k]\n\nReloads the current module."
+				print_line @@reload_opts.usage
+				return
+			end
+		}
+
+		reload(stop_existing)
+	end
+
+	#
+	# Reloads an exploit module and launches an exploit.
+	#
+	def cmd_rexploit(*args)
+		# Stop existing job and reload the module
+		reload(true)
+		# Delegate to the exploit command
 		cmd_exploit(*args)
 	end
-	
+
 	#
 	# Picks a reasonable payload and minimally configures it
 	#
 	def exploit_choose_payload(mod, target)
-	
+
 		# Choose either the real target or an invalid address
 		# This is used to determine the LHOST value
 		rhost = mod.datastore['RHOST'] || '50.50.50.50'
-		
+
 		# A list of preferred payloads in the best-first order
 		pref = [
 			'windows/meterpreter/reverse_tcp',
@@ -257,7 +289,7 @@ class Exploit
 			'windows/meterpreter/reverse_nonx_tcp',
 			'windows/meterpreter/reverse_ord_tcp',
 			'windows/shell/reverse_tcp',
-			'generic/shell_reverse_tcp'				
+			'generic/shell_reverse_tcp'
 		]
 		pset = mod.compatible_payloads.map{|x| x[0] }
 		pref.each do |n|
@@ -266,7 +298,7 @@ class Exploit
 				mod.datastore['LHOST']   = Rex::Socket.source_address(rhost)
 				return n
 			end
-		end		
+		end
 		return
 	end