#!/usr/bin/perl

$^W=1; #-w

use strict;
use Fcntl ':mode';

#act rather like ssh-agent(1), except:
#automagically do the right things if it reasonably appears we're already
#running an ssh-agent, namely:
#don't start another one, unless we're requesting to run a command under it
#if we're not starting a command under it and not (-k) killing it, output
#proper syntax to set the environment,
#if we're killing it, do so and output proper syntax to adjust environment

#This is how ssh-agent (from ssh-agent.c in
#http://ftp.debian.org/debian/pool/main/o/openssh/openssh_4.3p2.orig.tar.gz
#) checks whether or not it's running under a C-Shell or compatible shell:
#	shell = getenv("SHELL");
#	if (shell != NULL && strncmp(shell + strlen(shell) - 3, "csh", 3) == 0)
#		c_flag = 1;

#save our arguments, we may need to restore them identically later
my @original_argv=@ARGV;

#full pathnames of these - we may need them later
my $fuser='/bin/fuser';
my $ssh_agent='/usr/bin/ssh-agent';
my $ssh_add='/usr/bin/ssh-add';

#start figuring out how we'll likely need to behave

#Bourne-like or C-Shell-like environment flavor commands?
#like ssh-agent, we default to Bourne-like ($c_flag false)
my $c_flag=0;

#first we look at SHELL in the environment, using same logic as ssh-agent.c
$ENV{SHELL} =~ /csh$/o and $c_flag=1;

#we'll use our own simplistic argument processing (no module dependencies)
#ssh-agent [-a bind_address] [-c | -s] [-t life] [-d] [command [args ...]]
#ssh-agent [-c | -s] -k

my $kill=0; #did we get -k?
my $got_non_kill_args=0; #did we get something incompatible with -k?
my $got_command=0; #did we get a command?

while (@ARGV){
	$_=$ARGV[0];
	if(/^-a$/){
		$got_non_kill_args=1;
		shift @ARGV; shift @ARGV; next;
	}elsif(/^-c$/){
		#ssh-agent may treat -c and -s as mutually exclusive
		#for now we just use whichever we got last
		$c_flag=1;
		shift @ARGV; next;
	}elsif(/^-d$/){
		$got_non_kill_args=1;
		shift @ARGV; next;
	}elsif(/^-k$/){
		$kill=1;
		shift @ARGV; next;
	}elsif(/^-s$/){
		#ssh-agent may treat -c and -s as mutually exclusive
		#for now we just use whichever we got last
		$c_flag=0;
		shift @ARGV; next;
	}elsif(/^-t$/){
		$got_non_kill_args=1;
		shift @ARGV; shift @ARGV; next;
	}elsif(/^--$/){
		#end of options
		shift @ARGV;
		if(@ARGV){
			$got_command=1;
			$got_non_kill_args=1;
			last;
		};
	}elsif(/^-./){
		#rats, unrecognized option
		die "$0: unrecognized option $_, aborting";
	}else{
		#anything left must be a command
		$got_command=1;
		$got_non_kill_args=1;
		last;
	};
};

#sanity check what we've gotten so far
if($kill && $got_non_kill_args){
	#incompatible options
	die "$0: -k option is only compatible with -c or -s option, and no other options or arguments, aborting";
};

sub do_agent_exec{
	#do our ssh-agent exec
	exec {$ssh_agent} ('ssh-agent',@original_argv);
	die(
		"$0: exec {'$ssh_agent'} ('ssh-agent'",
		(
			@original_argv
			?
				(",'" . join("','",@original_argv) . "'")
			:
				()
		)
		,") failed, aborting"
	);
};

if($got_command){
	#we were given a command - just hand this off to ssh-agent
	&do_agent_exec;
};

#if we're down here, we need to figure out if we've got a running agent

#sanity check the environment first

#track if we found a legitimate ssh-agent PID of ours
my $SSH_AGENT_PID=undef;

#first, some subroutines for figuring out PID stuff

sub fuser_file_pids{
	#return PIDs found by fuser for file (e.g. ssh-agent binary or socket)
	$#_==0 or die "$0: &fuser_file_pids didn't get one argument, aborting";
	my $file=$_[0];

	my $chpid=open(FUSER, '-|');
	if(!defined($chpid)){
		#(implicit) fork failed
		die "$0: open(FUSER, '-|') failed, $!.\n";
	}elsif(!$chpid){
		#I'm child
		#for fuser, we want to discard STDERR
		close(STDERR) or die "$0: failed to close STDERR,";
		#so the diagnostics won't be seen ... whatever ...
		open(STDERR,'>>','/dev/null') or die "$0: failed to open STDERR as /dev/null, child aborting";
		exec {$fuser} ('fuser',$file); #fuser doesn't support -- (at least yet)
		die "$0: exec {'$fuser'} ('fuser','$file') failed, $!, child aborting";
	};
	#else #I'm parent
	my $fuser_data;
	{
		local $/; undef $/; #slurp
		while(<FUSER>){
			$fuser_data=$_;
		};
	};
	if(!close(FUSER) && $! != 0){
		warn "$0: close(FUSER) failed, \$!=$! \$?=$?";
		return ();
	}
	local $_=$fuser_data;
	chomp;
	s/\s+/ /g;
	s/^\s*//;
	s/\s*$//;

	return(split);
};

#once we grab them, track them and note that we did so
my @ssh_agent_pids=();
my $got_ssh_agent_pids=0;

sub ssh_agent_pids{
	#return ssh-agent PIDs found by fuser

	#if we already did it, just return what we got before
	if($got_ssh_agent_pids){
		return(@ssh_agent_pids);
	};

	#save this stuff and note that we did so
	$got_ssh_agent_pids=1;
	return(@ssh_agent_pids=&fuser_file_pids($ssh_agent));
};

sub pid_is_ssh_agent{
	#check if a PID is one of ssh-agent's
	$_[0] =~ /^\d+$/ or return undef;
	my $pid=$_[0];
	local @_=&ssh_agent_pids;
	for (@_){
		$_ == $pid and return 1;
	};
	return 0;
};

my $owner; #a place to note owner

sub is_valid_pid{
	#is $_[0] a valid ssh-agent pid of ours?
	$#_==0 or die "$0: &is_valid_pid didn't get one argument, aborting";
	local $_=$_[0];

	#one or more decimal digits?
	(/^\d+$/)
	or die "$0: &is_valid_pid argument didn't match (/^\d+\$/), aborting";

	kill(0,$_) #can we signal it?
	&& &pid_is_ssh_agent($_) #fuser reports it as ssh-agent's?
	&& (
		#check that we're not superuser (root),
		#as UID 0 can signal all PIDs
		$> != 0
		|| (
			#we're superuser -
			#we need to also check if it's our PID
			$owner=(stat("/proc/$_"))[4],
			defined($owner)
			&& $owner==$>
		)
	);
};

#try to find a legitimate PID
if(
	#first we try the environment
	exists($ENV{SSH_AGENT_PID}) && #in the environment?
	&is_valid_pid($ENV{SSH_AGENT_PID})
){
	#note it as legitimate
	$SSH_AGENT_PID=$ENV{SSH_AGENT_PID};
}else{
	#we didn't find a valid one via environment, try to find a legitimate PID
	for(&ssh_agent_pids){
		if(
			&is_valid_pid($_)
		){
			#note it as legitimate
			$SSH_AGENT_PID=$_;
			#no reason to check any other PIDs
			last;
			#while there could be more than one PID that would satisfy
			#our criteria, we wouldn't (easily) know how to pick which one
			#of them would be optimal anyway - e.g. we wouldn't want to
			#exclude ssh-agents that have a command - as that might be the
			#one we're looking for (e.g. as may be the case under X).
		};
	};
};

if(!defined($SSH_AGENT_PID)){
	#we didn't find a valid PID
	if($kill){
		#got a syntactically valid kill (-k) request,
		#but nothing to legitimately signal

		#we could just exec ssh-agent
		#but instead we choose to be a hair safer
		#it appears current ssh-agent implementations would
		#unconditionally signal SSH_AGENT_PID
		if(!$c_flag){
			print "unset SSH_AUTH_SOCK;\n";
			print "unset SSH_AGENT_PID;\n";
			if(exists($ENV{SSH_AGENT_PID})){
				#this might be a bit of a fib
				print "echo Agent pid $ENV{SSH_AGENT_PID} killed;\n";
			}else{
				#act a bit more ssh-agent-like here:
				warn "SSH_AGENT_PID not set, cannot kill agent\n";
				exit 1;
			};
		}else{
			print "unsetenv SSH_AUTH_SOCK;\n";
			print "unsetenv SSH_AGENT_PID;\n";
			if(exists($ENV{SSH_AGENT_PID})){
				#this might be a bit of a fib
				print "echo Agent pid $ENV{SSH_AGENT_PID} killed;\n";
			}else{
				#act a bit more ssh-agent-like here:
				warn "SSH_AGENT_PID not set, cannot kill agent\n";
				exit 1;
			};
		};
		exit 0; #we're done
	}else{
		#didn't get a valid PID, and not asked to kill (-k),
		#so, we need to exec ssh-agent
		&do_agent_exec;
	};
}elsif($kill){
	#we got a valid PID and asked to kill (-k)
	#we could hand this off to ssh-agent, but why bother?
	kill(15,$SSH_AGENT_PID);
	if(!$c_flag){
		print(
			"unset SSH_AUTH_SOCK;\n",
			"unset SSH_AGENT_PID;\n",
			"echo Agent pid $SSH_AGENT_PID killed;\n"
		);
	}else{
		print(
			"unsetenv SSH_AUTH_SOCK;\n",
			"unsetenv SSH_AGENT_PID;\n",
			"echo Agent pid $SSH_AGENT_PID killed;\n"
		);
	};
	exit 0; #we're done
};
#else (if we're still continuing down here) ...
#we got a valid PID
#not asked to kill (-k)
#not given a command to run (that case was handled earlier)
#we have good SSH_AGENT_PID, will need to put that in environment,
#need to verify or figure out SSH_AUTH_SOCK

sub is_valid_sock{
	#is $_[0] a valid ssh-agent socket for us?

	my $mode;

	#proper number of arguments?
	$#_==0 or return 0;

	$_=$_[0];

	#expected form?
	m'^/tmp/ssh-[^/]+/agent\.\d+$'

	#is it a socket of ours?
	&& (($mode,$owner)=(stat($_[0]))[2,4])

	&& defined($mode) #got $mode?
	&& defined($owner) #got $owner?
	&& $owner==$> #is it mine?
	&& S_IFSOCK($mode) #is it a socket?
	or return 0; #if we didn't pass one of our tests

	#there are various additional tests we could potentially do here
	#we go for the acid test - can ssh-add talk to it?

	#we're only interested in return value of ssh-add
	my $chpid=fork;
	if(!defined($chpid)){
		die "$0: fork failed, aborting";
	}elsif(!$chpid){
		#I'm child

		#we're only interested in return value, so ...
		close(STDIN) or die "$0: failed to close STDIN,";
		open(STDIN,'<','/dev/null')
			or die "$0: failed to open STDIN as /dev/null, child aborting";
		close(STDOUT)
			or die "$0: failed to close STDOUT,";
		open(STDOUT,'>>','/dev/null')
			or die "$0: failed to open STDOUT as /dev/null, child aborting";
		close(STDERR) or die "$0: failed to close STDERR,";
		open(STDERR,'>>','/dev/null')
			or die "$0: failed to open STDERR as /dev/null, child aborting";

		#at this point we're only doing this in the child
		$ENV{SSH_AUTH_SOCK}=$_[0];

		exec {$ssh_add} ('ssh-add','-l');
		die "$0: exec {'$ssh_add'} ('ssh-add','-l') failed, aborting";
	};

	#I'm parent, wait for child
	waitpid($chpid,0);

	($? & 255) == 0 #child didn't terminate from signal or dump core

	&& (
		#for 0 or 1: we successfully talked to agent
		$? >> 8 == 0 #and found identities
		||
		$? >> 8 == 1 #but found no identities
	)
	#anything not satisfying the above was unsuccessful,
	#e.g. $? >> 8 == 2 for ssh-add couldn't talk to SSH_AUTH_SOCK or
	#command not found or child die()ed, etc.
};

sub fuser_file_pids_matches_pid{
	#does a PID from &fuser_file_pids($_[0]) match $_[1]?

	#argument sanity checks
	$#_==1 or return undef;
	defined($_[0]) or return undef;
	defined($_[1]) or return undef;
	$_[1] =~ /^\d+$/ or return undef;

	for (&fuser_file_pids($_[0])){
		$_==$_[1] and return 1;
	};
	return 0;
};

#a reminder of where we're at:
#we got a valid PID
#not asked to kill (-k)
#not given a command to run (that case was handled earlier)
#we have good SSH_AGENT_PID, will need to put that in environment,
#need to verify or figure out SSH_AUTH_SOCK

#track if we found a legitimate SSH_AUTH_SOCK
my $SSH_AUTH_SOCK=undef;

#do we have a valid $ENV{SSH_AUTH_SOCK}?
if(
	#in the environment?
	defined($ENV{SSH_AUTH_SOCK})

	#we sanity check form before using fuser as our fuser doesn't support --
	&& $ENV{SSH_AUTH_SOCK} =~ m'^/tmp/ssh-[^/]+/agent\.\d+$' #expected form

	&& &is_valid_sock($ENV{SSH_AUTH_SOCK}) #is it legitimate?

	#PIDs match?
	&&(
		#unfortunately looks like we can only make this check if we're EUID 0
		$> != 0
		|| &fuser_file_pids_matches_pid($ENV{SSH_AUTH_SOCK},$SSH_AGENT_PID)
	)
){
	$SSH_AUTH_SOCK=$ENV{SSH_AUTH_SOCK}; #good, save it
}else{
	#didn't find a good SSH_AUTH_SOCK in the environment,
	#need to see if we can locate one

	#we're looking for something of the form:
	#/tmp/ssh-XXXXXXXX/agent.<ppid>
	opendir(DIR,'/tmp')
		or die "open(DIR,'/tmp') failed, aborting";
	while($_=readdir(DIR)){

		/^ssh-./ or next; #only considering expected names

		my $mode;
		my $sock_dir="/tmp/$_"; #potential path of directory containing socket

		#is it a directory of ours?
		($mode,$owner)=(stat($sock_dir))[2,4];

		if(
			defined($mode) #got $mode?
			&& defined($owner) #got $owner?
			&& $owner==$> #is it mine?
			&& S_IFDIR($mode) #is it a directory?
			&& opendir(DIR2,$sock_dir) #open, need to examine more closely
		){
			while($_=readdir(DIR2)){
				/^agent\.\d+$/ or next; #only considering expected names

				my $socket="$sock_dir/$_";

				#is it a valid socket for us?
				&is_valid_sock($socket) #is it legitimate?

				#PIDs match?
				&&(
					#unfortunately looks like we can only make this check if
					#we're EUID 0
					$> != 0
					|| &fuser_file_pids_matches_pid($socket,$SSH_AGENT_PID)
				)

				&& (
					$SSH_AUTH_SOCK=$socket,
					last #no need to check further
				)
			};
			closedir(DIR2);
		}
		if(defined($SSH_AUTH_SOCK)){
			#no need to check further
			last;
		};
	};
	closedir(DIR);
};

#exec ssh-agent if we didn't get good SSH_AGENT_PID and SSH_AUTH_SOCK
defined($SSH_AGENT_PID) or &do_agent_exec;
defined($SSH_AUTH_SOCK) or &do_agent_exec;

#got good SSH_AGENT_PID and SSH_AUTH_SOCK
if(!$c_flag){
	print (
		"SSH_AUTH_SOCK=$SSH_AUTH_SOCK; export SSH_AUTH_SOCK;\n",
		"SSH_AGENT_PID=$SSH_AGENT_PID; export SSH_AGENT_PID;\n",
		"echo Agent pid $SSH_AGENT_PID;\n"
	);
}else{
	print (
		"setenv SSH_AUTH_SOCK $SSH_AUTH_SOCK;\n",
		"setenv SSH_AGENT_PID $SSH_AGENT_PID;\n",
		"echo Agent pid $SSH_AGENT_PID;\n"
	);
};