lib and spec updates

lib/msf/core/exploit/remote/http/nifi/dbconnectionpool.rb

@@ -48,18 +48,23 @@ def delete_dbconnectionpool(token, db_con_pool, version = 0)
     opts = {
       'method' => 'DELETE',
       'uri' => normalize_uri(target_uri.path, 'nifi-api', 'controller-services', db_con_pool),
-      'vars_get' => {
-        'version' => version
-      }
+      'vars_get' => { 'version' => version }
     }
     opts['headers'] = { 'Authorization' => "Bearer #{token}" } if token
     res = send_request_cgi(opts)
-    while res.code == 400 && res.body.include?('is not the most up-to-date revision') && version <= 20
-      vprint_status("Found newer revision of #{db_con_pool}, attempting to delete version #{version}")
-      delete_dbconnectionpool(token, db_con_pool, version += 1)
-    end
+
     raise DBConnectionPoolError if res.nil?
 
+    while res.code == 400 && res.body.include?('is not the most up-to-date revision') && version <= 20
+      version += 1
+      opts['vars_get'] = { 'version' => version }
+
+      res = send_request_cgi(opts)
+      raise ProcessorError if res.nil?
+
+      vprint_status("Found newer revision of #{db_con_pool}, attempting to delete version #{version}") if res.code == 400 && res.body.include?('is not the most up-to-date revision')
+    end
+
     if version == 20
       print_bad("Aborting after attempting to delete 20 version of DB Connection Pool: #{db_con_pool}")
       raise DBConnectionPoolError

lib/msf/core/exploit/remote/http/nifi/processor.rb

@@ -1,173 +1,147 @@
-# -*- coding: binary -*-
+# -*- coding:binary -*-
 
-module Msf::Exploit::Remote::HTTP::Nifi::Processor
-  include Msf::Exploit::Remote::HttpClient
+require 'spec_helper'
 
-  class ProcessorError < StandardError
+RSpec.describe Msf::Exploit::Remote::HTTP::Nifi::Processor do
+  subject do
+    mod = ::Msf::Module.new
+    mod.extend described_class
+    mod
   end
 
-  # Start processor
-  #
-  # @param token [String] The bearer token from a valid login
-  # @param processor [String] UUID of the processes
-  def start_processor(token, processor)
-    vprint_status("Attempting to start Processor: #{processor}")
-    body = {
-      'state' => 'RUNNING',
-      'disconnectedNodeAcknowledged' => false,
-      'revision' => {
-        'clientId' => 'x',
-        'version' => 0
-      }
-    }
-    opts = {
-      'method' => 'PUT',
-      'uri' => normalize_uri(target_uri.path, 'nifi-api', 'processors', processor, 'run-status'),
-      'ctype' => 'application/json',
-      'data' => body.to_json
-    }
-    opts['headers'] = { 'Authorization' => "Bearer #{token}" } if token
-    res = send_request_cgi(opts)
-    raise ProcessorError if res.nil?
-
-    unless res.code == 200
-      print_bad("Unexpected response code: #{res.code}")
-      raise ProcessorError
-    end
-    print_good('Processor Start sent successfully')
+  let(:valid_code) do
+    200
   end
 
-  # Stop processor
-  #
-  # @param token [String] The bearer token from a valid login
-  # @param processor [String] UUID of the processes
-  def stop_processor(token, processor)
-    vprint_status("Attempting to stop Processor: #{processor}")
-    body = {
-      'revision' => {
-        'clientId' => 'x',
-        'version' => 1
-      },
-      'state' => 'STOPPED'
-    }
-    opts = {
-      'method' => 'PUT',
-      'uri' => normalize_uri(target_uri.path, 'nifi-api', 'processors', processor, 'run-status'),
-      'ctype' => 'application/json',
-      'data' => body.to_json
-    }
-    opts['headers'] = { 'Authorization' => "Bearer #{token}" } if token
-    res = send_request_cgi(opts)
-    raise ProcessorError if res.nil?
-
-    unless res.code == 200
-      print_bad("Unexpected response code: #{res.code}")
-      raise ProcessorError
+  describe '#nifi Processor start_processor' do
+    it 'raises error if page can not be reached' do
+      allow(subject).to receive(:send_request_cgi) do
+        res = Rex::Proto::Http::Response::E404.new
+        res
+      end
+
+      expect { subject.start_processor('a', 'a') }.to raise_error(Msf::Exploit::Remote::HTTP::Nifi::Processor::ProcessorError)
+    end
+
+    it 'raises error when unexpected response code is received' do
+      allow(subject).to receive(:send_request_cgi) do
+        res = Rex::Proto::Http::Response.new
+        res.code = 400
+        res.body = ''
+        res
+      end
+
+      expect { subject.start_processor('a', 'a') }.to raise_error(Msf::Exploit::Remote::HTTP::Nifi::Processor::ProcessorError)
     end
 
-    # Stop may not have worked (but must be done first). Terminate threads now
-    opts = {
-      'method' => 'DELETE',
-      'uri' => normalize_uri(target_uri.path, 'nifi-api', 'processors', processor, 'threads')
-    }
-    opts['headers'] = { 'Authorization' => "Bearer #{token}" } if token
-    res = send_request_cgi(opts)
-    raise ProcessorError if res.nil?
-
-    unless res.code == 200
-      print_bad("Unexpected response code: #{res.code}")
-      raise ProcessorError
+    it 'returns nil when successfull' do
+      allow(subject).to receive(:send_request_cgi) do
+        res = Rex::Proto::Http::Response.new
+        res.code = valid_code
+        res.body = ''
+        res
+      end
+
+      expect(subject.start_processor('a', 'a')).to be_nil
     end
-    print_good('Processor Start stop successfully')
   end
 
-  # Delete a processor
-  #
-  # @param token [String] The bearer token from a valid login
-  # @param processor [String] UUID of the processes
-  # @param version [Int] The version number to delete
-  def delete_processor(token, processor, version = 0)
-    vprint_status("Attempting to delete version #{version} of Processor: #{processor}")
-    opts = {
-      'method' => 'DELETE',
-      'uri' => normalize_uri(target_uri.path, 'nifi-api', 'processors', processor),
-      'vars_get' => { 'version' => version }
-    }
-    opts['headers'] = { 'Authorization' => "Bearer #{token}" } if token
-    res = send_request_cgi(opts)
-    # if we tried to delete the old revision, go ahead and delete the newer one
-    # arbitrary version limit of 20
-    while res.code == 400 && res.body.include?('is not the most up-to-date revision') && version <= 20
-      vprint_status("Found newer revision of #{processor}, attempting to delete version #{version}")
-      opts['vars_get'] = { 'version' => version }
-      delete_processor(token, processor, version += 1)
+  describe '#nifi Processor stop_processor' do
+    it 'raises error if page can not be reached' do
+      allow(subject).to receive(:send_request_cgi) do
+        res = Rex::Proto::Http::Response::E404.new
+        res
+      end
+
+      expect { subject.stop_processor('a', 'a') }.to raise_error(Msf::Exploit::Remote::HTTP::Nifi::Processor::ProcessorError)
     end
 
-    raise ProcessorError if res.nil?
+    it 'raises error when unexpected response code is received' do
+      allow(subject).to receive(:send_request_cgi) do
+        res = Rex::Proto::Http::Response.new
+        res.code = 400
+        res.body = ''
+        res
+      end
 
-    if version == 20
-      print_bad("Aborting after attempting to delete 20 version of Processor: #{processor}")
-      raise ProcessorError
+      expect { subject.stop_processor('a', 'a') }.to raise_error(Msf::Exploit::Remote::HTTP::Nifi::Processor::ProcessorError)
     end
 
-    unless res.code == 200
-      print_bad("Unexpected response code: #{res.code}")
-      raise ProcessorError
+    it 'returns nil when successfull' do
+      allow(subject).to receive(:send_request_cgi) do
+        res = Rex::Proto::Http::Response.new
+        res.code = valid_code
+        res.body = ''
+        res
+      end
+
+      expect(subject.stop_processor('a', 'a')).to be_nil
     end
-    print_good('Delete sent successfully')
   end
 
-  # Creates a processor in a process group
-  #
-  # @param token [String] The bearer token from a valid login
-  # @param process_group [String] UUID of a processor group
-  # @param type [String] What type of processor to create
-  # @return [String] The UUID of the root process group
-  def create_processor(token, process_group, type = 'org.apache.nifi.processors.standard.ExecuteProcess')
-    vprint_status("Attempting to create of processor in group: #{process_group} of type #{type}")
-    body = {
-      'component' => { 'type' => type },
-      'revision' => { 'version' => 0 }
-    }
-    opts = {
-      'method' => 'POST',
-      'uri' => normalize_uri(target_uri.path, 'nifi-api', 'process-groups', process_group, 'processors'),
-      'ctype' => 'application/json',
-      'data' => body.to_json
-    }
-    opts['headers'] = { 'Authorization' => "Bearer #{token}" } if token
-    res = send_request_cgi(opts)
-    return nil if res.nil?
-
-    unless res.code == 201
-      print_bad("Unexpected response code: #{res.code}")
-      raise ProcessorError
+  describe '#nifi Processor create_processor' do
+    it 'raises error if page can not be reached' do
+      allow(subject).to receive(:send_request_cgi) do
+        res = Rex::Proto::Http::Response::E404.new
+        res
+      end
+
+      expect { subject.create_processor('a', 'a') }.to raise_error(Msf::Exploit::Remote::HTTP::Nifi::Processor::ProcessorError)
+    end
+
+    it 'raises error when unexpected response code is received' do
+      allow(subject).to receive(:send_request_cgi) do
+        res = Rex::Proto::Http::Response.new
+        res.code = 400
+        res.body = ''
+        res
+      end
+
+      expect { subject.create_processor('a', 'a') }.to raise_error(Msf::Exploit::Remote::HTTP::Nifi::Processor::ProcessorError)
+    end
+
+    it 'returns UUID when successfull' do
+      allow(subject).to receive(:send_request_cgi) do
+        res = Rex::Proto::Http::Response.new
+        res.code = 201
+        res.body = '{"id":"628de124-3d0f-11ee-be56-0242ac120002"}'
+        res
+      end
+
+      expect(subject.create_processor('a', 'a')).to eq('628de124-3d0f-11ee-be56-0242ac120002')
     end
-    res.get_json_document['id']
   end
 
-  # Get a processor in a process group
-  #
-  # @param token [String] The bearer token from a valid login
-  # @param processor [String] UUID of a processoror
-  # @param field [String] the key from the JSON blob to return
-  # @return [String] THe value from the specified field
-  def get_processor_field(token, processor, field = 'id')
-    vprint_status("Attempting to get field #{field} of processor: #{processor}")
-    opts = {
-      'method' => 'GET',
-      'uri' => normalize_uri(target_uri.path, 'nifi-api', 'processors', processor)
-    }
-    opts['headers'] = { 'Authorization' => "Bearer #{token}" } if token
-    res = send_request_cgi(opts)
-
-    return nil if res.nil?
-
-    unless res.code == 200
-      print_bad("Unexpected response code: #{res.code}")
-      raise ProcessorError
+  describe '#nifi Processor get_processor_field' do
+    it 'raises error if page can not be reached' do
+      allow(subject).to receive(:send_request_cgi) do
+        res = Rex::Proto::Http::Response::E404.new
+        res
+      end
+
+      expect { subject.get_processor_field('a', 'a') }.to raise_error(Msf::Exploit::Remote::HTTP::Nifi::Processor::ProcessorError)
     end
 
-    res.get_json_document[field]
+    it 'raises error when unexpected response code is received' do
+      allow(subject).to receive(:send_request_cgi) do
+        res = Rex::Proto::Http::Response.new
+        res.code = 400
+        res.body = ''
+        res
+      end
+
+      expect { subject.get_processor_field('a', 'a') }.to raise_error(Msf::Exploit::Remote::HTTP::Nifi::Processor::ProcessorError)
+    end
+
+    it 'returns UUID when successfull' do
+      allow(subject).to receive(:send_request_cgi) do
+        res = Rex::Proto::Http::Response.new
+        res.code = 200
+        res.body = '{"id":"628de124-3d0f-11ee-be56-0242ac120002"}'
+        res
+      end
+
+      expect(subject.get_processor_field('a', 'a')).to eq('628de124-3d0f-11ee-be56-0242ac120002')
+    end
   end
 end