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