# Hook into ActionView rendering to inject Immunio's hooks.
require 'securerandom'

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

    def self.finalize_template(id)
      proc { CHECKSUM_CACHE.delete(id) if CHECKSUM_CACHE.has_key?(id) }
    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 has_source?
      @template.respond_to?(:source) && !@template.source.nil?
    end

    def is_text?
      @template.formats.first == :text
    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 template_sha
      CHECKSUM_CACHE[@template.object_id]
    end

    def compiled?
      @template.instance_variable_get :@compiled
    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 self.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 =~ /\{immunio-var:\d+:#{nonce}\}/ then
          # 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 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

    # Generate code injected in templates to wrap everything inside `<%= ... %>`.
    def self.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 = (#{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

    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)
    def self.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 self.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 self.current
      rendering_stack.last
    end

    def self.next_var_id
      rendering_stack.first.next_var_id
    end

    def self.vars
      rendering_stack.first.vars
    end

    def self.get_nonce
      rendering_stack.first.get_nonce
    end

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

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

    def self.wrap_code(code, handler, options = {})
      case
      when handler == 'ActionView::Template::Handlers::ERB'
        modifier = options[:escape] ? '=' : '=='
        "<%#{modifier} #{code} %>"
      when handler == 'Haml::Plugin'
        modifier = options[:escape] ? '=' : '!='
        "#{modifier} #{code}"
      end
    end

    private
      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)
        input.gsub!(/\{\/?immunio-(fragment|var):\d+:[a-zA-Z0-9]+\}/, "")
      end
  end

  # Regexp to test for blocks (... do) in the Ruby code of templates.
  BLOCK_EXPR = ActionView::Template::Handlers::Erubis::BLOCK_EXPR

  # Hooks for the ERB template engine.
  # (Default one used in Rails).
  module ErubisHooks
    extend ActiveSupport::Concern

    included do
      Immunio::Utils.alias_method_chain self, :add_expr, :immunio
    end

    def add_expr_with_immunio(src, code, indicator)
      # Wrap expressions in the templates to track their rendered value.
      # Do not wrap expressions with blocks, eg.: <%= form_tag do %>
      # TODO should we support blocks?
      Request.time "plugin", "#{Module.nesting[0]}::#{__method__}" do
        unless code =~ BLOCK_EXPR
          # escape unless we see the == indicator
          escape = !(indicator == '==')
          code = Immunio::Template.generate_render_var_code(code, escape)
        end
        Request.pause "plugin", "#{Module.nesting[0]}::#{__method__}" do
          add_expr_without_immunio(src, code, indicator)
        end
      end
    end
  end

  # Hooks for the HAML template engine.
  module HamlHooks
    extend ActiveSupport::Concern

    included do
      Immunio::Utils.alias_method_chain self, :push_script, :immunio
    end

    def push_script_with_immunio(code, opts = {}, &block)
      # Wrap expressions in the templates to track their rendered value.
      Request.time "plugin", "#{Module.nesting[0]}::#{__method__}" do
        if code !~ BLOCK_EXPR
          # escape if we're told to by HAML
          code = Immunio::Template.generate_render_var_code(code, opts[:escape_html])
        end
        Request.pause "plugin", "#{Module.nesting[0]}::#{__method__}" do
          push_script_without_immunio(code, opts, &block)
        end
      end
    end
  end

  # Hook for the `ActionView::TemplateRenderer`. These are called for root
  # templates.
  module TemplateRendererHooks
    extend ActiveSupport::Concern

    included do
      Immunio::Utils.alias_method_chain self, :render_template, :immunio
    end

    def render_template_with_immunio(template, *args)
      Request.time "plugin", "#{Module.nesting[0]}::#{__method__}" do
        renderer = Template.new(template)

        renderer.render @view do
          Request.pause "plugin", "#{Module.nesting[0]}::#{__method__}" do
            render_template_without_immunio(template, *args)
          end
        end
      end
    end
  end

  # Hook for the `ActionView::Template`. These are called for non-root
  # templates.
  module TemplateHooks
    extend ActiveSupport::Concern

    included do
      Immunio::Utils.alias_method_chain self, :render, :immunio
    end

    def render_with_immunio(context, *args, &block)
      Request.time "plugin", "#{Module.nesting[0]}::#{__method__}" do
        renderer = Template.new(self)

        renderer.render context do
          Request.pause "plugin", "#{Module.nesting[0]}::#{__method__}" do
            render_without_immunio(context, *args, &block)
          end
        end
      end
    end
  end

  # Hook for `ActionController::Caching::Fragments` responsible for handling the `<% cache do %>...` in templates.
  module FragmentCachingHooks
    extend ActiveSupport::Concern

    included do
      Immunio::Utils.alias_method_chain self, :write_fragment, :immunio
    end

    def write_fragment_with_immunio(key, content, options = nil)
      return content unless cache_configured?

      template = Template.current
      if template
        # We're rendering a template. Defer caching 'till we get the escaped content from the hook handler.
        content = Template.mark_and_defer_fragment_write(key, content, options)
      else
        # Not rendering a template. Ignore.
        # Shouldn't happen. But, just to be safe in case fragment caching is used in the controller for something else.
        content = write_fragment_without_immunio(key, content, options)
      end

      content
    end
  end

  # Hook for the `ActiveSupport::Hash#to_query`.
  # Use case: building a url within a decorator that renders a partial with an interpolation.
  module ActiveSupportHooks
    extend ActiveSupport::Concern

    included do
      Immunio::Utils.alias_method_chain self, :to_query, :immunio
    end

    def to_query_with_immunio(namespace = nil)
      escaped_string = to_query_without_immunio(namespace)
      Request.time "plugin", "#{Module.nesting[0]}::#{__method__}" do
        # Our markers got escaped, so un-unescaped them back.
        escaped_string.gsub!(
          Immunio::Template::ENCODED_IMMUNIO_TOKENS_RE,
          "{immunio-\\1:\\2:\\3}\\4{/immunio-\\1:\\2:\\3}")
      end
      escaped_string
    end
  end
end

# Add XSS hooks if enabled
if Immunio::agent.plugin_enabled?("xss") then
  # Hook into template engines.
  ActionView::Template::Handlers::Erubis.send :include, Immunio::ErubisHooks

  ActiveSupport.on_load(:after_initialize) do
    # Wait after Rails initialization to patch custom template engines.
    if defined? Haml::Compiler
      Haml::Compiler.send :include, Immunio::HamlHooks
    end

    Hash.send :include, Immunio::ActiveSupportHooks
  end

  # Hook into rendering process of Rails.
  ActionView::TemplateRenderer.send :include, Immunio::TemplateRendererHooks
  ActionView::Template.send :include, Immunio::TemplateHooks
  ActionController::Caching::Fragments.send :include, Immunio::FragmentCachingHooks
end