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

require 'msf/core'
require 'msf/core/exploit/powershell'

EXAMPLE_PATH = File.join(Msf::Config.data_directory, "exploits", "powershell", "powerdump.ps1")

describe Msf::Exploit::Powershell do
  let(:datastore) { { } }
  subject do
    mod = Module.new
    mod.extend described_class
    mod.stub(
      :datastore => datastore
    )

    mod
  end

  describe "::read_script" do
    it 'should read a sample script file' do
      script = subject.read_script(EXAMPLE_PATH)
      script.should be_kind_of(Rex::Exploitation::Powershell::Script)
    end
  end

  describe "::encode_script" do
    it 'should read and encode a sample script file' do
      script = subject.encode_script(EXAMPLE_PATH)
      script.should be
      script.length.should be > 0
    end
  end

  describe "::compress_script" do
    it 'should create a compress script' do
      script = File.read(EXAMPLE_PATH)
      compressed = subject.compress_script(script)
      compressed.length.should be < script.length
    end
  end

  describe "::generate_psh_command_line" do
    it 'should contain no full stop when :no_full_stop' do
      opts = {:no_full_stop => true}
      command = subject.generate_psh_command_line(opts)
      command.include?("powershell ").should be_true
    end

    it 'should contain full stop unless :no_full_stop' do
      opts = {}
      command = subject.generate_psh_command_line(opts)
      command.include?("powershell.exe ").should be_true

      opts = {:no_full_stop => false}
      command = subject.generate_psh_command_line(opts)
      command.include?("powershell.exe ").should be_true
    end

    it 'should ensure the path should always ends with \\' do
      opts = {:path => "test"}
      command = subject.generate_psh_command_line(opts)
      command.include?("test\\powershell.exe ").should be_true

      opts = {:path => "test\\"}
      command = subject.generate_psh_command_line(opts)
      command.include?("test\\powershell.exe ").should be_true
    end
  end

  describe "::generate_psh_args" do
    it 'should return empty string for nil opts' do
      subject.generate_psh_args(nil).should eql ""
    end

    command_args = [[:encodedcommand, "parp"],
                    [:executionpolicy, "bypass"],
                    [:inputformat, "xml"],
                    [:file, "x"],
                    [:noexit, true],
                    [:nologo, true],
                    [:noninteractive, true],
                    [:mta, true],
                    [:outputformat, 'xml'],
                    [:sta, true],
                    [:noprofile, true],
                    [:windowstyle, "hidden"],
                    [:command, "Z"]
    ]

    permutations = (0..command_args.length).to_a.combination(2).map{|i,j| command_args[i...j]}

    permutations.each do |perms|
      opts = {}
      perms.each do |k,v|
        opts[k] = v
        it "should generate correct arguments for #{opts}" do
          opts[:shorten] = true
          short_args = subject.generate_psh_args(opts)
          opts[:shorten] = false
          long_args = subject.generate_psh_args(opts)

          opt_length = opts.length - 1

          short_args.should_not be_nil
          long_args.should_not be_nil
          short_args.count('-').should eql opt_length
          long_args.count('-').should eql opt_length
          short_args[0].should_not eql " "
          long_args[0].should_not eql " "
          short_args[-1].should_not eql " "
          long_args[-1].should_not eql " "

          if opts[:command]
            long_args[-10..-1].should eql "-Command Z"
            short_args[-4..-1].should eql "-c Z"
          end
       end
      end
    end
  end

end