# Copyright (c) 2010, patrickHVE@googlemail.com
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
#     * Redistributions of source code must retain the above copyright
#       notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above copyright
#       notice, this list of conditions and the following disclaimer in the
#       documentation and/or other materials provided with the distribution.
#     * The names of the author may not be used to endorse or promote products
#       derived from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL patrickHVE@googlemail.com BE LIABLE FOR ANY
# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

module Rex
module Post
module Meterpreter
module Extensions
module Stdapi
module Railgun

#
# shared functions
#
# 
module DLLHelper

	# converts ruby string to zero-terminated ASCII string
	def str_to_ascii_z(str)
		return str+"\x00"
	end

	# converts 0-terminated ASCII string to ruby string
	def asciiz_to_str(asciiz)
		zero_byte_idx = asciiz.index("\x00")
		if zero_byte_idx != nil
			return asciiz[0, zero_byte_idx]
		else
			return asciiz
		end
	end

	# converts ruby string to zero-terminated WCHAR string
	def str_to_uni_z(str)
		enc = str.unpack("C*").pack("v*")
		enc += "\x00\x00"
		return enc
	end

	# converts 0-terminated UTF16 to ruby string
	def uniz_to_str(uniz)
		uniz.unpack("v*").pack("C*").unpack("A*")[0]
	end

	# parses a number param and returns the value
	# raises an exception if the param cannot be converted to a number
	# examples:
	#   nil => 0
	#   3 => 3
	#   "MB_OK" => 0
	#   "SOME_CONSTANT | OTHER_CONSTANT" => 17
	#   "tuna" => !!!!!!!!!!Exception
	#
	# Parameter "win_consts" is a WinConstantManager
	def param_to_number(v, win_consts = @win_consts)
		if v.class == NilClass then
			return 0
		elsif v.class == Fixnum then
			return v # ok, it's already a number
		elsif v.class == Bignum then
			return v # ok, it's already a number
		elsif v.class == String then
			dw = win_consts.parse(v) # might raise an exception
			if dw != nil
				return dw
			else
				raise ArgumentError, "Param #{v} (class #{v.class}) cannot be converted to a number. It's a string but matches no constants I know."
			end
		else
			raise "Param #{v} (class #{v.class}) should be a number but isn't"
		end
	end

	# assembles the buffers "in" and "inout"
	def assemble_buffer(direction, function, args)
		layout = {} # paramName => BufferItem
		blob = ""
		#puts " building buffer: #{direction}"
		function.params.each_with_index do |param_desc, param_idx|
			#puts "  processing #{param_desc[0]} #{param_desc[1]} #{param_desc[2]}"
			# we care only about inout buffers
			if param_desc[2] == direction
				buffer = nil
				# Special case:
				# The user can choose to supply a Null pointer instead of a buffer
				# in this case we don't need space in any heap buffer
				if param_desc[0][0,1] == 'P' # type is a pointer
					if args[param_idx] == nil
						next
					end
				end

				case param_desc[0] # required argument type
					when "PDWORD"
						dw = param_to_number(args[param_idx])
						buffer = [dw].pack('V')
					when "PWCHAR"
						raise "param #{param_desc[1]}: string expected" unless args[param_idx].class == String
						buffer = str_to_uni_z(args[param_idx])
					when "PCHAR"
						raise "param #{param_desc[1]}: string expected" unless args[param_idx].class == String
						buffer = str_to_ascii_z(args[param_idx])
					when "PBLOB"
						raise "param #{param_desc[1]}: please supply your BLOB as string!" unless args[param_idx].class == String
						buffer = args[param_idx]
					# other types (non-pointers) don't reference buffers
					# and don't need any treatment here
				end

				if buffer != nil
					#puts "   adding #{buffer.length} bytes to heap blob"
					layout[param_desc[1]] = BufferItem.new(param_idx, blob.length, buffer.length, param_desc[0])
					blob += buffer
					# sf: force 8 byte alignment to satisfy x64, wont matter on x86.
					while( blob.length % 8 != 0 )
						blob += "\x00" 
					end
					#puts "   heap blob size now #{blob.length}"
				end
			end
		end
		#puts "  built buffer: #{direction}"
		return [layout, blob]
	end
	
end

end; end; end; end; end; end;