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

require 'rex/rop_builder'
require 'rex/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
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