class AwsMfa

# File lib/aws_mfa.rb, line 9
def initialize
  validate_aws_installed
  @aws_config_dir = find_aws_config
end

# File lib/aws_mfa.rb, line 170
def execute
  profile = 'default'
  profile_index = ARGV.index('--profile')
  if (!profile_index.nil?)
    profile = ARGV.delete_at(profile_index + 1)
    ARGV.delete_at(profile_index)
  end
  arn = load_arn(profile)
  credentials = load_credentials(arn, profile)
  if ARGV.empty?
    print_credentials(credentials)
  else
    export_credentials(credentials)
    exec(*ARGV)
  end
end

# File lib/aws_mfa.rb, line 163
def export_credentials(credentials)
  ENV['AWS_SECRET_ACCESS_KEY'] = credentials['SecretAccessKey']
  ENV['AWS_ACCESS_KEY_ID'] = credentials['AccessKeyId']
  ENV['AWS_SESSION_TOKEN'] = credentials['SessionToken']
  ENV['AWS_SECURITY_TOKEN'] = credentials['SessionToken']
end

# File lib/aws_mfa.rb, line 18
def find_aws_config
  if ENV['AWS_CREDENTIAL_FILE']
    aws_config_file = ENV['AWS_CREDENTIAL_FILE']
    aws_config_dir = File.dirname(aws_config_file)
  else
    aws_config_dir = File.join(ENV['HOME'], '.aws')
    aws_config_file = File.join(aws_config_dir, 'config')
  end

  unless File.readable?(aws_config_file)
    raise Errors::ConfigurationNotFound, 'Aws configuration not found. You must run `aws configure`'
  end

  aws_config_dir
end

# File lib/aws_mfa.rb, line 46
def load_arn(profile='default')
  arn_file_name = 'mfa_device'
  if (! profile.eql? 'default')
    arn_file_name = "#{profile}_#{arn_file_name}"
  end
  arn_file = File.join(aws_config_dir, arn_file_name)

  if File.readable?(arn_file)
    arn = load_arn_from_file(arn_file)
  else
    arn = load_arn_from_aws(profile)
    write_arn_to_file(arn_file, arn)
  end

  arn
end

# File lib/aws_mfa.rb, line 67
def load_arn_from_aws(profile='default')
  puts 'Fetching MFA devices for your account...'
  if mfa_devices(profile).any?
    mfa_devices(profile).first.fetch('SerialNumber')
  else
    raise Errors::DeviceNotFound, 'No MFA devices were found for your account'
  end
end

# File lib/aws_mfa.rb, line 63
def load_arn_from_file(arn_file)
  File.read(arn_file)
end

# File lib/aws_mfa.rb, line 93
def load_credentials(arn, profile='default')
  credentials_file_name = 'mfa_credentials'
  if (! profile.eql? 'default')
    credentials_file_name = "#{profile}_#{credentials_file_name}"
  end
  credentials_file  = File.join(aws_config_dir, credentials_file_name)

  if File.readable?(credentials_file) && token_not_expired?(credentials_file)
    credentials = load_credentials_from_file(credentials_file)
  else
    credentials = load_credentials_from_aws(arn, profile)
    write_credentials_to_file(credentials_file, credentials)
  end

  JSON.parse(credentials).fetch('Credentials')
end

# File lib/aws_mfa.rb, line 114
def load_credentials_from_aws(arn, profile='default')
  code = request_code_from_user
  unset_environment
  credentials_command = "aws --profile #{profile} --output json sts get-session-token --serial-number #{arn} --token-code #{code}"
  result = AwsMfa::ShellCommand.new(credentials_command).call
  if result.succeeded?
    result.output
  else
    raise Errors::InvalidCode, 'There was a problem validating the MFA code with AWS'
  end
end

# File lib/aws_mfa.rb, line 110
def load_credentials_from_file(credentials_file)
  File.read(credentials_file)
end

# File lib/aws_mfa.rb, line 76
def mfa_devices(profile='default')
  @mfa_devices ||= begin
    list_mfa_devices_command = "aws --profile #{profile} --output json iam list-mfa-devices"
    result = AwsMfa::ShellCommand.new(list_mfa_devices_command).call
    if result.succeeded?
      JSON.parse(result.output).fetch('MFADevices')
    else
      raise Errors::Error, 'There was a problem fetching MFA devices from AWS'
    end
  end
end

# File lib/aws_mfa.rb, line 156
def print_credentials(credentials)
  puts "export AWS_SECRET_ACCESS_KEY='#{credentials['SecretAccessKey']}'"
  puts "export AWS_ACCESS_KEY_ID='#{credentials['AccessKeyId']}'"
  puts "export AWS_SESSION_TOKEN='#{credentials['SessionToken']}'"
  puts "export AWS_SECURITY_TOKEN='#{credentials['SessionToken']}'"
end

# File lib/aws_mfa.rb, line 130
def request_code_from_user
  puts 'Enter the 6-digit code from your MFA device:'
  code = $stdin.gets.chomp
  raise Errors::InvalidCode, 'That is an invalid MFA code' unless code =~ /^\d{6}$/
  code
end

# File lib/aws_mfa.rb, line 144
def token_not_expired?(credentials_file)
  # default is 12 hours
  expiration_period = 12 * 60 * 60
  mtime = File.stat(credentials_file).mtime
  now = Time.new
  if now - mtime < expiration_period
    true
  else
    false
  end
end

# File lib/aws_mfa.rb, line 137
def unset_environment
  ENV.delete('AWS_SECRET_ACCESS_KEY')
  ENV.delete('AWS_ACCESS_KEY_ID')
  ENV.delete('AWS_SESSION_TOKEN')
  ENV.delete('AWS_SECURITY_TOKEN')
end

# File lib/aws_mfa.rb, line 14
def validate_aws_installed
  raise Errors::CommandNotFound, 'Could not find the aws command' unless which('aws')
end

# File lib/aws_mfa.rb, line 35
def which(cmd)
  exts = ENV['PATHEXT'] ? ENV['PATHEXT'].split(';') : ['']
  ENV['PATH'].split(File::PATH_SEPARATOR).each do |path|
    exts.each { |ext|
      exe = File.join(path, "#{cmd}#{ext}")
      return exe if File.executable? exe
    }
  end
  return nil
end

# File lib/aws_mfa.rb, line 88
def write_arn_to_file(arn_file, arn)
  File.open(arn_file, 'w') { |f| f.print arn }
  puts "Using MFA device #{arn}. To change this in the future edit #{arn_file}."
end

# File lib/aws_mfa.rb, line 126
def write_credentials_to_file(credentials_file, credentials)
  File.open(credentials_file, 'w') { |f| f.print credentials }
end