require 'nokogiri'
require 'pstore'
 
 
#
#  Rack::DomainSprinkler
#
#  Modifies outgoing HTML markup so that common static assets like script source
#  files, images, and stylesheets will be served from one of a number of domains
#  rather than just a single one in order to improve parallelization of resource
#  downloads during page loads.
#
module Rack
  class DomainSprinkler
 
    #
    #  Options:
    #    A list of domains across which to spread requests
    #    The name of the file to be used for caching URL/path mappings
    #
    def initialize(app, domains=[], cache_loc="tmp/sprinklins.pstore")
      @app = app
      @@domains = domains
      @@url_cache = PStore.new(cache_loc)
    end
 
 
    def call(env)
      dup._call(env)
    end
 
 
    #
    #  Currently only looks for <link>, <script>, and <img> tags.
    #  Should be refactored to provide more generic mapping.
    #
    def _call(env)
      status, headers, response = @app.call(env)
      if headers["Content-Type"] =~ /^text\/html/
        document = Nokogiri::HTML(response.body)
        
        sprinkle(env, document, 'link', 'href')
        sprinkle(env, document, 'script', 'src')
        sprinkle(env, document, 'img', 'src')
        
        response.body = document.to_html
      end
 
      [status, headers, response]
    end
 
 
    def sprinkle(env, document, selector, property)
      document.search(selector).each do |node|
        url = node[property]
        if !url.empty? and url !~ /^\#/ and url !~ /^http\:\/\//
          node[property] = domainify(url, env)
        end
      end
    end
 
 
    def domainify(path, env)
      if path !~ /\//
        path = absolutify(path)
      end
      
      full_url = nil
      @@url_cache.transaction do
        unless @@url_cache[path]
          domain = @@domains.shift and @@domains.push(domain)
          @@url_cache[path] = "http://#{ domain }#{ path }"
        end
        full_url = @@url_cache[path]
      end
 
      full_url
    end
 
 
    def absolutify(path, env)
      path_prefix = env['REQUEST_PATH']
      unless path_prefix =~ /\/$/
        path_prefix.sub!(/\/[^\/]+$/, '')
      end
 
      "#{ tmp_path }/#{ path }"
    end
 
  end
end
 

module Rack
  class MemoryBloat
    def initialize(app, logger)
      @app = app
      @logger = logger
    end
 
    def call(env)
      memory_usage_before = memory_usage
      result = @app.call(env)
      memory_usage_after = memory_usage
      @logger.info "MemoryBloat: #{memory_usage_after - memory_usage_before} URL: #{Rack::Request.new(env).url}"
      result
    end
 
    private
    def memory_usage
      `ps -o rss= -p #{$$}`.to_i
    end
  end
end

# Full code is at http://github.com/cwninja/rack-cache-buster
 
require 'digest/md5'
 
module Rack
  class CacheBuster
    def initialize(app, key, target_time = nil)
      @app, @target_time = app, target_time
      @key = "-"+Digest::MD5.hexdigest(key || "blank-key").freeze
      @key_regexp = /#{@key}/.freeze
    end
 
    def call(env)
      status, headers, body = app.call(unpatch_etag(env))
      [status, patch_etag(limit_cache(headers)), body]
    end
 
  protected
    QUOTE_STRIPPER=/^"|"$/.freeze
    ETAGGY_HEADERS = ["HTTP_IF_NONE_MATCH", "HTTP_IF_MATCH", "HTTP_IF_RANGE"].freeze
    ETag = "ETag".freeze
 
    attr_reader :app
    attr_reader :key
 
 
    def limit_cache(headers)
      return headers if @target_time.nil?
      cache_control_header = CacheControlHeader.new(headers)
      if cache_control_header.expired?
        headers
      else
        cache_control_header.expire_time = [@target_time, cache_control_header.expire_time].min
        headers.merge(CacheControlHeader::CacheControl => cache_control_header.to_s)
      end
    end
 
    def unpatch_etag(headers)
      ETAGGY_HEADERS.inject(headers){|memo, k|
        memo.has_key?(k) ? memo.merge(k => strip_etag(memo[k])) : memo
      }
    end
 
    def strip_etag(s)
      s.gsub(@key_regexp, "")
    end
 
    def modify_etag(s)
      s = s.to_s.gsub(QUOTE_STRIPPER, "")
      s.empty? ? s : %Q{"#{s}#{key}"}
    end
 
    def patch_etag(headers)
      headers.merge(ETag => modify_etag(headers[ETag]))
    end
 
    autoload :CacheControlHeader, "rack/cache_buster/cache_control_header"
    autoload :Rails, "rack/cache_buster/rails"
  end
end
 

require "rack/cache_buster"
 
class Rack::CacheBuster::CacheControlHeader
  CacheControl = "Cache-Control".freeze
  Age = "Age".freeze
  MaxAge = "max-age".freeze
 
  def initialize(env)
    @age = env[Age].to_i
    @max_age = 0
    if env[CacheControl]
      parts = env[CacheControl].split(/ *, */)
      settings, options = parts.partition{|part| part =~ /=/ }
      settings = settings.inject({}){|acc, part|
        k, v = part.split(/ *= */, 2)
        acc.merge(k => v)
      }
      @max_age = settings.delete(MaxAge).to_i
      @other_parts = options + settings.map{|k,v| "#{k}=#{v}"}
    end
  end
 
  def expires_in
    @max_age - @age
  end
 
  def expired?
    expires_in <= 0
  end
 
  def expire_time
    Time.now + expires_in
  end
 
  def expire_time=(t)
    @max_age = [t.to_i - Time.now.to_i + @age, @age].max
  end
 
  def update_env(env)
    env[CacheControl] = to_s
  end
 
  def to_s
    ["max-age=#{@max_age}", *@other_parts].join(", ")
  end
end
 

require "rack/cache_buster"
 
class Rack::CacheBuster::Rails < Rack::CacheBuster
  def initialize(app)
    app_version    = File.read(File.join(::Rails.root, "REVISION")) rescue nil
    wind_down_time = Time.parse(File.read(File.join(::Rails.root, "WIND_DOWN"))) rescue nil
    super(app, app_version, wind_down_time)
  end
end
 

class LiveStats
  
  def initialize(app, cache, options={})
    @app = app
    @cache = cache
    @period = options[:period] || 3600
    @precision = options[:precision] || 10
  end
 
  def inc(key, value=1, timeout=nil)
    if @cache.incr(key, value).nil?
      @cache.add(key, value, timeout || @period)
    end
  end
 
  def call(env) 
    env["rack.livestats.track"] = true
    now = Time::now
    
    response = @app.call(env)
    
    if env["rack.livestats.track"]
      ts = ((Time::now - now) * 1000).to_i # in ms
      key = Time::now.to_i / @precision
      
      # time spent
      inc("t"+key.to_s, ts)
      # visits
      inc("v"+key.to_s)
    end
 
    response
  end
end

module Rack
  
  # A rack middleware for validating HTML via w3c validator
  class Validate
    
    def initialize( app )
      @app = app
    end
        
    def call( env )
      status, headers, response = @app.call( env )
      
      request = Rack::Request.new( env )
      if !request.params['rack-validate'].blank?
        if headers['Content-Type'] =~ /text\/html|application\/xhtml\+xml/
          body = response.body
        
          issues = Validator.validate( body )
        
          body.insert( 0, Validator.generate_report( issues ) )
        
          headers["Content-Length"] = body.length.to_s
          response = [body]
        end
      end
      
      [status, headers, response]
    end
  
    class Validator
      include W3CValidators
      
      def self.validate( response )
        validator = MarkupValidator.new    
        validator.validate_text( response )        
      end
      
      def self.generate_report( issues )        
        report = ""
        report << STYLES
        report << SCRIPT
        
        report << "<div id='message_toolbar'>"
          report << '<div id="controls">'        
            summary = "<span>There were " + issues.errors.size.to_s + " errors and " + issues.warnings.size.to_s + " warnings</span><br/><br/>"
      		  report << summary  		  
      		  report << LINKS  		  
    		  report << '</div>'
        
          if !issues.errors.empty?
            report << "<table id='errors_table' style='display:none;'>"
            report << "<tr><td colspan='2' class='header_column'>Errors</td></tr>"
            issues.errors.each do |item|
              report << "<tr><td class='line_number_column'>Line #{item.line}</td><td class='message_column'>#{html_escape( item.message )}</td></tr>"
            end
            report << "</table>"
          end
        
          if !issues.warnings.empty?
            report << "<table id='warnings_table' style='display:none;'>"
            report << "<tr><td colspan='2' class='header_column'>Warnings</td></tr>"
            issues.warnings.each do |item|
              report << "<tr><td class='line_number_column'>--</td><td class='message_column'>#{html_escape( item.message )}</td></tr>"
            end
            report << "</table>"
          end
        report << "</div>"
      end
      
      private 
      
        # Stealing HTML escape method from rails
        def self.html_escape( string )
          string.to_s.gsub(/[&"><]/) { |special| HTML_ESCAPE[special] }
        end
      
        # Stealing HTML escape constant from rails
        HTML_ESCAPE	=	{ '&' => '&amp;', '>' => '&gt;', '<' => '&lt;', '"' => '&quot;' }
        
        LINKS = <<-HERE
  		    <a href="#" onclick="$( '#errors_table' ).toggle();updateText( this, 'Show Errors', 'Hide Errors' );">Show Errors</a>
  		    <a href="#" onclick="$( '#warnings_table' ).toggle();updateText( this, 'Show Warnings', 'Hide Warnings' );">Show Warnings</a>
  		    <a href="#" onclick="$( '#message_toolbar' ).toggle();">Close Toolbar</a>
  		  HERE
        
        SCRIPT = <<-HERE
          <script src="http://www.google.com/jsapi"></script>
          <script type="text/javascript">
            google.load("jquery", "1.3.2");
            google.load("jqueryui", "1.7.2");
 
            function updateText( element, textOne, textTwo ) {
              if( element.innerHTML == textOne ) {
                element.innerHTML = textTwo;
              }
              else { 
                element.innerHTML = textOne;
              }
            }
          </script>        
        HERE
                        
        STYLES = <<-HERE
    		<style type='text/css'>
          #message_toolbar {
            width: 100%; 
            min-height: 50px;
            background: black;
            position: absolute;			  
            top:0;
            left:0;
            -moz-opacity:.80; 
            filter:alpha(opacity=80); 
            opacity:.80;
            font-size: smaller;
            font-family: "Courier New";
            color: white;
            text-align: center;
          }			
          #message_toolbar a{
            color: white;
            font-weight: bold;
          }
          .line_number_column {
            text-align: left;
            width: 100px;
          }			
          .error_message_column {
            text-align: left;
          }		
          #errors_table,
          #warnings_table {
            width: 75%;
            margin: 0 auto;
            border: none;
            color: white;
            margin-top: 10px;
          }
          #errors_table .header_column,
          #warnings_table .header_column {
            text-align: center;
            border-bottom: 2px dashed;
          }
          #errors_table .header_column {
            border-color: red;
          }
          #warnings_table .header_column {
            border-color: yellow;
          }
    		</style>
 
        HERE
    end        
  end
end

module Rack
  class Superlogger
    module LogProcessor
      def self.[](type)
        case type
        when Class
          type
        when String, Symbol
          const_get type.to_s.capitalize
        else
          raise ArgumentError, "Unexpected type class #{type.class}"
        end
      end
      
      class Base
        attr_reader :logger
        
        def initialize(options)
          @options = options
        end
        
        def process(env)
          raise "Not implemented"
        end
      end
 
      class Templated < Base
        def initialize(options)
          @logger   = options.delete(:logger)   or raise ArgumentError, "You must specify a logger"
          @template = options.delete(:template) or raise ArgumentError, "You must specify a template"
          super options
        end
        
        def process(env)
          request = Rack::Request.new(env)
          message = @template.dup
 
          Rack::Superlogger::REQUEST_METHODS.each do |method_name|
            env["rack.superlogger.data"][method_name.to_sym] = request.send(method_name.to_sym) if message.include?(":#{method_name}")
          end
 
          env["rack.superlogger.data"].each do |k, v|
            message.gsub! ":#{k}", v.to_s
          end
 
          @logger.info message
        end
      end
    end
      
    REQUEST_METHODS = Rack::Request.public_instance_methods(false).
                        reject { |method_name| method_name =~ /[=\[]|content_length/ }.freeze
 
    def initialize(app, options)
      @app, @processor = app, LogProcessor[options.delete(:type)].new(options)
    end
     
    def call(env)
      env["rack.superlogger.data"], env["rack.superlogger.raw_logger"] = {}, @processor.logger
            
      before = Time.now.to_f
      status, headers, body = @app.call(env)
      duration = ((Time.now.to_f - before.to_f) * 1000).floor 
      
      env["rack.superlogger.data"][:duration]       = duration.to_s
      env["rack.superlogger.data"][:status]         = status.to_s
      env["rack.superlogger.data"][:content_length] = headers["Content-length"]
      
      @processor.process env
      
      [status, headers, body]
    end
  end
end

require 'rubygems'
require 'rack'
require 'redis_request_limiter'
 
# Give 5 requests a minute to make testing really easy
use RedisAPILimiter, :per_user_limit => 5, :reset_requests_every => 60
run proc { |env| [200, {"Content-Type" => "text/html", "Content-Length" => "5"}, ["Hello"]]}
 

require 'redis'
 
class RedisRequestLimiter
  
  DEFAULT_OPTIONS = {
    :per_user_limit       => 100,
    :reset_requests_every => 300,
    :override_limit_proc  => proc { |env| [nil, nil] }, # limit, every
    :redis_rate_key       => "api-requests".freeze,
    :header_status_prefix => "X-API-Limit".freeze,
    :env_to_api_key_proc  => proc { |env| env["REMOTE_ADDR"] },
    :limiter_app          => proc do |env|
      headers = {"Content-Type" => "text/html"}
      resets_in = env['x-rack.rate-limiter.resets-at'] - Time.now.to_i
      body = "You are currently limited and will be reset in appx. %s seconds." % resets_in
      headers["Content-Length"] = body.length.to_s
      return 503, headers, [body]
    end
  }
  
  attr_accessor :per_user_limit, :reset_requests_every, :redis_rate_key, :override_limit_proc,
                :header_status_prefix, :limiter_app, :redis, :env_to_api_key_proc
  
  def initialize(app, opts = {})
    @app = app
    # Setup options w/ defaults
    real_opts = DEFAULT_OPTIONS.merge(opts)
    real_opts.each_pair do |key, value|
      send(:"#{key}=", value) if respond_to?(:"#{key}=")
    end
  end
  
  def call(env)
    self.dup._call(env)
  end
  
  def _call(env)
    self.preload_request_info(env)
    env = self.env_with_limit_information(env)
    if self.under_request_limit?(env)
      count_request!(env)
      status, headers, body = @app.call(env)
      self.add_headers(headers)
      [status, headers, body]
    else
      status, headers, body = @limiter_app.call(env)
      self.add_headers(headers)
      [status, headers, body]
    end
  end
  
  def self.redis
    @@redis ||= Redis.new
  end
  
  def self.redis=(value)
    @@redis = value
  end
  
  protected
  
  def env_with_limit_information(env)
    env = env.dup
    env['x-rack.rate-limiter.current-count'] = self.current_request_count
    env['x-rack.rate-limiter.last-period']   = self.last_period
    env['x-rack.rate-limiter.resets-at']     = self.last_period + @reset_requests_every
    env['x-rack.rate-limiter.request-limit'] = @per_user_limit
    return env
  end
  
  def add_headers(header_hash)
    remaining_count = [(@per_user_limit - self.current_request_count), 0].max
    header_hash["#{@header_status_prefix}-Requests-Used"]      = @current_request_count.to_s
    header_hash["#{@header_status_prefix}-Requests-Remaining"] = remaining_count.to_s
    header_hash["#{@header_status_prefix}-Period"]             = self.last_period.to_s
    header_hash["#{@header_status_prefix}-Seconds-Remaining"]  = ((self.last_period + @reset_requests_every) - Time.now.to_i).to_s
    header_hash["#{@header_status_prefix}-Is-Limited"]         = (remaining_count > 0 ? "No" : "Yes")
  end
  
  def under_request_limit?(env)
    if self.last_period == self.current_period
      self.current_request_count < @per_user_limit
    else
      @current_request_count = 0
      @last_period = current_period
      @redis.set(@count_env_key, "0")
      @redis.set(@period_env_key, @current_period)
      true
    end
  end
  
  def count_request!(env)
    @current_request_count = @redis.incr(@count_env_key)
  end
  
  def env_rate_limit_key(env, key = "count")
    "#{@redis_rate_key}:#{@env_to_api_key_proc.call(env)}:#{key}"
  end
  
  def preload_request_info(env)
    current_request_limit, current_reset_every = @override_limit_proc.call(env)
    # Since we dup, we can change ivars
    @reset_requests_every = current_reset_every   if !current_reset_every.nil?
    @per_user_limit       = current_request_limit if !current_request_limit.nil?
    @count_env_key        = env_rate_limit_key(env)
    @period_env_key       = env_rate_limit_key(env, "current-limit-period")
    @redis              ||= self.class.redis
  end
  
  def last_period
    @last_period ||= begin
      raw_stored_period = @redis.get(@period_env_key).to_i
      (raw_stored_period / @reset_requests_every) * @reset_requests_every
    end
  end
  
  def current_period
    @current_period ||= (Time.now.to_i / @reset_requests_every) * @reset_requests_every 
  end
  
  def current_request_count
    @current_request_count ||= @redis.get(@count_env_key).to_i
  end
  
end
 

require 'geoip'
 
module Rack
  # Rack::GeoIPCountry uses the geoip gem and the GeoIP database to lookup the country of a request by its IP address
  # The database can be downloaded from:
  # http://geolite.maxmind.com/download/geoip/database/GeoLiteCountry/GeoIP.dat.gz
  #
  # Usage:
  # use Rack::GeoIPCountry, :db => "path/to/GeoIP.dat"
  #
  # By default all requests are looked up and the X_GEOIP_* headers are added to the request
  # The headers can then be read in the application
  # The country name is added to the request header as X_GEOIP_COUNTRY, eg:
  # X_GEOIP_COUNTRY: United Kingdom
  #
  # The full set of GEOIP request headers is below:
  # X_GEOIP_COUNTRY_ID - The GeoIP country-ID as an integer, if not found set to 0
  # X_GEOIP_COUNTRY_CODE - The ISO3166-1 two-character country code, if not found set to --
  # X_GEOIP_COUNTRY_CODE3 - The ISO3166-2 three-character country code, if not found set to --
  # X_GEOIP_COUNTRY - The ISO3166 English-language name of the country, if not found set to N/A
  # X_GEOIP_CONTINENT - The two-character continent code, if not found set to --
  #
  #
  # You can use the included Mapping class to trigger lookup only for certain requests by specifying matching path prefix in options, eg:
  # use Rack::GeoIPCountry::Mapping, :prefix => '/video_tracking'
  # The above will lookup IP addresses only for requests matching /video_tracking etc.
  #
  # MIT License - Karol Hosiawa ( http://twitter.com/hosiawak )
  #
 
  class GeoIPCountry
    def initialize(app, options = {})
      options[:db] ||= 'GeoIP.dat'
      @db = GeoIP.new(options[:db])
      @app = app
    end
 
    def call(env)
      res = @db.country(env['REMOTE_ADDR'])
      env['X_GEOIP_COUNTRY_ID'] = res[2]
      env['X_GEOIP_COUNTRY_CODE'] = res[3]
      env['X_GEOIP_COUNTRY_CODE3'] = res[4]
      env['X_GEOIP_COUNTRY'] = res[5]
      env['X_GEOIP_CONTINENT'] = res[6]
      @app.call(env)
    end
 
    class Mapping
      def initialize(app, options = {})
        @app, @prefix = app, /^#{options.delete(:prefix)}/
        @geoip_country = GeoIPCountry.new(app, options)
      end
 
      def call(env)
        if env['PATH_INFO'] =~ @prefix
          @geoip_country.call(env)
        else
          @app.call(env)
        end
      end
    end
  end
end
 

# See also http://github.com/simonjefford/rack_firebug_logger
# for this middleware + tests + a rails plugin
class FirebugLogger
  def initialize(app, options = {})
    @app = app
    @options = options
  end
 
  def call(env)
    dup._call(env)
  end
 
  def _call(env)
    status, headers, body = @app.call(env)
    return [status, headers, body] unless (headers["Content-Type"] =~ /html/ && env['firebug.logs'])
    response = Rack::Response.new([], status, headers)
    js = generate_js(env['firebug.logs'])
    body.each do |line|
      line.gsub!("</body>", js)
      response.write(line)
    end
    response.finish
  end
 
  private
 
  def generate_js(logs)
    js = ["<script type=\"text/javascript\">"]
    start_group(js)
    logs.each do |level, log|
      level = sanitise_level(level)
      log.gsub!('"', '\"')
      js << "console.#{level.to_s}(\"#{log}\");"
    end
    end_group(js)
    js << "</script>"
    js << "</body>"
    js.join("\n")
  end
 
  def start_group(js)
    if @options[:group]
      js << "console.group(\"#{@options[:group]}\");"
    end
  end
 
  def sanitise_level(level)
    if [:info, :debug, :warn, :error].include?(level)
      level
    else
      :debug
    end
  end
 
  def end_group(js)
    if @options[:group]
      js << "console.groupEnd();"
    end
  end
end

class RackProctitle
  def initialize(app, options = nil)
    @app = app
    @prefix = options.delete(:prefix) if options
    
    @revision = APPLICATION_VERSION if defined?(APPLICATION_VERSION)
    @mutex = Mutex.new
    @titles = []
    @request_threads = []
    @queue_length = 0
    @request_count = 0
 
    @updater_thread = Thread.new do
      while true
        @mutex.synchronize do
          set_request_list_title
        end
        sleep 0.5
      end
    end
  end
 
  def call(env)
    Thread.current[:request_str] = ((env["REQUEST_URI"].nil? || env["REQUEST_URI"].empty?) ? "/" : env["REQUEST_URI"]).split("?", 2)[0]
    Thread.current[:arrived_at] = Time.now.to_f
    @mutex.synchronize do
      @request_threads.push(Thread.current)
      @queue_length += 1
      set_request_list_title
    end
    
    begin
      @app.call(env)
    ensure
      @mutex.synchronize do
        @queue_length -= 1
        @request_count += 1
 
        @last_time = Time.now.to_f - Thread.current[:arrived_at].to_f
        @last_request_str = Thread.current[:request_str].to_s
 
        @request_threads.delete(Thread.current)
        set_request_list_title
      end
    end
  end
 
  protected
  def set_request_list_title
    @title = if @request_threads.empty?
      idle_message
    else
      now = Time.now.to_f
      list = @request_threads.inject([]) do |str, thread|
        str << "#{time_delta_abbriv(now - thread[:arrived_at])} #{thread[:request_str]}"
      end.join(" | ")
      "handling #{list}"
    end
    update_process_title
  end
 
  def idle_message
    str = "idle#{' '*20}"
    str << "[last #{time_delta_abbriv(@last_time)} #{@last_request_str}]" if @last_time && @last_request_str
    str
  end
 
  def update_process_title
    title = "#{@prefix || 'rack'}"
    title << "/#{@revision}" if @revision
    title << " ["
    title << "#{@queue_length}/"
    title << "#{@request_count}"
    title << "]: #{@title}"
    $0 = title
  end
 
  def time_delta_abbriv(delta)
    if delta < 60
      "%.1fs" % delta
    elsif delta < 3600
      "#{delta.to_i / 60}m#{delta.to_i % 60}s"
    elsif delta < 86400
      "#{delta.to_i / 3600}h#{(delta.to_i % 3600) / 60}m"
    else
      "#{delta.to_i / 86400}d#{(delta.to_i % 86400) / 3600}h"
    end
  end
end