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 :daybreak
 231         write_daybreak(file, hash)
 232       end
 233     end
 234     puts  'Import successful!                        '
 235   end
 236
 237   def write_dbm(file, data)
 238     db = DBM.open(file, 0666, DBM::WRCREAT)
 239     data.each_pair do |key, value|
 240       db[key] = value
 241     end
 242     db.close
 243   end
 244
 245   def write_daybreak(file, data)
 246     db = Daybreak::DB.new(file + '.db')
 247     data.each_pair do |key, value|
 248       db[key] = value
 249     end
 250     db.close
 251   end
 252
 253   def create_folder
 254     if @registry
 255       folder = @registry
 256     else
 257       folder = File.join(@profile, 'registry_%s' % [@type.to_s])
 258     end
 259     Dir.mkdir(folder) unless File.directory?(folder)
 260     if File.directory?(folder) and Dir.glob(File.join(folder, '**')).select{|f|File.file? f}.length>0
 261       puts 'ERROR: Unable to import!'
 262       puts 'Import folder exists and is not empty: ' + folder
 263       exit
 264     end
 265     folder
 266   end
 267
 268   # used to create subregistry folders
 269   def create_subdir(path)
 270     dirs = File.dirname(path).split('/')
 271     dirs.length.times { |i|
 272       dir = dirs[0,i+1].join("/")+"/"
 273       unless File.exist?(dir)
 274         Dir.mkdir(dir)
 275       end
 276     }
 277   end
 278 end
 279
 280 opt_parser.parse!
 281
 282 case ARGV[0]
 283 when 'export'
 284   if File.exists? options[:dbfile]
 285     puts 'Export file already exists.'
 286     exit
 287   end
 288
 289   reg = ExportRegistry.new(options[:profile], options[:type], options[:registry])
 290
 291   data = reg.export
 292
 293   if not data.empty?
 294     File.open(options[:dbfile], 'w') do |f|
 295       f.write(Marshal.dump(data))
 296     end
 297     puts 'Written registry to ' + options[:dbfile]
 298   end
 299
 300 when 'import'
 301   unless File.exists? options[:dbfile]
 302     puts 'Import file does not exist.'
 303     exit
 304   end
 305
 306   reg = ImportRegistry.new(options[:profile], options[:type], options[:registry])
 307   data = Marshal.load File.read(options[:dbfile])
 308
 309   puts 'Read %d registry files from import file.' % data.length
 310   reg.import data
 311
 312 else
 313   puts opt_parser
 314
 315 end
 316