# -*- 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']),
        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', false]),
        OptBool.new('Powershell::sub_funcs', [true, 'Substitute function names', false]),
        OptEnum.new('Powershell::method', [true, 'Payload delivery method', 'reflection', %w(net reflection old msil)]),
      ], self.class)
  end

  #
  # Return a script from path or string
  #
  def read_script(script_path)
    return 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? or 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.select { |k, v| k =~ /^Powershell::(strip|sub)/ && v }.keys.map do |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 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.select { |k, v| k =~ /^Powershell::(strip|sub)/ && v }.keys.map do |k|
      mod_method = k.split('::').last.intern
      opts[mod_method.to_sym] = true
    end

    Rex::Powershell::Command.compress_script(script_in, eof, opts)
  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 [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] :use_single_quotes Wraps the -Command
  #   argument in single quotes unless :encode_final_payload
  #
  # @return [String] Powershell command line with 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']

    unless opts.key? :shorten
      opts[:shorten] = (datastore['Powershell::method'] != 'old')
    end
    template_path = File.join(Msf::Config.data_directory,
                              "templates",
                              "scripts")

    command = Rex::Powershell::Command.cmd_psh_payload(pay,
                                                                     payload_arch,
                                                                     template_path,
                                                                     opts)
    vprint_status("Powershell command length: #{command.length}")

    command
  end


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