lib/cms_scanner.rb in cms_scanner-0.0.32 vs lib/cms_scanner.rb in cms_scanner-0.0.33

@@ -15,10 +15,11 @@
 require 'cms_scanner/typhoeus/response'
 require 'cms_scanner/typhoeus/hydra'
 require 'cms_scanner/public_suffix/domain'
 # Custom Libs
 require 'cms_scanner/helper'
+require 'cms_scanner/exit_code'
 require 'cms_scanner/errors/http'
 require 'cms_scanner/cache/typhoeus'
 require 'cms_scanner/target'
 require 'cms_scanner/browser'
 require 'cms_scanner/version'
@@ -54,13 +55,17 @@
     @@total_requests = value
   end
 
   # Scan
   class Scan
+    attr_reader :run_error
+
     def initialize
       controllers << NS::Controller::Core.new
 
+      exit_hook
+
       yield self if block_given?
     end
 
     # @return [ Controllers ]
     def controllers
@@ -68,12 +73,16 @@
     end
 
     def run
       controllers.run
     rescue OptParseValidator::NoRequiredOption => e
+      @run_error = e
+
       formatter.output('@usage', msg: e.message)
     rescue => e
+      @run_error = e
+
       formatter.output('@scan_aborted',
                        reason: e.message,
                        trace: e.backtrace,
                        verbose: controllers.first.parsed_options[:verbose])
     ensure
@@ -89,9 +98,29 @@
     end
 
     # @return [ Hash ]
     def datastore
       controllers.first.datastore
+    end
+
+    # Hook to be able to have an exit code returned
+    # depending on the findings / errors
+    def exit_hook
+      at_exit do
+        if run_error
+          exit(NS::ExitCode::CLI_OPTION_ERROR) if run_error.is_a?(OptParseValidator::Error) ||
+                                                  run_error.is_a?(OptionParser::ParseError)
+
+          exit(NS::ExitCode::INTERRUPTED) if run_error.is_a?(Interrupt)
+          exit(NS::ExitCode::ERROR)
+        end
+
+        controller = controllers.first
+
+        # The parsed_option[:url] must be checked to avoid raising erros when only -h/-v are given
+        exit(NS::ExitCode::VULNERABLE) if controller.parsed_options[:url] && controller.target.vulnerable?
+        exit(NS::ExitCode::OK)
+      end
     end
   end
 end
 
 require "#{CMSScanner::APP_DIR}/app"