# -*- coding: binary -*-
require 'rex/text'
require 'rex/arch'
require 'msf/core/framework'

module Rex
module Exploitation

###
#
# This class provides an interface to generating cmdstagers.
#
###

class CmdStagerBase

  def initialize(exe)
    @linemax     = 2047 # covers most likely cases
    @exe         = exe
  end

  #
  # Generates the cmd payload including the h2bv2 decoder and encoded payload.
  # The resulting commands also perform cleanup, removing any left over files
  #
  def generate(opts = {})
    # Allow temporary directory override
    @tempdir = opts[:temp]
    @tempdir ||= "%TEMP%\\"
    if (@tempdir == '.')
      @tempdir = ''
    end

    opts[:linemax] ||= @linemax

    generate_cmds(opts)
  end


  #
  # This does the work of actually building an array of commands that
  # when executed will create and run an executable payload.
  #
  def generate_cmds(opts)

    # Initialize an arry of commands to execute
    cmds = []

    # Add the exe building commands
    cmds += generate_cmds_payload(opts)

    # Add the decoder script building commands
    cmds += generate_cmds_decoder(opts)

    compress_commands(cmds, opts)
  end


  #
  # Generate the commands to create an encoded version of the
  # payload file
  #
  def generate_cmds_payload(opts)

    # First encode the payload
    encoded = encode_payload(opts)

    # Now split it up into usable pieces
    parts = slice_up_payload(encoded, opts)

    # Turn each part into a valid command
    parts_to_commands(parts, opts)
  end

  #
  # This method is intended to be override by the child class
  #
  def encode_payload(opts)
    # Defaults to nothing
    ""
  end

  #
  # We take a string of data and turn it into an array of parts.
  #
  # We save opts[:extra] bytes out of every opts[:linemax] for the parts
  # appended and prepended to the resulting elements.
  #
  def slice_up_payload(encoded, opts)
    tmp = encoded.dup

    parts = []
    xtra_len = opts[:extra]
    xtra_len ||= 0
    while (tmp.length > 0)
      parts << tmp.slice!(0, (opts[:linemax] - xtra_len))
    end

    parts
  end

  #
  # Combine the parts of the encoded file with the stuff that goes
  # before / after it -- example "echo " and " >>file"
  #
  def parts_to_commands(parts, opts)
    # Return as-is
    parts
  end



  #
  # Generate the commands that will decode the file we just created
  #
  def generate_cmds_decoder(opts)
    # Defaults to no commands.
    []
  end



  #
  # Compress commands into as few lines as possible. Minimizes the number of
  # commands to execute while maximizing the number of commands per execution.
  #
  def compress_commands(cmds, opts)
    new_cmds = []
    line = ''

    concat = opts[:concat_operator] || cmd_concat_operator

    # We cannot compress commands if there is no way to combine commands on
    # a single line.
    return cmds unless concat

    cmds.each { |cmd|

      # If this command will fit, concat it and move on.
      if ((line.length + cmd.length + concat.length) < opts[:linemax])
        line << concat if line.length > 0
        line << cmd
        next
      end

      # The command wont fit concat'd to this line, if we have something,
      # we have to add it to the array now.
      if (line.length > 0)
        new_cmds << line
        line = ''
      end

      # If it won't fit even after emptying the current line, error out..
      if (cmd.length > opts[:linemax])
        raise RuntimeError, 'Line too long - %u bytes, max %u' % [cmd.length, opts[:linemax]]
      end

      # It will indeed fit by itself, lets add it.
      line << cmd

    }
    new_cmds << line if (line.length > 0)

    # Return the final array.
    new_cmds
  end

  #
  # Can be overriden.  For exmaple, use for unix use ";" instead
  #
  def cmd_concat_operator
    nil
  end

  # Should be overriden if the cmd stager needs to setup anything
  # before it's executed
  def setup(mod = nil)

  end

  #
  # Should be overriden if the cmd stager needs to do any clenaup
  #
  def teardown(mod = nil)

  end

end
end
end
