diff --git a/utils/vmp_utils b/utils/vmp_utils new file mode 100755 index 0000000..613abd7 --- /dev/null +++ b/utils/vmp_utils @@ -0,0 +1,86 @@ +#!/usr/bin/env ruby +# frozen_string_literal: true + +require 'date' +require 'json' +require 'redis' +require 'thor' + +class VMPUtils < Thor + desc 'export_tokens', 'Export VMPooler tokens' + long_desc <<-LONGDESC + vmp_utils export_tokens will export the tokens stored in redis as JSON + to the specified file. By default it will export all tokens used within + the last year. This can be overridden by setting the "since-date" option. + LONGDESC + option :since_date, default: (DateTime.now - 365).iso8601 + option :redis_host, required: true + option :redis_password + option :redis_port, type: :numeric, default: 6379 + option :output_file, required: true + def export_tokens + redis = Redis.new(host: options[:redis_host], port: options[:redis_port], password: options[:redis_password]) + begin + redis.ping + puts "connection successful to #{options[:redis_host]}" + tokens_to_export = {} + date_filter = DateTime.parse(options[:since_date]) + redis.scan_each(match: 'vmpooler__token__*') do |key| + print('.') + key_data = redis.hgetall(key) + export = false + if key_data.key?('last') && date_filter < DateTime.parse(key_data['last']) + export = true + elsif date_filter < DateTime.parse(key_data['created']) + export = true + end + tokens_to_export[key] = key_data if export + end + puts '' + File.open(options[:output_file], 'w') do |output| + output.puts JSON.pretty_generate(tokens_to_export) + end + rescue Redis::CannotConnectError => e + raise Thor::Error, e + rescue Redis::CommandError => e + raise Thor::Error, e + end + end + + desc 'import_tokens', 'Import previously exported VMPooler tokens' + long_desc <<-LONGDESC + vmp_utils import_tokens will load tokens from a previous export into the + specified redis instance. It expects a JSON file like the one generated by + export_tokens to be provided as input. + LONGDESC + option :redis_host, required: true + option :redis_password + option :redis_port, type: :numeric, default: 6379 + option :input_file, required: true + def import_tokens + raise Thor::Error, "input-file '#{options[:input_file]}' doesn't exist" unless File.exist?(options[:input_file]) + + begin + data_file = File.read(options[:input_file]) + tokens_to_import = JSON.parse(data_file) + rescue JSON::ParserError + raise Thor::Error, "Unable to parse please verify that #{options[:input_file]}. Please make sure it is a valid JSON file" + end + + redis = Redis.new(host: options[:redis_host], port: options[:redis_port], password: options[:redis_password]) + begin + redis.ping + puts "connection successful to #{options[:redis_host]}" + tokens_to_import.each do |key, value| + puts "importing token for #{value['user']}" + redis.hset(key, value) + end + rescue Redis::CannotConnectError => e + raise Thor::Error, e + rescue Redis::CommandError => e + raise Thor::Error, e + end + end +end + +VMPUtils.start(ARGV) diff --git a/vmpooler.gemspec b/vmpooler.gemspec index 5a3dd1e..20e35fe 100644 --- a/vmpooler.gemspec +++ b/vmpooler.gemspec @@ -46,5 +46,6 @@ Gem::Specification.new do |s| s.add_development_dependency 'rspec', '>= 3.2' s.add_development_dependency 'rubocop', '~> 1.1.0' s.add_development_dependency 'simplecov', '>= 0.11.2' + s.add_development_dependency 'thor', '~> 1.0', '>= 1.0.1' s.add_development_dependency 'yarjuf', '>= 2.0' end