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

data/lib/aikido/zen/sinks/curb.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 Curl
+      def self.load_sinks!
+        if Gem.loaded_specs["curb"]
+          require "curb"
+
+          ::Curl::Easy.prepend(Curl::EasyExtensions)
+        end
+      end
+
       SINK = Sinks.add("curb", scanners: [
-        Aikido::Zen::Scanners::SSRFScanner,
-        Aikido::Zen::OutboundConnectionMonitor
+        Scanners::SSRFScanner,
+        OutboundConnectionMonitor
       ])
 
-      module Extensions
+      module Helpers
         def self.wrap_request(curl, url: curl.url)
-          Aikido::Zen::Scanners::SSRFScanner::Request.new(
+          Scanners::SSRFScanner::Request.new(
             verb: nil, # Curb hides this by directly setting an option in C
             uri: URI(url),
             headers: curl.headers
@@ -33,29 +41,40 @@ module Aikido::Zen
             status = curl.status.to_i
           end
 
-          Aikido::Zen::Scanners::SSRFScanner::Response.new(status: status, headers: headers)
+          Scanners::SSRFScanner::Response.new(status: status, headers: headers)
         end
 
-        def perform
-          wrapped_request = Extensions.wrap_request(self)
+        def self.scan(request, connection, operation)
+          SINK.scan(
+            request: request,
+            connection: connection,
+            operation: operation
+          )
+        end
+      end
+
+      module EasyExtensions
+        extend Sinks::DSL
+
+        sink_around :perform do |super_call|
+          wrapped_request = Helpers.wrap_request(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: Aikido::Zen::OutboundConnection.from_uri(URI(url)),
-            request: wrapped_request,
-            operation: "request"
-          )
+          connection = OutboundConnection.from_uri(URI(url))
 
-          response = super
+          Helpers.scan(wrapped_request, connection, "request")
 
-          Aikido::Zen::Scanners::SSRFScanner.track_redirects(
+          response = super_call.call
+
+          Scanners::SSRFScanner.track_redirects(
             request: wrapped_request,
-            response: Extensions.wrap_response(self)
+            response: Helpers.wrap_response(self)
           )
 
           # When libcurl has follow_location set, it will handle redirections
@@ -67,14 +86,21 @@ 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 url != last_effective_url
-            last_effective_request = Extensions.wrap_request(self, url: last_effective_url)
-            context["ssrf.request"] = last_effective_request if context
-
-            SINK.scan(
-              connection: Aikido::Zen::OutboundConnection.from_uri(URI(last_effective_url)),
-              request: last_effective_request,
-              operation: "request"
-            )
+            last_effective_request = Helpers.wrap_request(self, url: last_effective_url)
+
+            # Code coverage is disabled here because the else clause is a no-op,
+            # so there is nothing to cover.
+            # :nocov:
+            if context
+              context["ssrf.request"] = last_effective_request
+            else
+              # empty
+            end
+            # :nocov:
+
+            connection = OutboundConnection.from_uri(URI(last_effective_url))
+
+            Helpers.scan(last_effective_request, connection, "request")
           end
 
           response
@@ -86,4 +112,4 @@ module Aikido::Zen
   end
 end
 
-::Curl::Easy.prepend(Aikido::Zen::Sinks::Curl::Extensions)
+Aikido::Zen::Sinks::Curl.load_sinks!

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

@@ -1,20 +1,45 @@
 # frozen_string_literal: true
 
-require_relative "../sink"
+require_relative "../scanners/ssrf_scanner"
 require_relative "../outbound_connection_monitor"
 
 module Aikido::Zen
   module Sinks
     module EventMachine
       module HttpRequest
+        def self.load_sinks!
+          if Gem.loaded_specs["em-http-request"]
+            require "em-http-request"
+
+            ::EventMachine::HttpRequest.use(EventMachine::HttpRequest::Middleware)
+
+            # NOTE: We can't use middleware to intercept requests as we want to ensure any
+            # modifications to the request from user-supplied middleware are already applied
+            # before we scan the request.
+            ::EventMachine::HttpClient.prepend(EventMachine::HttpRequest::HttpClientExtensions)
+          end
+        end
+
         SINK = Sinks.add("em-http-request", scanners: [
-          Aikido::Zen::Scanners::SSRFScanner,
-          Aikido::Zen::OutboundConnectionMonitor
+          Scanners::SSRFScanner,
+          OutboundConnectionMonitor
         ])
 
-        module Extensions
-          def send_request(*)
-            wrapped_request = Aikido::Zen::Scanners::SSRFScanner::Request.new(
+        module Helpers
+          def self.scan(request, connection, operation)
+            SINK.scan(
+              request: request,
+              connection: connection,
+              operation: operation
+            )
+          end
+        end
+
+        module HttpClientExtensions
+          extend Sinks::DSL
+
+          sink_before :send_request do
+            wrapped_request = Scanners::SSRFScanner::Request.new(
               verb: req.method.to_s,
               uri: URI(req.uri),
               headers: req.headers
@@ -24,16 +49,12 @@ module Aikido::Zen
             context = Aikido::Zen.current_context
             context["ssrf.request"] = wrapped_request if context
 
-            SINK.scan(
-              connection: Aikido::Zen::OutboundConnection.new(
-                host: req.host,
-                port: req.port
-              ),
-              request: wrapped_request,
-              operation: "request"
+            connection = OutboundConnection.new(
+              host: req.host,
+              port: req.port
             )
 
-            super
+            Helpers.scan(wrapped_request, connection, "request")
           end
         end
 
@@ -43,13 +64,13 @@ module Aikido::Zen
             context = Aikido::Zen.current_context
             context["ssrf.request"] = nil if context
 
-            Aikido::Zen::Scanners::SSRFScanner.track_redirects(
-              request: Aikido::Zen::Scanners::SSRFScanner::Request.new(
+            Scanners::SSRFScanner.track_redirects(
+              request: Scanners::SSRFScanner::Request.new(
                 verb: client.req.method,
                 uri: URI(client.req.uri),
                 headers: client.req.headers
               ),
-              response: Aikido::Zen::Scanners::SSRFScanner::Response.new(
+              response: Scanners::SSRFScanner::Response.new(
                 status: client.response_header.status,
                 headers: client.response_header.to_h
               )
@@ -61,11 +82,4 @@ module Aikido::Zen
   end
 end
 
-::EventMachine::HttpRequest
-  .use(Aikido::Zen::Sinks::EventMachine::HttpRequest::Middleware)
-
-# NOTE: We can't use middleware to intercept requests as we want to ensure any
-# modifications to the request from user-supplied middleware are already applied
-# before we scan the request.
-::EventMachine::HttpClient
-  .prepend(Aikido::Zen::Sinks::EventMachine::HttpRequest::Extensions)
+Aikido::Zen::Sinks::EventMachine::HttpRequest.load_sinks!

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

@@ -1,31 +1,26 @@
 # frozen_string_literal: true
 
-require_relative "../sink"
+require_relative "../scanners/ssrf_scanner"
 require_relative "../outbound_connection_monitor"
 
 module Aikido::Zen
   module Sinks
     module Excon
-      SINK = Sinks.add("excon", scanners: [
-        Aikido::Zen::Scanners::SSRFScanner,
-        Aikido::Zen::OutboundConnectionMonitor
-      ])
+      def self.load_sinks!
+        if Gem.loaded_specs["excon"]
+          require "excon"
 
-      module Extensions
-        # Maps Excon request params to an Aikido OutboundConnection.
-        #
-        # @param connection [Hash<Symbol, Object>] the data set in the connection.
-        # @param request [Hash<Symbol, Object>] the data overrides sent for each
-        #   request.
-        #
-        # @return [Aikido::Zen::OutboundConnection]
-        def self.build_outbound(connection, request)
-          Aikido::Zen::OutboundConnection.new(
-            host: request.fetch(:hostname) { connection[:hostname] },
-            port: request.fetch(:port) { connection[:port] }
-          )
+          ::Excon::Connection.prepend(ConnectionExtensions)
+          ::Excon::Middleware::RedirectFollower.prepend(RedirectFollowerExtensions)
         end
+      end
+
+      SINK = Sinks.add("excon", scanners: [
+        Scanners::SSRFScanner,
+        OutboundConnectionMonitor
+      ])
 
+      module Helpers
         def self.build_request(connection, request)
           uri = URI(format("%<scheme>s://%<host>s:%<port>i%<path>s", {
             scheme: request.fetch(:scheme) { connection[:scheme] },
@@ -35,45 +30,61 @@ module Aikido::Zen
           }))
           uri.query = request.fetch(:query) { connection[:query] }
 
-          Aikido::Zen::Scanners::SSRFScanner::Request.new(
+          Scanners::SSRFScanner::Request.new(
             verb: request.fetch(:method) { connection[:method] },
             uri: uri,
             headers: connection[:headers].to_h.merge(request[:headers].to_h)
           )
         end
 
-        def request(params = {}, *)
-          request = Extensions.build_request(@data, params)
+        def self.scan(request, connection, operation)
+          SINK.scan(
+            request: request,
+            connection: connection,
+            operation: operation
+          )
+        end
+      end
+
+      module ConnectionExtensions
+        extend Sinks::DSL
+
+        sink_around :request do |super_call, params = {}|
+          request = Helpers.build_request(@data, params)
 
           # 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"] = request
           end
 
-          SINK.scan(
-            connection: Aikido::Zen::OutboundConnection.from_uri(request.uri),
-            request: request,
-            operation: "request"
-          )
+          connection = OutboundConnection.from_uri(request.uri)
 
-          response = super
+          Helpers.scan(request, connection, "request")
 
-          Aikido::Zen::Scanners::SSRFScanner.track_redirects(
+          response = super_call.call
+
+          Scanners::SSRFScanner.track_redirects(
             request: request,
-            response: Aikido::Zen::Scanners::SSRFScanner::Response.new(
+            response: Scanners::SSRFScanner::Response.new(
               status: response.status,
               headers: response.headers.to_h
             )
           )
 
           response
-        rescue ::Excon::Error::Socket => err
-          # Excon wraps errors inside the lower level layer. This only happens
-          # to our scanning exceptions when a request is using RedirectFollower,
-          # so we unwrap them when it happens so host apps can handle errors
-          # consistently.
-          raise err.cause if err.cause.is_a?(Aikido::Zen::UnderAttackError)
+        rescue Sinks::DSL::PresafeError => err
+          outer_cause = err.cause
+          case outer_cause
+          when ::Excon::Error::Socket
+            inner_cause = outer_cause.cause
+            # Excon wraps errors inside the lower level layer. This only happens
+            # to our scanning exceptions when a request is using RedirectFollower,
+            # so we unwrap them when it happens so host apps can handle errors
+            # consistently.
+            raise inner_cause if inner_cause.is_a?(Aikido::Zen::UnderAttackError)
+          end
           raise
         ensure
           context["ssrf.request"] = prev_request if context
@@ -81,23 +92,30 @@ module Aikido::Zen
       end
 
       module RedirectFollowerExtensions
-        def response_call(data)
-          if (response = data[:response])
-            Aikido::Zen::Scanners::SSRFScanner.track_redirects(
-              request: Extensions.build_request(data, {}),
-              response: Aikido::Zen::Scanners::SSRFScanner::Response.new(
+        extend Sinks::DSL
+
+        sink_before :response_call do |datum|
+          response = datum[:response]
+
+          # Code coverage is disabled here because the else clause is a no-op,
+          # so there is nothing to cover.
+          # :nocov:
+          if !response.nil?
+            Scanners::SSRFScanner.track_redirects(
+              request: Helpers.build_request(datum, {}),
+              response: Scanners::SSRFScanner::Response.new(
                 status: response[:status],
                 headers: response[:headers]
               )
             )
+          else
+            # empty
           end
-
-          super
+          # :nocov:
         end
       end
     end
   end
 end
 
-::Excon::Connection.prepend(Aikido::Zen::Sinks::Excon::Extensions)
-::Excon::Middleware::RedirectFollower.prepend(Aikido::Zen::Sinks::Excon::RedirectFollowerExtensions)
+Aikido::Zen::Sinks::Excon.load_sinks!

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

@@ -3,118 +3,114 @@
 module Aikido::Zen
   module Sinks
     module File
-      SINK = Sinks.add("File", scanners: [
-        Aikido::Zen::Scanners::PathTraversalScanner
-      ])
+      def self.load_sinks!
+        # Create a copy of the original method for internal use only to prevent
+        # recursion in PathTraversalScanner.
+        #
+        # IMPORTANT: The alias must be created before the method is overridden,
+        # when the extensions are prepended.
+        ::File.singleton_class.alias_method(:expand_path__internal_for_aikido_zen, :expand_path)
+
+        ::File.singleton_class.prepend(FileClassExtensions)
+        ::File.prepend(FileExtensions)
+      end
+
+      SINK = Sinks.add("File", scanners: [Scanners::PathTraversalScanner])
 
-      module Extensions
-        def self.scan_path(filepath, operation)
+      module Helpers
+        def self.scan(filepath, operation)
           SINK.scan(
             filepath: filepath,
             operation: operation
           )
         end
+      end
 
-        # Module to extend only the initializer method of `File` (`File.new`)
-        module Initiliazer
-          def initialize(filename, *, **)
-            Extensions.scan_path(filename, "new")
-            super
-          end
-        end
+      module FileClassExtensions
+        extend Sinks::DSL
 
-        def open(filename, *, **)
-          Extensions.scan_path(filename, "open")
-          super
+        sink_before :open do |path|
+          Helpers.scan(path, "open")
         end
 
-        def read(filename, *)
-          Extensions.scan_path(filename, "read")
-          super
+        sink_before :read do |path|
+          Helpers.scan(path, "read")
         end
 
-        def write(filename, *, **)
-          Extensions.scan_path(filename, "write")
-          super
+        sink_before :write do |path|
+          Helpers.scan(path, "write")
         end
 
-        def join(*)
-          joined = super
-          Extensions.scan_path(joined, "join")
-          joined
+        sink_before :truncate do |file_name|
+          Helpers.scan(file_name, "truncate")
         end
 
-        def chmod(mode, *paths)
-          paths.each { |path| Extensions.scan_path(path, "chmod") }
-          super
+        sink_before :rename do |old_name, new_name|
+          Helpers.scan(old_name, "rename")
+          Helpers.scan(new_name, "rename")
         end
 
-        def chown(user, group, *paths)
-          paths.each { |path| Extensions.scan_path(path, "chown") }
-          super
+        sink_before :unlink do |*file_names|
+          file_names.each do |file_name|
+            Helpers.scan(file_name, "unlink")
+          end
         end
 
-        def rename(from, to)
-          Extensions.scan_path(from, "rename")
-          Extensions.scan_path(to, "rename")
-          super
+        sink_before :delete do |*file_names|
+          file_names.each do |file_name|
+            Helpers.scan(file_name, "delete")
+          end
         end
 
-        def symlink(from, to)
-          Extensions.scan_path(from, "symlink")
-          Extensions.scan_path(to, "symlink")
-          super
+        sink_before :symlink do |old_name, new_name|
+          Helpers.scan(old_name, "symlink")
+          Helpers.scan(new_name, "symlink")
         end
 
-        def truncate(file_name, *)
-          Extensions.scan_path(file_name, "truncate")
-          super
+        sink_before :chmod do |_mode_int, *file_names|
+          file_names.each do |file_name|
+            Helpers.scan(file_name, "chmod")
+          end
         end
 
-        def unlink(*args)
-          args.each do |arg|
-            Extensions.scan_path(arg, "unlink")
+        sink_before :chown do |_owner_int, group_int, *file_names|
+          file_names.each do |file_name|
+            Helpers.scan(file_name, "chown")
           end
-          super
         end
 
-        def delete(*args)
-          args.each do |arg|
-            Extensions.scan_path(arg, "delete")
+        sink_before :utime do |_atime, _mtime, *file_names|
+          file_names.each do |file_name|
+            Helpers.scan(file_name, "utime")
           end
-          super
         end
 
-        def utime(atime, mtime, *args)
-          args.each do |arg|
-            Extensions.scan_path(arg, "utime")
-          end
-          super
+        sink_after :join do |result|
+          Helpers.scan(result, "join")
+        end
+
+        sink_before :expand_path do |file_name|
+          Helpers.scan(file_name, "expand_path")
         end
 
-        def expand_path(filename, *)
-          Extensions.scan_path(filename, "expand_path")
-          super
+        sink_before :realpath do |file_name|
+          Helpers.scan(file_name, "realpath")
         end
 
-        def realpath(filename, *)
-          Extensions.scan_path(filename, "realpath")
-          super
+        sink_before :realdirpath do |file_name|
+          Helpers.scan(file_name, "realdirpath")
         end
+      end
+
+      module FileExtensions
+        extend Sinks::DSL
 
-        def realdirpath(filename, *)
-          Extensions.scan_path(filename, "realdirpath")
-          super
+        sink_before :initialize do |path|
+          Helpers.scan(path, "new")
         end
       end
     end
   end
 end
 
-# Internally, Path Traversal's scanner logic uses `expand_path`, in order to avoid recursion issues we keep
-# a copy of the original method, only to be used internally.
-# It's important to keep this line before prepend the Extensions module, otherwise the alias will call
-# the extended method.
-::File.singleton_class.alias_method :expand_path__internal_for_aikido_zen, :expand_path
-::File.singleton_class.prepend(Aikido::Zen::Sinks::File::Extensions)
-::File.prepend Aikido::Zen::Sinks::File::Extensions::Initiliazer
+Aikido::Zen::Sinks::File.load_sinks!

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

@@ -1,23 +1,31 @@
 # frozen_string_literal: true
 
-require_relative "../sink"
+require_relative "../scanners/ssrf_scanner"
 require_relative "../outbound_connection_monitor"
 
 module Aikido::Zen
   module Sinks
     module HTTP
+      def self.load_sinks!
+        if Gem.loaded_specs["http"]
+          require "http"
+
+          ::HTTP::Client.prepend(ClientExtensions)
+        end
+      end
+
       SINK = Sinks.add("http", scanners: [
-        Aikido::Zen::Scanners::SSRFScanner,
-        Aikido::Zen::OutboundConnectionMonitor
+        Scanners::SSRFScanner,
+        OutboundConnectionMonitor
       ])
 
-      module Extensions
+      module Helpers
         # Maps an HTTP Request to an Aikido OutboundConnection.
         #
         # @param req [HTTP::Request]
         # @return [Aikido::Zen::OutboundConnection]
         def self.build_outbound(req)
-          Aikido::Zen::OutboundConnection.new(
+          OutboundConnection.new(
             host: req.socket_host,
             port: req.socket_port
           )
@@ -28,7 +36,7 @@ module Aikido::Zen
         # @param req [HTTP::Request]
         # @return [Aikido::Zen::Scanners::SSRFScanner::Request]
         def self.wrap_request(req)
-          Aikido::Zen::Scanners::SSRFScanner::Request.new(
+          Scanners::SSRFScanner::Request.new(
             verb: req.verb,
             uri: URI(req.uri.to_s),
             headers: req.headers.to_h
@@ -36,32 +44,43 @@ module Aikido::Zen
         end
 
         def self.wrap_response(resp)
-          Aikido::Zen::Scanners::SSRFScanner::Response.new(
+          Scanners::SSRFScanner::Response.new(
             status: resp.status,
             headers: resp.headers.to_h
           )
         end
 
-        def perform(req, *)
-          wrapped_request = Extensions.wrap_request(req)
+        def self.scan(request, connection, operation)
+          SINK.scan(
+            request: request,
+            connection: connection,
+            operation: operation
+          )
+        end
+      end
+
+      module ClientExtensions
+        extend Sinks::DSL
+
+        sink_around :perform do |super_call, req|
+          wrapped_request = Helpers.wrap_request(req)
 
           # 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(
-            request: wrapped_request,
-            connection: Extensions.build_outbound(req),
-            operation: "request"
-          )
+          connection = Helpers.build_outbound(req)
+
+          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
@@ -73,4 +92,4 @@ module Aikido::Zen
   end
 end
 
-::HTTP::Client.prepend(Aikido::Zen::Sinks::HTTP::Extensions)
+Aikido::Zen::Sinks::HTTP.load_sinks!