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