vmpooler/utils/vmp_utils
Gene Liverman dd1b8167a9
(DIO-2186) Add vmp_utils cli tool
This adds a thor-based cli under the utils folder. Its initial function
is to make it easier to migrate tokens from one redis instance to
another.
2021-09-16 16:02:02 -04:00

86 lines
3 KiB
Ruby
Executable file

#!/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)