# Hook into ActionView rendering to inject Immunio's hooks.
require 'securerandom'
require 'ripper'
require 'digest/sha1'

module Immunio
  # Renders templates by filtering them through Immunio's hook handlers.
  class Template
    CHECKSUM_CACHE = Hash.new do |cache, template_id|
      template = ObjectSpace._id2ref(template_id)

      if template.respond_to?(:source) && !template.source.nil?
        finalizer = Immunio::Template.finalize_template(template_id)
        ObjectSpace.define_finalizer(template, finalizer)
        cache[template_id] = Digest::SHA1.hexdigest(template.source).freeze
      end
    end

    attr_accessor :vars

    def initialize(template)
      @template = template
      @next_var_id = 0
      @next_template_id = 0
      @vars = {}
      @scheduled_fragments_writes = []
    end

    def id
      (@template.respond_to?(:virtual_path) && @template.virtual_path) ||
        (@template.respond_to?(:source) && @template.source)
    end

    def ==(other)
      self.class === other && id == other.id
    end

    def is_text?
      @template.formats.first == :text
    end

    def template_sha
      CHECKSUM_CACHE[@template.object_id]
    end

    # Generate the next var unique ID to be used in a template.
    def next_var_id
      id = @next_var_id
      @next_var_id += 1
      id
    end

    def next_template_id
      id = @next_template_id
      @next_template_id += 1
      id
    end

    def get_nonce
      # Generate a two byte CSRNG nonce to make our substitutions unpreictable
      # Why only 2 bytes? The nonce is per render, so the odds of guessing it are very low
      # and entropy is finite so we don't want to drain the random pool unnecessarily
      @nonce ||= SecureRandom.hex(2)
    end

    def mark_and_defer_fragment_write(key, content, options)
      id = @scheduled_fragments_writes.size
      nonce = Template.get_nonce
      @scheduled_fragments_writes << [key, content, options]
      "{immunio-fragment:#{id}:#{nonce}}#{content}{/immunio-fragment:#{id}:#{nonce}}"
    end

    def render(context)
      load_source context
      # Don't handle templates with no source (inline text templates).
      unless has_source?
        rendered = yield
        rendered.instance_variable_set("@__immunio_processed", true) unless rendered.frozen?
        return rendered
      end

      begin
        root = true if rendering_stack.length == 0

        rendering_stack.push self
        # Calculate SHA1 of this template.
        template_sha
        Immunio.logger.debug {"ActionView rendering template with sha #{@template_sha}, root: #{root}"}
        rendered = yield
        rendered.instance_variable_set("@__immunio_processed", true) unless rendered.frozen?

        if root
          # This is the root template. Let ActionView render it, and then look
          # for XSS.

          # If the rendered result isn't a string, or a string-like, then let's
          # skip it for safety sake.
          unless rendered.respond_to? :to_str
            unless $__immunio_av_rendered_non_string
              Immunio.logger.warn { "ActionView rendered #{@template.inspect} to a non-string-like value: #{rendered.inspect}. This rendering will not be analyzed for XSS. Further warnings will be suppressed." }
              $__immunio_av_rendered_non_string = true
            end
            return rendered
          end

          rendered = rendered.to_str

          result = run_hook!("template_render_done", {
            content_type: Mime::Type.lookup_by_extension(@template.formats.first).to_s,
            rendered: rendered,
            vars: @vars
          })

          # We use the return value from the hook handler if present.
          rendered = result["rendered"] || rendered.dup

          remove_var_markers! rendered

          # If some fragments were marked to be cached, commit their content to cache.
          write_and_remove_fragments! context, rendered

          rendered.html_safe
        else
          # This is a partial template. Just render it.
          rendered
        end
      ensure
        top_template = rendering_stack.pop
        unless top_template == self
          raise Error, "Unexpected Immunio::Template on rendering stack. Expected #{id}, got #{top_template.try :id}."
        end
      end
    end

    ENCODED_IMMUNIO_TOKENS_RE = Regexp.compile(/(?:{|%7b)immunio-(var|fragment)(?::|%3a)(\d+)(?::|%3a)([0-9a-f]{1,4})(?:}|%7d)(.*?)(?:{|%7b)(?:\/|%2f)immunio-\1(?::|%3a)\2(?::|%3a)\3(?:}|%7d)/i)

    private

    class << self
      def finalize_template(id)
        proc { CHECKSUM_CACHE.delete(id) if CHECKSUM_CACHE.has_key?(id) }
      end

      def current
        rendering_stack.last
      end

      def next_var_id
        rendering_stack.first.next_var_id
      end

      def vars
        rendering_stack.first.vars
      end

      def get_nonce
        rendering_stack.first.get_nonce
      end

      # Save fragment info to the root template only
      def mark_and_defer_fragment_write(*args)
        rendering_stack.first.mark_and_defer_fragment_write(*args)
      end

      # Stack of the templates currently being rendered.
      def rendering_stack
        Thread.current["immunio.rendering_stack"] ||= []
      end

      def wrap_code(code, handler, options = {})
        wrap_method = {
          'ActionView::Template::Handlers::ERB' => :wrap_code_for_erb,
          'Haml::Plugin' => :wrap_code_for_haml,
          'Slim::RailsTemplate' => :wrap_code_for_slim
        }[handler]

        send wrap_method, code, options
      end

      def wrap_code_for_erb(code, options)
        modifier = options[:escape] ? '=' : '=='
        "<%#{modifier} #{code} %>"
      end

      def wrap_code_for_haml(code, options)
        modifier = options[:escape] ? '=' : '!='
        "#{modifier} #{code}"
      end

      def wrap_code_for_slim(code, options)
        modifier = options[:escape] ? '=' : '=='
        "#{modifier} #{code}"
      end

      def decode_immunio_tokens(rendered)
        # Look for URI or JS encoded immunio tokens in the rendering and decode them
        # WebConsole incompatibility: `rendered` can be of type `Mime::Type` Which
        # doesn't respond to `gsub!`.
        if rendered.respond_to?(:gsub!)
          was_html_safe = rendered.html_safe?
          was_frozen = rendered.frozen?

          if was_frozen
            # This is not an airtight solution. Object#dup does not copy methods
            # defined on the instance, and may be overridden by subclasses to do
            # things that would cause problems for us. But most likely there is no
            # problem with using dup. We can't use Object#clone because the clone
            # retains the frozen status of the original, preventing us from
            # modifying the string contents.
            rendered = rendered.dup
          end

          rendered.gsub! ENCODED_IMMUNIO_TOKENS_RE, "{immunio-\\1:\\2:\\3}\\4{/immunio-\\1:\\2:\\3}"

          rendered.instance_variable_set(:@html_safe, true) if was_html_safe
          rendered.freeze if was_frozen
        end

        rendered
      end

      def mark_var(content, code, template_id, template_sha, file, line, escape, is_text, handler)
        id = Template.next_var_id
        nonce = Template.get_nonce

        # NOTE: What happens here is pretty funky to preserve the html_safe SafeBuffer behaviour in ruby.
        # If escaped is true we directly concatenate the content between two SafeBuffers. This will cause
        # escaping if content is not itself a SafeBuffer.
        # Otherwise we explicitly convert to a string, and convert that to a SafeBuffer to ensure that
        # for instance no escaping is performed on the contents of a <%== %> Erubis interpolation.
        rendering =
          if escape && !is_text
            # explicitly convert (w/ escapes) and mark safe things that aren't String (SafeBuffer is_a String also)
            # `to_s` is used to render any object passed to a template.
            # It is called internally when appending to ActionView::OutputBuffer.
            # We force rendering to get the actual string.
            # This has no impact if `rendered` is already a string.
            content = content.to_s.html_safe unless content.is_a? String

            # As a failsafe, just return the content if it already contains our markers. This can occur when
            # a helper calls render partial to generate a component of a page. Both render calls are root level
            # templates from our perspective.
            if content =~ /\{immunio-var:\d+:#{nonce}\}/ then
              # don't add markers.
              Immunio.logger.debug {"WARNING: ActionView not marking interpolation which already contains markers: \"#{content}\""}
              return content
            end

            "{immunio-var:#{id}:#{nonce}}".html_safe + content + "{/immunio-var:#{id}:#{nonce}}".html_safe
          else
            content = "" if content.nil?

            # See comment above
            if (content.respond_to? :=~) &&
               (content =~ /\{immunio-var:\d+:#{nonce}\}/)
              # don't add markers.
              Immunio.logger.debug {"WARNING: ActionView not marking interpolation which already contains markers: \"#{content}\""}
              return content.html_safe
            end

            "{immunio-var:#{id}:#{nonce}}".html_safe + content.to_s.html_safe + "{/immunio-var:#{id}:#{nonce}}".html_safe
          end

        # If we got here, the interpolation has been wrapped in our markers and we
        # need to record send data about it to the hook
        Template.vars[id.to_s] = {
          template_sha: template_sha,
          template_id: template_id.to_s,
          nonce: nonce,
          code: wrap_code(code, handler, escape: escape),
          file: file,
          line: line
        }

        rendering
      end

      def render_var(code, rendered, template_id, template_sha, file, line, escape, is_text, handler)
        rendered = decode_immunio_tokens rendered

        if rendered.instance_variable_get("@__immunio_processed") then
          # Ignore buffers marked as __immunio_processed in render as these are full templates or partials
          return rendered
        elsif code =~ /yield( .*)?/
          # Ignore yielded blocks inside layouts
          return rendered
        end
        rendered = mark_var rendered, code, template_id, template_sha, file, line, escape, is_text, handler
        rendered.html_safe
      end

      def remove_comment(code)
        comment = Ripper.slice(code, "comment")

        if comment
          code = code.sub(Regexp.new(Regexp.escape(comment) + "\\Z"), "")
        end

        code
      end

      # Generate code injected in templates to wrap everything inside `<%= ... %>`.
      def generate_render_var_code(code, escape)
        template = Template.current
        if template
          template_id = template.next_template_id

          handler = template.instance_variable_get(:@template).handler
          handler_name = if handler.is_a? Class
                           handler.name
                         else
                           handler.class.name
                         end

          "(__immunio_result = (#{remove_comment(code)}); Immunio::Template.render_var(#{code.strip.inspect}, __immunio_result, #{template_id}, '#{template.template_sha}', __FILE__, __LINE__, #{!!escape}, #{template.is_text?}, '#{handler_name}'))"
        else
          code
        end
      end

      def remove_all_markers!(input)
        input.gsub!(/\{\/?immunio-(fragment|var):\d+:[a-zA-Z0-9]+\}/, "")
      end
    end

    # End of private class methods

    def has_source?
      @template.respond_to?(:source) && !@template.source.nil?
    end

    def compiled?
      @template.instance_variable_get :@compiled
    end

    def load_source(context)
      return if !@template.respond_to?(:source) || !@template.source.nil?

      # @template is a virtual template that doesn't contain the source. We need
      # to try to load the source. But, the virtual template doesn't know the
      # original format of the source template file.
      #
      # First, try to load it using the Rails defaults (usually "html" and
      # "txt"). If that doesn't work, try to use the original format from the
      # virtual template.
      #
      # Though one might think the format from the virtual template would always
      # work, unfortunately the format from the template refers to the "type" of
      # the template, which may or may not be the same as the format of the
      # lookup context, which specifies the file extension of the template. Ugh,
      # naming... For example, the lookup context format may be "txt" while the
      # template format is "text".
      #
      # Astute readers may note that there's the possibility the template
      # extension is not in the Rails default list of lookup formats, and also
      # does not match the template type. We are just going to leave that for
      # another day, and hope that day never comes...
      begin
        refreshed = Immunio::IOHooks.paused { @template.refresh(context) }
      rescue
        begin
          old_formats = context.lookup_context.formats
          context.lookup_context.formats = @template.formats
          refreshed = Immunio::IOHooks.paused { @template.refresh(context) }
        rescue
          Immunio.logger.warn { "Failed to refresh template source from #{@template} using contexts #{old_formats} and #{@template.formats}" }
        ensure
          context.lookup_context.formats = old_formats
        end
      end

      return if refreshed.nil?

      @template.instance_variable_set :@source, refreshed.source
    end

    def rendering_stack
      self.class.rendering_stack
    end

    def run_hook!(name, meta={})
      default_meta = {
        template_sha: template_sha,
        name: (@template.respond_to?(:virtual_path) && @template.virtual_path) || nil,
        origin: @template.identifier,
        nonce: Template.get_nonce
      }
      Immunio.run_hook! "action_view", name, default_meta.merge(meta)
    end

    def write_and_remove_fragments!(context, content)
      # Rails tests do use the context as the view context sometimes.
      if context.is_a? ActionController::Base
        controller = context
      elsif context.respond_to? :controller
        controller = context.controller
      else
        # Some rails unit tests don't have a controller...
        remove_all_markers! content
        return
      end

      # Iterate to handle nested fragments. Child fragments have lower ids than their parents.
      nonce = Template.get_nonce
      @scheduled_fragments_writes.each_with_index do |(key, _, options), id|
        # Remove the markers ...
        content.sub!(/\{immunio-fragment:#{id}:#{nonce}\}(.*)\{\/immunio-fragment:#{id}:#{nonce}\}/m) do
          # The escaped content inside the markers ($1), is written to cache.
          output = $1
          remove_all_markers! output
          controller.write_fragment_without_immunio key, output, options
          output
        end
      end
      # To be extra safe strip all markers from content
      remove_all_markers! content
    end

    def remove_var_markers!(input)
      nonce = Template.get_nonce
      # TODO is this the fastest way to remove the markers? Needs benchmarking ...
      input.gsub!(/\{\/?immunio-var:\d+:#{nonce}\}/, "")
    end

    def remove_all_markers!(input)
      self.class.remove_all_markers!(input)
    end
  end
end