#!/usr/bin/perl -w
# -*-perl-*-
#
# ssh-find-agent:
#
#   behaves like ssh-agent but doesn't start a
#     new one if it can find an old one.
#
#   by default, it cleans up stale agent sockets too.
#
#   it doesn't do anything to clean up stale agent processes,
#     but if you always use ssh-find-agent instead of ssh-agent
#     you should be ok?
#
# usage:   eval `ssh-find-agent` - should work with sh or csh
#
#   if that doesn't work, try setting $verbose=1.
#   sorry, no -v flag yet.
#
use Socket;
use strict;

#######################################################################
#
#              you may wish to change these settings
#
my $username = $ENV{USER};
#
# this is the identity i usually use, so i prefer to find an
#  agent that already has it loaded.
my $magic_id = "";
if ($username eq "xxx") { $magic_id = "xxx\@xxx.xxx"; };
#
# this may vary with your installation of ssh
my $sockdir = "/tmp/ssh-$username";
#
#  $cleanup=1 removes old sockets from $sockdir if no corresponding
#             agent process can be found.  should be a cmdline option.
my $cleanup = 1;
#
#  $verbose=1 prints out debugging info.  should be a cmdline option.
my $verbose = 1;
#
#######################################################################

# globals
my %agent_socks = ();
my %agent_ids = ();
my %agent_pids = ();

# cleanup - used when retrying
sub zap_arrays {
    my $k;
    foreach $k (keys %agent_socks) { delete $agent_socks{$k}; }
    foreach $k (keys %agent_ids) { delete $agent_ids{$k}; }
    foreach $k (keys %agent_pids) { delete $agent_pids{$k}; }
}

# print if $verbose is set.
# wrap up the print command in an echo to make it through a shell eval
sub vprint {
    $verbose || return;
    # XXX should quote @_
    print "echo \"@_\" ;\n";
}

# build the list of available unix-domain sockets
sub build_socket_list {
    my $file;
    foreach $file ( glob("$sockdir/*") ) {
	my $sockid = "";
	if ($file =~ /ssh-([0-9]+)-agent$/) { $sockid = $1; }
	else { warn "$file doesn't match pattern"; next; }
	if (!-S $file) { warn "$file is not a socket"; next; }
	
	$agent_socks{$sockid} = $file;
	$agent_ids{$sockid} = [];
    }
}

# connect to each one in turn and build a list of
# the identities it has
sub build_identity_list {
    my $sockid;
    foreach $sockid (keys %agent_socks) {
	my $file = $agent_socks{$sockid};
	
	$ENV{SSH_AUTH_SOCK} = $file;
	# pid doesn't seem to be needed here
	# $ENV("SSH_AGENT_PID") = "";
	
	if (! open(AGENT_IDS, "ssh-add -l 2>&1 |")) {
	    warn "couldn't run ssh-add -l for agent $sockid: $!";
	    delete $agent_socks{$sockid};
	} else {
	    while (<AGENT_IDS>) {
		if (/Could not connect to authentication server/) {
		    vprint "dead agent socket $file\n";
		    if ($cleanup) {
			unlink $file || warn "unlinking $file failed: $!";
		    }
		    delete $agent_socks{$sockid};
		    delete $agent_ids{$sockid};
		} elsif (/[0-9]+\s[0-9]+\s[0-9]+\s(.+)/) {
		    my $idname = $1;
		    my $idlist = $agent_ids{$sockid};
		    # print "  identity: $1\n";
		    push @$idlist, $idname;
		} elsif (/The agent has no identities./) {
		} else {
		    warn "unknown ssh-agent output line: $_";
		}
	    }
	    if (!close(AGENT_IDS)) {
		warn "error running ssh-add -l for agent $sockid: $!";
		delete $agent_socks{$sockid};
	    }
	}

	delete $ENV{SSH_AUTH_SOCK};
	#    delete $ENV{SSH_AGENT_PID};
    }
}

# find process ids for all known agents.
sub build_pid_list {
    my $sockid;
    foreach $sockid (keys %agent_ids) {
	my $pid = $sockid;
	if (kill(0, $pid) == 1) {
	    $agent_pids{$sockid} = $pid;
	} elsif (kill(0, 1+$pid) == 1) {
	    # ugh.  who knows why, but on linux the agent
	    # forks itself or something
	    $agent_pids{$sockid} = 1 + $pid;
	}
    }
}

# print out all the good agents found
sub choose_agent {
    my $sockid;
    my $best_sockid = "";
    my $best_count = -1;

    # XXX check pid too.

    foreach $sockid (keys %agent_ids) {
	exists($agent_pids{$sockid}) || next;

	my $pid = $agent_pids{$sockid};
	my $idlist = $agent_ids{$sockid};
	my $ok = 0;
	my $count = 0;

	vprint "found agent: pid $pid\n";

	my $id;
	foreach $id (@$idlist) {
	    if ($id eq $magic_id) { $ok = 1; }
	    vprint "  identity $id\n";
	    $count++;
	}

	if ($count > $best_count) {
	    if ($ok || $best_count < 1) {
		$best_count = $count;
		$best_sockid = $sockid;
	    }
	}
    }
    return $best_sockid;
}

# based on the SHELL environment variable, print out similar
# goop to what ssh-agent does, so you can just call
# eval `ssh-find-agent` and it should work.
sub print_shell_stuff {
    my ($sockid) = @_;
    my $file = "$sockdir/ssh-$sockid-agent";

    # ugh
    my $pid = 1 + $sockid;

    if ($ENV{SHELL} =~ /csh$/) {
	print "setenv SSH_AUTH_SOCK $file\n";
	print "setenv SSH_AGENT_PID $pid\n";
	# print "echo Agent pid $pid\n";
    } else {
	print "SSH_AUTH_SOCK=$file; export SSH_AUTH_SOCK;\n";
	print "SSH_AGENT_PID=$pid; export SSH_AGENT_PID;\n";
	# print "echo Agent pid $pid;\n";
    }
}

# build a complete list of agents and find the best
sub find_agent {
    zap_arrays;
    build_socket_list;
    build_identity_list;
    build_pid_list;

    return &choose_agent;
}

sub find_or_start {
    my $agent_id = &find_agent;
    if ($agent_id eq "") {
	# couldn't find one, fork off a new one.
	# we send the output to /dev/null and look for it using find_agent
	#  again - we could just feed the output of ssh-agent back to the
	#  user instead.
	system("ssh-agent >& /dev/null") && die "unable to start new agent: $!";
	$agent_id = &find_agent;
    }

    # something is wrong.
    $agent_id ne "" || die "still unable to find agent after starting one";

    if ($verbose) {
	my $idlist = $agent_ids{$agent_id};
	my $id;
	my $idnames = "";
	foreach $id (@$idlist) {
	    $idnames .= "<$id> "
	}
	vprint "chose agent $agent_id\n";
	vprint " identities: $idnames\n";
    }

    return $agent_id;
}

#######################################################################

my $agent_id = &find_or_start;
print_shell_stuff $agent_id;