#!/usr/bin/env ruby
# -*- coding: binary -*-

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 'msfenv'

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

require 'rex'
require 'msf/ui'
require 'msf/base'
require 'msf/core/payload_generator'


  class MsfVenomError < StandardError; end
  class UsageError < MsfVenomError; end
  class NoTemplateError < MsfVenomError; end
  class IncompatibleError < MsfVenomError; end


  require 'optparse'


  # Creates a new framework object.
  #
  # @note Ignores any previously cached value.
  # @param (see ::Msf::Simple::Framework.create)
  # @return [Msf::Framework]
  def init_framework(create_opts={})
    create_opts[:module_types] ||= [
      ::Msf::MODULE_PAYLOAD, ::Msf::MODULE_ENCODER, ::Msf::MODULE_NOP
    ]
    @framework = ::Msf::Simple::Framework.create(create_opts.merge('DisableDatabase' => true))
  end

  # Cached framework object
  #
  # @return [Msf::Framework]
  def framework
    return @framework if @framework

    init_framework

    @framework
  end


  def parse_args(args)
    opts = {}
    datastore = {}
    opt = OptionParser.new
    opt.banner = "Usage: #{$0} [options] <var=val>"
    opt.separator('')
    opt.separator('Options:')

    opt.on('-p', '--payload       <payload>', String,
      'Payload to use. Specify a \'-\' or stdin to use custom payloads') do |p|
      if p == '-'
        opts[:payload] = 'stdin'
      else
        opts[:payload] = p
      end
    end

    opt.on('--payload-options', "List the payload's standard options") do
      opts[:list_options] = true
    end

    opt.on('-l', '--list          [type]', Array, 'List a module type. Options are: payloads, encoders, nops, all') do |l|
      if l.nil? or l.empty?
        l = ["all"]
      end
      opts[:list] = l
    end

    opt.on('-n', '--nopsled       <length>', Integer, 'Prepend a nopsled of [length] size on to the payload') do |n|
      opts[:nops] = n.to_i
    end

    opt.on('-f', '--format        <format>', String, "Output format (use --help-formats for a list)") do |f|
      opts[:format] = f
    end

    opt.on('--help-formats', String, "List available formats") do
      init_framework(:module_types => [])
      msg = "Executable formats\n" +
            "\t" + ::Msf::Util::EXE.to_executable_fmt_formats.join(", ") + "\n" +
            "Transform formats\n" +
            "\t" + ::Msf::Simple::Buffer.transform_formats.join(", ")
      raise UsageError, msg
    end

    opt.on('-e', '--encoder       <encoder>', String, 'The encoder to use') do |e|
      opts[:encoder] = e
    end

    opt.on('-a', '--arch          <arch>', String, 'The architecture to use') do |a|
      opts[:arch] = a
    end

    opt.on('--platform      <platform>', String, 'The platform of the payload') do |l|
      opts[:platform] = l
    end

    opt.on('-s', '--space         <length>', Integer, 'The maximum size of the resulting payload') do |s|
      opts[:space] = s
    end

    opt.on('--encoder-space <length>', Integer, 'The maximum size of the encoded payload (defaults to the -s value)') do |s|
      opts[:encoder_space] = s
    end

    opt.on('-b', '--bad-chars     <list>', String, 'The list of characters to avoid example: \'\x00\xff\'') do |b|
      opts[:badchars] = Rex::Text.hex_to_raw(b)
    end

    opt.on('-i', '--iterations    <count>', Integer, 'The number of times to encode the payload') do |i|
      opts[:iterations] = i
    end

    opt.on('-c', '--add-code      <path>', String, 'Specify an additional win32 shellcode file to include') do |x|
      opts[:add_code] = x
    end

    opt.on('-x', '--template      <path>', String, 'Specify a custom executable file to use as a template') do |x|
      opts[:template] = x
    end

    opt.on('-k', '--keep', 'Preserve the template behavior and inject the payload as a new thread') do
      opts[:keep] = true
    end

    opt.on('-o', '--out           <path>', 'Save the payload') do |x|
      opts[:out] = x
    end

    opt.on('-v', '--var-name      <name>', String, 'Specify a custom variable name to use for certain output formats') do |x|
      opts[:var_name] = x
    end

    opt.on('--smallest', 'Generate the smallest possible payload') do
      opts[:smallest] = true
    end

    opt.on_tail('-h', '--help', 'Show this message') do
      raise UsageError, "#{opt}"
    end

    begin
      opt.parse!(args)
    rescue OptionParser::InvalidOption => e
      raise UsageError, "Invalid option\n#{opt}"
    rescue OptionParser::MissingArgument => e
      raise UsageError, "Missing required argument for option\n#{opt}"
    end

    if opts.empty?
      raise UsageError, "No options\n#{opt}"
    end

    if args
      args.each do |x|
        k,v = x.split('=', 2)
        datastore[k.upcase] = v.to_s
      end
      if opts[:payload].to_s =~ /[\_\/]reverse/ and datastore['LHOST'].nil?
        datastore['LHOST'] = Rex::Socket.source_address
      end
    end

    if opts[:payload].nil? # if no payload option is selected assume we are reading it from stdin
      opts[:payload] = "stdin"
    end

    if opts[:payload] == 'stdin' and not opts[:list]
      $stderr.puts "Attempting to read payload from STDIN..."
      begin
        ::Timeout.timeout(30) do
          opts[:stdin] = payload_stdin
        end
      rescue Timeout::Error
        opts[:stdin] = ''
      end
    end

    opts[:datastore] = datastore

    opts
  end


  # Read a raw payload from stdin (or whatever IO object we're currently
  # using as stdin, see {#initialize})
  #
  # @return [String]
  def payload_stdin
    @in = $stdin
    @in.binmode
    payload = @in.read
    payload
  end

  def dump_payloads
    init_framework(:module_types => [ ::Msf::MODULE_PAYLOAD ])
    tbl = Rex::Ui::Text::Table.new(
        'Indent'  => 4,
        'Header'  => "Framework Payloads (#{framework.stats.num_payloads} total)",
        'Columns' =>
            [
                "Name",
                "Description"
            ])

    framework.payloads.each_module { |name, mod|
      tbl << [ name, mod.new.description.split.join(' ') ]
    }

    "\n" + tbl.to_s + "\n"
  end

  def dump_encoders(arch = nil)
    init_framework(:module_types => [ ::Msf::MODULE_ENCODER ])
    tbl = Rex::Ui::Text::Table.new(
        'Indent'  => 4,
        'Header'  => "Framework Encoders" + ((arch) ? " (architectures: #{arch})" : ""),
        'Columns' =>
            [
                "Name",
                "Rank",
                "Description"
            ])
    cnt = 0

    framework.encoders.each_module(
        'Arch' => arch ? arch.split(',') : nil) { |name, mod|
      tbl << [ name, mod.rank_to_s, mod.new.name ]

      cnt += 1
    }

    (cnt > 0) ? "\n" + tbl.to_s + "\n" : "\nNo compatible encoders found.\n\n"
  end

  def dump_nops
    init_framework(:module_types => [ ::Msf::MODULE_NOP ])
    tbl = Rex::Ui::Text::Table.new(
        'Indent'  => 4,
        'Header'  => "Framework NOPs (#{framework.stats.num_nops} total)",
        'Columns' =>
            [
                "Name",
                "Description"
            ])

    framework.nops.each_module { |name, mod|
      tbl << [ name, mod.new.description.split.join(' ') ]
    }

    "\n" + tbl.to_s + "\n"
  end



if __FILE__ == $0

  begin
    generator_opts = parse_args(ARGV)
  rescue MsfVenomError, Msf::OptionValidateError => e
    $stderr.puts e.message
    exit(-1)
  end

  if generator_opts[:list]
    generator_opts[:list].each do |mod|
      case mod.downcase
        when "payloads"
          $stdout.puts dump_payloads
        when "encoders"
          $stdout.puts dump_encoders(generator_opts[:arch])
        when "nops"
          $stdout.puts dump_nops
        when "all"
          # Init here so #dump_payloads doesn't create a framework with
          # only payloads, etc.
          init_framework
          $stdout.puts dump_payloads
          $stdout.puts dump_encoders
          $stdout.puts dump_nops
        else
          $stderr.puts "Invalid module type"
      end
    end
    exit(0)
  end

  if generator_opts[:list_options]
    payload_mod = framework.payloads.create(generator_opts[:payload])

    if payload_mod.nil?
      $stderr.puts "Invalid payload: #{generator_opts[:payload]}"
      exit
    end

    $stderr.puts "Options for #{payload_mod.fullname}:\n\n"
    $stdout.puts ::Msf::Serializer::ReadableText.dump_module(payload_mod, '    ')

    $stderr.puts "Advanced options for #{payload_mod.fullname}:\n\n"
    $stdout.puts ::Msf::Serializer::ReadableText.dump_advanced_options(payload_mod, '    ')

    $stderr.puts "Evasion options for #{payload_mod.fullname}:\n\n"
    $stdout.puts ::Msf::Serializer::ReadableText.dump_evasion_options(payload_mod, '    ')
    exit(0)
  end

  generator_opts[:framework] = framework
  generator_opts[:cli] = true

  begin
    venom_generator =  Msf::PayloadGenerator.new(generator_opts)
    payload = venom_generator.generate_payload
  rescue ::Exception => e
    elog("#{e.class} : #{e.message}\n#{e.backtrace * "\n"}")
    $stderr.puts e.message
  end

  # No payload generated, no point to go on
  exit(-1) unless payload

  if generator_opts[:out]
    begin
      ::File.open(generator_opts[:out], 'wb') do |f|
        f.write(payload)
      end
      $stderr.puts "Saved as: #{generator_opts[:out]}"
    rescue ::Exception => e
      # If I can't save it, then I can't save it. I don't think it matters what error.
      elog("#{e.class} : #{e.message}\n#{e.backtrace * "\n"}")
      $stderr.puts e.message
    end
  else
    output_stream = $stdout
    output_stream.binmode
    output_stream.write payload
    # trailing newline for pretty output
    $stderr.puts unless payload =~ /\n$/
  end

end