aikido-zen 1.0.0.pre.beta.1 → 1.0.1.beta.2

data/lib/aikido/zen/sinks/httpclient.rb

@@ -1,19 +1,27 @@
 # frozen_string_literal: true
 
-require_relative "../sink"
+require_relative "../scanners/ssrf_scanner"
 require_relative "../outbound_connection_monitor"
 
 module Aikido::Zen
   module Sinks
     module HTTPClient
+      def self.load_sinks!
+        if Gem.loaded_specs["httpclient"]
+          require "httpclient"
+
+          ::HTTPClient.prepend(HTTPClient::HTTPClientExtensions)
+        end
+      end
+
       SINK = Sinks.add("httpclient", scanners: [
-        Aikido::Zen::Scanners::SSRFScanner,
-        Aikido::Zen::OutboundConnectionMonitor
+        Scanners::SSRFScanner,
+        OutboundConnectionMonitor
       ])
 
-      module Extensions
+      module Helpers
         def self.wrap_request(req)
-          Aikido::Zen::Scanners::SSRFScanner::Request.new(
+          Scanners::SSRFScanner::Request.new(
             verb: req.http_header.request_method,
             uri: req.http_header.request_uri,
             headers: req.headers
@@ -21,48 +29,69 @@ module Aikido::Zen
         end
 
         def self.wrap_response(resp)
-          Aikido::Zen::Scanners::SSRFScanner::Response.new(
+          # Code coverage is disabled here because `do_get_header` is not called,
+          # because WebMock does not mock it.
+          # :nocov:
+          Scanners::SSRFScanner::Response.new(
             status: resp.http_header.status_code,
             headers: resp.headers
           )
+          # :nocov:
+        end
+
+        def self.scan(request, connection, operation)
+          SINK.scan(
+            request: request,
+            connection: connection,
+            operation: operation
+          )
         end
 
-        def self.perform_scan(req, &block)
+        def self.sink(req, &block)
           wrapped_request = wrap_request(req)
-          connection = Aikido::Zen::OutboundConnection.from_uri(req.http_header.request_uri)
+          connection = OutboundConnection.from_uri(req.http_header.request_uri)
 
           # Store the request information so the DNS sinks can pick it up.
-          if (context = Aikido::Zen.current_context)
+          context = Aikido::Zen.current_context
+          if context
             prev_request = context["ssrf.request"]
             context["ssrf.request"] = wrapped_request
           end
 
-          SINK.scan(connection: connection, request: wrapped_request, operation: "request")
+          scan(wrapped_request, connection, "request")
 
           yield
         ensure
           context["ssrf.request"] = prev_request if context
         end
+      end
+
+      module HTTPClientExtensions
+        extend Sinks::DSL
 
-        def do_get_block(req, *)
-          Extensions.perform_scan(req) { super }
+        private
+
+        sink_around :do_get_block do |super_call, req|
+          Helpers.sink(req, &super_call)
         end
 
-        def do_get_stream(req, *)
-          Extensions.perform_scan(req) { super }
+        sink_around :do_get_stream do |super_call, req|
+          Helpers.sink(req, &super_call)
         end
 
-        def do_get_header(req, res, *)
-          super.tap do
-            Aikido::Zen::Scanners::SSRFScanner.track_redirects(
-              request: Extensions.wrap_request(req),
-              response: Extensions.wrap_response(res)
-            )
-          end
+        sink_after :do_get_header do |_result, req, res, _sess|
+          # Code coverage is disabled here because `do_get_header` is not called,
+          # because WebMock does not mock it.
+          # :nocov:
+          Scanners::SSRFScanner.track_redirects(
+            request: Helpers.wrap_request(req),
+            response: Helpers.wrap_response(res)
+          )
+          # :nocov:
         end
       end
     end
   end
 end
 
-::HTTPClient.prepend(Aikido::Zen::Sinks::HTTPClient::Extensions)
+Aikido::Zen::Sinks::HTTPClient.load_sinks!

data/lib/aikido/zen/sinks/httpx.rb

@@ -1,19 +1,27 @@
 # frozen_string_literal: true
 
-require_relative "../sink"
+require_relative "../scanners/ssrf_scanner"
 require_relative "../outbound_connection_monitor"
 
 module Aikido::Zen
   module Sinks
     module HTTPX
+      def self.load_sinks!
+        if Gem.loaded_specs["httpx"]
+          require "httpx"
+
+          ::HTTPX::Session.prepend(HTTPX::SessionExtensions)
+        end
+      end
+
       SINK = Sinks.add("httpx", scanners: [
-        Aikido::Zen::Scanners::SSRFScanner,
-        Aikido::Zen::OutboundConnectionMonitor
+        Scanners::SSRFScanner,
+        OutboundConnectionMonitor
       ])
 
-      module Extensions
+      module Helpers
         def self.wrap_request(request)
-          Aikido::Zen::Scanners::SSRFScanner::Request.new(
+          Scanners::SSRFScanner::Request.new(
             verb: request.verb,
             uri: request.uri,
             headers: request.headers.to_hash
@@ -21,35 +29,46 @@ module Aikido::Zen
         end
 
         def self.wrap_response(response)
-          Aikido::Zen::Scanners::SSRFScanner::Response.new(
+          Scanners::SSRFScanner::Response.new(
             status: response.status,
             headers: response.headers.to_hash
           )
         end
 
-        def send_request(request, *)
-          wrapped_request = Extensions.wrap_request(request)
+        def self.scan(request, connection, operation)
+          SINK.scan(
+            request: request,
+            connection: connection,
+            operation: operation
+          )
+        end
+      end
+
+      module SessionExtensions
+        extend Sinks::DSL
+
+        sink_around :send_request do |super_call, request|
+          wrapped_request = Helpers.wrap_request(request)
 
           # Store the request information so the DNS sinks can pick it up.
-          if (context = Aikido::Zen.current_context)
+          context = Aikido::Zen.current_context
+          if context
             prev_request = context["ssrf.request"]
             context["ssrf.request"] = wrapped_request
           end
 
-          SINK.scan(
-            connection: Aikido::Zen::OutboundConnection.from_uri(request.uri),
-            request: wrapped_request,
-            operation: "request"
-          )
+          connection = OutboundConnection.from_uri(request.uri)
+
+          Helpers.scan(wrapped_request, connection, "request")
 
           request.on(:response) do |response|
-            Aikido::Zen::Scanners::SSRFScanner.track_redirects(
+            Scanners::SSRFScanner.track_redirects(
               request: wrapped_request,
-              response: Extensions.wrap_response(response)
+              response: Helpers.wrap_response(response)
             )
           end
 
-          super
+          super_call.call
         ensure
           context["ssrf.request"] = prev_request if context
         end
@@ -58,4 +77,4 @@ module Aikido::Zen
   end
 end
 
-::HTTPX::Session.prepend(Aikido::Zen::Sinks::HTTPX::Extensions)
+Aikido::Zen::Sinks::HTTPX.load_sinks!

data/lib/aikido/zen/sinks/kernel.rb

@@ -3,71 +3,32 @@
 module Aikido::Zen
   module Sinks
     module Kernel
-      SINK = Sinks.add("Kernel", scanners: [
-        Aikido::Zen::Scanners::ShellInjectionScanner
-      ])
-
-      module Extensions
-        # Checks if the user introduced input is trying to execute other commands
-        # using Shell Injection kind of attacks.
-        #
-        # @param command [String] the _full command_ that will be executed.
-        # @param context [Aikido::Zen::Context]
-        # @param sink [Aikido::Zen::Sink] the Sink that is running the scan.
-        # @param operation [Symbol, String] name of the method being scanned.
-        #
-        # @return [Aikido::Zen::Attacks::ShellInjectionAttack, nil] an Attack if any
-        # user input is detected as part of a Shell Injection Attack, or +nil+ if it's safe.
-        def self.scan_command(command, operation)
-          SINK.scan(
-            command: command,
-            operation: operation
-          )
-        end
+      def self.load_sinks!
+        ::Kernel.singleton_class.prepend(KernelExtensions)
+        ::Kernel.prepend(KernelExtensions)
+      end
 
-        # `system, spawn` functions can be invoked in several ways. For more details,
-        # see [the documentation](https://ruby-doc.org/3.4.1/Kernel.html#method-i-spawn)
-        #
-        # In our context, we care primarily about two common scenarios:
-        #   - one argument (String)
-        #       e.g.: system("ls"), system("echo something")
-        #   - two arguments (Hash, String)
-        #       e.g.: system({"foo" => "bar"}, "ls"), system({"foo" => "bar"}, "echo something")
-        #
-        # In all other cases, we do not protect against shell argument injections. Specifically:
-        #
-        # If a user input contains something like $(whoami) and is passed as part of the command
-        # arguments (e.g., user_input = "$(whoami)"):
-        #
-        #   system("echo", user_input)   This is safe because Ruby automatically escapes arguments
-        #                                passed to system/spawn in this form.
-        #
-        #   system("echo #{user_input}") This is not safe because Ruby interpolates the user_input
-        #                                into the command string, resulting in a potentially harmful
-        #                                command like `echo $(whoami)`.
-        def send_arg_to_scan(args, operation)
-          if args.size == 1 && args[0].is_a?(String)
-            Extensions.scan_command(args[0], operation)
-          end
+      SINK = Sinks.add("Kernel", scanners: [Scanners::ShellInjectionScanner])
 
-          if args.size == 2 && args[0].is_a?(Hash)
-            Extensions.scan_command(args[1], operation)
-          end
+      module Helpers
+        def self.scan(command, operation)
+          SINK.scan(command: command, operation: operation)
         end
+      end
 
-        def system(*args, **)
-          send_arg_to_scan(args, "system")
-          super
-        end
+      module KernelExtensions
+        extend Sinks::DSL
 
-        def spawn(*args, **)
-          send_arg_to_scan(args, "spawn")
-          super
+        %i[system spawn].each do |method_name|
+          sink_before method_name do |*args|
+            # Remove the optional environment argument before the command-line.
+            args.shift if args.first.is_a?(Hash)
+            Helpers.scan(args.first, method_name)
+          end
         end
       end
     end
   end
 end
 
-::Kernel.singleton_class.prepend Aikido::Zen::Sinks::Kernel::Extensions
-::Kernel.prepend Aikido::Zen::Sinks::Kernel::Extensions
+Aikido::Zen::Sinks::Kernel.load_sinks!

data/lib/aikido/zen/sinks/mysql2.rb

@@ -1,21 +1,33 @@
 # frozen_string_literal: true
 
-require_relative "../sink"
-
 module Aikido::Zen
   module Sinks
     module Mysql2
+      def self.load_sinks!
+        if Gem.loaded_specs["mysql2"]
+          require "mysql2"
+
+          ::Mysql2::Client.prepend(ClientExtensions)
+        end
+      end
+
       SINK = Sinks.add("mysql2", scanners: [Scanners::SQLInjectionScanner])
 
-      module Extensions
-        def query(query, *)
-          SINK.scan(query: query, dialect: :mysql, operation: "query")
+      module Helpers
+        def self.scan(query, operation)
+          SINK.scan(query: query, dialect: :mysql, operation: operation)
+        end
+      end
+
+      module ClientExtensions
+        extend Sinks::DSL
 
-          super
+        sink_before :query do |sql|
+          Helpers.scan(sql, "query")
         end
       end
     end
   end
 end
 
-::Mysql2::Client.prepend(Aikido::Zen::Sinks::Mysql2::Extensions)
+Aikido::Zen::Sinks::Mysql2.load_sinks!

data/lib/aikido/zen/sinks/net_http.rb

@@ -1,25 +1,32 @@
 # frozen_string_literal: true
 
-require_relative "../sink"
+require_relative "../scanners/ssrf_scanner"
 require_relative "../outbound_connection_monitor"
 
 module Aikido::Zen
   module Sinks
     module Net
       module HTTP
+        def self.load_sinks!
+          # In stdlib but not always required
+          require "net/http"
+
+          ::Net::HTTP.prepend(Net::HTTP::HTTPExtensions)
+        end
+
         SINK = Sinks.add("net-http", scanners: [
-          Aikido::Zen::Scanners::SSRFScanner,
-          Aikido::Zen::OutboundConnectionMonitor
+          Scanners::SSRFScanner,
+          OutboundConnectionMonitor
         ])
 
-        module Extensions
+        module Helpers
           # Maps a Net::HTTP connection to an Aikido OutboundConnection,
           # which our tooling expects.
           #
           # @param http [Net::HTTP]
           # @return [Aikido::Zen::OutboundConnection]
           def self.build_outbound(http)
-            Aikido::Zen::OutboundConnection.new(
+            OutboundConnection.new(
               host: http.address,
               port: http.port
             )
@@ -34,7 +41,7 @@ module Aikido::Zen
               path: req.path
             }))
 
-            Aikido::Zen::Scanners::SSRFScanner::Request.new(
+            Scanners::SSRFScanner::Request.new(
               verb: req.method,
               uri: uri,
               headers: req.to_hash,
@@ -43,33 +50,44 @@ module Aikido::Zen
           end
 
           def self.wrap_response(response)
-            Aikido::Zen::Scanners::SSRFScanner::Response.new(
+            Scanners::SSRFScanner::Response.new(
               status: response.code.to_i,
               headers: response.to_hash,
               header_normalizer: ->(val) { Array(val).join(", ") }
             )
           end
 
-          def request(req, *)
-            wrapped_request = Extensions.wrap_request(req, self)
+          def self.scan(request, connection, operation)
+            SINK.scan(
+              request: request,
+              connection: connection,
+              operation: operation
+            )
+          end
+        end
+
+        module HTTPExtensions
+          extend Sinks::DSL
+
+          sink_around :request do |super_call, req|
+            wrapped_request = Helpers.wrap_request(req, self)
 
             # Store the request information so the DNS sinks can pick it up.
-            if (context = Aikido::Zen.current_context)
+            context = Aikido::Zen.current_context
+            if context
               prev_request = context["ssrf.request"]
               context["ssrf.request"] = wrapped_request
             end
 
-            SINK.scan(
-              connection: Extensions.build_outbound(self),
-              request: wrapped_request,
-              operation: "request"
-            )
+            connection = Helpers.build_outbound(self)
+
+            Helpers.scan(wrapped_request, connection, "request")
 
-            response = super
+            response = super_call.call
 
-            Aikido::Zen::Scanners::SSRFScanner.track_redirects(
+            Scanners::SSRFScanner.track_redirects(
               request: wrapped_request,
-              response: Extensions.wrap_response(response)
+              response: Helpers.wrap_response(response)
             )
 
             response
@@ -82,4 +100,4 @@ module Aikido::Zen
   end
 end
 
-::Net::HTTP.prepend(Aikido::Zen::Sinks::Net::HTTP::Extensions)
+Aikido::Zen::Sinks::Net::HTTP.load_sinks!

data/lib/aikido/zen/sinks/patron.rb

@@ -1,56 +1,75 @@
 # frozen_string_literal: true
 
-require_relative "../sink"
+require_relative "../scanners/ssrf_scanner"
 require_relative "../outbound_connection_monitor"
 
 module Aikido::Zen
   module Sinks
     module Patron
+      def self.load_sinks!
+        if Gem.loaded_specs["patron"]
+          require "patron"
+
+          ::Patron::Session.prepend(SessionExtensions)
+        end
+      end
+
       SINK = Sinks.add("patron", scanners: [
-        Aikido::Zen::Scanners::SSRFScanner,
-        Aikido::Zen::OutboundConnectionMonitor
+        Scanners::SSRFScanner,
+        OutboundConnectionMonitor
       ])
 
-      module Extensions
+      module Helpers
         def self.wrap_response(request, response)
           # In this case, automatic redirection happened inside libcurl.
           if response.url != request.url && !response.url.to_s.empty?
-            Aikido::Zen::Scanners::SSRFScanner::Response.new(
+            Scanners::SSRFScanner::Response.new(
               status: 302, # We can't know what the actual status was, but we just need a 3XX
               headers: response.headers.merge("Location" => response.url)
             )
           else
-            Aikido::Zen::Scanners::SSRFScanner::Response.new(
+            Scanners::SSRFScanner::Response.new(
               status: response.status,
               headers: response.headers
             )
           end
         end
 
-        def handle_request(request)
-          wrapped_request = Aikido::Zen::Scanners::SSRFScanner::Request.new(
+        def self.scan(request, connection, operation)
+          SINK.scan(
+            request: request,
+            connection: connection,
+            operation: operation
+          )
+        end
+      end
+
+      module SessionExtensions
+        extend Sinks::DSL
+
+        sink_around :handle_request do |super_call, request|
+          wrapped_request = Scanners::SSRFScanner::Request.new(
             verb: request.action,
             uri: URI(request.url),
             headers: request.headers
           )
 
           # Store the request information so the DNS sinks can pick it up.
-          if (context = Aikido::Zen.current_context)
+          context = Aikido::Zen.current_context
+          if context
             prev_request = context["ssrf.request"]
             context["ssrf.request"] = wrapped_request
           end
 
-          SINK.scan(
-            connection: Aikido::Zen::OutboundConnection.from_uri(URI(request.url)),
-            request: wrapped_request,
-            operation: "request"
-          )
+          connection = OutboundConnection.from_uri(URI(request.url))
+
+          Helpers.scan(wrapped_request, connection, "request")
 
-          response = super
+          response = super_call.call
 
-          Aikido::Zen::Scanners::SSRFScanner.track_redirects(
+          Scanners::SSRFScanner.track_redirects(
             request: wrapped_request,
-            response: Extensions.wrap_response(request, response)
+            response: Helpers.wrap_response(request, response)
           )
 
           # When libcurl has follow_location set, it will handle redirections
@@ -62,18 +81,16 @@ module Aikido::Zen
           # stop the response from being exposed to the user. This downgrades
           # the SSRF into a blind SSRF, which is better than doing nothing.
           if request.url != response.url && !response.url.to_s.empty?
-            last_effective_request = Aikido::Zen::Scanners::SSRFScanner::Request.new(
+            last_effective_request = Scanners::SSRFScanner::Request.new(
               verb: request.action,
               uri: URI(response.url),
               headers: request.headers
             )
             context["ssrf.request"] = last_effective_request if context
 
-            SINK.scan(
-              connection: Aikido::Zen::OutboundConnection.from_uri(URI(response.url)),
-              request: last_effective_request,
-              operation: "request"
-            )
+            connection = OutboundConnection.from_uri(URI(response.url))
+
+            Helpers.scan(last_effective_request, connection, "request")
           end
 
           response
@@ -85,4 +102,4 @@ module Aikido::Zen
   end
 end
 
-::Patron::Session.prepend(Aikido::Zen::Sinks::Patron::Extensions)
+Aikido::Zen::Sinks::Patron.load_sinks!

data/lib/aikido/zen/sinks/pg.rb

@@ -1,51 +1,74 @@
 # frozen_string_literal: true
 
-require_relative "../sink"
-
 module Aikido::Zen
   module Sinks
     module PG
+      def self.load_sinks!
+        if Gem.loaded_specs["pg"]
+          require "pg"
+
+          ::PG::Connection.prepend(PG::ConnectionExtensions)
+        end
+      end
+
       SINK = Sinks.add("pg", scanners: [Scanners::SQLInjectionScanner])
 
-      # For some reason, the ActiveRecord pg adapter does not wrap exceptions in
-      # StatementInvalid, which leads to inconsistent handling. This guarantees
-      # that all Zen errors are wrapped in a StatementInvalid, so documentation
-      # can be consistent.
-      WRAP_EXCEPTIONS = if defined?(ActiveRecord::StatementInvalid)
-        <<~RUBY
-          rescue Aikido::Zen::SQLInjectionError
-            raise ActiveRecord::StatementInvalid
-        RUBY
+      module Helpers
+        # For some reason, the ActiveRecord pg adaptor does not wrap exceptions
+        # in ActiveRecord::StatementInvalid, leading to inconsistent handling.
+        # This guarantees that Aikido::Zen::SQLInjectionErrors are wrapped in
+        # an ActiveRecord::StatementInvalid.
+        def self.safe(&block)
+          # Code coverage is disabled here because this ActiveRecord behavior is
+          # exercised in end-to-end tests, which are not covered by SimpleCov.
+          # :nocov:
+          if !defined?(ActiveRecord::StatementInvalid)
+            Sinks::DSL.safe(&block)
+          else
+            begin
+              Sinks::DSL.safe(&block)
+            rescue Aikido::Zen::SQLInjectionError => err
+              raise ActiveRecord::StatementInvalid, cause: err
+            end
+          end
+          # :nocov:
+        end
+
+        def self.scan(query, operation)
+          SINK.scan(
+            query: query,
+            dialect: :postgresql,
+            operation: operation
+          )
+        end
       end
 
-      module Extensions
+      module ConnectionExtensions
+        extend Sinks::DSL
+
         %i[
           send_query exec sync_exec async_exec
           send_query_params exec_params sync_exec_params async_exec_params
-        ].each do |method|
-          module_eval <<~RUBY, __FILE__, __LINE__ + 1
-            def #{method}(query, *)
-              SINK.scan(query: query, dialect: :postgresql, operation: :#{method})
-              super
-            #{WRAP_EXCEPTIONS}
+        ].each do |method_name|
+          presafe_sink_before method_name do |query|
+            Helpers.safe do
+              Helpers.scan(query, method_name)
             end
-          RUBY
+          end
         end
 
         %i[
           send_prepare prepare async_prepare sync_prepare
-        ].each do |method|
-          module_eval <<~RUBY, __FILE__, __LINE__ + 1
-            def #{method}(_, query, *)
-              SINK.scan(query: query, dialect: :postgresql, operation: :#{method})
-              super
-            #{WRAP_EXCEPTIONS}
+        ].each do |method_name|
+          presafe_sink_before method_name do |_, query|
+            Helpers.safe do
+              Helpers.scan(query, method_name)
             end
-          RUBY
+          end
         end
       end
     end
   end
 end
 
-::PG::Connection.prepend(Aikido::Zen::Sinks::PG::Extensions)
+Aikido::Zen::Sinks::PG.load_sinks!