bin/rbotdb

   1 #!/usr/bin/env ruby
   2 #-- vim:sw=2:et
   3 #++
   4 #
   5 # :title: RBot Registry Export, Import and Migration Script.
   6 #
   7 # You can use this script to,
   8 #   - export the rbot registry in a format that is platform/engine independent
   9 #   - import 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('export_%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 '     export: store rbot registry platform-independently in a file.'
  51   opt.separator '     import: restore rbot registry from such a file.'
  52   opt.separator ''
  53   opt.separator 'Options:'
  54
  55   opt.on('-t', '--type TYPE', TYPES, 'format to export/import. 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 export to/import from. Defaults to: %s.' % options[:dbfile]) do |dbfile|
  68     options[:dbfile] = dbfile
  69   end
  70
  71   opt.separator ''
  72 end
  73
  74 class ExportRegistry
  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 export
  84     listings = search
  85     puts listings.inspect
  86     puts 'Found registry types: bdb=%d tc=%d dbm=%d daybreak=%d sqlite=%d' % [
  87       listings[:bdb].length, listings[:tc].length,
  88       listings[:dbm].length, listings[:daybreak].length, listings[:sqlite].length
  89     ]
  90     if listings[@type].empty?
  91       puts 'No suitable registry found!'
  92       exit
  93     end
  94     puts 'Using registry type: %s' % @type
  95     read(listings[@type])
  96   end
  97
  98   def read(listing)
  99     print "~Reading... (this might take a moment)\r"
 100     data = {}
 101     count = 0
 102     listing.each do |file|
 103       begin
 104         data[file.key] = case @type
 105         when :tc
 106           read_tc(file)
 107         when :bdb
 108           read_bdb(file)
 109         when :dbm
 110           read_dbm(file)
 111         when :daybreak
 112           read_daybreak(file)
 113         when :sqlite
 114           read_sqlite(file)
 115         end
 116         count += data[file.key].length
 117       rescue
 118         puts 'ERROR: <%s> %s' % [$!.class, $!]
 119         puts $@.join("\n")
 120         puts 'Keep in mind that, even minor version differences of'
 121         puts 'Barkeley DB or Tokyocabinet make files unreadable. Use this'
 122         puts 'script on the exact same platform rbot was running!'
 123         exit
 124       end
 125     end
 126     puts 'Read %d registry files, with %d entries.' % [data.length, count]
 127     data
 128   end
 129
 130   def read_bdb(file)
 131     data = {}
 132     db = BDB::Hash.open(file.abs, nil, 'r')
 133     db.each do |key, value|
 134       data[key] = value
 135     end
 136     db.close
 137     data
 138   end
 139
 140   def read_tc(file)
 141     data = {}
 142     db = TokyoCabinet::BDB.new
 143     db.open(file.abs, TokyoCabinet::BDB::OREADER)
 144     db.each do |key, value|
 145       data[key] = value
 146     end
 147     db.close
 148     data
 149   end
 150
 151   def read_dbm(file)
 152     db = DBM.open(file.abs.gsub(/\.[^\.]+$/,''), 0666, DBM::READER)
 153     data = db.to_hash
 154     db.close
 155     data
 156   end
 157
 158   def read_daybreak(file)
 159     data = {}
 160     db = Daybreak::DB.new(file.abs)
 161     db.each do |key, value|
 162       data[key] = value
 163     end
 164     db.close
 165     data
 166   end
 167
 168   def read_sqlite(file)
 169     data = {}
 170     db = SQLite3::Database.new(file.abs)
 171     res = db.execute('SELECT key, value FROM data')
 172     res.each do |row|
 173       key, value = row
 174       data[key] = value
 175     end
 176     db.close
 177     data
 178   end
 179
 180   # searches in profile directory for existing registry formats
 181   def search
 182     {
 183       :bdb => list(get_registry, '*.db'),
 184       :tc => list(get_registry('_tc'), '*.tdb'),
 185       :dbm => list(get_registry('_dbm'), '*.*'),
 186       :daybreak => list(get_registry('_daybreak'), '*.db'),
 187       :sqlite => list(get_registry('_sqlite'), '*.db'),
 188     }
 189   end
 190
 191   def get_registry(suffix='')
 192     if @registry
 193       File.expand_path(@registry)
 194     else
 195       File.join(@profile, 'registry'+suffix)
 196     end
 197   end
 198
 199   class RegistryFile
 200     def initialize(folder, name)
 201       @folder = folder
 202       @name = name
 203       @key = name.gsub(/\.[^\.]+$/,'')
 204     end
 205     attr_reader :folder, :name, :key
 206     def abs
 207       File.expand_path(File.join(@folder, @name))
 208     end
 209     def ext
 210       File.extname(@name)
 211     end
 212   end
 213
 214   def list(folder, ext='*.db')
 215     return [] if not File.directory? folder
 216     Dir.chdir(folder) do
 217       Dir.glob(File.join('**', ext)).map do |name|
 218         RegistryFile.new(folder, name) if File.exists?(name)
 219       end
 220     end
 221   end
 222 end
 223
 224 class ImportRegistry
 225   def initialize(profile, type, registry)
 226     @profile = File.expand_path profile
 227     @registry = registry ? File.expand_path(registry) : nil
 228     @type = type
 229     puts 'Using type=%s profile=%s' % [@type, @profile]
 230   end
 231
 232   def import(data)
 233     puts 'Using registry type: %s' % @type
 234     folder = create_folder
 235     print "~Importing... (this might take a moment)\r"
 236     data.each do |file, hash|
 237       file = File.join(folder, file)
 238       create_subdir(file)
 239       case @type
 240       when :dbm
 241         write_dbm(file, hash)
 242       when :tc
 243         write_tc(file, hash)
 244       when :daybreak
 245         write_daybreak(file, hash)
 246       when :sqlite
 247         write_sqlite(file, hash)
 248       end
 249     end
 250     puts  'Import successful!                        '
 251   end
 252
 253   def write_dbm(file, data)
 254     db = DBM.open(file, 0666, DBM::WRCREAT)
 255     data.each_pair do |key, value|
 256       db[key] = value
 257     end
 258     db.close
 259   end
 260
 261   def write_tc(file, data)
 262     db = TokyoCabinet::BDB.new
 263     db.open(file + '.tdb',
 264           TokyoCabinet::BDB::OREADER |
 265           TokyoCabinet::BDB::OCREAT |
 266           TokyoCabinet::BDB::OWRITER)
 267     data.each_pair do |key, value|
 268       db[key] = value
 269     end
 270     db.optimize
 271     db.close
 272   end
 273
 274   def write_daybreak(file, data)
 275     db = Daybreak::DB.new(file + '.db')
 276     data.each_pair do |key, value|
 277       db[key] = value
 278     end
 279     db.close
 280   end
 281
 282   def write_sqlite(file, data)
 283     db = SQLite3::Database.new(file + '.db')
 284     db.execute('CREATE TABLE data (key string, value blob)')
 285     data.each_pair do |key, value|
 286       db.execute('INSERT INTO data VALUES (?, ?)',
 287             key, value)
 288     end
 289     db.close
 290   end
 291
 292   def create_folder
 293     if @registry
 294       folder = @registry
 295     else
 296       folder = File.join(@profile, 'registry_%s' % [@type.to_s])
 297     end
 298     Dir.mkdir(folder) unless File.directory?(folder)
 299     if File.directory?(folder) and Dir.glob(File.join(folder, '**')).select{|f|File.file? f}.length>0
 300       puts 'ERROR: Unable to import!'
 301       puts 'Import folder exists and is not empty: ' + folder
 302       exit
 303     end
 304     folder
 305   end
 306
 307   # used to create subregistry folders
 308   def create_subdir(path)
 309     dirs = File.dirname(path).split('/')
 310     dirs.length.times { |i|
 311       dir = dirs[0,i+1].join("/")+"/"
 312       unless File.exist?(dir)
 313         Dir.mkdir(dir)
 314       end
 315     }
 316   end
 317 end
 318
 319 opt_parser.parse!
 320 if options[:type].nil?
 321   puts 'Missing Argument: -t [type]'
 322   exit
 323 end
 324
 325 case ARGV[0]
 326 when 'export'
 327   if File.exists? options[:dbfile]
 328     puts 'Export file already exists.'
 329     exit
 330   end
 331
 332   reg = ExportRegistry.new(options[:profile], options[:type], options[:registry])
 333
 334   data = reg.export
 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 'import'
 344   unless File.exists? options[:dbfile]
 345     puts 'Import file does not exist.'
 346     exit
 347   end
 348
 349   reg = ImportRegistry.new(options[:profile], options[:type], options[:registry])
 350   data = Marshal.load File.read(options[:dbfile])
 351
 352   puts 'Read %d registry files from import file.' % data.length
 353   reg.import data
 354
 355 else
 356   puts opt_parser
 357
 358 end
 359