sshsa.rb

#!/usr/bin/env ruby

# sshsa (SSH Switch Agent)
#
# sshsa manages multiple ssh-agents with different sets of keys. It
# lets you switch between them with one command. It was written on OS X
# with Keychain stuff in mind, so probably won't work unmodified on
# Linux.
#
# To use:
#
# ## Install this script
#
# Put this script (sshsa) somewhere in your $PATH and make it executable.
#
#
# ## Set up Profile
#
# Add the following lines to your bash .profile:
#
# sshsa() { eval $(sshsa "$@"); }
# ssh-add -D 2> /dev/null
#
# The first line is the wrapper for this script which will take the output
# and source it into your current shell, so we can inject new environment
# variables.
#
# The second line removes all keys from the default agent.
#
# (The way the ssh-agent and Keychain integration works is that when
# an ssh-agent is started, it will add all keys to it that have passphrases
# in the keychain. That would bypass everything we're going for here, so
# we want to neuter that agent when a shell is opened.)
#
#
# ## Create your identities
#
# Create a folder: ~/.ssh/identities
#
# In that folder, create a new folder for each identity you want. Then move
# your keys (public and private) into the appropriate identity folders.
#
# Note: MOVE your keys there, don't leave them in ~/.ssh
#
#
# ## Update Bash prompt
#
# The script exports an env variable called SSH_AGENT_IDENTITY_NAME when
# switching between identities, so that can be used to decorate your bash
# prompt.
#
#
# ## Try switching to one of your identities
#
# The first time you switch to an identity, it will tell you the keys were
# not added. Just use the supplied command to re-enter the passphrase and
# store in the Keychain. In the future, you won't need to reenter the
# password.
#

require 'fileutils'
require 'open3'

identity = ARGV[0]

identities_path = File.expand_path("~/.ssh/identities")
identity_path = File.join(identities_path, identity)
agents_path = File.expand_path("~/.ssh/agents")
agent_path = File.join(agents_path, identity)

def run_command_in_agent(agentfile, command)
  stdout, stderr, status = Open3.capture3("/usr/bin/env", "-i", "/bin/sh", "-c", ". #{agentfile} > /dev/null 2>/dev/null; #{command}")
  [ status, stdout, stderr ]
end

def keys_for_identity(identity_path)
  Dir[identity_path + "/*"].select do |identity_file|
    identity_file !~ /\.pub$/
  end
end

def create_agent_for_identity(identity_path, agent_path)
  `ssh-agent > #{agent_path}`

  keys = keys_for_identity(identity_path)
  keys.each do |key|
    status, stdout, stderr = run_command_in_agent(agent_path, "ssh-add -K #{key}")
    if !status.success?
      STDERR.puts "could not add identity #{key} to agent: #{status} #{stdout} #{stderr}"
      STDERR.puts "You probably need to add the passphrase to the keychain with the following:"
      STDERR.puts "  ssh-add -K #{key}"
      exit 1
    end
  end

  status, stdout, stderr = run_command_in_agent(agent_path, 'ssh-add -l')
  if !status.success?
    STDERR.puts "could not create new agent for #{identity_path}: #{stdout} #{stderr}"
    exit 1
  end

  keys = keys_for_identity(identity_path)

  # Pre-Sierra, the agent would get all keys from the keychain, so we need to
  # remove any unexpected ones.

  found = []
  stdout.each_line do |line|
    bits, fingerprint, path = line.split(/\s+/)
    if !keys.include?(path)
      run_command_in_agent(agent_path, "ssh-add -d #{path}")
    else
      found << path
    end
  end

  # now, make sure the keys we do want *are* there. It's possible they
  # aren't in the keychain, in which case the user will need to be prompted
  # for passphrases (when implemented).
  notfound = keys - found
  if notfound.length > 0
    STDERR.puts "Some keys were not added to the agent, likely because they're not in your keychain."
    STDERR.puts "Use the following command to add then to your agent and the keychain:"
    notfound.each do |key|
      STDERR.puts "  ssh-add -K #{key}"
    end
  end
end

unless File.exists?(identity_path)
  STDERR.puts "can't find #{identity} in #{identities_path}"
  exit 1
end

FileUtils.mkdir_p(agents_path)

valid_agent = false

if File.exists? agent_path
  # test that the existing agent works
  status, stdout, stderr = run_command_in_agent(agent_path, "ssh-add -l")

  # `ssh-add -l` return 1 when the agent is running but has no identities. :/
  valid_agent = status.success? || stdout =~ /The agent has no identities/
end

unless valid_agent
  create_agent_for_identity(identity_path, agent_path)
end

output = File.read(agent_path)
output += "SSH_AGENT_IDENTITY_NAME=#{identity}; export SSH_AGENT_IDENTITY_NAME;"

puts output