#!/usr/bin/env ruby
#
# $Id$
#
# This tool will collect, export, and import ROP gadgets
# from various file formats (PE, ELF, Macho)
# $Revision$
#

msfbase = __FILE__
while File.symlink?(msfbase)
	msfbase = File.expand_path(File.readlink(msfbase), File.dirname(msfbase))
end

$:.unshift(File.expand_path(File.join(File.dirname(msfbase), 'lib')))
require 'fastlib'
require 'msfenv'



$:.unshift(ENV['MSF_LOCAL_LIB']) if ENV['MSF_LOCAL_LIB']

require 'rex'
require 'rex/ropbuilder'
require 'rex/ui/text/output/stdio'
require 'rex/ui/text/color'
require 'optparse'

def opt2i(o)
	o.index("0x")==0 ? o.hex : o.to_i
end

opts = {}
color = true

opt = OptionParser.new
opt.banner = "Usage #{$PROGRAM_NAME} <option> [targets]"
opt.separator('')
opt.separator('Options:')

opt.on('-d', '--depth [size]', 'Number of maximum bytes to backwards disassemble from return instructions') do |d|
	opts[:depth] = opt2i(d)
end

opt.on('-s', '--search [regex]', 'Search for gadgets matching a regex, match intel syntax or raw bytes') do |regex|
	opts[:pattern] = regex
end

opt.on('-n', '--nocolor', 'Disable color. Useful for piping to other tools like the less and more commands') do
	color = false
end

opt.on('-x', '--export [filename]', 'Export gadgets to CSV format') do |csv|
	opts[:export] = csv
end

opt.on('-i', '--import [filename]', 'Import gadgets from previous collections') do |csv|
	opts[:import] = csv
end

opt.on('-v', '--verbose', 'Output very verbosely') do
	opts[:verbose] = true
end

opt.on_tail('-h', '--help', 'Show this message') do
	puts opt
	exit(1)
end

begin
	opt.parse!
rescue OptionParser::InvalidOption
	puts "Invalid option, try -h for usage"
	exit(1)
end

if opts.empty? and (ARGV.empty? or ARGV.nil?)
	puts "no options"
	puts opt
	exit(1)
end

# set defaults
opts[:depth] ||= 5

gadgets = []

if opts[:import].nil?
	files = []
	ARGV.each do |file|
		if(File.directory?(file))
			dir = Dir.open(file)
			dir.entries.each do |ent|
				path = File.join(file, ent)
				next if not File.file?(path)
				files << File.join(path)
			end
		else
			files << file
		end
	end
	
	ropbuilder = Rex::RopBuilder::RopCollect.new

	files.each do |file|
		ret, retn = []
		ropbuilder = Rex::RopBuilder::RopCollect.new(file)
		ropbuilder.print_msg("Collecting gadgets from %bld%cya#{file}%clr\n", color)
		retn = ropbuilder.collect(opts[:depth], "\xc2") # retn
		ret = ropbuilder.collect(opts[:depth], "\xc3") # ret
		ropbuilder.print_msg("Found %grn#{ret.count + retn.count}%clr gadgets\n\n", color)

		# compile a list of all gadgets from all files
		ret.each do |gadget|
			gadgets << gadget
			if opts[:verbose]
				ropbuilder.print_msg("#{gadget[:file]} gadget: %bld%grn#{gadget[:address]}%clr\n", color)
				ropbuilder.print_msg("#{gadget[:disasm]}\n", color)
			end
		end

		retn.each do |gadget|
			gadgets << gadget
			if opts[:verbose]
				ropbuilder.print_msg("#{gadget[:file]} gadget: %bld%grn#{gadget[:address]}%clr\n", color)
				ropbuilder.print_msg("#{gadget[:disasm]}\n", color)
			end
		end
		
	end

	ropbuilder.print_msg("Found %bld%grn#{gadgets.count}%clr gadgets total\n\n", color)
end

if opts[:import]

	ropbuilder = Rex::RopBuilder::RopCollect.new()
	ropbuilder.print_msg("Importing gadgets from %bld%cya#{opts[:import]}\n", color)
	gadgets = ropbuilder.import(opts[:import])

	gadgets.each do |gadget|
		ropbuilder.print_msg("gadget: %bld%cya#{gadget[:address]}%clr\n", color)
		ropbuilder.print_msg(gadget[:disasm] + "\n", color)
	end

	ropbuilder.print_msg("Imported %grn#{gadgets.count}%clr gadgets\n", color)
end

if opts[:pattern]
	matches = ropbuilder.pattern_search(opts[:pattern])
	if opts[:verbose]
		ropbuilder.print_msg("Found %grn#{matches.count}%clr matches\n", color)
	end
end

if opts[:export]
	ropbuilder.print_msg("Exporting %grn#{gadgets.count}%clr gadgets to %bld%cya#{opts[:export]}%clr\n", color)
	csv = ropbuilder.to_csv(gadgets)

	if csv.nil?
		exit(1)
	end

	begin
		fd = File.new(opts[:export], 'w')
		fd.puts csv
		fd.close
	rescue
		puts "Error writing #{opts[:export]} file"
		exit(1)
	end
	ropbuilder.print_msg("%bld%redSuccess!%clr gadgets exported to %bld%cya#{opts[:export]}%clr\n", color)
end