database.rb in ronin-1.0.0.pre1

- old
+ new

@@ -1,9 +1,9 @@
 #
 # Ronin - A Ruby platform for exploit development and security research.
 #
-# Copyright (c) 2006-2009 Hal Brodigan (postmodern.mod3 at gmail.com)
+# Copyright (c) 2006-2010 Hal Brodigan (postmodern.mod3 at gmail.com)
 #
 # This program is free software; you can redistribute it and/or modify
 # it under the terms of the GNU General Public License as published by
 # the Free Software Foundation; either version 2 of the License, or
 # (at your option) any later version.
@@ -17,77 +17,116 @@
 # along with this program; if not, write to the Free Software
 # Foundation, Inc., 51 Franklin St, Fifth Floor, Boston, MA  02110-1301  USA
 #
 
 require 'ronin/database/exceptions/invalid_config'
+require 'ronin/database/exceptions/unknown_repository'
+require 'ronin/database/migrations/migrations'
 require 'ronin/model'
-require 'ronin/arch'
-require 'ronin/os'
-require 'ronin/author'
-require 'ronin/license'
-require 'ronin/product'
 require 'ronin/config'
+require 'ronin/installation'
 
+require 'addressable/uri'
 require 'yaml'
 require 'dm-core'
 
 module Ronin
+  #
+  # Manages the {Database} configuration and the defined repositories.
+  # Also provides a simple wrapper around DataMapper, for initializing,
+  # auto-upgrading and querying {Database} repositories.
+  #
   module Database
     # Database configuration file
     CONFIG_FILE = File.join(Config::PATH,'database.yml')
 
     # Database log file
     DEFAULT_LOG_PATH = File.join(Config::PATH,'database.log')
 
     # Database log level
     DEFAULT_LOG_LEVEL = :info
 
-    # Default configuration of the database
-    DEFAULT_CONFIG = "sqlite3://" + File.join(Config::PATH,'database.sqlite3')
+    # Default database repository
+    DEFAULT_REPOSITORY = Addressable::URI.new(
+      :scheme => 'sqlite3',
+      :path => File.join(Config::PATH,'database.sqlite3')
+    )
 
+    @repositories = {}
+    @log = nil
+
     #
-    # Returns the Database configuration that is stored in the
-    # +CONFIG_FILE+. Defaults to +DEFAULT_CONFIG+ if +CONFIG_FILE+ does not
-    # exist.
+    # Returns the Database repositories to use.
     #
+    # @return [Hash{Symbol => Addressable::URI}]
+    #   The database repository names and URIs.
+    #
     # @raise [InvalidConfig]
-    #   The config file did not contain a YAML Hash or String.
+    #   The config file did not contain a YAML Hash.
     #
-    def Database.config
-      unless (class_variable_defined?('@@ronin_database_config'))
-        @@ronin_database_config = DEFAULT_CONFIG
+    # @since 1.0.0
+    #
+    def Database.repositories
+      if @repositories.empty?
+        @repositories[:default] = DEFAULT_REPOSITORY
 
         if File.file?(CONFIG_FILE)
-          conf = YAML.load(CONFIG_FILE)
+          conf = YAML.load_file(CONFIG_FILE)
 
-          unless (conf.kind_of?(Hash) || conf.kind_of?(String))
-            raise(InvalidConfig,"#{CONFIG_FILE} must contain either a Hash or a String",caller)
+          unless conf.kind_of?(Hash)
+            raise(InvalidConfig,"#{CONFIG_FILE} must contain a YAML Hash of repositories")
           end
 
-          @@ronin_database_config = conf
+          conf.each do |name,uri|
+            @repositories[name.to_sym] = Addressable::URI.parse(uri)
+          end
         end
       end
 
-      return @@ronin_database_config ||= DEFAULT_CONFIG
+      return @repositories
     end
 
     #
-    # Sets the Database configuration.
+    # Determines if the Database provides a specific repository.
     #
-    # @param [String, Hash] configuration
-    #   The DataMapper configuration to set the Ronin::Database with.
+    # @param [String, Symbol] name
+    #   Name of the repository.
     #
-    def Database.config=(configuration)
-      @@ronin_database_config = configuration
+    # @return [Boolean]
+    #   Specifies if the Database provides the repository.
+    #
+    # @since 1.0.0
+    #
+    def Database.repository?(name)
+      Database.repositories.has_key?(name.to_sym)
     end
 
     #
-    # @return [DataMapper::Logger, nil]
-    #   The current Database log.
+    # Saves the Database configuration to `CONFIG_FILE`.
     #
-    def Database.log
-      @@ronin_database_log ||= nil
+    # @yield []
+    #   If a block is given, it will be called before the database
+    #   configuration is saved.
+    #
+    # @return [true]
+    #
+    # @since 1.0.0
+    #
+    def Database.save
+      yield if block_given?
+
+      File.open(CONFIG_FILE,'w') do |file|
+        hash = {}
+        
+        Database.repositories.each do |name,value|
+          hash[name.to_s] = value.to_s
+        end
+
+        YAML.dump(hash,file)
+      end
+
+      return true
     end
 
     #
     # Setup the Database log.
     #
@@ -100,64 +139,153 @@
     # @option options [IO] :stream
     #   The stream to use for the log.
     #
     # @option options [Symbol] :level
     #   The level of messages to log.
-    #   May be either +:fatal+, +:error+, +:warn+, +:info+ or +:debug+.
+    #   May be either `:fatal`, `:error`, `:warn`, `:info` or `:debug`.
     #
-    # @return [DataMapper::Logger]
-    #   The new Database log.
+    # @return [true]
+    #   Specifies that the log has been setup.
     #
     def Database.setup_log(options={})
       path = (options[:path] || DEFAULT_LOG_PATH)
       stream = (options[:stream] || File.new(path,'w+'))
       level = (options[:level] || DEFAULT_LOG_LEVEL)
 
-      return @@ronin_database_log = DataMapper::Logger.new(stream,level)
+      @log = DataMapper::Logger.new(stream,level)
+      return true
     end
 
     #
+    # Determines if a specific database repository is setup.
+    #
+    # @param [Symbol] name
+    #   The database repository name.
+    #
     # @return [Boolean]
     #   Specifies wether or not the Database is setup.
     #
-    def Database.setup?
-      repository = DataMapper.repository(Model::REPOSITORY_NAME)
+    def Database.setup?(name=:default)
+      repository = DataMapper.repository(name)
 
       return repository.class.adapters.has_key?(repository.name)
     end
 
     #
-    # Updates the Database, by running auto-upgrades, but only if the
-    # Database is already setup.
+    # Upgrades the Database, by running migrations for a given
+    # ronin library, but only if the Database has been setup.
     #
-    # @yield []
-    #   The block to call before the Database is updated.
+    # @return [Boolean]
+    #   Specifies whether the Database was migrated or is currently
+    #   not setup.
     #
-    def Database.update!(&block)
-      block.call if block
-
-      DataMapper.auto_upgrade!(Model::REPOSITORY_NAME) if Database.setup?
-      return nil
+    def Database.upgrade!
+      if Database.setup?
+        Migrations.migrate_up!
+      else
+        false
+      end
     end
 
     #
     # Sets up the Database.
     #
-    # @param [String, Hash] configuration
-    #   The DataMapper configuration to use to setup the Database.
-    # 
     # @yield []
     #   The block to call after the Database has been setup, but before
     #   it is updated.
     #
-    def Database.setup(configuration=Database.config,&block)
+    # @see Database.upgrade!
+    #
+    def Database.setup(&block)
       # setup the database log
-      Database.setup_log unless Database.log
+      Database.setup_log unless @log
 
-      # setup the database repository
-      DataMapper.setup(Model::REPOSITORY_NAME, configuration)
+      # setup the database repositories
+      Database.repositories.each do |name,uri|
+        DataMapper.setup(name,uri)
+      end
 
-      Database.update!(&block)
+      # auto-upgrade the database repository
+      Database.upgrade!(&block)
+    end
+
+    #
+    # Performs Database transactions within a given repository.
+    #
+    # @param [String, Symbol] name
+    #   The name of the repository to access.
+    #
+    # @return [DataMapper::Repository]
+    #   The Database repository.
+    #
+    # @raise [UnknownRepository]
+    #   The specified Database repository is unknown.
+    #
+    # @since 1.0.0
+    #
+    def Database.repository(name,&block)
+      name = name.to_sym
+
+      unless Database.repository?(name)
+        raise(UnknownRepository,"unknown database repository #{name}")
+      end
+
+      return DataMapper.repository(name,&block)
+    end
+
+    #
+    # Clears the Database, by running destructive auto-migrations.
+    #
+    # @param [String, Symbol] name
+    #   The name of the Database repository to clear.
+    #
+    # @yield []
+    #   If a block is given, it will be called after the Database
+    #   repository has been cleared.
+    #
+    # @return [nil]
+    #
+    # @raise [UnknownRepository]
+    #   The specified Database repository is unknown.
+    #
+    # @since 1.0.0
+    #
+    def Database.clear(name)
+      name = name.to_sym
+
+      unless Database.repository?(name)
+        raise(UnknownRepository,"unknown database repository #{name}")
+      end
+
+      DataMapper.auto_migrate!(name)
+
+      yield if block_given?
       return nil
+    end
+
+    #
+    # Performs Database transactions in each of the Database
+    # repositories.
+    #
+    # @yield []
+    #   The given block will be ran within the context of each Database
+    #   repository.
+    #
+    # @return [Array]
+    #   The results from each database transaction.
+    #
+    # @since 1.0.0
+    #
+    def Database.map(&block)
+      results = []
+
+      Database.repositories.each_key do |name|
+        DataMapper.repository(name) do
+          result = block.call()
+          results << result unless result.nil?
+        end
+      end
+
+      return results
     end
   end
 end