arachni 0.2.2.1

data/lib/anemone/page_store.rb

@@ -0,0 +1,160 @@
+require 'forwardable'
+
+module Anemone
+  class PageStore
+    extend Forwardable
+
+    def_delegators :@storage, :keys, :values, :size, :each
+
+    def initialize(storage = {})
+      @storage = storage
+    end
+
+    # We typically index the hash with a URI,
+    # but convert it to a String for easier retrieval
+    def [](index)
+      @storage[index.to_s]
+    end
+
+    def []=(index, other)
+      @storage[index.to_s] = other
+    end
+
+    def delete(key)
+      @storage.delete key.to_s
+    end
+
+    def has_key?(key)
+      @storage.has_key? key.to_s
+    end
+
+    def each_value
+      each { |key, value| yield value }
+    end
+
+    def values
+      result = []
+      each { |key, value| result << value }
+      result
+    end
+
+    def touch_key(key)
+      self[key] = Page.new(key)
+    end
+
+    def touch_keys(keys)
+      @storage.merge! keys.inject({}) { |h, k| h[k.to_s] = Page.new(k); h }
+    end
+
+    # Does this PageStore contain the specified URL?
+    # HTTP and HTTPS versions of a URL are considered to be the same page.
+    def has_page?(url)
+      schemes = %w(http https)
+      if schemes.include? url.scheme
+        u = url.dup
+        return schemes.any? { |s| u.scheme = s; has_key?(u) }
+      end
+
+      has_key? url
+    end
+
+    #
+    # Use a breadth-first search to calculate the single-source
+    # shortest paths from *root* to all pages in the PageStore
+    #
+    def shortest_paths!(root)
+      root = URI(root) if root.is_a?(String)
+      raise "Root node not found" if !has_key?(root)
+
+      q = Queue.new
+
+      q.enq root
+      root_page = self[root]
+      root_page.depth = 0
+      root_page.visited = true
+      self[root] = root_page
+      while !q.empty?
+        page = self[q.deq]
+        page.links.each do |u|
+          begin
+            link = self[u]
+            next if link.nil? || !link.fetched? || link.visited
+
+            q << u unless link.redirect?
+            link.visited = true
+            link.depth = page.depth + 1
+            self[u] = link
+
+            if link.redirect?
+              u = link.redirect_to
+              redo
+            end
+          end
+        end
+      end
+
+      self
+    end
+
+    #
+    # Removes all Pages from storage where redirect? is true
+    #
+    def uniq!
+      each_value { |page| delete page.url if page.redirect? }
+      self
+    end
+
+    #
+    # If given a single URL (as a String or URI), returns an Array of Pages which link to that URL
+    # If given an Array of URLs, returns a Hash (URI => [Page, Page...]) of Pages linking to those URLs
+    #
+    def pages_linking_to(urls)
+      unless urls.is_a?(Array)
+        urls = [urls]
+        single = true
+      end
+
+      urls.map! do |url|
+        unless url.is_a?(URI)
+          URI(url) rescue nil
+        else
+          url
+        end
+      end
+      urls.compact
+
+      links = {}
+      urls.each { |url| links[url] = [] }
+      values.each do |page|
+        urls.each { |url| links[url] << page if page.links.include?(url) }
+      end
+
+      if single and !links.empty?
+        return links[urls.first]
+      else
+        return links
+      end
+    end
+
+    #
+    # If given a single URL (as a String or URI), returns an Array of URLs which link to that URL
+    # If given an Array of URLs, returns a Hash (URI => [URI, URI...]) of URLs linking to those URLs
+    #
+    def urls_linking_to(urls)
+      unless urls.is_a?(Array)
+        urls = [urls] unless urls.is_a?(Array)
+        single = true
+      end
+
+      links = pages_linking_to(urls)
+      links.each { |url, pages| links[url] = pages.map{|p| p.url} }
+
+      if single and !links.empty?
+        return links[urls.first]
+      else
+        return links
+      end
+    end
+
+  end
+end

data/lib/anemone/storage.rb

@@ -0,0 +1,34 @@
+module Anemone
+  module Storage
+
+    def self.Hash(*args)
+      hash = Hash.new(*args)
+      # add close method for compatibility with Storage::Base
+      class << hash; def close; end; end
+      hash
+    end
+
+    def self.PStore(*args)
+      require Arachni::Options.instance.dir['lib'] + 'anemone/storage/pstore'
+      self::PStore.new(*args)
+    end
+
+    def self.TokyoCabinet(file = 'anemone.tch')
+      require Arachni::Options.instance.dir['lib'] + 'anemone/storage/tokyo_cabinet'
+      self::TokyoCabinet.new(file)
+    end
+
+    def self.MongoDB(mongo_db = nil, collection_name = 'pages')
+      require Arachni::Options.instance.dir['lib'] + 'anemone/storage/mongodb'
+      mongo_db ||= Mongo::Connection.new.db('anemone')
+      raise "First argument must be an instance of Mongo::DB" unless mongo_db.is_a?(Mongo::DB)
+      self::MongoDB.new(mongo_db, collection_name)
+    end
+
+    def self.Redis(opts = {})
+      require Arachni::Options.instance.dir['lib'] + 'anemone/storage/redis'
+      self::Redis.new(opts)
+    end
+
+  end
+end

data/lib/anemone/storage/base.rb

@@ -0,0 +1,75 @@
+require Arachni::Options.instance.dir['lib'] + 'anemone/storage/exceptions'
+
+module Anemone
+  module Storage
+    class Base
+
+      def initialize(adapter)
+        @adap = adapter
+
+        # verify adapter conforms to this class's methods
+        methods.each do |method|
+          if !@adap.respond_to?(method.to_sym)
+            raise "Storage adapter must support method #{method}"
+          end
+        end
+      end
+
+      def [](key)
+        @adap[key]
+        rescue
+          puts key
+          raise RetrievalError, $!
+      end
+
+      def []=(key, value)
+        @adap[key] = value
+        rescue
+          raise InsertionError, $!
+      end
+
+      def delete(key)
+        @adap.delete(key)
+        rescue
+          raise DeletionError, $!
+      end
+
+      def each
+        @adap.each { |k, v| yield k, v }
+        rescue
+          raise GenericError, $!
+      end
+
+      def merge!(hash)
+        @adap.merge!(hash)
+        rescue
+          raise GenericError, $!
+      end
+
+      def close
+        @adap.close
+        rescue
+          raise CloseError, $!
+      end
+
+      def size
+        @adap.size
+        rescue
+          raise GenericError, $!
+      end
+
+      def keys
+        @adap.keys
+        rescue
+          raise GenericError, $!
+      end
+
+      def has_key?(key)
+        @adap.has_key?(key)
+        rescue
+          raise GenericError, $!
+      end
+
+    end
+  end
+end

data/lib/anemone/storage/exceptions.rb

@@ -0,0 +1,15 @@
+module Anemone
+  module Storage
+
+    class GenericError < Error; end;
+
+    class ConnectionError < Error; end
+
+    class RetrievalError < Error; end
+
+    class InsertionError < Error; end
+
+    class CloseError < Error; end
+
+  end
+end

data/lib/anemone/storage/mongodb.rb

@@ -0,0 +1,89 @@
+begin
+  require 'mongo'
+rescue LoadError
+  puts "You need the mongo gem to use Anemone::Storage::MongoDB"
+  exit
+end
+
+module Anemone
+  module Storage
+    class MongoDB 
+
+      BINARY_FIELDS = %w(body headers data)
+
+      def initialize(mongo_db, collection_name)
+        @db = mongo_db
+        @collection = @db[collection_name]
+        @collection.remove
+        @collection.create_index 'url'
+      end
+
+      def [](url)
+        if value = @collection.find_one('url' => url.to_s)
+          load_page(value)
+        end
+      end
+
+      def []=(url, page)
+        hash = page.to_hash
+        BINARY_FIELDS.each do |field|
+          hash[field] = BSON::Binary.new(hash[field]) unless hash[field].nil?
+        end
+        @collection.update(
+          {'url' => page.url.to_s},
+          hash,
+          :upsert => true
+        )
+      end
+
+      def delete(url)
+        page = self[url]
+        @collection.remove('url' => url.to_s)
+        page
+      end
+
+      def each
+        @collection.find do |cursor|
+          cursor.each do |doc|
+            page = load_page(doc)
+            yield page.url.to_s, page 
+          end
+        end
+      end
+
+      def merge!(hash)
+        hash.each { |key, value| self[key] = value }
+        self
+      end
+
+      def size
+        @collection.count
+      end
+
+      def keys
+        keys = []
+        self.each { |k, v| keys << k.to_s }
+        keys
+      end
+
+      def has_key?(url)
+        !!@collection.find_one('url' => url.to_s)
+      end
+
+      def close
+        @db.connection.close
+      end
+
+      private
+
+      def load_page(hash)
+        BINARY_FIELDS.each do |field|
+          hash[field] = hash[field].to_s
+        end
+        Page.from_hash(hash)
+      end
+
+    end
+  end
+end
+

data/lib/anemone/storage/pstore.rb

@@ -0,0 +1,50 @@
+require 'pstore'
+require 'forwardable'
+
+module Anemone
+  module Storage
+    class PStore
+      extend Forwardable
+
+      def_delegators :@keys, :has_key?, :keys, :size
+
+      def initialize(file)
+        File.delete(file) if File.exists?(file)
+        @store = ::PStore.new(file)
+        @keys = {}
+      end
+
+      def [](key)
+        @store.transaction { |s| s[key] }
+      end
+
+      def []=(key,value)
+        @keys[key] = nil
+        @store.transaction { |s| s[key] = value }
+      end
+
+      def delete(key)
+        @keys.delete(key)
+        @store.transaction { |s| s.delete key}
+      end
+
+      def each
+        @keys.each_key do |key|
+          value = nil
+          @store.transaction { |s| value = s[key] }
+          yield key, value
+        end
+      end
+
+      def merge!(hash)
+        @store.transaction do |s|
+          hash.each { |key, value| s[key] = value; @keys[key] = nil }
+        end
+        self
+      end
+
+      def close; end
+
+    end
+  end
+end

data/lib/anemone/storage/redis.rb

@@ -0,0 +1,90 @@
+require 'redis'
+
+module Anemone
+  module Storage
+    class Redis
+
+      MARSHAL_FIELDS = %w(links visited fetched)
+
+      def initialize(opts = {})
+        @redis = ::Redis.new(opts)
+        @key_prefix = opts[:key_prefix] || 'anemone'
+        keys.each { |key| delete(key) }
+      end
+
+      def [](key)
+        rkey = "#{@key_prefix}:pages:#{key.to_s}"
+        rget(rkey)
+      end
+
+      def []=(key, value)
+        rkey = "#{@key_prefix}:pages:#{key.to_s}"
+        hash = value.to_hash
+        MARSHAL_FIELDS.each do |field|
+          hash[field] = Marshal.dump(hash[field])
+        end
+        hash.each do |field, value|
+          @redis.hset(rkey, field, value)
+        end
+      end
+
+      def delete(key)
+        rkey = "#{@key_prefix}:pages:#{key.to_s}"
+        page = self[key]
+        @redis.del(rkey)
+        page
+      end
+
+      def each
+        rkeys = @redis.keys("#{@key_prefix}:pages:*")
+        rkeys.each do |rkey|
+          page = rget(rkey)
+          yield page.url.to_s, page
+        end
+      end
+
+      def merge!(hash)
+        hash.each { |key, value| self[key] = value }
+        self
+      end
+
+      def size
+        @redis.keys("#{@key_prefix}:pages:*").size
+      end
+
+      def keys
+        keys = []
+        self.each { |k, v| keys << k.to_s }
+        keys
+      end
+
+      def has_key?(key)
+        rkey = "#{@key_prefix}:pages:#{key.to_s}"
+        @redis.exists(rkey)
+      end
+
+      def close
+        @redis.quit
+      end
+
+      private
+
+      def load_value(hash)
+        MARSHAL_FIELDS.each do |field|
+          unless hash[field].nil? || hash[field] == ''
+            hash[field] = Marshal.load(hash[field]) 
+          end
+        end
+        Page.from_hash(hash)
+      end
+
+      def rget(rkey)
+        hash = @redis.hgetall(rkey)
+        if !!hash
+          load_value(hash)
+        end
+      end
+
+    end
+  end
+end