bin/rbotdb

   1 #!/usr/bin/env ruby
   2 #-- vim:sw=2:et
   3 #++
   4 #
   5 # :title: RBot Registry Backup, Restore and Migration Script.
   6 #
   7 # You can use this script to,
   8 #   - backup the rbot registry in a format that is platform/engine independent
   9 #   - restore these backups in supported formats (dbm, daybreak)
  10 #   - migrate old rbot registries bdb (ruby 1.8) and tokyocabinet.
  11 #
  12 # For more information, just execute the script without any arguments!
  13 #
  14 # Author:: apoc (Matthias Hecker) <apoc@geekosphere.org>
  15 # Copyright:: (C) 2014 Matthias Hecker
  16 # License:: GPLv3
  17
  18 begin; require 'rubygems'; rescue Exception; end
  19
  20 # load registry formats:
  21 begin; require 'bdb'; rescue Exception; end
  22 begin; require 'tokyocabinet'; rescue Exception; end
  23 begin; require 'dbm'; rescue Exception; end
  24 begin; require 'daybreak'; rescue Exception; end
  25 begin; require 'sqlite3'; rescue Exception; end
  26
  27 puts 'RBot Registry Backup/Restore/Migrate'
  28 puts '[%s]' % ['Ruby: ' + RUBY_VERSION,
  29                'DBM: ' + (DBM::VERSION rescue '-'),
  30                'BDB: ' + (BDB::VERSION rescue '-'),
  31                'TokyoCabinet: ' + (TokyoCabinet::VERSION rescue '-'),
  32                'Daybreak: ' + (Daybreak::VERSION rescue '-'),
  33                'SQLite: ' + (SQLite3::VERSION rescue '-'),
  34               ].join(' | ')
  35
  36 require 'date'
  37 require 'optparse'
  38
  39 TYPES = [:bdb, :tc, :dbm, :daybreak, :sqlite]
  40 options = {
  41   :profile => '~/.rbot',
  42   :registry => nil,
  43   :dbfile => './%s.rbot' % DateTime.now.strftime('backup_%Y-%m-%d_%H%M%S'),
  44   :type => nil
  45 }
  46 opt_parser = OptionParser.new do |opt|
  47   opt.banner = 'Usage: rbotdb COMMAND [OPTIONS]'
  48   opt.separator ''
  49   opt.separator 'Commands:'
  50   opt.separator '     backup: store rbot registry platform-independently in a file.'
  51   opt.separator '     restore: restore rbot registry from such a file.'
  52   opt.separator ''
  53   opt.separator 'Options:'
  54
  55   opt.on('-t', '--type TYPE', TYPES, 'format to backup/restore. Values: %s.' % [TYPES.join(', ')]) do |type|
  56     options[:type] = type
  57   end
  58
  59   opt.on('-p', '--profile [PROFILE]', 'rbot profile directory. Defaults to: %s.' % options[:profile]) do |profile|
  60     options[:profile] = profile
  61   end
  62
  63   opt.on('-r', '--registry [REGISTRY]', 'registry-path to read/write, Optional, defaults to: <PROFILE>/registry_<TYPE>.') do |profile|
  64     options[:registry] = profile
  65   end
  66
  67   opt.on('-f', '--file [DBFILE]', 'cross-platform file to backup to/restore from. Defaults to: %s.' % options[:dbfile]) do |dbfile|
  68     options[:dbfile] = dbfile
  69   end
  70
  71   opt.separator ''
  72 end
  73
  74 class BackupRegistry
  75   def initialize(profile, type, registry)
  76     @profile = File.expand_path profile
  77     @type = type
  78     @registry = registry
  79     puts 'Using type=%s profile=%s registry=%s' % [@type, @profile, @registry.inspect]
  80   end
  81
  82   # returns a hash with the complete registry data
  83   def backup
  84     listings = search
  85     puts 'Found registry types: bdb=%d tc=%d dbm=%d daybreak=%d sqlite=%d' % [
  86       listings[:bdb].length, listings[:tc].length,
  87       listings[:dbm].length, listings[:daybreak].length, listings[:sqlite].length
  88     ]
  89     if listings[@type].empty?
  90       puts 'No suitable registry found!'
  91       exit
  92     end
  93     puts 'Using registry type: %s' % @type
  94     read(listings[@type])
  95   end
  96
  97   def read(listing)
  98     print "~Reading... (this might take a moment)\r"
  99     data = {}
 100     count = 0
 101     listing.each do |file|
 102       begin
 103         data[file.key] = case @type
 104         when :tc
 105           read_tc(file)
 106         when :bdb
 107           read_bdb(file)
 108         when :dbm
 109           read_dbm(file)
 110         when :daybreak
 111           read_daybreak(file)
 112         when :sqlite
 113           read_sqlite(file)
 114         end
 115         count += data[file.key].length
 116       rescue
 117         puts 'ERROR: <%s> %s' % [$!.class, $!]
 118         puts $@.join("\n")
 119         puts 'Keep in mind that, even minor version differences of'
 120         puts 'Barkeley DB or Tokyocabinet make files unreadable. Use this'
 121         puts 'script on the exact same platform rbot was running!'
 122         exit
 123       end
 124     end
 125     puts 'Read %d registry files, with %d entries.' % [data.length, count]
 126     data
 127   end
 128
 129   def read_bdb(file)
 130     data = {}
 131     begin
 132       db = BDB::Hash.open(file.abs, nil, 'r')
 133     rescue BDB::Fatal
 134       db = BDB::Btree.open(file.abs, nil, 'r')
 135     end
 136     db.each do |key, value|
 137       data[key] = value
 138     end
 139     db.close
 140     data
 141   end
 142
 143   def read_tc(file)
 144     data = {}
 145     db = TokyoCabinet::BDB.new
 146     db.open(file.abs, TokyoCabinet::BDB::OREADER)
 147     db.each do |key, value|
 148       data[key] = value
 149     end
 150     db.close
 151     data
 152   end
 153
 154   def read_dbm(file)
 155     db = DBM.open(file.abs.gsub(/\.[^\.]+$/,''), 0666, DBM::READER)
 156     data = db.to_hash
 157     db.close
 158     data
 159   end
 160
 161   def read_daybreak(file)
 162     data = {}
 163     db = Daybreak::DB.new(file.abs)
 164     db.each do |key, value|
 165       data[key] = value
 166     end
 167     db.close
 168     data
 169   end
 170
 171   def read_sqlite(file)
 172     data = {}
 173     db = SQLite3::Database.new(file.abs)
 174     res = db.execute('SELECT key, value FROM data')
 175     res.each do |row|
 176       key, value = row
 177       data[key] = value
 178     end
 179     db.close
 180     data
 181   end
 182
 183   # searches in profile directory for existing registry formats
 184   def search
 185     {
 186       :bdb => list(get_registry, '*.db'),
 187       :tc => list(get_registry('_tc'), '*.tdb'),
 188       :dbm => list(get_registry('_dbm'), '*.*'),
 189       :daybreak => list(get_registry('_daybreak'), '*.db'),
 190       :sqlite => list(get_registry('_sqlite'), '*.db'),
 191     }
 192   end
 193
 194   def get_registry(suffix='')
 195     if @registry
 196       File.expand_path(@registry)
 197     else
 198       File.join(@profile, 'registry'+suffix)
 199     end
 200   end
 201
 202   class RegistryFile
 203     def initialize(folder, name)
 204       @folder = folder
 205       @name = name
 206       @key = name.gsub(/\.[^\.]+$/,'')
 207     end
 208     attr_reader :folder, :name, :key
 209     def abs
 210       File.expand_path(File.join(@folder, @name))
 211     end
 212     def ext
 213       File.extname(@name)
 214     end
 215   end
 216
 217   def list(folder, ext='*.db')
 218     return [] if not File.directory? folder
 219     Dir.chdir(folder) do
 220       Dir.glob(File.join('**', ext)).map do |name|
 221         RegistryFile.new(folder, name) if File.exists?(name)
 222       end
 223     end
 224   end
 225 end
 226
 227 class RestoreRegistry
 228   def initialize(profile, type, registry)
 229     @profile = File.expand_path profile
 230     @registry = registry ? File.expand_path(registry) : nil
 231     @type = type
 232     puts 'Using type=%s profile=%s' % [@type, @profile]
 233   end
 234
 235   def restore(data)
 236     puts 'Using registry type: %s' % @type
 237     folder = create_folder
 238     print "~Restoring... (this might take a moment)\r"
 239     data.each do |file, hash|
 240       file = File.join(folder, file)
 241       create_subdir(file)
 242       case @type
 243       when :dbm
 244         write_dbm(file, hash)
 245       when :tc
 246         write_tc(file, hash)
 247       when :daybreak
 248         write_daybreak(file, hash)
 249       when :sqlite
 250         write_sqlite(file, hash)
 251       end
 252     end
 253     puts  'Restore successful!                        '
 254   end
 255
 256   def write_dbm(file, data)
 257     db = DBM.open(file, 0666, DBM::WRCREAT)
 258     data.each_pair do |key, value|
 259       db[key] = value
 260     end
 261     db.close
 262   end
 263
 264   def write_tc(file, data)
 265     db = TokyoCabinet::BDB.new
 266     db.open(file + '.tdb',
 267           TokyoCabinet::BDB::OREADER |
 268           TokyoCabinet::BDB::OCREAT |
 269           TokyoCabinet::BDB::OWRITER)
 270     data.each_pair do |key, value|
 271       db[key] = value
 272     end
 273     db.optimize
 274     db.close
 275   end
 276
 277   def write_daybreak(file, data)
 278     db = Daybreak::DB.new(file + '.db')
 279     data.each_pair do |key, value|
 280       db[key] = value
 281     end
 282     db.close
 283   end
 284
 285   def write_sqlite(file, data)
 286     db = SQLite3::Database.new(file + '.db')
 287     db.execute('CREATE TABLE data (key PRIMARY_KEY, value)')
 288     data.each_pair do |key, value|
 289       db.execute('INSERT INTO data VALUES (?, ?)',
 290             key, value)
 291     end
 292     db.close
 293   end
 294
 295   def create_folder
 296     if @registry
 297       folder = @registry
 298     else
 299       folder = File.join(@profile, 'registry_%s' % [@type.to_s])
 300     end
 301     Dir.mkdir(folder) unless File.directory?(folder)
 302     if File.directory?(folder) and Dir.glob(File.join(folder, '**')).select{|f|File.file? f}.length>0
 303       puts 'ERROR: Unable to restore!'
 304       puts 'Restore folder exists and is not empty: ' + folder
 305       exit
 306     end
 307     folder
 308   end
 309
 310   # used to create subregistry folders
 311   def create_subdir(path)
 312     dirs = File.dirname(path).split('/')
 313     dirs.length.times { |i|
 314       dir = dirs[0,i+1].join("/")+"/"
 315       unless File.exist?(dir)
 316         Dir.mkdir(dir)
 317       end
 318     }
 319   end
 320 end
 321
 322 opt_parser.parse!
 323 if ARGV.length > 0 and options[:type].nil?
 324   puts opt_parser
 325   puts 'Missing Argument: -t [type]'
 326   exit
 327 end
 328
 329 case ARGV[0]
 330 when 'backup'
 331   if File.exists? options[:dbfile]
 332     puts 'Backup file already exists.'
 333     exit
 334   end
 335
 336   reg = BackupRegistry.new(options[:profile], options[:type], options[:registry])
 337
 338   data = reg.backup
 339
 340   if not data.empty?
 341     File.open(options[:dbfile], 'w') do |f|
 342       f.write(Marshal.dump(data))
 343     end
 344     puts 'Written registry to ' + options[:dbfile]
 345   end
 346
 347 when 'restore'
 348   unless File.exists? options[:dbfile]
 349     puts 'Backup file does not exist.'
 350     exit
 351   end
 352
 353   reg = RestoreRegistry.new(options[:profile], options[:type], options[:registry])
 354   data = Marshal.load File.read(options[:dbfile])
 355
 356   puts 'Read %d registry files from backup file.' % data.length
 357   reg.restore data
 358
 359 else
 360   puts opt_parser
 361
 362 end
 363