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

@@ -32,29 +32,46 @@
 # Module
 module CMSScanner
   APP_DIR = Pathname.new(__FILE__).dirname.join('..', 'app').expand_path
   NS      = self
 
-  def self.included(base)
-    remove_const(:NS)
-    const_set(:NS, base)
-    super(base)
-  end
-
   # Number of requests performed to display at the end of the scan
   Typhoeus.on_complete do |response|
     self.total_requests += 1 unless response.cached?
   end
 
-  # @return [ Integer ]
-  def self.total_requests
-    @@total_requests ||= 0
+  # Module to be able to use these class methods when the CMSScanner
+  # is included in another module
+  module ClassMethods
+    # @return [ Integer ]
+    def total_requests
+      @@total_requests ||= 0
+    end
+
+    # @param [ Integer ]
+    def total_requests=(value)
+      @@total_requests = value
+    end
+
+    # The lowercase name of the scanner
+    # Mainly used in directory paths like the default cookie-jar file and
+    # path to load the cli options from files
+    #
+    # @return [ String ]
+    def app_name
+      to_s.underscore
+    end
   end
 
-  # @param [ Integer ]
-  def self.total_requests=(value)
-    @@total_requests = value
+  extend ClassMethods
+
+  def self.included(base)
+    remove_const(:NS)
+    const_set(:NS, base)
+
+    base.extend(ClassMethods)
+    super(base)
   end
 
   # Scan
   class Scan
     attr_reader :run_error
@@ -104,23 +121,27 @@
 
     # 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(run_error_exit_code) if run_error
 
-          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
+
+    # @return [ Integer ] The exit code related to the run_error
+    def run_error_exit_code
+      return NS::ExitCode::CLI_OPTION_ERROR if run_error.is_a?(OptParseValidator::Error) ||
+                                               run_error.is_a?(OptionParser::ParseError)
+
+      return NS::ExitCode::INTERRUPTED if run_error.is_a?(Interrupt)
+
+      NS::ExitCode::ERROR
     end
   end
 end
 
 require "#{CMSScanner::APP_DIR}/app"