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

require 'parser/current'

module Contrast
  module Utils
    # This utility allows us to parse and rewrite the AST in Ruby 2.5,
    # allowing us to track String interpolation propagation by replacing those
    # events with String#+ events instead.
    #
    # TODO: RUBY-714 remove w/ EOL of 2.5
    # @deprecated Changes to this class are discouraged as this approach is
    #   being phased out with support for those language versions.
    class RubyAstRewriter < Parser::TreeRewriter
      VARIABLES = %i[ivar cvar gvar].cs__freeze

      # Rewrite the given source_string, converting the interpolation action
      # therein to a concatenation action
      #
      # @param source_string [String] the String to rewrite
      # @return [String] either the rewritten string or the original on error.
      def rewrite source_string
        ast = Parser::CurrentRuby.parse(source_string)
        buffer = Parser::Source::Buffer.new('rewrite')
        buffer.source = source_string
        begin
          super(buffer, ast)
        rescue StandardError => _e
          source_string
        end
      end

      # This overloads the on_dstr method in Parser::AST::Processor to handle
      # the replace within the given node.
      def on_dstr node
        return if node.children.all? { |child_node| child_node.type == :str }

        new_content = +'('
        node.children.each_with_index do |child_node, index|
          # A begin node looks like #{some_code} in ruby, we do a substring
          # from [2...-1] to get rid of the #{ & trailing }.
          if child_node.type == :begin
            code = child_node.
                location.
                expression.
                source_buffer.
                source[child_node.location.begin.begin_pos...child_node.location.end.end_pos]
            code = code[2...-1]
            new_content += "((#{ code }).to_s)"

          # For interpolations that use class, instance, or global variables,
          # those are NOT within a begin block, but instead are a ivar, cvar,
          # or gvar node, not stripping of interpolation markers required.
          elsif VARIABLES.include?(child_node.type)
            variable = child_node.children.first
            new_content << "((#{ variable }).to_s)"

          # When interpolation includes strings before or after an
          # interpolation they are simple :str nodes, in order to preserve
          # escaping we need to do a dump of the string value.
          elsif child_node.type == :str
            literal_value = child_node.children.first
            literal_value = literal_value.dump[1...-1]
            new_content <<  "\"#{ literal_value }\""
          end
          new_content << ' + ' unless index == node.children.length - 1
        end
        new_content << ')'
        if node.location.cs__respond_to?(:heredoc_body)
          replace(node.location.heredoc_body, new_content)
        else
          replace(node.location.expression, new_content)
        end
      end
    end
  end
end