#!/usr/bin/expect --
# IBM(c) 2007 EPL license http://www.eclipse.org/legal/epl-v10.html
#
#
# sets up ssh on the input node list
# called from by xdsh  <nodelist> -K command 
# Environment Variables:
#
#   DSH_REMOTE_CMD  set to path to remote shell (ssh)
#     root password must agree on all the nodes
#
#   XCAT_ROOT set to root of xCAT install
#
#   DSH_REMOTE_PASSWORD  - user password for -s option 
#
#   SSH_SETUP_COMMAND - Command to be sent to the IB switch to setup SSH. 
#
#   DSH_FROM_USERID_HOME - The home directory of the userid from 
#                  where the ssh keys will be obtained
#                  to send
#
#   DSH_FROM_USERID - The userid from where the ssh keys will be obtained
#                  to send
#                  to the node,  or generated and then obtained to send to the
#                  node.
#   DSH_TO_USERID - The userid on the node where the ssh keys will be updated. 
#
# Usage: remoteshell.expect 
#   [-t remote_shell hostname]  test the remoteshell on the host
#   [-k] Generates the ssh keys needed 
#   [-s node_list]  copies the ssh keys to the nodes
#   [-h ] usage
#
#    exit 0 - good
#    exit 1 - abort
#    exit 2 - usage error
#
#
################################################################################

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

set timeout 80
log_user 0

#Clear xdsh environment variables

set empty ""

# path to remote command
if { [info exists env(DSH_REMOTE_CMD)] } {
   set remoteshell $env(DSH_REMOTE_CMD)
} else {
   set remoteshell "/usr/bin/ssh"
}
if { [info exists env(XCATROOT)] } {
	set xcatroot $env(XCATROOT)
} else {
	set xcatroot "/opt/xcat" 
}

if { [info exists env(SSH_SETUP_COMMAND)] } {
    set ssh_setup_cmd $env(SSH_SETUP_COMMAND)
} else {
    set ssh_setup_cmd ""
}
# User on the Management node that has the ssh keys
if { [info exists env(DSH_FROM_USERID)] } {
    set from_userid $env(DSH_FROM_USERID)
} else {
    set from_userid "root"
}

# User on the node where we will send the ssh keys
if { [info exists env(DSH_TO_USERID)] } {
    set to_userid $env(DSH_TO_USERID)
} else {
    set to_userid "root"
}

# get the from userid and its home directory to get the ssh keys

#set home [exec /usr/bin/perl -e {$user = $ENV{'DSH_FROM_USERID'} ? $ENV{'DSH_FROM_USERID'} : "root"; @info = getpwnam($user); print $info[7] . "\n";}]

#if { [string compare $home $empty] ==0 } {
#set home $env(HOME)
#}
set home $env(DSH_FROM_USERID_HOME)

#
# check input arguments
#
if { [llength $argv] != 0 } {
   if { [string compare "-h" [lindex $argv 0]] ==0 } {  # usage

     puts "Usage: remoteshell.expect" 
     puts "  -t remote_shell hostname - test the remoteshell on the host"
     puts "  -k Generates the ssh keys needed" 
     puts "  -s node_list  - copies keys to the nodes"
     puts "  -h  usage"
     exit 0
    }
	# -t means test to see if the shell is already setup
	if { [string compare "-t" [lindex $argv 0]] ==0 } {
		set env(LC_ALL) "C"

		set pid [ spawn [lindex $argv 1] [lindex $argv 2] -l $to_userid echo test.success ]
		expect {
			timeout { exit 1 }
			"Are you sure you want to continue connecting (yes/no)?" {
				exec /bin/kill $pid
				exit 1
			}
			"*ssword*" {
				exec /bin/kill $pid               		
				exit 1
			}
			"Permission denied." {
				exit 1
			}
			"test.success" { 
				exit 0 
			}
		}
		exit 1
	} elseif { [string compare "-k" [lindex $argv 0]] ==0 } {
		# -k  means generate the keys to the nodes

		# check to see if ssh rsa keys exists, if not generate it
		if {![file exists "$home/.ssh/id_rsa"]} {
			set env(LC_ALL) "C"
			spawn /usr/bin/ssh-keygen -t rsa
			expect "Generating public/private rsa"
			expect -re "Enter file.*:"
			send "\r"
			expect -re "Enter passphrase.*:"
			send "\r"
			expect -re "Enter same passphrase.*:"
			send "\r"
			expect eof
		} elseif {[file exists "$home/.ssh/id_rsa"] &&([file size "$home/.ssh/id_rsa"] == 0)} {
			set env(LC_ALL) "C"
			spawn /usr/bin/ssh-keygen -t rsa
			expect "Generating public/private rsa"
			expect -re "Enter file.*:"
			send "\r"
			expect -re "Overwrite.*\(y/n\)\?"
			send "y\r"
			expect -re "Enter passphrase.*:"
			send "\r"
			expect -re "Enter same passphrase.*:"
			send "\r"
			expect eof
		}

	} elseif { [string compare "-s" [lindex $argv 0]] ==0 } {
                # copy the keys to the nodes
		set nodes [lindex $argv 1]
		set empty ""
		if { [string compare $nodes $empty] ==0 } {
			#no target
			exit 1
		}
		set nodelist [split $nodes ',']
		set nodelist2 $nodelist
		set printlist [ join $nodelist ", " ]
		set scp "/usr/bin/scp"
		set directory "$home/.ssh"
      set fh_auth_keys  [ open "$home/.ssh/tmp/authorized_keys" "r"]
      set auth_keys [read $fh_auth_keys]
      close $fh_auth_keys
	   if { [info exists env(DSH_REMOTE_PASSWORD)] } {
		    set word $env(DSH_REMOTE_PASSWORD)
		} else {
		  exec /bin/stty -echo
		  set timeout 360
		  puts "Enter the password for the userid:$to_userid to access the following target nodes: $printlist"
		  expect_user "*\n"
		  set timeout 80
		  exec /bin/stty echo
		  set word $expect_out(buffer)
		  set expect_out(buffer) ""
		  send_user "\n"
		}
                
		foreach node $nodelist {
                       puts $node
			if { [string compare $ssh_setup_cmd $empty] ==0 } {
				set env(LC_ALL) "C"
				set pid [ spawn $remoteshell $node -l $to_userid /bin/mkdir -p /tmp/$to_userid/.ssh ]
				expect {
					"Are you sure you want to continue connecting (yes/no)?" {
						send "yes\r"
						exp_continue
					}
					"*ssword*" {
						send "$word\r"
						exp_continue
					}
					"Permission denied*" {
						exec /bin/kill $pid
					}
				}
				set pid [ spawn $scp  $directory/tmp/authorized_keys $to_userid\@$node:/tmp/$to_userid/.ssh ] 
				expect {
					"Are you sure you want to continue connecting (yes/no)?" {
						send "yes\r"
						exp_continue
					}
					"*ssword*" {
						send "$word\r"
						exp_continue
					}
					"Permission denied*" {
						exec /bin/kill $pid
					}
				}
				set pid [ spawn $scp  $directory/id_rsa $to_userid\@$node:/tmp/$to_userid/.ssh ] 
				expect {
					"Are you sure you want to continue connecting (yes/no)?" {
						send "yes\r"
						exp_continue
					}
					"*ssword*" {
						send "$word\r"
						exp_continue
					}
					"Permission denied*" {
						exec /bin/kill $pid
					}
				}
				set pid [ spawn $scp  $directory/copy.sh $to_userid\@$node:/tmp/$to_userid/.ssh ] 
				expect {
					"Are you sure you want to continue connecting (yes/no)?" {
						send "yes\r"
						exp_continue
					}
					"*ssword*" {
						send "$word\r"
						exp_continue
					}
					"Permission denied*" {
						exec /bin/kill $pid
					}
				}
				set pid [ spawn $remoteshell $node -l $to_userid /tmp/$to_userid/.ssh/copy.sh ]
				expect {
					"Are you sure you want to continue connecting (yes/no)?" {
						send "yes\r"
						exp_continue
					}
					"*ssword*" {
						send "$word\r"
						exp_continue
					}
					"Permission denied*" {
						exec /bin/kill $pid
					}
				}
			} else {
				# Loop all keys in $auth_keys and generate ssh setup command one by one
				# Remove '\n'
				set auth_keys [string range $auth_keys 0 end-1]
				set all_keys [split $auth_keys "\n"]
				set ssh_setup_string ""
				foreach each_key $all_keys {
					# Skip blank line
					if { [string compare each_key $empty] == 0 } {
						continue;
					}
					set ssh_setup_string "$ssh_setup_string;$ssh_setup_cmd \"$each_key\""
				}
				# Trim semicolon at the tail
				set ssh_setup_string [string range $ssh_setup_string 1 end]
				set pid [ spawn $remoteshell $node -l $to_userid $ssh_setup_string ]
				expect {
					"Are you sure you want to continue connecting (yes/no)?" {
						send "yes\r"
						exp_continue
					}
					"*ssword*" {
						send "$word\r"
						exp_continue
					}
					"Permission denied*" {
						exec /bin/kill $pid	
					}
				}
			}
		}
		set word ""
		set failed "0"
		foreach node $nodelist2 {
			set hostname $node
			set hlist [split $hostname '.']
			set name ""
			set nlist ""
			set j [llength $hlist]
			for { set i 0 } { $i<$j} {incr i } {
				set element [lindex $hlist $i]
				if { $i == 0 } {
					set name $element
				} else {
					set name "$name.$element"
				}
				lappend nlist $name
			}
			foreach hn $nlist {
				set env(LC_ALL) "C"
				set pid [ spawn $remoteshell $hn -l $to_userid echo test.success ]
				expect {
					"Are you sure you want to continue connecting (yes/no)?" {
						send "yes\r"
						exp_continue
					}
					"*ssword*" {
						exec /bin/kill $pid
						set failed "1"
					}
					"*refused." {
						exec /bin/kill $pid
						set failed "1"
					}
					"Permission denied*" {
						exec /bin/kill $pid
						set failed "1"
					}
					timeout {
						exec /bin/kill $pid
						set failed "1"
					}
					"test.success" { continue }
					}
			}
		}
		if { [string compare "1" $failed ] ==0 } {
			exit 1
		} else {
			exit 0
		}
   } else {  #  usage  error 
              puts "flag entered is not valid" 
              puts "Usage: remoteshell.expect" 
              puts "  -t remote_shell hostname - test the remoteshell on the host"
              puts "  -k Generates the ssh keys needed" 
              puts "  -s node_list  - copies keys to the nodes"
	      exit  2
	    }
} else {   # usage error
              puts "Command requires a flag" 
              puts "Usage: remoteshell.expect" 
              puts "  -t remote_shell hostname - test the remoteshell on the host"
              puts "  -k Generates the ssh keys needed" 
              puts "  -s node_list  - copies keys to the nodes"
	      exit  2
}