Add health check functionality

lib/msf/core/rpc/json/v1_0/rpc_command.rb

@@ -148,4 +148,4 @@ def post_process_result(result, method, params)
       end
     end
   end
-end
+end

lib/msf/core/rpc/v10/client.rb

@@ -91,7 +91,7 @@ def call(meth, *args)
       do_logout_cleanup
     end
 
-    unless meth == "auth.login"
+    if meth != 'auth.login' && meth != 'health.check'
       unless self.token
         raise RuntimeError, "client not authenticated"
       end

lib/msf/core/rpc/v10/health.rb

@@ -0,0 +1,35 @@
+# -*- coding: binary -*-
+module Msf
+  module RPC
+    class Health
+
+      # Returns whether the framework object is currently healthy and ready to accept
+      # requests
+      #
+      # @return [Hash]
+      #
+      def self.check(framework)
+        # A couple of rudimentary checks to ensure that nothing breaks when interacting
+        # with framework object
+        is_healthy = (
+          !framework.version.to_s.empty? &&
+          # Ensure that the db method can be invoked and returns a truthy value as
+          # the rpc clients interact with framework's database object which raises can
+          # raise an exception
+          framework.db
+        )
+
+        unless is_healthy
+          return { status: 'DOWN' }
+        end
+
+        { status: 'UP' }
+      rescue => e
+        elog('Health status failing', error: e)
+
+        { status: 'DOWN' }
+      end
+
+    end
+  end
+end

lib/msf/core/rpc/v10/rpc_health.rb

@@ -0,0 +1,18 @@
+# -*- coding: binary -*-
+module Msf
+module RPC
+class RPC_Health < RPC_Base
+
+  # Returns whether the service is currently healthy and ready to accept
+  # requests. This endpoint is not authenticated.
+  #
+  # @return [Hash]
+  # @example Here's how you would use this from the client:
+  #  rpc.call('health.check')
+  def rpc_check_noauth
+    Msf::RPC::Health.check(framework)
+  end
+
+end
+end
+end

lib/msf/core/rpc/v10/service.rb

@@ -36,6 +36,7 @@ def initialize(framework, options={})
     self.users              = self.options[:users] || []
     self.job_status_tracker = Msf::RPC::RpcJobStatusTracker.new
 
+    add_handler("health",  Msf::RPC::RPC_Health.new(self))
     add_handler("core",    Msf::RPC::RPC_Core.new(self))
     add_handler("auth",    Msf::RPC::RPC_Auth.new(self))
     add_handler("console", Msf::RPC::RPC_Console.new(self))

lib/msf/core/web_services/json_rpc_app.rb

@@ -14,6 +14,7 @@ class JsonRpcApp < Sinatra::Base
 
     # Servlet registration
     register AuthServlet
+    register HealthServlet
     register JsonRpcServlet
 
     # Custom error handling

lib/msf/core/web_services/servlet/health_servlet.rb

@@ -0,0 +1,22 @@
+module Msf::WebServices::HealthServlet
+
+  def self.api_path
+    '/api/v1/health'
+  end
+
+  def self.registered(app)
+    app.get self.api_path, &health_check
+  end
+
+  #######
+  private
+  #######
+
+  def self.health_check
+    lambda {
+      health_check = Msf::RPC::Health.check(framework)
+      is_success = health_check[:status] == 'UP'
+      set_json_data_response(response: health_check, code: is_success ? 200 : 503)
+    }
+  end
+end

lib/msf_autoload.rb

@@ -181,6 +181,7 @@ def custom_inflections
       'rpc_auth' => 'RPC_Auth',
       'rpc_job' => 'RPC_Job',
       'rpc_core' => 'RPC_Core',
+      'rpc_health' => 'RPC_Health',
       'rpc_module' => 'RPC_Module',
       'cli' => 'CLI',
       'sqlitei' => 'SQLitei',

msf-json-rpc.ru

@@ -24,29 +24,17 @@ run Msf::WebServices::JsonRpcApp
 #
 warmup do |app|
   client = Rack::MockRequest.new(app)
-  response = client.post(
-    '/api/v1/json-rpc',
-    input: {
-      jsonrpc: '2.0',
-      method: 'core.version',
-      id: 1,
-      params: []
-    }.to_json
-  )
-
-  warmup_error_message = "Metasploit JSON RPC did not successfully start up. Unexpected response returned: #{response.body}"
+  response = client.get('/api/v1/health')
+
+  warmup_error_message = "Metasploit JSON RPC did not successfully start up. Unexpected response returned: '#{response.body}'"
   begin
     parsed_response = JSON.parse(response.body)
   rescue JSON::ParserError => e
     raise warmup_error_message, e
   end
 
-  is_valid_response = (
-    parsed_response['jsonrpc'] == '2.0' &&
-      parsed_response['id'] == 1 &&
-      !parsed_response.dig('result', 'version').to_s.empty? &&
-      !parsed_response.dig('result', 'ruby').to_s.empty?
-  )
+  expected_response = { 'data' => { 'status' => 'UP' } }
+  is_valid_response = parsed_response == expected_response
 
   unless is_valid_response
     raise warmup_error_message

spec/api/json_rpc_spec.rb

@@ -13,7 +13,8 @@
   include_context 'Msf::Framework#threads cleaner'
 
   let(:app) { subject }
-  let(:api_url) { '/api/v1/json-rpc' }
+  let(:health_check_url) { '/api/v1/health' }
+  let(:rpc_url) { '/api/v1/json-rpc' }
   let(:framework) { app.settings.framework }
   let(:module_name) { 'scanner/ssl/openssl_heartbleed' }
   let(:a_valid_result_uuid) { { 'result' => hash_including({ 'uuid' => match(/\w+/) }) } }
@@ -27,7 +28,7 @@
   end
 
   def create_job
-    post api_url, {
+    post rpc_url, {
       'jsonrpc': '2.0',
       'method': 'module.check',
       'id': 1,
@@ -42,7 +43,7 @@ def create_job
   end
 
   def get_job_results(uuid)
-    post api_url, {
+    post rpc_url, {
       'jsonrpc': '2.0',
       'method': 'module.results',
       'id': 1,
@@ -52,6 +53,19 @@ def get_job_results(uuid)
     }.to_json
   end
 
+  def get_rpc_health_check
+    post rpc_url, {
+      'jsonrpc': '2.0',
+      'method': 'health.check',
+      'id': 1,
+      'params': []
+    }.to_json
+  end
+
+  def get_rest_health_check
+    get health_check_url
+  end
+
   def last_json_response
     JSON.parse(last_response.body)
   end
@@ -96,6 +110,80 @@ def wait_for_expect(retry_count = 20, sleep_duration = 0.5)
     end
   end
 
+  describe 'health status' do
+    context 'when using the REST health check functionality' do
+      it 'passes the health check' do
+        expected_response = {
+          "data" => {
+            "status"=>"UP"
+          }
+        }
+
+        get_rest_health_check
+        expect(last_response).to be_ok
+        expect(last_json_response).to eq(expected_response)
+      end
+    end
+
+    context 'when there is an issue' do
+      before(:each) do
+        allow(framework).to receive(:version).and_raise 'Mock error'
+      end
+
+      it 'fails the health check' do
+        expected_response = {
+          "data" => {
+            "status"=>"DOWN"
+          }
+        }
+
+        get_rest_health_check
+
+        expect(last_response.status).to be 503
+        expect(last_json_response).to eq(expected_response)
+      end
+    end
+
+    context 'when using the RPC health check functionality' do
+      context 'when the service is healthy' do
+        it 'passes the health check' do
+          expected_response = {
+            "id"=>1,
+            "jsonrpc"=>"2.0",
+            "result"=> {
+              "status"=>"UP"
+            }
+          }
+
+          get_rpc_health_check
+          expect(last_response).to be_ok
+          expect(last_json_response).to eq(expected_response)
+        end
+      end
+
+      context 'when there is an issue' do
+        before(:each) do
+          allow(framework).to receive(:version).and_raise 'Mock error'
+        end
+
+        it 'fails the health check' do
+          expected_response = {
+            "id"=>1,
+            "jsonrpc"=>"2.0",
+            "result"=> {
+              "status"=>"DOWN"
+            }
+          }
+
+          get_rpc_health_check
+
+          expect(last_response).to be_ok
+          expect(last_json_response).to eq(expected_response)
+        end
+      end
+    end
+  end
+
   describe 'Running a check job and verifying results' do
     context 'when the module returns check code safe' do
       before(:each) do