[registry] improved export/import/migrate script.

1 files changed, 245 insertions, 134 deletions

diff --git a/bin/rbotdb b/bin/rbotdb
index 4b0279ff..5f7491a7 100755
--- a/bin/rbotdb
+++ b/bin/rbotdb
@@ -2,194 +2,305 @@
 #-- vim:sw=2:et
 #++
 #
-# :title: Registry import/export and migration.
+# :title: RBot Registry Export, Import and Migration Script.
 #
 # You can use this script to,
-#   - backup the rbot registry in a format that is platform independent
-#   - restore backups
+#   - export the rbot registry in a format that is platform/engine independent
+#   - import these backups in supported formats (dbm, daybreak)
 #   - migrate old rbot registries bdb (ruby 1.8) and tokyocabinet.
 #
+# For more information, just execute the script without any arguments!
+#
 # Author:: apoc (Matthias Hecker) <apoc@geekosphere.org>
 # Copyright:: (C) 2014 Matthias Hecker
 # License:: GPLv3
 
-begin; require 'rubygems'; rescue Exception; puts "[#{$!}]"; end
-begin; require 'dbm'; rescue Exception; puts "[#{$!}]"; end
-begin; require 'bdb'; rescue Exception; puts "[#{$!}]"; end
-begin; require 'tokyocabinet'; rescue Exception; puts "[#{$!}]"; end
-
-puts 'RBot registry backup/import script.'
-puts 'Ruby: %s | DBM: %s | BDB: %s | TC: %s' % [RUBY_VERSION, 
-                           (DBM::VERSION rescue '-'), 
-                           (BDB::VERSION rescue '-'),
-                           (TokyoCabinet::VERSION rescue '-')]
-
-if ARGV.length > 3 or ARGV.length < 2
-  puts """
-  Usage rbotdb [backup|restore] <BACKUP/RESTORE-FILE> [<PROFILE>] 
-
-  Examples:
-    rbotdb backup ~/rbot_db_backup.yaml
-    rbotdb backup ~/rbot_db_backup.yaml.gz ~/.rbot_two
-    rbotdb restore ~/rbot_db_backup.yaml
-  """
-  exit
-end
+begin; require 'rubygems'; rescue Exception; end
 
-mode = ARGV[0] if %w{backup restore}.include? ARGV[0]
-file = File.expand_path ARGV[1]
-profile = File.expand_path(ARGV[2] ? ARGV[2] : '~/.rbot')
-$last_error = ''
+# old registry formats:
+begin; require 'bdb'; rescue Exception; end
+begin; require 'tokyocabinet'; rescue Exception; end
 
-class Backup
-  class RegistryFile
-    def initialize(registry, path)
-      @registry = registry
-      @path = path
-    end
-    def path
-      @path
-    end
-    def abs
-      File.expand_path(File.join(@registry, @path))
-    end
-    def ext
-      File.extname(@path)
-    end
-    def valid?
-      File.file?(abs) and %w{.db .tdb}.include? ext
-    end
+# new formats:
+begin; require 'dbm'; rescue Exception; end
+begin; require 'daybreak'; rescue Exception; end
+
+puts 'RBot Registry Backup/Restore/Migrate'
+puts '[%s]' % ['Ruby: ' + RUBY_VERSION,
+               'DBM: ' + (DBM::VERSION rescue '-'),
+               'BDB: ' + (BDB::VERSION rescue '-'),
+               'TokyoCabinet: ' + (TokyoCabinet::VERSION rescue '-'),
+               'Daybreak: ' + (Daybreak::VERSION rescue '-'),
+              ].join(' | ')
+
+require 'date'
+require 'optparse'
+
+TYPES = [:bdb, :tc, :dbm, :daybreak, :auto]
+options = {
+  :profile => '~/.rbot',
+  :dbfile => './%s.rbot' % DateTime.now.strftime('export_%Y-%m-%d_%H%M%S'),
+  :type => :auto
+}
+opt_parser = OptionParser.new do |opt|
+  opt.banner = 'Usage: rbotdb COMMAND [OPTIONS]'
+  opt.separator ''
+  opt.separator 'Commands:'
+  opt.separator '     export: store rbot registry platform-independently in a file.'
+  opt.separator '     import: restore rbot registry from such a file.'
+  opt.separator ''
+  opt.separator 'Options:'
+
+  opt.on('-p', '--profile [PROFILE]', 'rbot profile directory. Defaults to: %s.' % options[:profile]) do |profile|
+    options[:profile] = profile
+  end
+
+  opt.on('-f', '--file [DBFILE]', 'cross-platform file to export to/import from. Defaults to: %s.' % options[:dbfile]) do |dbfile|
+    options[:dbfile] = dbfile
   end
 
-  def initialize(profile)
-    @profile = profile
-    @registry = File.join(profile, './registry')
+  opt.on('-t', '--type TYPE', TYPES, 'format to export/import. Values: %s. Defaults to %s.' % [TYPES.join(', '), options[:type]]) do |type|
+    options[:type] = type
   end
 
-  # list all database files:
-  def list
-    return nil if not File.directory? @registry
-    Dir.chdir @registry
-    Dir.glob(File.join('**', '*')).map do |name|
-      RegistryFile.new(@registry, name)
+  opt.separator ''
+end
+
+class ExportRegistry
+  def initialize(profile, type)
+    @profile = File.expand_path profile
+    @type = type
+    puts 'Using type=%s profile=%s' % [@type, @profile]
+  end
+
+  # returns a hash with the complete registry data
+  def export
+    listings = search
+    puts 'Found registry types: bdb=%d tc=%d dbm=%d daybreak=%d' % [
+      listings[:bdb].length, listings[:tc].length,
+      listings[:dbm].length, listings[:daybreak].length
+    ]
+    if @type == :auto
+      @type = :bdb if listings[:bdb].length > 0
+      @type = :tc if listings[:tc].length > 0
+      @type = :dbm if listings[:dbm].length > 0
+      @type = :daybreak if listings[:daybreak].length > 0
+    end
+    if @type == :auto or listings[@type].empty?
+      puts 'No suitable registry found!'
+      return
     end
+    puts 'Using registry type: %s' % @type
+    read(listings[@type])
   end
 
-  def load(ext='.db')
-    @data = {}
-    list.each do |file|
-      next unless file.ext == ext
-      db = loadDBM(file)
-      db = loadBDB(file) if not db
-      db = loadTC(file) if not db
-      if not db
-        puts 'ERROR: unable to load db: %s, last error: %s' % [file.abs, $last_error]
-      else
-        puts 'Loaded: %s [%d values]' % [file.abs, db.length]
-        @data[file.path] = db
+  def read(listing)
+    data = {}
+    count = 0
+    listing.each do |file|
+      begin
+        data[file.key] = case @type
+        when :tc
+          read_tc(file)
+        when :bdb
+          read_bdb(file)
+        when :dbm
+          read_dbm(file)
+        when :daybreak
+          read_daybreak(file)
+        end
+        count += data[file.key].length
+      rescue
+        puts 'ERROR: <%s> %s' % [$!.class, $!]
+        puts $@.join("\n")
+        puts 'Keep in mind that, even minor version differences of'
+        puts 'Barkeley DB or Tokyocabinet make files unreadable. Use this'
+        puts 'script on the exact same platform rbot was running!'
+        exit
       end
     end
+    puts 'Read %d registry files, with %d entries.' % [data.length, count]
+    data
   end
 
-  def write(file)
-    File.open(file, 'w') do |f|
-      f.write(Marshal.dump(@data))
+  def read_bdb(file)
+    data = {}
+    db = BDB::Hash.open(file.abs, nil, 'r')
+    db.each do |key, value|
+      data[key] = value
     end
+    db.close
+    data
   end
 
-  private
-
-  def loadDBM(file)
-    path = file.abs
-    begin
-      dbm = DBM.open(path.gsub(/\.[^\.]+$/,''), 0666, DBM::READER)
-      data = dbm.to_hash
-      dbm.close
-    rescue
-      $last_error = "[%s]\n%s" % [$!, $@.join("\n")]
+  def read_tc(file)
+    data = {}
+    db = TokyoCabinet::BDB.new
+    db.open(file.abs, TokyoCabinet::BDB::OREADER)
+    db.each do |key, value|
+      data[key] = value
     end
+    db.close
     data
   end
 
-  def loadBDB(file)
-    path = file.abs
-    begin
-      db = BDB::Hash.open(path, nil, 'r')
-      data = {}
-      db.each do |key, value|
-        data[key] = value
-      end
-      db.close
-    rescue
-      $last_error = "[%s]\n%s" % [$!, $@.join("\n")]
-    end
+  def read_dbm(file)
+    db = DBM.open(file.abs.gsub(/\.[^\.]+$/,''), 0666, DBM::READER)
+    data = db.to_hash
+    db.close
     data
   end
 
-  def loadTC(file)
-    path = file.abs
-    begin
-      db = TokyoCabinet::BDB.new
-      db.open(path, TokyoCabinet::BDB::OREADER)
-      data = {}
-      db.each do |key, value|
-        data[key] = value
-      end
-      db.close
-    rescue
-      $last_error = "[%s]\n%s" % [$!, $@.join("\n")]
+  def read_daybreak(file)
+    data = {}
+    db = Daybreak::DB.new(file.abs)
+    db.each do |key, value|
+      data[key] = value
     end
+    db.close
     data
   end
-end
 
-puts 'mode = ' + mode
-puts 'profile= ' + profile
-puts 'file = ' + file
+  # searches in profile directory for existing registry formats
+  def search
+    {
+      :tc => list(File.join(@profile, 'registry'), '*.tdb'),
+      :bdb => list(File.join(@profile, 'registry'), '*.db'),
+      :dbm => list(File.join(@profile, 'registry_dbm'), '*.*'),
+      :daybreak => list(File.join(@profile, 'registry_daybreak'), '*.db'),
+    }
+  end
 
-if mode == 'backup'
+  class RegistryFile
+    def initialize(folder, name)
+      @folder = folder
+      @name = name
+      @key = name.gsub(/\.[^\.]+$/,'')
+    end
+    attr_reader :folder, :name, :key
+    def abs
+      File.expand_path(File.join(@folder, @name))
+    end
+    def ext
+      File.extname(@name)
+    end
+  end
 
-  backup = Backup.new(profile)
+  def list(folder, ext='*.db')
+    return [] if not File.directory? folder
+    Dir.chdir(folder) do
+      Dir.glob(File.join('**', ext)).map do |name|
+        RegistryFile.new(folder, name) if File.exists?(name)
+      end
+    end
+  end
+end
 
-  if File.exists? file
-    puts 'ERROR! Backup file already exists!'
-    exit
+class ImportRegistry
+  def initialize(profile, type)
+    @profile = File.expand_path profile
+    @type = (type == :auto) ? :dbm : type
+    puts 'Using type=%s profile=%s' % [@type, @profile]
   end
 
-  backup.load '.tdb'
-  backup.write(file)
+  def import(data)
+    puts 'Using registry type: %s' % @type
+    folder = create_folder
+    data.each do |file, hash|
+      file = File.join(folder, file)
+      create_subdir(file)
+      case @type
+      when :dbm
+        write_dbm(file, hash)
+      when :daybreak
+        write_daybreak(file, hash)
+      end
+    end
+    puts 'Import completed.'
+  end
 
-else
+  def write_dbm(file, data)
+    db = DBM.open(file, 0666, DBM::WRCREAT)
+    data.each_pair do |key, value|
+      db[key] = value
+    end
+    db.close
+  end
 
-  registry = File.join(profile, './registry')
+  def write_daybreak(file, data)
+    db = Daybreak::DB.new(file + '.db')
+    data.each_pair do |key, value|
+      db[key] = value
+    end
+    db.close
+  end
 
-  data = Marshal.load File.read(file)
-  data.each_pair do |path, db|
-    path = File.expand_path(File.join(registry, path))
+  def create_folder
+    folder = @profile
+    case @type
+    when :dbm
+      folder = File.join(folder, 'registry_dbm')
+    when :daybreak
+      folder = File.join(folder, 'registry_daybreak')
+    else
+      puts 'ERROR: Unsupported import type: %s' % @type
+      exit
+    end
+    Dir.mkdir(folder) unless File.directory?(folder)
+    if File.directory?(folder) and Dir.glob(File.join(folder, '**')).select{|f|File.file? f}.length>0
+      puts 'ERROR: Unable to import!'
+      puts 'Import folder exists and is not empty: ' + folder
+      exit
+    end
+    folder
+  end
 
-    # create directories:
-    dirs = File.dirname(path).split("/")
+  # used to create subregistry folders
+  def create_subdir(path)
+    dirs = File.dirname(path).split('/')
     dirs.length.times { |i|
       dir = dirs[0,i+1].join("/")+"/"
       unless File.exist?(dir)
-        puts 'create subdir:'+dir
         Dir.mkdir(dir)
       end
     }
-    path.gsub!(/\.([^\.]+)$/,'')
+  end
+end
 
-    if File.exists? path+'.db' or File.exists? path+'.tdb'
-      raise 'error, unable to restore to existing db'
-    end
+opt_parser.parse!
+
+case ARGV[0]
+when 'export'
+  if File.exists? options[:dbfile]
+    puts 'Export file already exists.'
+    exit 
+  end
+
+  reg = ExportRegistry.new(options[:profile], options[:type])
 
-    puts 'restore to: '+path
-    dbm = DBM.open(path, 0666, DBM::WRCREAT)
-    db.each_pair do |key, value|
-      dbm[key] = value
+  data = reg.export
+
+  if not data.empty?
+    File.open(options[:dbfile], 'w') do |f|
+      f.write(Marshal.dump(data))
     end
-    dbm.close
+    puts 'Written registry to ' + options[:dbfile]
   end
 
-end
+when 'import'
+  unless File.exists? options[:dbfile]
+    puts 'Import file does not exist.'
+    exit 
+  end
+
+  reg = ImportRegistry.new(options[:profile], options[:type])
+  data = Marshal.load File.read(options[:dbfile])
 
+  puts 'Read %d registry files from import file.' % data.length
+  reg.import data
+
+else
+  puts opt_parser
+
+end