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