git.oblomov.eu Git - rbot/blob - 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     Dir.mkdir(@profile) unless File.directory?(@profile)
 297     if @registry
 298       folder = @registry
 299     else
 300       folder = File.join(@profile, 'registry_%s' % [@type.to_s])
 301     end
 302     Dir.mkdir(folder) unless File.directory?(folder)
 303     if File.directory?(folder) and Dir.glob(File.join(folder, '**')).select{|f|File.file? f}.length>0
 304       puts 'ERROR: Unable to restore!'
 305       puts 'Restore folder exists and is not empty: ' + folder
 306       exit
 307     end
 308     folder
 309   end
 310
 311   # used to create subregistry folders
 312   def create_subdir(path)
 313     dirs = File.dirname(path).split('/')
 314     dirs.length.times { |i|
 315       dir = dirs[0,i+1].join("/")+"/"
 316       unless File.exist?(dir)
 317         Dir.mkdir(dir)
 318       end
 319     }
 320   end
 321 end
 322
 323 opt_parser.parse!
 324 if ARGV.length > 0 and options[:type].nil?
 325   puts opt_parser
 326   puts 'Missing Argument: -t [type]'
 327   exit
 328 end
 329
 330 case ARGV[0]
 331 when 'backup'
 332   if File.exists? options[:dbfile]
 333     puts 'Backup file already exists.'
 334     exit
 335   end
 336
 337   reg = BackupRegistry.new(options[:profile], options[:type], options[:registry])
 338
 339   data = reg.backup
 340
 341   if not data.empty?
 342     File.open(options[:dbfile], 'w') do |f|
 343       f.write(Marshal.dump(data))
 344     end
 345     puts 'Written registry to ' + options[:dbfile]
 346   end
 347
 348 when 'restore'
 349   unless File.exists? options[:dbfile]
 350     puts 'Backup file does not exist.'
 351     exit
 352   end
 353
 354   reg = RestoreRegistry.new(options[:profile], options[:type], options[:registry])
 355   data = Marshal.load File.read(options[:dbfile])
 356
 357   puts 'Read %d registry files from backup file.' % data.length
 358   reg.restore data
 359
 360 else
 361   puts opt_parser
 362
 363 end
 364