#! /usr/bin/python

#+
# NAME:
#	ssh_to
# PURPOSE:
#	Simplify login between accounts on subnet
# CATEGORY:
#	ucsd/gen/python
# CALLING SEQUENCE:
#	ssh_to <-init> <user@host> <user2@host2> <...>
# INPUTS:
#	user@host	account for which ssh access is needed
#			(if the user name on the remote host is
#			the same as on the account from which
#			ssh_to is executed then 'user@ can be
#			omitted.
# OPTIONAL INPUT PARAMETERS:
#	-init		if specified it must be the first command
#			line argument. Initializes ssh access
#			to the remote account (see PROCEDURE).
# OUTPUTS:
#	(private/public key pair; updates of local 'identification'
#	file and remote 'authorization' file(s))
# OPTIONAL OUTPUT PARAMETERS:
# CALLS:
#	tiny.is_there, tiny.args, tiny.start, tiny.run_cmd, tiny.keys
# RESTRICTIONS:
# >	Assumes that ssh2 creates 2048 bit dsa keys
#	with key name id_dsa_2048_a for the private key
#	and id_dsa_2048_a.pub for the public key.
# >	Assumes that ~/.ssh2 already exists on the remote
#	machines.
# >	For the soft links to work ~/bin must be in the path.
# EXAMPLE:
#	ssh_to -init user@host		(initializes user@host)
#	user@host			(logs into user@host)
# PROCEDURE:
#	If -init is set:
#	1. Check for private key of type id_dsa_2048_a_*.
#	2. If the private key does not exist, then run
#	   ssh-keygen2 to generate private key id_dsa_2048_a
#	   and public key id_dsa_2048_a.pub. Rename the key
#	   pair to id_dsa_2048_a_$RANDOM(.pub).
#	3. Put private key in .ssh2/identification
#	4. Loop over all hosts, and for each
#	   a. put the public key in remote ~/.ssh2 directory,
#	   b. put the name of the public key in remote ~/.ssh2/authorization.
#	   c. put a link to ssh2_to in ~/bin 
# MODIFICATION HISTORY:
#	JUL-2004, Paul Hick (UCSD/CASS)
#	JAN-2004, Paul Hick (UCSD/CASS; pphick@ucsd.edu)
#		Converted from bash to Python.
#-

# If -init is not set then call ssh

import os, sys, random
from tiny import is_there, args, start, run_cmd, keys

say = 'ssh_to, '

verbose = is_there('-verbose', sys.argv)
init	= is_there('-init'   , sys.argv)

arg = args(sys.argv)
ssh = 'ssh'

bin = os.path.join(os.environ['HOME'],'bin')
txt = os.path.join( bin, 'ssh_to.txt' )

ssh_bin = ''
if not init:
	host = ((os.path.split(sys.argv[0])))[1]

	if os.path.isfile(txt):
		tmp = (open(txt,'r')).read()
		pbeg = tmp.find(host)
		if pbeg != -1:
			pbeg = tmp.find('=' ,pbeg)
			pend = tmp.find('\n',pbeg)
			ssh_bin = tmp[pbeg+1:pend]

if ssh_bin == '':
	ssh_bin = start('-bin=', sys.argv)

if ssh_bin == '':
	ssh = run_cmd( 'which '+ssh+' 2> /dev/null', verbose )
	if ssh == '':
		print say+'ssh not found'
		sys.exit()
	ssh = (ssh.split('\n'))[0]
else:
	ssh = os.path.join( ssh_bin, ssh )
	if not os.path.isfile(ssh) :
		print say+ssh+'not found'

if not init:
	tmp = [ssh, host]

	key = keys(sys.argv)
	if len(key ) > 0:
		tmp.extend(key)

	if len(arg) > 1:
		tmp.extend(arg[1:])

	print say+ssh+' '+host
	os.execvp(ssh, tmp)			# Transfer control (process does not return)



tmp = run_cmd( ssh+' -V 2>&1', verbose )
open_ssh = tmp.find('OpenSSH') == 0

# Check for ~/.ssh2 directory.

if open_ssh:
	ssh_dir = os.path.join( os.environ['HOME'], '.ssh' )
	prefix  = 'id_dsa'
else:
	ssh_dir = os.path.join( os.environ['HOME'], '.ssh2' )
	prefix  = 'id_dsa_2048_a'

if not os.path.isdir( ssh_dir ):
	print say+ssh_dir+' does not exist'
	sys.exit()

os.chdir( ssh_dir )

private = []
for key in os.listdir('.'):
	if key.find(prefix) == 0:
		tmp = os.path.splitext(key)
		if tmp[1] != '.pub':
			private.append(key)

# Generate private key if it doesn't exist.

if len(private) == 0:

	keygen = os.path.join((os.path.split(ssh))[0], 'ssh-keygen')
	print say+'generating key with '+keygen
	tmp = run_cmd( keygen+' -t dsa', verbose )

	if open_ssh:
		private = prefix
		public  = private+'.pub'

	else:
		random.seed()
		private = prefix+'_'+'%05d'%random.randint(1,99999)
		public  = private+'.pub'

		os.rename(prefix 	, private)
		os.rename(prefix+'.pub' , public )

else:
	private = private[0]
	public = private+'.pub'

print say+'using '+ssh+' with key '+private

if open_ssh:

	auth = '.ssh/authorized_keys'
	val  = (open(public,'r')).read()
	val  = (val.split('\n'))[0]

else:

	# Check wheter private key is already in .ssh2/identification
	# If not then add it.

	key = 'IdKey '+private

	if os.path.isfile( 'identification' ):
		tmp = (open('identification','r')).read()
		if tmp.find(key) == -1:
			tmp = tmp+key+'\n'
	else:
		tmp = '\n'

	(open('identification','w')).write(tmp)

	auth = '.ssh2/authorization'
	val  = (open(public,'r')).read()
	key  = 'Key '+public

hosts = arg[1:]

for host in hosts:

	print '\n\nSet up: '+host+'\n\n'

	# Copy public key into remote .ssh2 directory
	# Add line to .ssh2/authentication (if necessary)
	# Note that we assume that the remote .ssh2 directory exists.

	if open_ssh:
		cmd = ssh+' '+host	+ \
			' "'		  	+ \
			'touch '+auth+'; '	+ \
			'grep \\"'+val+'\\" '+auth+' > /dev/null; '+ \
			'if [ \\$? == 0 ]; then exit; fi; '	+ \
			'echo \\"'+val+'\\" >> '+auth		 	+ \
			'"'

	else:
		cmd =   'cat '+public+' | '	+ \
			 ssh+' '+host		+ \
			' "'		  	+ \
			'cat - > .ssh2/'+public+'; '		+ \
			'touch '+auth+'; '	+ \
			'grep \\"'+key+'\\" '+auth+'; '+ \
			'if [ \\$? == 0 ]; then exit; fi; '	 + \
			'echo \\"'+key+'\\" >> '+auth		 + \
			'"'

	run_cmd( cmd, verbose )
	
	# Add soft link to $HOME/bin (if necessary)

	if not os.path.isdir( bin ):
		os.mkdir( bin )

	tmp = os.path.join( bin, host )

	if os.path.islink( tmp ):
		if not os.path.exists( tmp ):	# Should remove broken link
			os.remove (tmp)

	if not os.path.islink( tmp ):
		os.symlink( sys.argv[0], tmp )

	if ssh_bin != '':
		if os.path.isfile( txt ):
			tmp = (open(txt,'r')).read()
			pbeg = tmp.find(host)
			if pbeg != -1:
				pend = tmp.find('\n',pbeg)
				tmp = tmp[0:pbeg]+tmp[pend+1:]
		else:
			tmp = ''

		tmp = tmp+host+'='+ssh_bin+'\n'
		(open(txt,'w')).write(tmp)
			
sys.exit()