aikido-zen 0.2.0-x86_64-linux → 1.0.1.beta.2-x86_64-linux

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!

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!