From d36e38fca61b2eceeaa208304047bbf348ceebcd Mon Sep 17 00:00:00 2001
From: James Lee <egypt@metasploit.com>
Date: Tue, 15 Jan 2013 10:34:31 -0600
Subject: [PATCH] Move encoding into handle_connection

* Allows payloads that override generate_stage to still take advantage
  of stage encoding
* Also adds doc comments for a few methods
---
 lib/msf/core/encoded_payload.rb |  6 +++--
 lib/msf/core/payload/stager.rb  | 41 ++++++++++++++++++++++++++++-----
 2 files changed, 39 insertions(+), 8 deletions(-)

diff --git a/lib/msf/core/encoded_payload.rb b/lib/msf/core/encoded_payload.rb
index 413d1e521f..d5c481e7ce 100755
--- a/lib/msf/core/encoded_payload.rb
+++ b/lib/msf/core/encoded_payload.rb
@@ -41,6 +41,7 @@ class EncodedPayload
 	# This method generates the full encoded payload and returns the encoded
 	# payload buffer.
 	#
+	# @return [String] The encoded payload.
 	def generate(raw = nil)
 		self.raw           = raw
 		self.encoded       = nil
@@ -86,8 +87,9 @@ class EncodedPayload
 
 	#
 	# Generates the raw payload from the payload instance.  This populates the
-	# raw attribute.
+	# {#raw} attribute.
 	#
+	# @return [String] The raw, unencoded payload.
 	def generate_raw
 		self.raw = (reqs['Prepend'] || '') + pinst.generate + (reqs['Append'] || '')
 
@@ -216,7 +218,7 @@ class EncodedPayload
 			# If the encoded payload is nil, raise an exception saying that we
 			# suck at life.
 			if (self.encoded == nil)
-				encoder = nil
+				self.encoder = nil
 
 				raise NoEncodersSucceededError,
 					"#{pinst.refname}: All encoders failed to encode.",
diff --git a/lib/msf/core/payload/stager.rb b/lib/msf/core/payload/stager.rb
index 7a66ae3203..60f46a0e7c 100644
--- a/lib/msf/core/payload/stager.rb
+++ b/lib/msf/core/payload/stager.rb
@@ -30,6 +30,9 @@ module Msf::Payload::Stager
 	#
 	# Return the stager payload's raw payload.
 	#
+	# Can be nil if the stager is not pre-assembled.
+	#
+	# @return [String,nil]
 	def payload
 		return module_info['Stager']['Payload']
 	end
@@ -37,6 +40,7 @@ module Msf::Payload::Stager
 	#
 	# Return the stager payload's assembly text, if any.
 	#
+	# @return [String,nil]
 	def assembly
 		return module_info['Stager']['Assembly']
 	end
@@ -44,6 +48,9 @@ module Msf::Payload::Stager
 	#
 	# Return the stager payload's offsets.
 	#
+	# These will be used for substitutions during stager generation.
+	#
+	# @return [Hash]
 	def offsets
 		return module_info['Stager']['Offsets']
 	end
@@ -51,6 +58,9 @@ module Msf::Payload::Stager
 	#
 	# Returns the raw stage payload.
 	#
+	# Can be nil if the final stage is not pre-assembled.
+	#
+	# @return [String,nil]
 	def stage_payload
 		return module_info['Stage']['Payload']
 	end
@@ -58,6 +68,7 @@ module Msf::Payload::Stager
 	#
 	# Returns the assembly text of the stage payload.
 	#
+	# @return [String]
 	def stage_assembly
 		return module_info['Stage']['Assembly']
 	end
@@ -65,6 +76,10 @@ module Msf::Payload::Stager
 	#
 	# Returns variable offsets within the stage payload.
 	#
+	# These will be used for substitutions during generation of the final
+	# stage.
+	#
+	# @return [Hash]
 	def stage_offsets
 		return module_info['Stage']['Offsets']
 	end
@@ -77,6 +92,11 @@ module Msf::Payload::Stager
 		true
 	end
 
+
+	#
+	# Whether to use an Encoder on the second stage
+	#
+	# @return [Boolean]
 	def encode_stage?
 		# Convert to string in case it hasn't been normalized
 		!!(datastore['EnableStageEncoding'].to_s == "true")
@@ -85,6 +105,7 @@ module Msf::Payload::Stager
 	#
 	# Generates the stage payload and substitutes all offsets.
 	#
+	# @return [String] The generated payload stage, as a string.
 	def generate_stage
 		# Compile the stage as necessary
 		p = build(stage_payload, stage_assembly, stage_offsets, '-stg1')
@@ -92,21 +113,23 @@ module Msf::Payload::Stager
 		# Substitute variables in the stage
 		substitute_vars(p, stage_offsets) if (stage_offsets)
 
-		# Encode the stage if stage encoding is enabled
-		p = encode_stage(p)
-
 		return p
 	end
 
 	#
 	# Transmit the associated stage.
 	#
+	# @param (see handle_connection_stage)
+	# @return (see handle_connection_stage)
 	def handle_connection(conn, opts={})
 		# If the stage should be sent over the client connection that is
 		# established (which is the default), then go ahead and transmit it.
 		if (stage_over_connection?)
 			p = generate_stage
 
+			# Encode the stage if stage encoding is enabled
+			p = encode_stage(p)
+
 			# Give derived classes an opportunity to an intermediate state before
 			# the stage is sent.  This gives derived classes an opportunity to
 			# augment the stage and the process through which it is read on the
@@ -146,10 +169,15 @@ module Msf::Payload::Stager
 	end
 
 	#
-	# Called by handle_connection to allow the stage to process
-	# whatever it is it needs to process.  The default is to simply attempt to
-	# create a session.
+	# Allow the stage to process whatever it is it needs to process.
 	#
+	# Override to deal with sending the final stage in cases where
+	# {#generate_stage} is not the whole picture, such as when uploading
+	# an executable. The default is to simply attempt to create a session
+	# on the given +conn+ socket with {Handler#create_session}.
+	#
+	# @param (see Handler#create_session)
+	# @return (see Handler#create_session)
 	def handle_connection_stage(conn, opts={})
 		create_session(conn, opts)
 	end
@@ -180,6 +208,7 @@ module Msf::Payload::Stager
 			'Encoder'       => stage_enc_mod,
 			'SaveRegisters' => ['edi'],
 			'ForceEncode'   => true)
+		print_status("Encoded stage with #{encp.encoder.refname}")
 
 		# If the encoding succeeded, use the encoded buffer.  Otherwise, fall
 		# back to using the non-encoded stage