require 'spec_helper'

RSpec.describe Msf::Ui::Console::CommandDispatcher::Exploit do
  include_context 'Msf::DBManager'
  include_context 'Msf::UIDriver'
  include_context 'Rex::Job#start run inline'
  include_context 'Msf::Framework#threads cleaner', verify_cleanup_required: false

  let(:remote_exploit_mod) do
    mod_klass = Class.new(Msf::Exploit) do
      def initialize
        super(
          'Name' => 'mock module',
          'Description' => 'mock module',
          'Author' => ['Unknown'],
          'License' => MSF_LICENSE,
          'Arch' => ARCH_CMD,
          'Platform' => ['unix'],
          'Targets' => [['Automatic', {}]],
          'DefaultTarget' => 0,
        )

        register_options(
          [
            Msf::Opt::RHOSTS,
            Msf::Opt::RPORT(3000),
            Msf::OptFloat.new('FloatValue', [false, 'A FloatValue which should be normalized before framework runs this module', 3.5])
          ]
        )
      end

      def check
        print_status("Checking for target #{datastore['RHOSTS']}:#{datastore['RPORT']} with normalized datastore value #{datastore['FloatValue'].inspect}")
      end

      def run
        print_status("Running for target #{datastore['RHOSTS']}:#{datastore['RPORT']} with normalized datastore value #{datastore['FloatValue'].inspect}")
      end

      alias_method :exploit, :run

      def cleanup
        print_status("Cleanup for target #{datastore['RHOSTS']}:#{datastore['RPORT']}")
      end
    end

    mod = mod_klass.new
    datastore = Msf::ModuleDataStore.new(mod)
    allow(mod).to receive(:framework).and_return(framework)
    mod.send(:datastore=, datastore)
    datastore.import_options(mod.options)
    Msf::Simple::Framework.simplify_module(mod)
    mod
  end

  let(:non_remote_exploit_mod) do
    mod_klass = Class.new(Msf::Exploit) do
      def initialize
        super(
          'Name' => 'mock module',
          'Description' => 'mock module',
          'Author' => ['Unknown'],
          'License' => MSF_LICENSE,
          'Arch' => ARCH_CMD,
          'Platform' => ['unix'],
          'Targets' => [['Automatic', {}]],
          'DefaultTarget' => 0,
          )

        register_options(
          [
            Msf::OptFloat.new('FloatValue', [false, 'A FloatValue which should be normalized before framework runs this module', 3.5])
          ]
        )

        deregister_options('RHOSTS')
      end

      def run
        print_status("Running with normalized datastore value #{datastore['FloatValue'].inspect}")
      end

      alias_method :exploit, :run

      def cleanup
        print_status('Cleanup')
      end
    end

    mod = mod_klass.new
    datastore = Msf::ModuleDataStore.new(mod)
    allow(mod).to receive(:framework).and_return(framework)
    mod.send(:datastore=, datastore)
    datastore.import_options(mod.options)
    Msf::Simple::Framework.simplify_module(mod)
    mod
  end

  subject do
    instance = described_class.new(driver)
    instance
  end

  def set_default_payload(mod)
    mod.datastore['PAYLOAD'] = 'generic/no_session_payload'
    mod.datastore['LHOST'] = '127.0.0.1'
  end

  before do
    run_rex_jobs_inline!

    allow(driver).to receive(:input).and_return(driver_input)
    allow(driver).to receive(:output).and_return(driver_output)
    current_mod.init_ui(driver_input, driver_output)
    allow(subject).to receive(:mod).and_return(current_mod)

    framework.modules.add_module_path(File.join(FILE_FIXTURES_PATH, 'modules'))
  end

  describe '#cmd_check' do
    context 'when checking a remote exploit module' do
      let(:current_mod) { remote_exploit_mod }

      it 'reports missing RHOST values' do
        allow(current_mod).to receive(:run).and_call_original
        current_mod.datastore['RHOSTS'] = ''
        subject.cmd_check
        expected_output = [
          'Msf::OptionValidateError One or more options failed to validate: RHOSTS.'
        ]

        expect(@combined_output).to match_array(expected_output)
        expect(subject.mod).not_to have_received(:run)
      end

      it 'runs a single RHOST value' do
        current_mod.datastore['RHOSTS'] = '192.0.2.1'
        subject.cmd_check
        expected_output = [
          'Checking for target 192.0.2.1:3000 with normalized datastore value 3.5',
          'Cleanup for target 192.0.2.1:3000',
          '192.0.2.1:3000 - Check failed: The state could not be determined.'
        ]

        expect(@combined_output).to match_array(expected_output)
      end

      it 'runs multiple RHOST values' do
        current_mod.datastore['RHOSTS'] = '192.0.2.1 192.0.2.2'
        subject.cmd_check
        expected_output = [
          'Checking for target 192.0.2.1:3000 with normalized datastore value 3.5',
          'Cleanup for target 192.0.2.1:3000',
          '192.0.2.1:3000 - Check failed: The state could not be determined.',
          'Checking for target 192.0.2.2:3000 with normalized datastore value 3.5',
          'Cleanup for target 192.0.2.2:3000',
          '192.0.2.2:3000 - Check failed: The state could not be determined.'
        ]

        expect(@combined_output).to match_array(expected_output)
      end

      it 'normalizes the datastore before running' do
        current_mod.datastore['RHOSTS'] = '192.0.2.1 192.0.2.2'
        current_mod.datastore.store('FloatValue', '5.0')
        subject.cmd_check
        expected_output = [
          'Checking for target 192.0.2.1:3000 with normalized datastore value 5.0',
          'Cleanup for target 192.0.2.1:3000',
          '192.0.2.1:3000 - Check failed: The state could not be determined.',
          'Checking for target 192.0.2.2:3000 with normalized datastore value 5.0',
          'Cleanup for target 192.0.2.2:3000',
          '192.0.2.2:3000 - Check failed: The state could not be determined.'
        ]

        expect(@combined_output).to match_array(expected_output)
      end

      it 'supports inline options' do
        current_mod.datastore.store('FloatValue', '5.0')
        subject.cmd_check('RHOSTS=192.0.2.5', 'FloatValue=10.0')
        expected_output = [
          'Checking for target 192.0.2.5:3000 with normalized datastore value 10.0',
          'Cleanup for target 192.0.2.5:3000',
          '192.0.2.5:3000 - Check failed: The state could not be determined.'
        ]

        expect(@combined_output).to match_array(expected_output)
      end

      it 'supports multiple inlined RHOST values' do
        current_mod.datastore.store('FloatValue', '5.0')
        subject.cmd_check('RHOSTS=192.0.2.5 192.0.2.6', 'FloatValue=10.0')
        expected_output = [
          'Checking for target 192.0.2.5:3000 with normalized datastore value 10.0',
          'Cleanup for target 192.0.2.5:3000',
          '192.0.2.5:3000 - Check failed: The state could not be determined.',
          'Checking for target 192.0.2.6:3000 with normalized datastore value 10.0',
          'Cleanup for target 192.0.2.6:3000',
          '192.0.2.6:3000 - Check failed: The state could not be determined.'
        ]

        expect(@combined_output).to match_array(expected_output)
      end

      it 'incorrectly handles unknown flags, and inadvertently runs the module with the old rhosts value' do
        current_mod.datastore['RHOSTS'] = '192.0.2.1'
        subject.cmd_check('-j')
        expected_output = [
          'Checking for target 192.0.2.1:3000 with normalized datastore value 3.5',
          'Cleanup for target 192.0.2.1:3000',
          '192.0.2.1:3000 - Check failed: The state could not be determined.'
        ]

        expect(@combined_output).to match_array(expected_output)
      end
    end

    context 'when checking a non remote exploit module' do
      let(:current_mod) { non_remote_exploit_mod }

      it 'notifies the user that this module does not support check' do
        subject.cmd_check
        expected_output = [
          'This module does not support check.'
        ]

        expect(@combined_output).to match_array(expected_output)
      end
    end
  end

  describe '#cmd_run' do
    before do
      set_default_payload(current_mod)
    end

    context 'when running a remote exploit module' do
      let(:current_mod) { remote_exploit_mod }

      it 'reports missing RHOST values' do
        allow(current_mod).to receive(:run).and_call_original
        current_mod.datastore['RHOSTS'] = nil
        subject.cmd_run
        expected_output = [
          'Msf::OptionValidateError One or more options failed to validate: RHOSTS.'
        ]

        expect(@combined_output).to match_array(expected_output)
        expect(subject.mod).not_to have_received(:run)
      end

      it 'attempts to run modules with blank RHOSTS' do
        allow(current_mod).to receive(:run).and_call_original
        current_mod.datastore['RHOSTS'] = ''
        subject.cmd_run
        expected_output = [
          'Msf::OptionValidateError One or more options failed to validate: RHOSTS.'
        ]

        expect(@combined_output).to match_array(expected_output)
        expect(subject.mod).not_to have_received(:run)
      end

      it 'reports a missing payload value' do
        allow(current_mod).to receive(:run).and_call_original
        current_mod.datastore['PAYLOAD'] = nil
        current_mod.datastore['RHOSTS'] = '192.0.2.1'
        subject.cmd_run
        expected_output = [
          'Exploit failed: A payload has not been selected.',
          'Exploit completed, but no session was created.'
        ]

        expect(@combined_output).to match_array(expected_output)
        expect(subject.mod).not_to have_received(:run)
      end

      it 'validates payload options' do
        set_default_payload(current_mod)
        allow(current_mod).to receive(:run).and_call_original
        current_mod.datastore['REQUIRED_PAYLOAD_OPTION'] = 'foo'
        current_mod.datastore['RHOSTS'] = '192.0.2.1'
        subject.cmd_run
        expected_output = [
          'Exploit completed, but no session was created.',
          'Msf::OptionValidateError One or more options failed to validate: REQUIRED_PAYLOAD_OPTION.'
        ]

        expect(@combined_output).to match_array(expected_output)
        expect(subject.mod).not_to have_received(:run)
      end

      it 'runs a single RHOST value' do
        set_default_payload(current_mod)
        current_mod.datastore['RHOSTS'] = '192.0.2.1'
        subject.cmd_run
        expected_output = [
          'Running for target 192.0.2.1:3000 with normalized datastore value 3.5',
          'Cleanup for target 192.0.2.1:3000',
          'Exploit completed, but no session was created.'
        ]

        expect(@combined_output).to match_array(expected_output)
      end

      it 'runs multiple RHOST values' do
        set_default_payload(current_mod)
        current_mod.datastore['RHOSTS'] = '192.0.2.1 192.0.2.2'
        subject.cmd_run
        expected_output = [
          'Exploiting target 192.0.2.1',
          'Running for target 192.0.2.1:3000 with normalized datastore value 3.5',
          'Cleanup for target 192.0.2.1:3000',
          'Exploiting target 192.0.2.2',
          'Running for target 192.0.2.2:3000 with normalized datastore value 3.5',
          'Cleanup for target 192.0.2.2:3000',
          'Exploit completed, but no session was created.'
        ]

        expect(@combined_output).to match_array(expected_output)
      end

      it 'normalizes the datastore before running' do
        set_default_payload(current_mod)
        current_mod.datastore['RHOSTS'] = '192.0.2.1 192.0.2.2'
        current_mod.datastore.store('FloatValue', '5.0')
        subject.cmd_run
        expected_output = [
          'Exploiting target 192.0.2.1',
          'Running for target 192.0.2.1:3000 with normalized datastore value 5.0',
          'Cleanup for target 192.0.2.1:3000',
          'Exploiting target 192.0.2.2',
          'Running for target 192.0.2.2:3000 with normalized datastore value 5.0',
          'Cleanup for target 192.0.2.2:3000',
          'Exploit completed, but no session was created.'
        ]

        expect(@combined_output).to match_array(expected_output)
      end

      it 'supports inline options' do
        set_default_payload(current_mod)
        current_mod.datastore.store('FloatValue', '5.0')
        subject.cmd_run('RHOSTS=192.0.2.5', 'FloatValue=10.0')
        expected_output = [
          'Running for target 192.0.2.5:3000 with normalized datastore value 10.0',
          'Cleanup for target 192.0.2.5:3000',
          'Exploit completed, but no session was created.'
        ]

        expect(@combined_output).to match_array(expected_output)
      end

      it 'supports multiple inlined RHOST values' do
        set_default_payload(current_mod)
        current_mod.datastore.store('FloatValue', '5.0')
        subject.cmd_run('RHOSTS=192.0.2.5 192.0.2.6', 'FloatValue=10.0')
        expected_output = [
          'Exploiting target 192.0.2.5',
          'Running for target 192.0.2.5:3000 with normalized datastore value 10.0',
          'Cleanup for target 192.0.2.5:3000',
          'Exploiting target 192.0.2.6',
          'Running for target 192.0.2.6:3000 with normalized datastore value 10.0',
          'Cleanup for target 192.0.2.6:3000',
          'Exploit completed, but no session was created.'
        ]

        expect(@combined_output).to match_array(expected_output)
      end

      it 'supports rhosts as arguments' do
        set_default_payload(current_mod)
        current_mod.datastore.store('FloatValue', '5.0')
        subject.cmd_run('192.0.2.5', '192.0.2.6')
        expected_output = [
          'Exploiting target 192.0.2.5',
          'Running for target 192.0.2.5:3000 with normalized datastore value 5.0',
          'Cleanup for target 192.0.2.5:3000',
          'Exploiting target 192.0.2.6',
          'Running for target 192.0.2.6:3000 with normalized datastore value 5.0',
          'Cleanup for target 192.0.2.6:3000',
          'Exploit completed, but no session was created.'
        ]

        expect(@combined_output).to match_array(expected_output)
      end

      it 'honors the -j flag, and the module is run as a job' do
        set_default_payload(current_mod)
        current_mod.datastore['RHOSTS'] = '192.0.2.1'
        subject.cmd_run('-j')
        expected_output = [
          'Running rex job 0 inline',
          'Running for target 192.0.2.1:3000 with normalized datastore value 3.5',
          'Exploit running as background job 0.',
          'Exploit completed, but no session was created.'
        ]

        expect(@combined_output).to match_array(expected_output)
      end

      it 'honors the -j flag, and the module is run as a job when there are multiple hosts' do
        set_default_payload(current_mod)
        current_mod.datastore['RHOSTS'] = '192.0.2.1 192.0.2.2'
        subject.cmd_run('-j')
        expected_output = [
          'Exploiting target 192.0.2.1',
          'Running rex job 0 inline',
          'Running for target 192.0.2.1:3000 with normalized datastore value 3.5',
          'Exploiting target 192.0.2.2',
          'Running rex job 1 inline',
          'Running for target 192.0.2.2:3000 with normalized datastore value 3.5',
          'Exploit completed, but no session was created.'
        ]

        expect(@combined_output).to match_array(expected_output)
      end
    end

    context 'when running a non remote exploit module' do
      let(:current_mod) { non_remote_exploit_mod }

      it 'reports a missing payload value' do
        allow(current_mod).to receive(:run).and_call_original
        current_mod.datastore['PAYLOAD'] = nil
        current_mod.datastore['RHOSTS'] = '192.0.2.1'
        subject.cmd_run
        expected_output = [
          'Exploit failed: A payload has not been selected.',
          'Exploit completed, but no session was created.'
        ]

        expect(@combined_output).to match_array(expected_output)
        expect(subject.mod).not_to have_received(:run)
      end

      it 'runs when a payload is set' do
        set_default_payload(current_mod)
        subject.cmd_run
        expected_output = [
          'Running with normalized datastore value 3.5',
          'Cleanup',
          'Exploit completed, but no session was created.'
        ]

        expect(@combined_output).to match_array(expected_output)
      end

      it 'runs the module once, even if multiple rhost values are set' do
        set_default_payload(current_mod)
        current_mod.datastore.store('FloatValue', '10.0')
        current_mod.datastore['RHOSTS'] = '192.0.2.1 192.0.2.2 192.0.2.3'
        subject.cmd_run
        expected_output = [
          'Running with normalized datastore value 10.0',
          'Cleanup',
          'Exploit completed, but no session was created.'
        ]

        expect(@combined_output).to match_array(expected_output)
      end

      it 'normalized the datastore before running' do
        set_default_payload(current_mod)
        current_mod.datastore.store('FloatValue', '5.0')
        subject.cmd_run
        expected_output = [
          'Running with normalized datastore value 5.0',
          'Cleanup',
          'Exploit completed, but no session was created.'
        ]

        expect(@combined_output).to match_array(expected_output)
      end

      it 'supports inline options' do
        set_default_payload(current_mod)
        subject.cmd_run('FloatValue=10.0')
        expected_output = [
          'Running with normalized datastore value 10.0',
          'Cleanup',
          'Exploit completed, but no session was created.'
        ]

        expect(@combined_output).to match_array(expected_output)
      end

      it 'supports payloads being set via an option' do
        set_default_payload(current_mod)
        subject.cmd_run('-p', 'foo/bar/baz')
        expected_output = [
          'Exploit failed: You specified an invalid payload: foo/bar/baz',
          'Exploit completed, but no session was created.'
        ]

        expect(@combined_output).to match_array(expected_output)
      end

      it 'supports payloads being set via a datastore option' do
        set_default_payload(current_mod)
        subject.cmd_run('payload=foo/bar/baz')
        expected_output = [
          'Exploit failed: You specified an invalid payload: foo/bar/baz',
          'Exploit completed, but no session was created.'
        ]

        expect(@combined_output).to match_array(expected_output)
      end
    end
  end

  describe '#cmd_rerun' do
  end

  describe '#cmd_exploit' do
  end

  describe '#cmd_reload' do
  end
end