# Author: Luke Gane
# Version: 1.0.0
# Last updated: 2015-06-18

#from __future__ import print_function # Uncomment for Python 2.X

# === Notes ===

# This Python script generates a shell script that (optionally) prompts for a 
# password and then opens a terminal window with multiple tabs, each containing 
# an SSH session. This script can generate two shell scripts: one that assumes 
# all connections are using the same password and connects automatically, and 
# one that simply initiates the connections but requires the user to manually 
# enter a password for each one. The former is useful if the user regularly 
# needs to connect to multiple servers that share a password (e.g., in a test 
# environment) while the latter is convenient if authentication keys are in use 
# or the user simply wishes to avoid having to open tabs and enter multiple SSH 
# commands on a regular basis.

# Each tab will be assigned a title (specified by the strings in the "titles" 
# list below), as tabs with titles make it easier to keep track of which session 
# is in which tab. However, these titles will only persist if the tabs are 
# opened using a profile that retains the titles passed to gnome-terminal. To 
# create such a profile if you don't already have one:
# 	1) Go to Edit > Profiles > New in a terminal window.
# 	2) Name the new profile "Keep init titles" (or whatever you want as long 
# 	   as you update the profileName variable accordingly).
# 	3) Under the "Title and Command" tab, set "When terminal commands set 
# 	   their own titles:" to "Keep initial title".

# The shell script generated by this script has only been tested on Ubuntu, but 
# it should work anywhere GNOME Terminal is installed.

# In order to use the shell script generated when samePassword is set to True, 
# sshpass must be installed (on Ubuntu, run "sudo apt-get install sshpass").

# Note that if an SSH connection is refused, its tab will close immediately. If 
# all the connections are refused (which might occur if, for example, an 
# incorrect password is entered), all the tabs will close, as will the parent 
# window.

# If every server uses a different password or you prefer to use authentication 
# keys, set samePassword to False to generate a simpler shell script that just 
# opens a terminal window with a tab for each session.

# For maximum convenience, it's nice to just double-click on a link to the shell 
# script to run it. On Ubuntu, this requires two simple steps:
# 	1) Make the script executable (chmod +x NameOfScript.sh).
# 	2) Tell Ubuntu to run executable text files instead of opening them for 
# 	   editing. In the file browser, go to Edit > Preferences; on the 
# 	   "Behaviour" tab, under the "Executable Text Files" heading, select 
# 	   "Run executable text files when they are opened".

# === Set these values appropriately ===

# Specify session and tab information here
profileName = 'Keep init titles' # Name of a profile that keeps the initial titles assigned to terminal windows
titles = ['Server 1', 'Server 2', 'Server 3'] # For the tabs
users = ['root', 'root', 'root']
addresses = ['127.0.0.1', '127.0.0.1', '127.0.0.1']

# Specify a name for the script to be generated
scriptName = 'ConnectToServersTest.sh'

# Specify whether or not all servers use the same password
samePassword = True;

# === Check data ===

if (not (len(titles) == len(users) == len(addresses))):
	raise Exception('Inconsistent number of titles, users, and addresses specified.')

if (len(titles) == 0):
	raise Exception('No session data specified.')

# === Parts of the command ===

if (samePassword):
	commandStart = '#!/bin/bash\n\ngnome-terminal -e "bash -c \'read -s -p \\"Password for test servers: \\" inputPassword; gnome-terminal '

	# Open a single window first...
	commandMiddleFirst = '--window-with-profile=\\"{profile}\\" --title=\\"{title}\\" -e \\"sshpass -p \\$inputPassword ssh -o StrictHostKeyChecking=no {user}@{address}\\" '

	# ...then open as many tabs as necessary
	commandMiddle = '--tab-with-profile=\\"{profile}\\" --title=\\"{title}\\" -e \\"sshpass -p \\$inputPassword ssh -o StrictHostKeyChecking=no {user}@{address}\\" '

	commandEnd = '\'"'
	
else:
	# Quite a bit simpler in this case
	commandStart = '#!/bin/bash\n\ngnome-terminal '
	commandMiddleFirst = '--window-with-profile="{profile}" --title="{title}" -e "ssh {user}@{address}" '
	commandMiddle = '--tab-with-profile="{profile}" --title="{title}" -e "ssh {user}@{address}" '
	commandEnd = ''

# === Create the script ===

with open(scriptName, 'w') as scriptFile:
	
	commandString = commandStart + commandMiddleFirst.format(profile = profileName, title=titles[0], user=users[0], address=addresses[0]);
	
	for cur in range(1,len(titles)):
		commandString += commandMiddle.format(profile = profileName, title=titles[cur], user=users[cur], address=addresses[cur])
		
	commandString += commandEnd
	
	scriptFile.write(commandString)

# === Notification and reminder ===

print('Script "' + scriptName + '" created; remember to run "chmod +x ' + scriptName + '" before using.')