# -*- coding: binary -*-

require 'msf/base'

module Msf
module Simple

###
#
# Simple payload wrapper class for performing generation.
#
###
module Payload

  include Module

  #
  # Generate a payload with the mad skillz.  The payload can be generated in
  # a number of ways.
  #
  # opts can have:
  #
  #   Encoder     => A encoder module name.
  #   BadChars    => A string of bad characters.
  #   Format      => The format to represent the data as: ruby, perl, c, raw
  #   Options     => A hash of options to set.
  #   OptionStr   => A string of options in VAR=VAL form separated by
  #                  whitespace.
  #   NoComment   => Disables prepention of a comment
  #   NopSledSize => The number of NOPs to use
  #   MaxSize     => The maximum size of the payload.
  #   Iterations  => Number of times to encode.
  #   Force       => Force encoding.
  #
  # raises:
  #
  #   BadcharError => If the supplied encoder fails to encode the payload
  #   NoKeyError => No valid encoder key could be found
  #   ArgumentParseError => Options were supplied improperly
  #
  def self.generate_simple(payload, opts, &block)

    # Clone the module to prevent changes to the original instance
    payload = payload.replicant
    Msf::Simple::Framework.simplify_module(payload)
    yield(payload) if block_given?

    # Import any options we may need
    payload._import_extra_options(opts)
    framework = payload.framework

    # Generate the payload
    e = EncodedPayload.create(payload,
        'BadChars' => opts['BadChars'],
        'MinNops'  => opts['NopSledSize'],
        'Encoder'  => opts['Encoder'],
        'Iterations'  => opts['Iterations'],
        'ForceEncode' => opts['ForceEncode'],
        'Space'    => opts['MaxSize'])

    fmt = opts['Format'] || 'raw'

    exeopts = {
      :inject => opts['KeepTemplateWorking'],
      :template => opts['Template'],
      :template_path => opts['ExeDir']
    }

    arch = payload.arch
    plat = opts['Platform'] || payload.platform

    # Save off the original payload length
    len = e.encoded.length


    if arch.index(ARCH_JAVA) and fmt == 'war'
      return e.encoded_war.pack
    end

    output = Msf::Util::EXE.to_executable_fmt(framework, arch, plat, e.encoded, fmt, exeopts)

    if not output
      # Generate jar if necessary
      if fmt == 'jar'
        return e.encoded_jar.pack
      end

      # Serialize the generated payload to some sort of format
      fmt ||= "ruby"
      output = Buffer.transform(e.encoded, fmt)

      # Prepend a comment
      if (fmt != 'raw' and opts['NoComment'] != true)
        ((ou = payload.options.options_used_to_s(payload.datastore)) and ou.length > 0) ? ou += "\n" : ou = ''
        output =
          Buffer.comment(
            "#{payload.refname} - #{len} bytes#{payload.staged? ? " (stage 1)" : ""}\n" +
            "http://www.metasploit.com\n" +
            ((e.encoder) ? "Encoder: #{e.encoder.refname}\n" : '') +
            ((e.nop) ?     "NOP gen: #{e.nop.refname}\n" : '') +
            "#{ou}",
            fmt) +
          output

        # If it's multistage, include the second stage too
        if payload.staged?
          stage = payload.generate_stage

          # If a stage was generated, then display it
          if stage and stage.length > 0
            output +=
              "\n" +
              Buffer.comment(
                "#{payload.refname} - #{stage.length} bytes (stage 2)\n" +
                "http://www.metasploit.com\n",
                fmt) +
              Buffer.transform(stage, fmt)
          end
        end

      end

    end

    # How to warn?
    #if exeopts[:fellback]
    #	$stderr.puts(OutError + "Warning: Falling back to default template: #{exeopts[:fellback]}")
    #end

    return output
  end

  #
  # Calls the class method.
  #
  def generate_simple(opts, &block)
    Msf::Simple::Payload.generate_simple(self, opts, &block)
  end

end

end
end