# Copyright (c) 2023 Contrast Security, Inc. See https://www.contrastsecurity.com/enduser-terms-0317a for more details.
# frozen_string_literal: true

require 'contrast/extension/assess/exec_trigger'
require 'contrast/components/logger'
require 'contrast/agent/assess/events/event_data'
require 'contrast/agent/patching/policy/patch'

module Contrast
  module Extension
    module Assess
      # This module provides us with a way to invoke Kernel propagation for those
      # methods which are too complex to fit into one of the standard
      # Contrast::Agent::Assess::Policy::Propagator molds without cluttering up the
      # Kernel Module or exposing our methods there.
      module KernelPropagator
        class << self
          extend Contrast::Components::Logger::InstanceMethods
          include Contrast::Components::Logger::InstanceMethods
          include Contrast::Extension::Assess::ExecTrigger

          # We're 'tracking' sprintf now, meaning if anything is tracked on the way
          # in, the entire result will be tracked out. We're going to take this
          # approach for now b/c it's fast and easy. I don't super love it, and by
          # that I mean I hate it.
          #
          # To actually track this, we'd have to find the index of the new things
          # being added, then remove the tags at the range of the format marker,
          # which is some arbitrary length thing, and add the new tags from the
          # inserted string, shifted down by the length of the aforementioned
          # marker.
          #
          # marker is in the format %[flags][width][.precision]type, type being a
          # single character. We could regexp this with %.+[bBdiouxXeEfgGaAcps%]
          #
          # also, b/c Ruby hates us, there are things called absolute markers,
          # (digit)$, that go in the flags section. These cannot be mixed w/ the
          # order assumed type
          #
          # oh, and there's also %<name>type and %{name}... b/c of course there is
          # -HM
          def sprintf_tagger patcher, preshift, ret, _block
            return unless (properties = Contrast::Agent::Assess::Tracker.properties!(ret))

            format_string = preshift.args[0]
            args = preshift.args[1]

            parent_events = []
            track_sprintf(ret, format_string, args, parent_events)
            event_data = Contrast::Agent::Assess::Events::EventData.new(patcher,
                                                                        ret,
                                                                        preshift.object,
                                                                        ret,
                                                                        preshift.args)
            properties.build_event(event_data, 1)

            properties.event.instance_variable_set(:@_parent_events, parent_events)
            ret
          end

          def track_sprintf result, format_string, args, parent_events
            handle_sprintf_value(format_string, result, parent_events)
            case args
            when String
              handle_sprintf_value(args, result, parent_events)
            when Hash
              handle_sprintf_hash(args, result, parent_events)
            when Array
              handle_sprintf_array(args, result, parent_events)
            end
          rescue StandardError => e
            logger.error('Unable to track dataflow through sprintf', e)
          end

          private

          def handle_sprintf_value value, result, parent_events
            properties = Contrast::Agent::Assess::Tracker.properties(result)
            return unless properties

            value_properties = Contrast::Agent::Assess::Tracker.properties(value)
            return unless value_properties

            parent_events << value_properties.event if value_properties.event
            properties.splat_from(value, result)
          end

          def handle_sprintf_array args, result, parent_events
            args.each do |value|
              handle_sprintf_value(value, result, parent_events)
            end
          end

          def handle_sprintf_hash args, result, parent_events
            args.each_value do |value|
              handle_sprintf_value(value, result, parent_events)
            end
          end
        end
      end

      # Used for Kernel exec aliasing since we have no other way of accessing the
      # Kernel#exec under C if we alias it there.
      module ContrastKernel
        # Method to replace the call to Kernel#exec when applying alias patch.
        # It will invoke  Contrast::Extension::Assess::KernelPropagator before
        # calling the original method.
        #
        # @param source [String, Proc] Potentially untrusted shell command to execute.
        # @return nil - This method will invoke the Kernel#exec which will
        def cs__kernel_exec source
          # Check if in contrast scope and we have source.
          unless Contrast::Agent::Patching::Policy::Patch.in_contrast_scope? || source.nil?
            Contrast::Agent::Patching::Policy::Patch.with_contrast_scope do
              Contrast::Extension::Assess::KernelPropagator.apply_trigger(source)
            end
          end
          # Call this in the end else any code bellow this call won't be executed.
          Kernel.exec(source)
        end
      end
    end
  end
end