# coding:utf-8
# AutoSSH.py: run a shell script in all machines of a certain LCC
# author: Thúlio Carvalho

import paramiko
import sys
import threading
import os
import time
import getpass

user = ""
passwd = ""
machines = []
script = ""

def setupLCC(lab):
    global machines

    if (lab == 1):
        invalid = [1, 5, 11, 22, 24, 27]
        machines = range(41)
        for i in invalid: machines.remove(i)

        for i in range(len(machines)):
            machines[i] = "%02d" % (machines[i])

    elif (lab == 2):
        invalid = [2, 7, 8, 35]
        machines = range(41)
        for i in invalid: machines.remove(i)

        for i in range(len(machines)):
            machines[i] = "%02d" % (machines[i])

    elif (lab == 3):
        invalid = [120]
        machines = range(121)
        for i in invalid: machines.remove(i)

        for i in range(len(machines)):
            machines[i] = "%03d" % (machines[i])



def login(self, host):
    global user, passwd, script
    ssh = paramiko.SSHClient()
    ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())

    try:
        ssh.connect(host, port=23456, username=user, password=passwd)
    except paramiko.AuthenticationException:
        print "[E] Authentication Exception"
    except paramiko.SSHException:
        print "[E] SSH Exception"

    try:
        (_, stdout, stderr) = ssh.exec_command("./" + script)
        return (stdout.read(), stderr.read())
    except Exception, e:
        print "[E] Script Exception"
    finally:
        ssh.close()


def configureUser():
    global user, passwd

    user = raw_input("Username: ")
    pprompt = lambda: (getpass.getpass(), getpass.getpass('Retype password: '))

    p1, p2 = pprompt()
    while p1 != p2:
        print('Passwords do not match. Try again')
        p1, p2 = pprompt()

    lab = int(raw_input("Choose the LCC:\n" +
                    "1 - LCC1\n" +
                    "2 - LCC2\n" +
                    "3 - LCC3\n"))

    return lab

def usage():
    print "You must parse the script file!"
    print "Example: $python AutoSSH.py script.sh"
    print "Exiting..."

def main(argv):
    if (len(argv) < 2):
        usage()
        return

    lab = configureUser()
    setupLCC(lab)

    try:
        for i in machines:
            threading.Thread(target=login,args=(i)).start()
            time.sleep(0.7)

    except Exception, e:
        print '[-] General Exception'

if __name__ == "__main__":
    main(sys.argv)