hedra 1.0.1 → 2.0.0

data/lib/hedra/cli.rb

@@ -42,7 +42,111 @@ module Hedra
     end
   end
 
-  class CLI < Thor
+  class BaselineCLI < Thor
+    desc 'list', 'List saved baselines'
+    def list
+      baseline = Baseline.new
+      baselines = baseline.list
+
+      if baselines.empty?
+        puts 'No baselines saved.'
+      else
+        puts 'Saved baselines:'
+        baselines.each do |b|
+          puts "  - #{b[:name]} (#{b[:url_count]} URLs, created: #{b[:created_at]})"
+        end
+      end
+    end
+
+    desc 'compare NAME URL_OR_FILE', 'Compare current results against baseline'
+    option :file, type: :boolean, aliases: '-f', desc: 'Treat argument as file with URLs'
+    option :output, type: :string, aliases: '-o', desc: 'Output file'
+    def compare(name, target)
+      baseline = Baseline.new
+      urls = options[:file] ? File.readlines(target).map(&:strip).reject(&:empty?) : [target]
+
+      client = HttpClient.new
+      analyzer = Analyzer.new
+      current_results = []
+
+      urls.each do |url|
+        response = client.get(url)
+        result = analyzer.analyze(url, response.headers.to_h)
+        current_results << result
+      rescue StandardError => e
+        warn "Failed to scan #{url}: #{e.message}"
+      end
+
+      comparisons = baseline.compare(name, current_results)
+      print_comparisons(comparisons)
+
+      if options[:output]
+        File.write(options[:output], JSON.pretty_generate(comparisons))
+        puts "Comparison saved to #{options[:output]}"
+      end
+    rescue StandardError => e
+      warn "Comparison failed: #{e.message}"
+      exit 1
+    end
+
+    desc 'delete NAME', 'Delete a baseline'
+    def delete(name)
+      baseline = Baseline.new
+      baseline.delete(name)
+      puts "Baseline deleted: #{name}"
+    rescue StandardError => e
+      warn "Failed to delete baseline: #{e.message}"
+      exit 1
+    end
+
+    private
+
+    def print_comparisons(comparisons)
+      pastel = Pastel.new
+
+      comparisons.each do |comp|
+        puts "\n#{pastel.bold(comp[:url])}"
+        puts "Baseline Score: #{comp[:baseline_score]} | Current Score: #{comp[:current_score]}"
+
+        change = comp[:score_change]
+        if change.positive?
+          puts pastel.green("Score improved by #{change} points")
+        elsif change.negative?
+          puts pastel.red("Score decreased by #{change.abs} points")
+        else
+          puts 'Score unchanged'
+        end
+
+        if comp[:new_findings].any?
+          puts pastel.yellow("\nNew findings:")
+          comp[:new_findings].each { |f| puts "  - #{f[:header]}: #{f[:issue]}" }
+        end
+
+        if comp[:resolved_findings].any?
+          puts pastel.green("\nResolved findings:")
+          comp[:resolved_findings].each { |f| puts "  - #{f[:header]}: #{f[:issue]}" }
+        end
+      end
+    end
+  end
+
+  class CacheCLI < Thor
+    desc 'clear', 'Clear response cache'
+    def clear
+      cache = Cache.new
+      cache.clear
+      puts 'Cache cleared.'
+    end
+
+    desc 'clear-expired', 'Clear expired cache entries'
+    def clear_expired
+      cache = Cache.new
+      cache.clear_expired
+      puts 'Expired cache entries cleared.'
+    end
+  end
+
+  class CLI < Thor # rubocop:disable Metrics/ClassLength
     class_option :verbose, type: :boolean, aliases: '-v', desc: 'Verbose output'
     class_option :quiet, type: :boolean, aliases: '-q', desc: 'Quiet mode'
     class_option :debug, type: :boolean, desc: 'Enable debug logging'
@@ -58,24 +162,69 @@ module Hedra
     option :rate, type: :string, desc: 'Rate limit (e.g., 5/s)'
     option :proxy, type: :string, desc: 'HTTP/SOCKS proxy URL'
     option :user_agent, type: :string, desc: 'Custom User-Agent header'
-    option :follow_redirects, type: :boolean, default: false, desc: 'Follow redirects'
+    option :follow_redirects, type: :boolean, default: true, desc: 'Follow redirects'
     option :output, type: :string, aliases: '-o', desc: 'Output file'
-    option :format, type: :string, default: 'table', desc: 'Output format (table, json, csv)'
-    def scan(target)
+    option :format, type: :string, default: 'table', desc: 'Output format (table, json, csv, html)'
+    option :cache, type: :boolean, default: false, desc: 'Enable response caching'
+    option :cache_ttl, type: :numeric, default: 3600, desc: 'Cache TTL in seconds'
+    option :check_certificates, type: :boolean, default: true, desc: 'Check SSL certificates'
+    option :check_security_txt, type: :boolean, default: false, desc: 'Check for security.txt'
+    option :save_baseline, type: :string, desc: 'Save results as baseline'
+    option :progress, type: :boolean, default: true, desc: 'Show progress bar'
+    def scan(target) # rubocop:disable Metrics/AbcSize, Metrics/CyclomaticComplexity, Metrics/PerceivedComplexity
       setup_logging
       urls = options[:file] ? read_urls_from_file(target) : [target]
 
       client = build_http_client
-      analyzer = Analyzer.new
+      analyzer = Analyzer.new(
+        check_certificates: options[:check_certificates],
+        check_security_txt: options[:check_security_txt]
+      )
+      cache = options[:cache] ? Cache.new(ttl: options[:cache_ttl]) : nil
+      rate_limiter = options[:rate] ? RateLimiter.new(options[:rate]) : nil
+      circuit_breakers = {}
       results = []
 
+      progress = options[:progress] && !options[:quiet] ? ProgressTracker.new(urls.length, quiet: options[:quiet]) : nil
+
       with_concurrency(urls, options[:concurrency]) do |url|
-        response = client.get(url)
-        result = analyzer.analyze(url, response.headers.to_h)
-        results << result
-        print_result(result) unless options[:quiet] || options[:output]
-      rescue StandardError => e
-        log_error("Failed to scan #{url}: #{e.message}")
+        rate_limiter&.acquire
+
+        # Get or create circuit breaker for this domain
+        domain = URI.parse(url).host
+        circuit_breakers[domain] ||= CircuitBreaker.new
+
+        begin
+          circuit_breakers[domain].call do
+            # Check cache first
+            cached = cache&.get(url)
+            if cached
+              result = cached
+              log_info("Cache hit: #{url}") if @verbose
+            else
+              response = client.get(url)
+              result = analyzer.analyze(url, response.headers.to_h, http_client: client)
+              cache&.set(url, result)
+            end
+
+            results << result
+            print_result(result) unless options[:quiet] || options[:output]
+          end
+        rescue CircuitOpenError
+          log_error("Circuit breaker open for #{domain}, skipping #{url}")
+        rescue StandardError => e
+          log_error("Failed to scan #{url}: #{e.message}")
+        ensure
+          progress&.increment
+        end
+      end
+
+      progress&.finish
+
+      if options[:save_baseline]
+        baseline = Baseline.new
+        baseline.save(options[:save_baseline], results)
+        say "Baseline saved: #{options[:save_baseline]}", :green unless options[:quiet]
       end
 
       export_results(results) if options[:output]
@@ -87,14 +236,19 @@ module Hedra
     option :proxy, type: :string, desc: 'HTTP/SOCKS proxy URL'
     option :user_agent, type: :string, desc: 'Custom User-Agent header'
     option :timeout, type: :numeric, aliases: '-t', default: 10, desc: 'Request timeout'
+    option :check_certificates, type: :boolean, default: true, desc: 'Check SSL certificates'
+    option :check_security_txt, type: :boolean, default: true, desc: 'Check for security.txt'
     def audit(url)
       setup_logging
       client = build_http_client
-      analyzer = Analyzer.new
+      analyzer = Analyzer.new(
+        check_certificates: options[:check_certificates],
+        check_security_txt: options[:check_security_txt]
+      )
 
       begin
         response = client.get(url)
-        result = analyzer.analyze(url, response.headers.to_h)
+        result = analyzer.analyze(url, response.headers.to_h, http_client: client)
 
         if options[:json]
           output = JSON.pretty_generate(result)
@@ -163,8 +317,8 @@ module Hedra
     option :output, type: :string, aliases: '-o', required: true, desc: 'Output file'
     option :input, type: :string, aliases: '-i', desc: 'Input JSON file with results'
     def export(format)
-      unless %w[json csv].include?(format)
-        say "Invalid format: #{format}. Use json or csv.", :red
+      unless %w[json csv html].include?(format)
+        say "Invalid format: #{format}. Use json, csv, or html.", :red
         exit 1
       end
 
@@ -177,6 +331,62 @@ module Hedra
     desc 'plugin SUBCOMMAND', 'Manage plugins'
     subcommand 'plugin', Hedra::PluginCLI
 
+    desc 'baseline SUBCOMMAND', 'Manage security baselines'
+    subcommand 'baseline', Hedra::BaselineCLI
+
+    desc 'cache SUBCOMMAND', 'Manage response cache'
+    subcommand 'cache', Hedra::CacheCLI
+
+    desc 'ci_check URL_OR_FILE', 'CI/CD friendly check (exit code based on score threshold)'
+    option :file, type: :boolean, aliases: '-f', desc: 'Treat argument as file with URLs'
+    option :threshold, type: :numeric, default: 80, desc: 'Minimum score threshold'
+    option :fail_on_critical, type: :boolean, default: true, desc: 'Fail if critical issues found'
+    option :output, type: :string, aliases: '-o', desc: 'Output file'
+    option :format, type: :string, default: 'json', desc: 'Output format (json, csv, html)'
+    def ci_check(target) # rubocop:disable Metrics/AbcSize
+      setup_logging
+      urls = options[:file] ? read_urls_from_file(target) : [target]
+
+      client = build_http_client
+      analyzer = Analyzer.new
+      results = []
+      failed = false
+
+      urls.each do |url|
+        response = client.get(url)
+        result = analyzer.analyze(url, response.headers.to_h, http_client: client)
+        results << result
+
+        if result[:score] < options[:threshold]
+          say "FAIL: #{url} - Score #{result[:score]} below threshold #{options[:threshold]}", :red
+          failed = true
+        end
+
+        if options[:fail_on_critical] && result[:findings].any? { |f| f[:severity] == :critical }
+          say "FAIL: #{url} - Critical security issues found", :red
+          failed = true
+        end
+      rescue StandardError => e
+        log_error("Failed to check #{url}: #{e.message}")
+        failed = true
+      end
+
+      # Export results if output specified
+      if options[:output]
+        exporter = Exporter.new
+        exporter.export(results, options[:format], options[:output])
+        say "Results exported to #{options[:output]}", :green unless options[:quiet]
+      end
+
+      if failed
+        say "\nCI check failed", :red
+        exit 1
+      else
+        say "\nCI check passed", :green
+        exit 0
+      end
+    end
+
     private
 
     def setup_logging
@@ -190,7 +400,7 @@ module Hedra
         timeout: options[:timeout] || 10,
         proxy: options[:proxy],
         user_agent: options[:user_agent],
-        follow_redirects: options[:follow_redirects] || false,
+        follow_redirects: options.key?(:follow_redirects) ? options[:follow_redirects] : true,
         verbose: @verbose
       )
     end
@@ -295,6 +505,12 @@ module Hedra
       say "Results exported to #{options[:output]}", :green unless options[:quiet]
     end
 
+    def log_info(message)
+      return unless @verbose
+
+      puts Pastel.new.cyan("INFO: #{message}")
+    end
+
     def severity_badge(severity)
       pastel = Pastel.new
       case severity.to_s

data/lib/hedra/config.rb

@@ -11,7 +11,7 @@ module Hedra
     DEFAULT_CONFIG = {
       'timeout' => 10,
       'concurrency' => 10,
-      'follow_redirects' => false,
+      'follow_redirects' => true,
       'user_agent' => "Hedra/#{VERSION}",
       'proxy' => nil,
       'output_format' => 'table'

data/lib/hedra/exporter.rb

@@ -11,6 +11,8 @@ module Hedra
         export_json(results, output_file)
       when 'csv'
         export_csv(results, output_file)
+      when 'html'
+        export_html(results, output_file)
       else
         raise Error, "Unsupported format: #{format}"
       end
@@ -45,5 +47,10 @@ module Hedra
         end
       end
     end
+
+    def export_html(results, output_file)
+      reporter = HtmlReporter.new
+      reporter.generate(results, output_file)
+    end
   end
 end

data/lib/hedra/html_reporter.rb

@@ -0,0 +1,136 @@
+# frozen_string_literal: true
+
+require 'erb'
+require 'time'
+
+module Hedra
+  # Generate HTML reports
+  class HtmlReporter
+    TEMPLATE = <<~HTML.freeze
+      <!DOCTYPE html>
+      <html lang="en">
+      <head>
+        <meta charset="UTF-8">
+        <meta name="viewport" content="width=device-width, initial-scale=1.0">
+        <title>Hedra Security Report</title>
+        <style>
+          * { margin: 0; padding: 0; box-sizing: border-box; }
+          body { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; background: #f5f5f5; padding: 20px; }
+          .container { max-width: 1200px; margin: 0 auto; background: white; border-radius: 8px; box-shadow: 0 2px 4px rgba(0,0,0,0.1); }
+          .header { background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); color: white; padding: 30px; border-radius: 8px 8px 0 0; }
+          .header h1 { font-size: 32px; margin-bottom: 10px; }
+          .header .meta { opacity: 0.9; font-size: 14px; }
+          .summary { display: grid; grid-template-columns: repeat(auto-fit, minmax(200px, 1fr)); gap: 20px; padding: 30px; border-bottom: 1px solid #eee; }
+          .summary-card { text-align: center; padding: 20px; background: #f9f9f9; border-radius: 8px; }
+          .summary-card .value { font-size: 36px; font-weight: bold; margin: 10px 0; }
+          .summary-card .label { color: #666; font-size: 14px; }
+          .score-a { color: #10b981; }
+          .score-b { color: #f59e0b; }
+          .score-c { color: #ef4444; }
+          .results { padding: 30px; }
+          .result-item { margin-bottom: 30px; border: 1px solid #eee; border-radius: 8px; overflow: hidden; }
+          .result-header { background: #f9f9f9; padding: 20px; border-bottom: 1px solid #eee; }
+          .result-header h2 { font-size: 18px; margin-bottom: 10px; word-break: break-all; }
+          .result-header .score { display: inline-block; padding: 5px 15px; border-radius: 20px; font-weight: bold; font-size: 14px; }
+          .result-body { padding: 20px; }
+          .finding { padding: 15px; margin-bottom: 10px; border-left: 4px solid; border-radius: 4px; background: #f9f9f9; }
+          .finding.critical { border-color: #ef4444; background: #fef2f2; }
+          .finding.warning { border-color: #f59e0b; background: #fffbeb; }
+          .finding.info { border-color: #3b82f6; background: #eff6ff; }
+          .finding-header { font-weight: bold; margin-bottom: 5px; }
+          .finding-issue { margin-bottom: 5px; }
+          .finding-fix { font-size: 14px; color: #666; font-style: italic; }
+          .badge { display: inline-block; padding: 3px 8px; border-radius: 3px; font-size: 12px; font-weight: bold; text-transform: uppercase; }
+          .badge.critical { background: #ef4444; color: white; }
+          .badge.warning { background: #f59e0b; color: white; }
+          .badge.info { background: #3b82f6; color: white; }
+          .no-findings { text-align: center; padding: 40px; color: #10b981; font-size: 18px; }
+          .footer { text-align: center; padding: 20px; color: #666; font-size: 14px; border-top: 1px solid #eee; }
+        </style>
+      </head>
+      <body>
+        <div class="container">
+          <div class="header">
+            <h1>🛡️ Hedra Security Report</h1>
+            <div class="meta">Generated: <%= Time.now.strftime('%Y-%m-%d %H:%M:%S %Z') %></div>
+          </div>
+      #{'    '}
+          <div class="summary">
+            <div class="summary-card">
+              <div class="label">URLs Scanned</div>
+              <div class="value"><%= results.length %></div>
+            </div>
+            <div class="summary-card">
+              <div class="label">Average Score</div>
+              <div class="value <%= score_class(avg_score) %>"><%= avg_score.round %>/100</div>
+            </div>
+            <div class="summary-card">
+              <div class="label">Total Findings</div>
+              <div class="value"><%= total_findings %></div>
+            </div>
+            <div class="summary-card">
+              <div class="label">Critical Issues</div>
+              <div class="value" style="color: #ef4444;"><%= critical_count %></div>
+            </div>
+          </div>
+      #{'    '}
+          <div class="results">
+            <% results.each do |result| %>
+              <div class="result-item">
+                <div class="result-header">
+                  <h2><%= result[:url] %></h2>
+                  <span class="score <%= score_class(result[:score]) %>">Score: <%= result[:score] %>/100</span>
+                  <span style="color: #666; margin-left: 15px; font-size: 14px;"><%= result[:timestamp] %></span>
+                </div>
+                <div class="result-body">
+                  <% if result[:findings].empty? %>
+                    <div class="no-findings">✓ All security headers properly configured</div>
+                  <% else %>
+                    <% result[:findings].each do |finding| %>
+                      <div class="finding <%= finding[:severity] %>">
+                        <div class="finding-header">
+                          <span class="badge <%= finding[:severity] %>"><%= finding[:severity] %></span>
+                          <%= finding[:header] %>
+                        </div>
+                        <div class="finding-issue"><%= finding[:issue] %></div>
+                        <% if finding[:recommended_fix] %>
+                          <div class="finding-fix">💡 <%= finding[:recommended_fix] %></div>
+                        <% end %>
+                      </div>
+                    <% end %>
+                  <% end %>
+                </div>
+              </div>
+            <% end %>
+          </div>
+      #{'    '}
+          <div class="footer">
+            Generated by Hedra <%= Hedra::VERSION %> | <a href="https://github.com/blackstack/hedra">GitHub</a>
+          </div>
+        </div>
+      </body>
+      </html>
+    HTML
+
+    def generate(results, output_file)
+      avg_score = results.sum { |r| r[:score] }.to_f / results.length
+      total_findings = results.sum { |r| r[:findings].length }
+      critical_count = results.sum { |r| r[:findings].count { |f| f[:severity] == :critical } }
+
+      html = ERB.new(TEMPLATE).result(binding)
+      File.write(output_file, html)
+    end
+
+    private
+
+    def score_class(score)
+      if score >= 80
+        'score-a'
+      elsif score >= 60
+        'score-b'
+      else
+        'score-c'
+      end
+    end
+  end
+end

data/lib/hedra/http_client.rb

@@ -8,16 +8,20 @@ module Hedra
     DEFAULT_USER_AGENT = "Hedra/#{VERSION} Security Header Analyzer".freeze
     MAX_RETRIES = 3
     RETRY_DELAY = 1
+    MAX_REDIRECTS = 10
 
-    def initialize(timeout: 10, proxy: nil, user_agent: nil, follow_redirects: false, verbose: false)
+    def initialize( # rubocop:disable Metrics/ParameterLists
+      timeout: 10, proxy: nil, user_agent: nil, follow_redirects: true, verbose: false, max_retries: MAX_RETRIES
+    )
       @timeout = timeout
       @proxy = proxy
       @user_agent = user_agent || DEFAULT_USER_AGENT
       @follow_redirects = follow_redirects
       @verbose = verbose
+      @max_retries = max_retries
     end
 
-    def get(url)
+    def get(url, redirect_count: 0)
       retries = 0
 
       begin
@@ -27,23 +31,28 @@ module Hedra
         response = client.get(url)
 
         if @follow_redirects && response.status.redirect?
+          raise NetworkError, "Too many redirects (#{MAX_REDIRECTS})" if redirect_count >= MAX_REDIRECTS
+
           location = response.headers['Location']
+          location = resolve_redirect_url(url, location)
           log "Following redirect to #{location}"
-          return get(location)
+          return get(location, redirect_count: redirect_count + 1)
         end
 
         raise NetworkError, "HTTP #{response.status}: #{response.status.reason}" unless response.status.success?
 
         log "Success: #{response.status}"
         response
-      rescue HTTP::Error, Errno::ECONNREFUSED, Errno::ETIMEDOUT => e
+      rescue HTTP::Error, Errno::ECONNREFUSED, Errno::ETIMEDOUT, Errno::EHOSTUNREACH => e
         retries += 1
-        raise NetworkError, "Failed after #{MAX_RETRIES} retries: #{e.message}" unless retries <= MAX_RETRIES
+        if retries <= @max_retries && retryable_error?(e)
+          delay = RETRY_DELAY * (2**(retries - 1))
+          log "Retry #{retries}/#{@max_retries} after #{delay}s: #{e.message}"
+          sleep delay
+          retry
+        end
 
-        delay = RETRY_DELAY * (2**(retries - 1))
-        log "Retry #{retries}/#{MAX_RETRIES} after #{delay}s: #{e.message}"
-        sleep delay
-        retry
+        raise NetworkError, "Failed after #{@max_retries} retries: #{e.message}"
       end
     end
 
@@ -62,6 +71,34 @@ module Hedra
       client
     end
 
+    def resolve_redirect_url(base_url, location)
+      # Handle relative redirects
+      return location if location.start_with?('http://', 'https://')
+
+      base_uri = URI.parse(base_url)
+      port_part = if base_uri.port && ![80, 443].include?(base_uri.port)
+                    ":#{base_uri.port}"
+                  else
+                    ''
+                  end
+
+      if location.start_with?('/')
+        "#{base_uri.scheme}://#{base_uri.host}#{port_part}#{location}"
+      else
+        # Relative to current path
+        base_path = base_uri.path.split('/')[0..-2].join('/')
+        "#{base_uri.scheme}://#{base_uri.host}#{port_part}#{base_path}/#{location}"
+      end
+    end
+
+    def retryable_error?(error)
+      # Retry on network errors, timeouts, but not on HTTP errors like 404
+      error.is_a?(HTTP::TimeoutError) ||
+        error.is_a?(Errno::ECONNREFUSED) ||
+        error.is_a?(Errno::ETIMEDOUT) ||
+        error.is_a?(Errno::EHOSTUNREACH)
+    end
+
     def log(message)
       puts "[HTTP] #{message}" if @verbose
     end

data/lib/hedra/progress_tracker.rb

@@ -0,0 +1,45 @@
+# frozen_string_literal: true
+
+module Hedra
+  # Track and display progress for batch operations
+  class ProgressTracker
+    def initialize(total, quiet: false)
+      @total = total
+      @current = 0
+      @quiet = quiet
+      @start_time = Time.now
+      @mutex = Mutex.new
+    end
+
+    def increment
+      @mutex.synchronize do
+        @current += 1
+        update_display unless @quiet
+      end
+    end
+
+    def finish
+      return if @quiet
+
+      puts "\n"
+      elapsed = Time.now - @start_time
+      puts "Completed #{@total} items in #{elapsed.round(2)}s"
+    end
+
+    private
+
+    def update_display
+      percentage = (@current.to_f / @total * 100).round(1)
+      bar_width = 40
+      filled = (bar_width * @current / @total).round
+      bar = ('█' * filled) + ('░' * (bar_width - filled))
+
+      elapsed = Time.now - @start_time
+      rate = @current / elapsed
+      eta = (@total - @current) / rate
+
+      print "\r[#{bar}] #{percentage}% (#{@current}/#{@total}) ETA: #{eta.round}s"
+      $stdout.flush
+    end
+  end
+end

data/lib/hedra/rate_limiter.rb

@@ -0,0 +1,60 @@
+# frozen_string_literal: true
+
+module Hedra
+  # Token bucket rate limiter
+  class RateLimiter
+    def initialize(rate_string)
+      @requests, period = parse_rate(rate_string)
+      @period = period_to_seconds(period)
+      @tokens = @requests
+      @last_refill = Time.now
+      @mutex = Mutex.new
+    end
+
+    def acquire
+      @mutex.synchronize do
+        refill_tokens
+
+        if @tokens >= 1
+          @tokens -= 1
+          return
+        end
+
+        # Wait until we have tokens
+        sleep_time = @period / @requests
+        sleep(sleep_time)
+        refill_tokens
+        @tokens -= 1
+      end
+    end
+
+    private
+
+    def parse_rate(rate_string)
+      # Format: "10/s", "100/m", "1000/h"
+      match = rate_string.match(%r{^(\d+)/([smh])$})
+      raise Error, "Invalid rate format: #{rate_string}" unless match
+
+      [match[1].to_i, match[2]]
+    end
+
+    def period_to_seconds(period)
+      case period
+      when 's' then 1
+      when 'm' then 60
+      when 'h' then 3600
+      else
+        raise Error, "Invalid period: #{period}"
+      end
+    end
+
+    def refill_tokens
+      now = Time.now
+      elapsed = now - @last_refill
+
+      tokens_to_add = (elapsed / @period) * @requests
+      @tokens = [@tokens + tokens_to_add, @requests].min
+      @last_refill = now
+    end
+  end
+end