require 'spec_helper'

require 'msf/core'
require 'msf/core/exploit/cmdstager'

RSpec.describe Msf::Exploit::CmdStager do

  def create_exploit(info ={})
    mod = Msf::Exploit.allocate
    mod.extend described_class
    mod.send(:initialize, info)
    mod
  end

  describe "#select_cmdstager" do

    subject do
      create_exploit
    end

    context "when no flavor" do

      it "raises ArgumentError" do
        expect { subject.select_cmdstager }.to raise_error(ArgumentError, /Unable to select CMD Stager/)
      end
    end

    context "when correct flavor" do

      context "with default decoder" do

        let(:flavor) do
          :vbs
        end

        before do
          subject.select_cmdstager(:flavor => flavor)
        end

        it "selects flavor" do
          expect(subject.flavor).to eq(flavor)
        end

        it "selects default decoder" do
          expect(subject.decoder).to eq(subject.default_decoder(flavor))
        end
      end

      context "without default decoder" do

        let(:flavor) do
          :tftp
        end

        before do
          subject.select_cmdstager(:flavor => flavor)
        end

        it "selects flavor" do
          expect(subject.flavor).to eq(flavor)
        end

        it "hasn't decoder" do
          expect(subject.decoder).to be_nil
        end
      end

      context "with incompatible target" do

        subject do
          create_exploit({
            'DefaultTarget' => 0,
            'Targets' =>
              [
                ['Linux',
                  {
                    'Platform' => 'linux',
                    'CmdStagerFlavor' => 'tftp'
                  }
                ]
              ]
          })
        end

        let(:flavor) do
          :vbs
        end

        it "raises ArgumentError" do
          expect { subject.select_cmdstager(:flavor => flavor) }.to raise_error(ArgumentError, /The CMD Stager '\w+' isn't compatible with the target/)
        end
      end
    end
  end

  describe "#default_decoder" do

    subject do
      create_exploit
    end

    context "when valid flavor as input" do

      context "with default decoder" do
        let(:flavor) do
          :vbs
        end

        let(:expected_decoder) do
          described_class::DECODERS[:vbs]
        end

        it "returns the decoder path" do
          expect(subject.default_decoder(flavor)).to eq(expected_decoder)
        end
      end

      context "without default decoder" do
        let(:flavor) do
          :bourne
        end

        it "returns nil" do
          expect(subject.default_decoder(flavor)).to be_nil
        end
      end
    end

    context "when invalid flavor as input" do
      let(:flavor) do
        :invalid_flavor
      end

      it "returns nil" do
        expect(subject.default_decoder(flavor)).to be_nil
      end
    end

    context "when nil flavor as input" do
      let(:flavor) do
        nil
      end

      it "should be nil" do
        expect(subject.default_decoder(flavor)).to be_nil
      end
    end
  end

  describe "#module_flavors" do

    context "when the module hasn't CmdStagerFlavor info" do

      context "neither the target" do

        subject do
          create_exploit
        end

        it "returns empty array" do
          expect(subject.module_flavors).to eq([])
        end
      end

      context "the target has CmdStagerFlavor info" do

        subject do
          create_exploit({
            'DefaultTarget' => 0,
            'Targets' =>
              [
                ['Windows',
                  {
                    'CmdStagerFlavor' => 'vbs'
                  }
                ]
              ]
          })
        end

        let(:expected_flavor) do
          ['vbs']
        end

        it "returns an array with the target flavor" do
          expect(subject.module_flavors).to eq(expected_flavor)
        end
      end
    end

    context "when the module has CmdStagerFlavor info" do

      context "but the target hasn't CmdStagerFlavor info" do

        subject do
          create_exploit('CmdStagerFlavor' => 'vbs')
        end

        let(:expected_flavor) do
          ['vbs']
        end

        it "returns an array with the module flavor" do
          expect(subject.module_flavors).to eq(expected_flavor)
        end
      end

      context "and the target has CmdStagerFlavor info" do

        subject do
          create_exploit({
            'CmdStagerFlavor' => 'vbs',
            'DefaultTarget'   => 0,
            'Targets' =>
              [
                ['Windows TFTP',
                  {
                    'CmdStagerFlavor' => 'tftp'
                  }
                ]
              ]
          })
        end

        let(:expected_flavor) do
          ['vbs', 'tftp']
        end

        it "returns an array with all the flavors available to the module" do
          expect(subject.module_flavors).to eq(expected_flavor)
        end
      end
    end
  end

  describe "#target_flavor" do

    context "when the module hasn't CmdStagerFlavor info" do

      context "neither the target" do

        subject do
          create_exploit
        end

        it "returns nil" do
          expect(subject.target_flavor).to be_nil
        end
      end

      context "the target has CmdStagerFlavor info" do

        subject do
          create_exploit({
            'DefaultTarget' => 0,
            'Targets' =>
              [
                ['Windows',
                  {
                    'CmdStagerFlavor' => 'vbs'
                  }
                ]
              ]
          })
        end

        let(:expected_flavor) do
          'vbs'
        end

        it "returns the target flavor" do
          expect(subject.target_flavor).to eq(expected_flavor)
        end
      end
    end

    context "when the module has CmdStagerFlavor info" do

      context "but the target hasn't CmdStagerFlavor info" do

        subject do
          create_exploit('CmdStagerFlavor' => 'vbs')
        end

        let(:expected_flavor) do
          'vbs'
        end

        it "returns the module flavor" do
          expect(subject.target_flavor).to eq(expected_flavor)
        end
      end

      context "and the target has CmdStagerFlavor info" do

        subject do
          create_exploit({
            'CmdStagerFlavor' => 'vbs',
            'DefaultTarget'   => 0,
            'Targets' =>
              [
                ['Windows TFTP',
                  {
                    'CmdStagerFlavor' => 'tftp'
                  }
                ]
              ]
          })
        end

        let(:expected_flavor) do
          'tftp'
        end

        it "returns the target flavor" do
          expect(subject.target_flavor).to eq(expected_flavor)
        end
      end
    end
  end

  describe "#compatible_flavor?" do

    context "when there isn't target flavor" do

      subject do
        create_exploit
      end

      let(:flavor) do
        :vbs
      end

      it "is compatible" do
        expect(subject.compatible_flavor?(flavor)).to be_truthy
      end
    end

    context "when the target flavor is a string" do

      subject do
        create_exploit('CmdStagerFlavor' => 'vbs')
      end

      context "and good flavor" do
        let(:flavor) do
          :vbs
        end

        it "is compatible" do
          expect(subject.compatible_flavor?(flavor)).to be_truthy
        end
      end

      context "and bad flavor" do
        let(:flavor) do
          :tftp
        end

        it "isn't compatible" do
          expect(subject.compatible_flavor?(flavor)).to be_falsey
        end
      end
    end

    context "when the target flavor is a symbol" do

      subject do
        create_exploit('CmdStagerFlavor' => :vbs)
      end

      context "and good flavor" do
        let(:flavor) do
          :vbs
        end

        it "is compatible" do
          expect(subject.compatible_flavor?(flavor)).to be_truthy
        end
      end

      context "and bad flavor" do
        let(:flavor) do
          :tftp
        end

        it "isn't compatible" do
          expect(subject.compatible_flavor?(flavor)).to be_falsey
        end
      end
    end

    context "when the target flavor is an Array" do

      subject do
        create_exploit('CmdStagerFlavor' => ['vbs', :tftp])
      end

      context "and good flavor" do
        let(:flavor) do
          :vbs
        end

        it "is compatible" do
          expect(subject.compatible_flavor?(flavor)).to be_truthy
        end
      end

      context "and bad flavor" do
        let(:flavor) do
          :echo
        end

        it "isn't compatible" do
          expect(subject.compatible_flavor?(flavor)).to be_falsey
        end
      end

    end
  end

  describe "#guess_flavor" do

    context "when the module hasn't targets" do

      context "neither platforms" do
        subject do
          create_exploit
        end

        it "doesn't guess" do
          expect(subject.guess_flavor).to be_nil
        end
      end

      context "but platforms" do

        context "one platform with default flavor" do
          let(:platform) do
            'win'
          end

          let(:expected_flavor) do
            :vbs
          end

          subject do
            create_exploit('Platform' => platform)
          end

          it "guess the platform defulat flavor" do
            expect(subject.guess_flavor).to eq(expected_flavor)
          end
        end

        context "one platform without default flavor" do
          let (:platform) do
            'java'
          end

          subject do
            create_exploit('Platform' => platform)
          end

          it "doesn't guess" do
            expect(subject.guess_flavor).to be_nil
          end
        end

        context "two platforms" do
          let(:platform) do
            ['unix', 'linux']
          end

          subject do
            create_exploit('Platform' => platform)
          end

          it "doesn't guess" do
            expect(subject.guess_flavor).to be_nil
          end
        end
      end
    end

    context "when the module has one target" do

      context "and the target has one platform" do

        context "with default flavor"do
          let (:expected_flavor) do
            :vbs
          end

          let (:platform) do
            'win'
          end

          subject do
            create_exploit({
              'DefaultTarget' => 0,
              'Targets' =>
                [
                  ['Windows',
                    {
                      'Platform' => platform
                    }
                  ]
                ]
            })
          end

          it "guess the target flavor" do
            expect(subject.guess_flavor).to eq(expected_flavor)
          end

        end

        context "without default flavor" do
          let (:platform) do
            'java'
          end

          subject do
            create_exploit({
              'DefaultTarget' => 0,
              'Targets' =>
                [
                  ['Java',
                    {
                      'Platform' => platform
                    }
                  ]
                ]
            })
          end

          it "doesn't guess" do
            expect(subject.guess_flavor).to be_nil
          end
        end
      end

      context "the target has two platforms" do
        subject do
          create_exploit({
            'DefaultTarget' => 0,
            'Targets' =>
              [
                ['MultiPlatform',
                  {
                    'Platform' => %w{ linux unix}
                  }
                ]
              ]
          })
        end

        it "doesn't guess" do
          expect(subject.guess_flavor).to be_nil
        end
      end
    end
  end

  describe "#select_flavor" do

    context "when flavor set in the datastore" do

      subject do
        create_exploit({
          'DefaultOptions' => {
            'CMDSTAGER::FLAVOR' => 'vbs'
          }
        })
      end

      let(:datastore_flavor) do
        :vbs
      end

      it "returns the datastore flavor" do
        expect(subject.select_flavor).to eq(datastore_flavor)
      end

      context "and flavor set in the opts" do

        let(:opts_flavor) do
          :bourne
        end

        it "returns the opts flavor" do
          expect(subject.select_flavor(:flavor => :bourne)).to eq(opts_flavor)
        end
      end
    end
  end

  describe "#select_decoder" do

    context "when decoder set in the datastore" do

      let(:decoder) do
        File.join(Rex::Exploitation::DATA_DIR, "exploits", "cmdstager", "vbs_b64")
      end

      subject do
        create_exploit({
          'DefaultOptions' => {
            'CMDSTAGER::DECODER' => decoder
          }
        })
      end

      it "returns datastore flavor" do
        expect(subject.select_decoder).to eq(decoder)
      end

      context "and decoder set in the opts" do

        let(:decoder_opts) do
          File.join(Rex::Exploitation::DATA_DIR, "exploits", "cmdstager", "vbs_b64_adodb")
        end

        it "returns the decoder_opts" do
          expect(subject.select_decoder(:decoder => decoder_opts)).to eq(decoder_opts)
        end
      end
    end
  end

  describe "#opts_with_decoder" do
    subject do
      create_exploit
    end

    context "with :decoder option" do

      let(:decoder) do
        File.join(Rex::Exploitation::DATA_DIR, "exploits", "cmdstager", "vbs_b64")
      end

      it "returns the :decoder option" do
        expect(subject.opts_with_decoder(:decoder => decoder)).to include(:decoder)
      end
    end

    context "without decoder option" do
      it ":hasn't decoder option" do
        expect(subject.opts_with_decoder).not_to include(:decoder)
      end
    end

  end

  describe "#create_stager" do
    subject do
      create_exploit
    end

    context "with correct flavor" do

      let(:flavor) do
        :vbs
      end

      let(:expected_class) do
        described_class::STAGERS[flavor]
      end

      before do
        subject.flavor = flavor
      end

      it "creates the correct instance" do
        expect(subject.create_stager.class).to eq(expected_class)
      end
    end

    context "with incorrect flavor" do
      let(:flavor) do
        :incorrect_flavor
      end

      let(:expected_class) do
        described_class::STAGERS[flavor]
      end

      before do
        subject.flavor = flavor
      end

      it "raises a NoMethodError" do
        expect { subject.create_stager }.to raise_error(NoMethodError)
      end
    end
  end
end
