class Net::SSH::Authentication::KeyManager

def add(key_file)
  key_files.push(File.expand_path(key_file)).uniq!
  self
end

def add_key_data(key_data_)
  key_data.push(key_data_).uniq!
  self
end

def add_keycert(keycert_file)
  keycert_files.push(File.expand_path(keycert_file)).uniq!
  self
end

def add_keycert_data(keycert_data_)
  keycert_data.push(keycert_data_).uniq!
  self
end

def agent
  return unless use_agent?
  @agent ||= Agent.connect(logger, options[:agent_socket_factory], options[:identity_agent])
rescue AgentNotAvailable
  @use_agent = false
  nil
end

def clear!
  key_files.clear
  key_data.clear
  keycert_data.clear
  known_identities.clear
  self
end

def each_identity
  prepared_identities = prepare_identities_from_files + prepare_identities_from_data
  user_identities = load_identities(prepared_identities, false, true)
  if agent
    agent.identities.each do |key|
      corresponding_user_identity = user_identities.detect { |identity|
        identity[:public_key] && identity[:public_key].to_pem == key.to_pem
      }
      user_identities.delete(corresponding_user_identity) if corresponding_user_identity
      if !options[:keys_only] || corresponding_user_identity
        known_identities[key] = { from: :agent, identity: key }
        yield key
      end
    end
  end
  user_identities = load_identities(user_identities, !options[:non_interactive], false)
  user_identities.each do |identity|
    key = identity.delete(:public_key)
    known_identities[key] = identity
    yield key
  end
  known_identity_blobs = known_identities.keys.map(&:to_blob)
  keycerts.each do |keycert|
    next if known_identity_blobs.include?(keycert.to_blob)
    (_, corresponding_identity) = known_identities.detect { |public_key, _|
      public_key.to_pem == keycert.to_pem
    }
    if corresponding_identity
      known_identities[keycert] = corresponding_identity
      yield keycert
    end
  end
  self
end

def finish
  @agent.close if @agent
  @agent = nil
end

def initialize(logger, options = {})
  self.logger = logger
  @key_files = []
  @key_data = []
  @keycert_files = []
  @keycert_data = []
  @use_agent = options[:use_agent] != false
  @known_identities = {}
  @agent = nil
  @options = options
end

def keycerts
  keycert_files.map { |keycert_file| KeyFactory.load_public_key(keycert_file) } +
    keycert_data.map { |data| KeyFactory.load_data_public_key(data) }
end

def load_identities(identities, ask_passphrase, ignore_decryption_errors)
  identities.map do |identity|
    case identity[:load_from]
    when :pubkey_file
      key = KeyFactory.load_public_key(identity[:pubkey_file])
      { public_key: key, from: :file, file: identity[:privkey_file] }
    when :privkey_file
      private_key = KeyFactory.load_private_key(
        identity[:privkey_file], options[:passphrase], ask_passphrase, options[:password_prompt]
      )
      key = private_key.send(:public_key)
      { public_key: key, from: :file, file: identity[:privkey_file], key: private_key }
    when :data
      private_key = KeyFactory.load_data_private_key(
        identity[:data], options[:passphrase], ask_passphrase, "<key in memory>", options[:password_prompt]
      )
      key = private_key.send(:public_key)
      { public_key: key, from: :key_data, data: identity[:data], key: private_key }
    else
      identity
    end
  rescue OpenSSL::PKey::RSAError, OpenSSL::PKey::DSAError, OpenSSL::PKey::ECError, OpenSSL::PKey::PKeyError, ArgumentError => e
    if ignore_decryption_errors
      identity
    else
      process_identity_loading_error(identity, e)
      nil
    end
  rescue Exception => e
    process_identity_loading_error(identity, e)
    nil
  end.compact
end

def no_keys?
  key_files.empty? && key_data.empty?
end

def prepare_identities_from_data
  key_data.map do |data|
    { load_from: :data, data: data }
  end
end

def prepare_identities_from_files
  key_files.map do |file|
    if readable_file?(file)
      identity = {}
      cert_file = file + "-cert.pub"
      public_key_file = file + ".pub"
      if readable_file?(cert_file)
        identity[:load_from] = :pubkey_file
        identity[:pubkey_file] = cert_file
      elsif readable_file?(public_key_file)
        identity[:load_from] = :pubkey_file
        identity[:pubkey_file] = public_key_file
      else
        identity[:load_from] = :privkey_file
      end
      identity.merge(privkey_file: file)
    end
  end.compact
end

def process_identity_loading_error(identity, e)
  case identity[:load_from]
  when :pubkey_file
    error { "could not load public key file `#{identity[:pubkey_file]}': #{e.class} (#{e.message})" }
  when :privkey_file
    error { "could not load private key file `#{identity[:privkey_file]}': #{e.class} (#{e.message})" }
  else
    raise e
  end
end

def readable_file?(path)
  File.file?(path) && File.readable?(path)
end

def sign(identity, data, sig_alg = nil)
  info = known_identities[identity] or raise KeyManagerError, "the given identity is unknown to the key manager"
  if info[:key].nil? && info[:from] == :file
    begin
      info[:key] = KeyFactory.load_private_key(info[:file], options[:passphrase], !options[:non_interactive], options[:password_prompt])
    rescue OpenSSL::OpenSSLError, Exception => e
      raise KeyManagerError, "the given identity is known, but the private key could not be loaded: #{e.class} (#{e.message})"
    end
  end
  if info[:key]
    if sig_alg.nil?
      signed = info[:key].ssh_do_sign(data.to_s)
      sig_alg = identity.ssh_signature_type
    else
      signed = info[:key].ssh_do_sign(data.to_s, sig_alg)
    end
    return Net::SSH::Buffer.from(:string, sig_alg,
                                 :mstring, signed).to_s
  end
  if info[:from] == :agent
    raise KeyManagerError, "the agent is no longer available" unless agent
    case sig_alg
    when "rsa-sha2-512"
      return agent.sign(info[:identity], data.to_s, Net::SSH::Authentication::Agent::SSH_AGENT_RSA_SHA2_512)
    when "rsa-sha2-256"
      return agent.sign(info[:identity], data.to_s, Net::SSH::Authentication::Agent::SSH_AGENT_RSA_SHA2_256)
    else
      return agent.sign(info[:identity], data.to_s)
    end
  end
  raise KeyManagerError, "[BUG] can't determine identity origin (#{info.inspect})"
end

def use_agent=(use_agent)
  finish if !use_agent
  @use_agent = use_agent
end

def use_agent?
  @use_agent
end