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

require 'contrast/utils/string_utils'

module Contrast
  module Api
    module Decorators
      # Used to decorate the {Contrast::Api::Dtm::TraceEvent} protobuf model to
      # convert our {Contrast::Agent::AssessContrastEvent} to the dtm.
      module TraceEvent
        def self.included klass
          klass.extend(ClassMethods)
        end

        # Wrapper around build_event_object for the args array. Handles
        # tainting the correct argument.
        def build_event_args contrast_event, taint_target
          contrast_event.args.each_index do |idx|
            truncate_arg = taint_target != idx
            event_arg = Contrast::Api::Dtm::TraceEventObject.build(contrast_event.args[idx], truncate_arg)
            args << event_arg
          end
        end

        # TeamServer only supports one object's taint ranges at a time.
        # We'll find the taint ranges for the target and return those
        def build_taint_ranges contrast_event, taint_target
          # If there's no taint_target, this isn't a dataflow trace, but a
          # trigger one
          return Contrast::Utils::ObjectShare::EMPTY_ARRAY unless taint_target

          properties = case taint_target
                       when Contrast::Utils::ObjectShare::OBJECT_KEY
                         Contrast::Agent::Assess::Tracker.properties(contrast_event.object)
                       when Contrast::Utils::ObjectShare::RETURN_KEY
                         Contrast::Agent::Assess::Tracker.properties(contrast_event.ret)
                       else
                         target = contrast_event.args[taint_target]
                         if target.is_a?(Hash)
                           if contrast_event.policy_node&.targets&.any?
                             Contrast::Agent::Assess::Tracker.properties(target[contrast_event.policy_node.targets[0]])
                           else
                             Contrast::Agent::Assess::Tracker.properties(target[contrast_event.policy_node.sources[0]])
                           end
                         else
                           Contrast::Agent::Assess::Tracker.properties(target)
                         end
                       end
          return unless properties.tracked?

          self.taint_ranges += properties.tags_to_dtm
        end

        def build_parent_ids contrast_event
          contrast_event&.parent_events&.each do |event|
            next unless event

            parent = Contrast::Api::Dtm::ParentObjectId.new
            parent.id = event.event_id.to_i
            parent_object_ids << parent
          end
        end

        def build_stack contrast_event
          # We delayed doing this as long as possible b/c it's expensive
          stack_dtms = Contrast::Utils::StackTraceUtils.build_assess_stack_array(contrast_event.stack_trace)
          self.stack += stack_dtms
        end

        # Class methods for TraceEvent
        module ClassMethods
          def build contrast_event
            event_dtm = new
            # Figure out what the target of this event was. It's a little
            # annoying for us since P can be named (thanks, Ruby) where
            # as for everyone else it is idx based.
            taint_target = contrast_event.determine_taint_target(event_dtm) # This can't be pulled into the decorator because SourceEvent has a custom impl :/

            event_dtm.type = contrast_event.policy_node.node_type
            event_dtm.action = contrast_event.policy_node.build_action
            event_dtm.timestamp_ms = contrast_event.time.to_i
            event_dtm.thread = Contrast::Utils::StringUtils.force_utf8(contrast_event.thread)
            truncate_obj = Contrast::Utils::ObjectShare::OBJECT_KEY != taint_target
            event_dtm.object = Contrast::Api::Dtm::TraceEventObject.build(contrast_event.object, truncate_obj)
            truncate_ret = Contrast::Utils::ObjectShare::RETURN_KEY != taint_target
            event_dtm.ret = Contrast::Api::Dtm::TraceEventObject.build(contrast_event.ret, truncate_ret)
            event_dtm.build_event_args(contrast_event, taint_target)
            event_dtm.build_parent_ids(contrast_event)
            event_dtm.build_taint_ranges(contrast_event, taint_target)
            event_dtm.build_stack(contrast_event)
            event_dtm.object_id = contrast_event.event_id.to_i
            event_dtm.signature = Contrast::Api::Dtm::TraceEventSignature.build(contrast_event.ret, contrast_event.policy_node, contrast_event.args)
            event_dtm
          end
        end
      end
    end
  end
end

Contrast::Api::Dtm::TraceEvent.include(Contrast::Api::Decorators::TraceEvent)