From 176e999f08c037ec2d11e4a50e6df3a58f8a48a0 Mon Sep 17 00:00:00 2001
From: HD Moore <hd_moore@rapid7.com>
Date: Mon, 9 Nov 2009 00:33:40 +0000
Subject: [PATCH] Merge patches from Ryan Linn. Fixes #438

git-svn-id: file:///home/svn/framework3/trunk@7413 4d416f70-5f16-0410-b530-b9f4589650da
---
 lib/msf/ui/console/command_dispatcher/core.rb |  5 +++++
 lib/rex/ui/text/shell.rb                      | 22 ++++++++++++++-----
 2 files changed, 21 insertions(+), 6 deletions(-)

diff --git a/lib/msf/ui/console/command_dispatcher/core.rb b/lib/msf/ui/console/command_dispatcher/core.rb
index 9358e78156..16d1e75007 100644
--- a/lib/msf/ui/console/command_dispatcher/core.rb
+++ b/lib/msf/ui/console/command_dispatcher/core.rb
@@ -366,6 +366,11 @@ class Core
 	# Instructs the driver to stop executing.
 	#
 	def cmd_exit(*args)
+		if(framework.sessions.count() > 0 and 
+			(args.length < 1 or (args[0] =~ /\-Y/i) == nil))
+				print_status("You have active sessions open, to exit anyway type \"exit -y\"")
+				return
+		end
 		driver.stop
 	end
 
diff --git a/lib/rex/ui/text/shell.rb b/lib/rex/ui/text/shell.rb
index 8920404630..692929fe17 100644
--- a/lib/rex/ui/text/shell.rb
+++ b/lib/rex/ui/text/shell.rb
@@ -40,6 +40,7 @@ module Shell
 		# Set the stop flag to false
 		self.stop_flag      = false
 		self.disable_output = false
+		self.stop_count	    = 0
 
 		# Initialize the prompt
 		self.init_prompt = prompt
@@ -113,21 +114,30 @@ module Shell
 
 		begin
 			
-			while (not self.stop_flag and (line = input.pgets))
+			while true
+				# If the stop flag was set or we've hit EOF, break out
+				break if (self.stop_flag or self.stop_count > 1)
+
+				line = input.pgets
 				log_output(input.prompt)
 
 				# If a block was passed in, pass the line to it.  If it returns true,
 				# break out of the shell loop.
 				if (block)
-					break if (block.call(line))
+					break if (line == nil or block.call(line))
+				elsif(input.eof? or line == nil)
+				# If you have sessions active, this will give you a shot to exit gravefully
+				# If you really are ambitious, 2 eofs will kick this out
+					self.stop_count += 1
+					next if(self.stop_count > 1)
+					run_single("quit")
+				else
 				# Otherwise, call what should be an overriden instance method to
 				# process the line.
-				else
 					run_single(line)
+					self.stop_count = 0
 				end
 
-				# If the stop flag was set or we've hit EOF, break out
-				break if (input.eof? or self.stop_flag)
 			end
 		# Prevent accidental console quits
 		rescue ::Interrupt
@@ -313,7 +323,7 @@ protected
 	attr_writer   :input, :output # :nodoc:
 	attr_accessor :stop_flag, :init_prompt # :nodoc:
 	attr_accessor :prompt_char, :tab_complete_proc # :nodoc:
-	attr_accessor :log_source # :nodoc:
+	attr_accessor :log_source, :stop_count # :nodoc:
 
 end