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

cs__scoped_require 'contrast/utils/assess/tracking_util'
cs__scoped_require 'contrast/utils/class_util'
cs__scoped_require 'contrast/utils/duck_utils'
cs__scoped_require 'contrast/utils/object_share'
cs__scoped_require 'contrast/utils/prevent_serialization'
cs__scoped_require 'contrast/utils/stack_trace_utils'
cs__scoped_require 'contrast/utils/string_utils'
cs__scoped_require 'contrast/utils/timer'

module Contrast
  module Agent
    module Assess
      # This class holds the data about an event in the application
      # We'll use it to build an event that TeamServer can consume if
      # the object to which this event belongs ends in a trigger.
      class ContrastEvent
        include Contrast::Utils::PreventSerialization

        class << self
          def safe_args_representation args
            return nil unless args
            return Contrast::Utils::ObjectShare::EMPTY_ARRAY if args.empty?

            rep = []
            args.each do |arg|
              # We have to handle named args
              rep <<  if arg.is_a?(Hash)
                        safe_arg_hash_representation(arg)
                      else
                        Contrast::Utils::ClassUtil.to_contrast_string(arg)
                      end
            end
            rep
          end

          def safe_arg_hash_representation hash
            # since this is the named hash for arguments, only the value is
            # suspect here
            hash.transform_values { |v| Contrast::Utils::ClassUtil.to_contrast_string(v) }
          end

          # if given an object that can be duped, duplicate it. otherwise just
          # return the original object. swallow all exceptions from
          # non-duplicable things.
          #
          # we can't just check respond_to? though b/c dup exists on the
          # base Object class
          def safe_dup original
            return nil unless original

            begin
              original.dup
            rescue StandardError
              original
            end
          end
        end

        attr_reader :event_id, :parent_ids

        # We need this to track the parent id's of events to build up a flow
        # chart of the finding
        @atomic_id = 0
        @atomic_mutex = Mutex.new
        def self.next_atomic_id
          @atomic_mutex.synchronize do
            @atomic_id += 1
            # Rollover things
          rescue StandardError
            @atomic_id = 1
          end
        end

        def initialize policy_node, tagged, object, ret, args
          @policy_node = policy_node
          # so long as this event is built in a factory, we know Contrast Code
          # will be the first three events
          @caller = caller(3, 20)
          @time = Contrast::Utils::Timer.now_ms
          @thread = Thread.current.object_id

          # These methods rely on the above being set. Don't move them!
          @event_id = Contrast::Agent::Assess::ContrastEvent.next_atomic_id
          @parent_ids = find_parent_ids(policy_node, object, ret, args)
          snapshot(tagged, object, ret, args)
        end

        # Parent IDs are the event ids of all the sources of this event which
        # were tracked prior to this event occurring
        def find_parent_ids policy_node, object, ret, args
          mapped = policy_node.sources.map do |source|
            value_of_source(source, object, ret, args)
          end
          selected = mapped.select do |source|
            source &&
                Contrast::Utils::DuckUtils.quacks_to?(source, :cs__properties) &&
                source.cs__properties.events &&
                source.cs__properties.events.last
          end
          selected.map do |source|
            source.cs__properties.events.last.event_id
          end
        end

        def snapshot tagged, object, ret, args
          target = @policy_node.target
          case target
            # If the target is nil, this rule was violated simply by a method
            # being called. We'll save all the information, but nothing will be
            # marked up, as nothing need be tracked
          when nil
            @object = Contrast::Utils::ClassUtil.to_contrast_string(object)
            @args = cs__class.safe_args_representation(args)
            @ret = Contrast::Utils::ClassUtil.to_contrast_string(ret)
            # If the target is O, then we dup the O and safely represent the rest
          when Contrast::Utils::ObjectShare::OBJECT_KEY
            @object = cs__class.safe_dup(tagged)
            @args = cs__class.safe_args_representation(args)
            @ret = Contrast::Utils::ClassUtil.to_contrast_string(ret)
            # If the target is R, then we dup the R and safely represent the rest
          when Contrast::Utils::ObjectShare::RETURN_KEY
            @object = Contrast::Utils::ClassUtil.to_contrast_string(object)
            @args = cs__class.safe_args_representation(args)
            @ret = cs__class.safe_dup(tagged)
            # If the target is P*, then we need to dup things a differently. We
            # need to find the true target inside so that we can mark it up
            # later, but the other args should be represented as their safe form.
          else
            @object = Contrast::Utils::ClassUtil.to_contrast_string(object)
            @args = cs__class.safe_args_representation(args)
            @ret = Contrast::Utils::ClassUtil.to_contrast_string(ret)
            save_target_arg(target, tagged)
          end
        end

        # I know we're creating an extra string here since we replace the safe
        # one w/ a dup, but good enough for now. Trying not to make this too
        # complicated. - HM 8/8/19
        def save_target_arg target, tagged
          return if @args.cs__frozen?

          if target.is_a?(Integer)
            @args[target] = cs__class.safe_dup(tagged)
            return
          end

          @args.each_with_index do |search, index|
            next unless search.is_a?(Hash)
            next unless search[target]

            search[target] = cs__class.safe_dup(tagged)
            @highlight = index
            break
          end
        end

        def value_of_source source, object, ret, args
          case source
          when Contrast::Utils::ObjectShare::OBJECT_KEY
            object
          when Contrast::Utils::ObjectShare::RETURN_KEY
            ret
          else
            if source.is_a?(Integer)
              args[source]
            else
              args.each do |search|
                next unless search.is_a?(Hash)

                s = search[source]
                return s if s
              end
            end
          end
        end

        # Convert this event into a DTM that TeamServer can consume
        def to_dtm_event
          event = Contrast::Api::Dtm::TraceEvent.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 = determine_taint_target(event)

          event.type = @policy_node.node_type
          event.action = @policy_node.build_action
          event.timestamp_ms = @time.to_i
          event.thread = Contrast::Utils::StringUtils.force_utf8(@thread)
          truncate_obj = Contrast::Utils::ObjectShare::OBJECT_KEY != taint_target
          event.object = build_event_object(@object, truncate_obj)
          truncate_ret = Contrast::Utils::ObjectShare::RETURN_KEY != taint_target
          event.ret = build_event_object(@ret, truncate_ret)
          built_args = build_event_args(taint_target)
          built_args.each do |arg|
            event.args << arg
          end
          taint_ranges = find_taint_ranges(@object, @args, @ret, taint_target)
          taint_ranges.each do |range|
            event.taint_ranges << range
          end

          # We delayed doing this as long as possible b/c it's expensive
          stack = Contrast::Utils::StackTraceUtils.build_assess_stack_array(@caller)
          event.stack += stack

          event.object_id = event_id.to_i
          parent_ids&.each do |id|
            parent = Contrast::Api::Dtm::ParentObjectId.new
            parent.id = id.to_i
            event.parent_object_ids << parent
          end

          # not to be confused w/ the partial signature
          build_complete_signature(event)

          event
        end

        # We're not going to build the signature string here, b/c we have all
        # the composite pieces of it. Instead, we're going to let TeamServer
        # render this for us.
        def build_complete_signature event
          signature = Contrast::Api::Dtm::TraceEventSignature.new
          event.signature = signature
          return_type = @ret ? @ret.cs__class.name : Contrast::Utils::ObjectShare::NIL_STRING
          signature.return_type = Contrast::Utils::StringUtils.force_utf8(return_type)
          signature.class_name = Contrast::Utils::StringUtils.force_utf8(@policy_node.class_name)
          signature.method_name = Contrast::Utils::StringUtils.force_utf8(@policy_node.method_name)
          if @args
            @args&.each do |arg|
              arg_type = arg ? arg.cs__class.name : Contrast::Utils::ObjectShare::NIL_STRING
              signature.arg_types << Contrast::Utils::StringUtils.force_utf8(arg_type)
            end
          end
          signature.constructor = @policy_node.method_name == :new
          # if there's a ret, then this method isn't nil. not 100% full proof since you can
          # return nil, but this is the best we've got currently.
          signature.void_method = @ret.nil?
          # 8 is STATIC in Java... we have to placate them for now
          # it has been requested that flags be removed since it isn't used
          signature.flags = 8 unless @policy_node.instance_method?
        end

        # Wrapper around build_event_object for the args array. Handles
        # tainting the correct argument.
        def build_event_args taint_target
          event_args = []
          @args.each_index do |idx|
            truncate_arg = taint_target != idx
            event_arg = build_event_object(@args[idx], truncate_arg)
            event_args << event_arg
          end
          event_args
        end

        # Build the event object. We were originally going to include taint on
        # each one, but TS doesn't accept / use that, so it is a waste of time.
        #
        # We'll truncate any object that isn't important to the taint ranges of
        # this event, so that we don't murder TeamServer by, for instance,
        # hypothetically sending the entire rendered HTML page >_>  <_<  >_>
        ELLIPSIS = '...'
        UNTRUNCATED_PORTION_LENGTH = 25
        TRUNCATION_LENGTH = (UNTRUNCATED_PORTION_LENGTH * 2) + 3 # ELLIPSIS length
        def build_event_object object, truncate
          event_object = Contrast::Api::Dtm::TraceEventObject.new
          obj_string = Contrast::Utils::StringUtils.force_utf8(object)
          if truncate && Contrast::Utils::StringUtils.ret_length(obj_string) > TRUNCATION_LENGTH
            tmp = []
            tmp << obj_string[0, UNTRUNCATED_PORTION_LENGTH]
            tmp << ELLIPSIS
            tmp << obj_string[
                obj_string.length - UNTRUNCATED_PORTION_LENGTH,
                UNTRUNCATED_PORTION_LENGTH]
            obj_string = tmp.join
          end
          event_object.value = Base64.encode64(obj_string)
          event_object.tracked = Contrast::Utils::Assess::TrackingUtil.tracked?(object)
          event_object
        end

        # We have to do a little work to figure out what our TS appropriate
        # target is. To break this down, the logic is as follows:
        # 1) If my policy_node has a target, work on targets. Else, work on sources.
        #    Per TS law, each policy_node must have at least a source or a target.
        #    The only type of policy_node w/o targets is a Trigger, but that may
        #    change.
        # 2) If I have a highlight, it means that I have a P target that is
        #    not in integer form (it was a named / keyword type for which I had
        #    to find the index). I need to address this so that TS can process
        #    it.
        # 3) I'll set the event's source and target to TS values.
        # 4) Return the highlight or the first source/target as the taint
        #    target.
        def determine_taint_target event
          if @policy_node&.targets&.any?
            event.source = @policy_node.source_string if @policy_node.source_string
            event.target = if @highlight
                             "P#{ @highlight }"
                           else
                             @policy_node.target_string
                           end
            @highlight || @policy_node.targets[0]
          elsif @policy_node&.sources&.any?
            event.source = if @highlight
                             "P#{ @highlight }"
                           else
                             @policy_node.source_string
                           end
            event.target = @policy_node.target_string if @policy_node.target_string
            @highlight || @policy_node.sources[0]
          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 find_taint_ranges object, args, ret, 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
                         object.cs__properties
                       when Contrast::Utils::ObjectShare::RETURN_KEY
                         ret.cs__properties
                       else
                         target = args[taint_target]
                         if target.is_a?(Hash)
                           if @policy_node&.targets&.any?
                             target[@policy_node.targets[0]].cs__properties
                           else
                             target[@policy_node.sources[0]].cs__properties
                           end
                         else
                           target.cs__properties
                         end
                       end

          return Contrast::Utils::ObjectShare::EMPTY_ARRAY unless properties.tracked?

          properties.tags_to_dtm
        end
      end
    end
  end
end