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     db = BDB::Hash.open(file.abs, nil, 'r')
 132     db.each do |key, value|
 133       data[key] = value
 134     end
 135     db.close
 136     data
 137   end
 138
 139   def read_tc(file)
 140     data = {}
 141     db = TokyoCabinet::BDB.new
 142     db.open(file.abs, TokyoCabinet::BDB::OREADER)
 143     db.each do |key, value|
 144       data[key] = value
 145     end
 146     db.close
 147     data
 148   end
 149
 150   def read_dbm(file)
 151     db = DBM.open(file.abs.gsub(/\.[^\.]+$/,''), 0666, DBM::READER)
 152     data = db.to_hash
 153     db.close
 154     data
 155   end
 156
 157   def read_daybreak(file)
 158     data = {}
 159     db = Daybreak::DB.new(file.abs)
 160     db.each do |key, value|
 161       data[key] = value
 162     end
 163     db.close
 164     data
 165   end
 166
 167   def read_sqlite(file)
 168     data = {}
 169     db = SQLite3::Database.new(file.abs)
 170     res = db.execute('SELECT key, value FROM data')
 171     res.each do |row|
 172       key, value = row
 173       data[key] = value
 174     end
 175     db.close
 176     data
 177   end
 178
 179   # searches in profile directory for existing registry formats
 180   def search
 181     {
 182       :bdb => list(get_registry, '*.db'),
 183       :tc => list(get_registry('_tc'), '*.tdb'),
 184       :dbm => list(get_registry('_dbm'), '*.*'),
 185       :daybreak => list(get_registry('_daybreak'), '*.db'),
 186       :sqlite => list(get_registry('_sqlite'), '*.db'),
 187     }
 188   end
 189
 190   def get_registry(suffix='')
 191     if @registry
 192       File.expand_path(@registry)
 193     else
 194       File.join(@profile, 'registry'+suffix)
 195     end
 196   end
 197
 198   class RegistryFile
 199     def initialize(folder, name)
 200       @folder = folder
 201       @name = name
 202       @key = name.gsub(/\.[^\.]+$/,'')
 203     end
 204     attr_reader :folder, :name, :key
 205     def abs
 206       File.expand_path(File.join(@folder, @name))
 207     end
 208     def ext
 209       File.extname(@name)
 210     end
 211   end
 212
 213   def list(folder, ext='*.db')
 214     return [] if not File.directory? folder
 215     Dir.chdir(folder) do
 216       Dir.glob(File.join('**', ext)).map do |name|
 217         RegistryFile.new(folder, name) if File.exists?(name)
 218       end
 219     end
 220   end
 221 end
 222
 223 class RestoreRegistry
 224   def initialize(profile, type, registry)
 225     @profile = File.expand_path profile
 226     @registry = registry ? File.expand_path(registry) : nil
 227     @type = type
 228     puts 'Using type=%s profile=%s' % [@type, @profile]
 229   end
 230
 231   def restore(data)
 232     puts 'Using registry type: %s' % @type
 233     folder = create_folder
 234     print "~Restoring... (this might take a moment)\r"
 235     data.each do |file, hash|
 236       file = File.join(folder, file)
 237       create_subdir(file)
 238       case @type
 239       when :dbm
 240         write_dbm(file, hash)
 241       when :tc
 242         write_tc(file, hash)
 243       when :daybreak
 244         write_daybreak(file, hash)
 245       when :sqlite
 246         write_sqlite(file, hash)
 247       end
 248     end
 249     puts  'Restore successful!                        '
 250   end
 251
 252   def write_dbm(file, data)
 253     db = DBM.open(file, 0666, DBM::WRCREAT)
 254     data.each_pair do |key, value|
 255       db[key] = value
 256     end
 257     db.close
 258   end
 259
 260   def write_tc(file, data)
 261     db = TokyoCabinet::BDB.new
 262     db.open(file + '.tdb',
 263           TokyoCabinet::BDB::OREADER |
 264           TokyoCabinet::BDB::OCREAT |
 265           TokyoCabinet::BDB::OWRITER)
 266     data.each_pair do |key, value|
 267       db[key] = value
 268     end
 269     db.optimize
 270     db.close
 271   end
 272
 273   def write_daybreak(file, data)
 274     db = Daybreak::DB.new(file + '.db')
 275     data.each_pair do |key, value|
 276       db[key] = value
 277     end
 278     db.close
 279   end
 280
 281   def write_sqlite(file, data)
 282     db = SQLite3::Database.new(file + '.db')
 283     db.execute('CREATE TABLE data (key string, value blob)')
 284     data.each_pair do |key, value|
 285       db.execute('INSERT INTO data VALUES (?, ?)',
 286             key, value)
 287     end
 288     db.close
 289   end
 290
 291   def create_folder
 292     if @registry
 293       folder = @registry
 294     else
 295       folder = File.join(@profile, 'registry_%s' % [@type.to_s])
 296     end
 297     Dir.mkdir(folder) unless File.directory?(folder)
 298     if File.directory?(folder) and Dir.glob(File.join(folder, '**')).select{|f|File.file? f}.length>0
 299       puts 'ERROR: Unable to restore!'
 300       puts 'Restore folder exists and is not empty: ' + folder
 301       exit
 302     end
 303     folder
 304   end
 305
 306   # used to create subregistry folders
 307   def create_subdir(path)
 308     dirs = File.dirname(path).split('/')
 309     dirs.length.times { |i|
 310       dir = dirs[0,i+1].join("/")+"/"
 311       unless File.exist?(dir)
 312         Dir.mkdir(dir)
 313       end
 314     }
 315   end
 316 end
 317
 318 opt_parser.parse!
 319 if ARGV.length > 0 and options[:type].nil?
 320   puts opt_parser
 321   puts 'Missing Argument: -t [type]'
 322   exit
 323 end
 324
 325 case ARGV[0]
 326 when 'backup'
 327   if File.exists? options[:dbfile]
 328     puts 'Backup file already exists.'
 329     exit
 330   end
 331
 332   reg = BackupRegistry.new(options[:profile], options[:type], options[:registry])
 333
 334   data = reg.backup
 335
 336   if not data.empty?
 337     File.open(options[:dbfile], 'w') do |f|
 338       f.write(Marshal.dump(data))
 339     end
 340     puts 'Written registry to ' + options[:dbfile]
 341   end
 342
 343 when 'restore'
 344   unless File.exists? options[:dbfile]
 345     puts 'Backup file does not exist.'
 346     exit
 347   end
 348
 349   reg = RestoreRegistry.new(options[:profile], options[:type], options[:registry])
 350   data = Marshal.load File.read(options[:dbfile])
 351
 352   puts 'Read %d registry files from backup file.' % data.length
 353   reg.restore data
 354
 355 else
 356   puts opt_parser
 357
 358 end
 359