#!/bin/sh
# ALL the continued lines following this one are interpreted
# by the bourne shell and ignored by Tcl, which picks up
# execution after the exec line.\
exec tclsh "$0" ${1+"$@"}

# Default agent file
set AGENTFILE $::env(HOME)/.ssh-agent-mgr.rc
set SHELL $::env(SHELL)

set ::exit_status(OK) 0
set ::exit_status(USAGE) 1
set ::exit_status(AGENTFILE) 2
set ::exit_status(STOP) 3
set ::exit_status(START) 4
set ::exit_status(START_AGENT) 5
set ::exit_status(AGENTFILE_CREATION) 6
set ::exit_status(ADD_KEY) 7
set ::exit_status(NO_AGENT) 8
set ::exit_status(KILL_AGENT) 9
set ::exit_status(AGENT_RUNNING) 10
set ::exit_status(UNKNOWN_SHELL) 11
set ::exit_status(NOT_OWNER_OF_AGENT_FILE) 12

#========================================================================
# This procedure gives the user information on how to run the command
#========================================================================
proc usage { } {
    puts stderr "Usage: $::argv0 \[--agent-file=<file>\] \[--shell=<shell>\] start|stop|restart|check|addkeys|help"
    puts stderr "\t--agent-file - user specified state file. (Default: $::env(HOME)/.ssh-agent-mgr.rc)"
    puts stderr "\t--shell - Currently valid values for <shell> are:"
    puts stderr "\t\t\tbash, sh, csh, tclsh, and wish"
    puts stderr "\tstart   - start a new ssh-agent if one does not exist"
    puts stderr "\tstop    - stop the currently running ssh-agent and "
    puts stderr "\t\t\tremoves the state file"
    puts stderr "\trestart - effectively does a stop then start in one command"
    puts stderr "\tcheck   - display information about currently running"
    puts stderr "\t\t\tssh-agent"
    puts stderr "\taddkeys - Add keys to the ssh agent"
    puts stderr "\thelp    - this usage message"
    exit $::exit_status(USAGE)
}

switch -exact -- $::tcl_platform(os) {
    SunOS {set ::SSH_PATH /usr/bin}
    Linux {set ::SSH_PATH /usr/bin}
    default {set ::SSH_PATH /usr/bin/}
}

#========================================================================
# This procedure will create a shell command line appropriate for
#   exporting the environment variable and its value into the parent
#   shell.
#========================================================================
proc export_env { var } {
    set shell [ file tail $::SHELL ]
    switch -exact -- $shell {
	sh -
	bash {
	    return "$var=$::env($var); export $var;"
	}
	tcsh -
	csh {
	    return "setenv $var $::env($var);"
	}
	tclsh -
	wish {
	    return "set ::env($var) $::env($var)"
	}
    }
    puti "Unknown shell: $::SHELL"
    exit $::exit_status(UNKNOWN_SHELL)
}

#========================================================================
# This procedure handles importing of variables. It reads the line and
#   dynamically determines the syntax shell syntax used. The variable
#   and its value is exported into the shell's environment space
#========================================================================
proc import_env { line } {
    if { [ regexp {^[ \t]*$} $line -> ] } {
	# Skip blank lines
	return ""
    }
    switch -regexp -- $line {
	setenv {
	    # csh derivative
	    regexp {setenv[ \t*]([^ \t]*)[ \t*](.*)} $line -> var value
	}
	{^([^=]*)=([^;]*);} {
	    # sh derivative
	    regexp {([^=]*)=([^;]*);} $line -> var value
	}
	::env {
	    # tclsh/wish derivative
	    regexp {set ::env\(([^)]*\)[ \t]+(.*)} $line -> var value
	}
    }
    if { [ info exists var ] && \
	     [ info exists value ] } {
	#----------------------------------------------------------------
	# only set the environment space if the information was gleaned
	#   from the input
	#----------------------------------------------------------------
	set ::env($var) $value
    }
}

#========================================================================
# This procedure will either return with the return status or exit
#   with the return status
#========================================================================
proc return_status { retval should_exit } {
    if { $should_exit } {
	exit $retval
    }
    return $retval
}

#========================================================================
# This procedure dumps the environment variables in the requested form
#========================================================================
proc dump_ssh_env { } {
    puts [ export_env SSH_AGENT_PID ]
    puts [ export_env SSH_AUTH_SOCK ]
}

#========================================================================
# This procedure gives instructions on how to run the script
#========================================================================
proc runAgent { Cmd } {
    return "$::argv0 [lrange $::argv 0 [ expr [ llength $::argv ] - 2 ] ] $Cmd"
}

#========================================================================
# This procedure will start the ssh-agent
#========================================================================
proc start_agent {args} {
    puti "Starting a new agent."
    if {[catch {exec $::SSH_PATH/ssh-agent -s} data]} {
        puti "Error while starting agent: $data"
        exit $::exit_status(START_AGENT)
    }

    ;## remove last line (echo)
    set data [join [lreplace [split $data "\n"] end end] "\n"]

    set top [ file dirname $::AGENTFILE ]
    if { ! [ file isdirectory ${top} ] } {
        ;## if there happens to be a file named ${::TOP}, rename it
        catch {file rename -- ${top} ${top}.bogus}
        catch {file mkdir ${top}}
    }

    if { [ catch {
        set fid [open $::AGENTFILE w]
        puts $fid $data
        close $fid
        file attributes $::AGENTFILE -permissions 400
    } err ] } {
        puti "Error while writing file ${::AGENTFILE}: $err"
	# Shut down the agent that was started since it will not be usable
	#   in the future.
	foreach line [ split $data "\n" ] {
	    import_env $line
	}
	catch { exec $::SSH_PATH/ssh-agent -k } err
        exit $::exit_status(AGENTFILE_CREATION)
    }

    source_file 1

    return ""
}

#========================================================================
# This takes care of adding the identity, rsa, and dsa keys to the agent
#========================================================================
proc addKey { } {
    if { [ catch { exec $::SSH_PATH/ssh-add < /dev/tty \
		       > /dev/null 2> /dev/tty } err ] } {
	puti "Error while adding key: $err"
	exit $::exit_status(ADD_KEY)
    }

    return ""
}

#========================================================================
# This procedure prints infomational messages in a way as not to
#   interfear with setting of environment variables.
#========================================================================
proc puti { message } {
    foreach line [ split $message "\n" ] {
	puts stderr "\# $line"
    }
}

#========================================================================
# Check the status of the ssh-agent
#========================================================================
proc sanity_check { conditions } {
    catch {exec $::SSH_PATH/ssh-add -l} err

    # These conditions everyone must pass
    switch -regexp -- $conditions {
	default {
	    switch -regexp -- $err {
		{not open a connection} {
		    append err "\n\# Use the '[ runAgent start ]' command to start an ssh agent."
		    return -code error $err
		}
		default {}
	    }
	}
    }
    switch -regexp -- $conditions {
	have_agent { }
	default {
	    switch -regexp -- $err {
		{agent has no identities} {
		    append err "\nUse the following command to add an identity to the agent:"
		    append err "\n[ runAgent addkeys ]"
		    return -code error $err
		}
		default {}
	    }
	}
    }
    return ""
}

#========================================================================
# Read environment variables from the state file
#========================================================================
proc source_file { should_exit } {
    if { [ catch {
        set fid [open $::AGENTFILE r]
        set data [read $fid]
        close $fid
    } err ] } {
	puti "Error while reading file ${::AGENTFILE}: $err"
	return [ return_status $::exit_status(AGENTFILE) $should_exit ]
    }

    foreach line [ split $data "\n" ] {
	import_env $line
    }

    return ""
}

#========================================================================
# This procedure gets an agent started and adds the keys
#========================================================================
proc start { should_exit } {
    #--------------------------------------------------------------------
    # Check to see if an agent is currently running
    #--------------------------------------------------------------------
    if { [ file exists $::AGENTFILE ] } {
	source_file 0
	if { ! [ catch { sanity_check have_agent } err ] } {
	    puti "ssh-agent is already running"
	    #------------------------------------------------------------
	    # Dump the info on how to set the environment variables
	    #------------------------------------------------------------
	    dump_ssh_env
	    return [ return_status $::exit_status(AGENT_RUNNING) $should_exit ]
	}
    }
    #--------------------------------------------------------------------
    # Make sure things are ready for starting
    #--------------------------------------------------------------------
    file delete -force -- $::AGENTFILE
    #--------------------------------------------------------------------
    # Start the agent
    #--------------------------------------------------------------------
    if { [ catch { start_agent } err ] } {
	puti "Unable to start agent"
	return [ return_status $::exit_status(START) $should_exit ]
    }
    #--------------------------------------------------------------------
    # Add the key
    #--------------------------------------------------------------------
    if { [ catch { addKey } err ] } {
	puti "Unable to add key to agent"
	return [ return_status $::exit_status(START) $should_exit ]
    } 
    #--------------------------------------------------------------------
    # Make sure it all looks good
    #--------------------------------------------------------------------
   if { [ catch {sanity_check all} err ] } {
	puti "Agent fails sanity check: $err"
	return [ return_status $::exit_status(START) $should_exit ]
    }
    #--------------------------------------------------------------------
    # Dump the info on how to set the environment variables
    #--------------------------------------------------------------------
    dump_ssh_env
    #--------------------------------------------------------------------
    # Let caller know what happened
    #--------------------------------------------------------------------
    return_status $::exit_status(OK) $should_exit
}

#========================================================================
# This procedure stops an active agent and removes the state file
#========================================================================
proc stop { should_exit } {
    #--------------------------------------------------------------------
    # 
    #--------------------------------------------------------------------
    if { [ file exists $::AGENTFILE ] } {
	source_file 1
	puti "Killing existing agent."
	#----------------------------------------------------------------
	# Actually try to kill the agent
	#----------------------------------------------------------------
	if { [ catch { exec $::SSH_PATH/ssh-agent -k } err ] } {
	    if { ! [ regexp {No such process} $err ] } {
		puti "Error while killing agent: $err"
		exit $::exit_status(KILL_AGENT)
	    }
	}
	#----------------------------------------------------------------
	# Remove the file containing the state information
	#----------------------------------------------------------------
	file delete -force -- $::AGENTFILE
    } else {
	if {[stop_via_ps] == 0} {
	    puti "Error: Could not find existing agent."
	    return [ return_status $::exit_status(NO_AGENT) $should_exit ]
        }
    }
    return [ return_status 0 $should_exit ]
}

#========================================================================
# This procedure stops an active agent and removes the state file
#========================================================================
proc stop_via_ps { } {
    set retval 0
    #--------------------------------------------------------------------
    # Open list of processes to kill
    #--------------------------------------------------------------------
    if [catch {set processes [exec ps -u $::tcl_platform(user) -o pid,comm]} err] {
	return $retval;
    }
    foreach line [split $processes '\n'] {
        foreach {pid cmd} $line {
            if { [regexp {^(.*/|)ssh-agent$} $cmd ->] } {
	        catch {exec kill -HUP $pid}
	        after 1000
	        catch {exec kill -9 $pid}
		incr retval 1
	    }
        }
    }
    return $retval
}

#========================================================================
;## MAIN ###
#========================================================================
for { set idx 0 } { $idx < $::argc } { incr idx } {
    switch -regexp -- [ lindex $::argv $idx ] {
	{^--} {
	    regexp {^--([^=]*)=?(.*)$} [ lindex $::argv $idx ] -> opt val
	    switch -exact -- $opt {
		shell {
		    if { [ string length $val ] > 0 } {
			set ::SHELL $val
		    }
		}
		agent-file {
		    # allow the user to specify a different file as
		    #  the agent resource file
		    if { [ string length $val ] > 0 } {
			set ::AGENTFILE $val
		    }
		}
		default {
		    puts stderr "ERROR: Unknown option: $opt"
		    usage
		}
	    }
	}
	{^addkeys$} -
        {^check$} -
	{^help$} -
	{^restart$} -
        {^start$} -
        {^stop$}  { set cmd [ lindex $::argv $idx ] }
	default {
	    puts stderr "ERROR: Unknown command: [ lindex $::argv $idx ]"
	    usage
	}
	
    }
}

if { ! [ info exists cmd ] } {
    usage
}

#------------------------------------------------------------------------
# Check to see if the AGENTFILE exists and if it is owned by the current
#   user.
#------------------------------------------------------------------------
if { [ file isfile $::AGENTFILE ] && ! [ file owned $AGENTFILE ] } {
    #--------------------------------------------------------------------
    # The file exists but is owned by someone else
    #--------------------------------------------------------------------
    exit $::exit_status(NOT_OWNER_OF_AGENT_FILE)
}
#------------------------------------------------------------------------
# Execute the command
#------------------------------------------------------------------------
switch -exact -- $cmd {
    addkeys {
	#----------------------------------------------------------------
	# Source the state file
	#----------------------------------------------------------------
	source_file 1
	#----------------------------------------------------------------
	# :TODO: Need to check to see if keys are already in the agent
	#----------------------------------------------------------------
	#----------------------------------------------------------------
	# Add the key
	#----------------------------------------------------------------
	if { [ catch { addKey } err ] } {
	    puti "Unable to add key to agent"
	    return [ return_status $::exit_status(START) $should_exit ]
	} 
    }

    check {
        if { [ file exists $::AGENTFILE ] } {
            source_file 1
            if { [ catch {sanity_check all} err ] } {
                puti "$err"
                exit $::exit_status(NO_AGENT)
            }
	    #------------------------------------------------------------
	    # Dump the keys that are known
	    #------------------------------------------------------------
	    catch {exec $::SSH_PATH/ssh-add -l} err
	    foreach line [ split $err "\n" ] {
		puti $line
	    }
	    #------------------------------------------------------------
	    # Dump the environment
	    #------------------------------------------------------------
	    dump_ssh_env
        } else {
            puti "An ssh agent does not appear to be running."
            puti "Use the 'start' option to start an ssh agent."
            exit $::exit_status(NO_AGENT)
        }
        exit 0
    }

    help {
	usage
    }
    restart {
	stop 0
	start 1
	exit 0
    }
    stop {
	stop 1
    }

    start {
	start 1
    }

    default { }
}

exit $::exit_status(OK)

