#!/usr/bin/python
'''campusssh.py -- SSH to a machine on UMich campus.
Optionally, use a machine in a particular lab. Requires patched
version of hostinfo.
'''

import os
import sys
from subprocess import Popen, PIPE

BLACKLIST = ("dc", "ugl")

def getmach(host=None):
    """Get a CAEN Linux machine to SSH into."""
    args = ["hostinfo", "-linux"]
    if host is not None:
        args.extend(["-host", host])
    output = Popen(args, stdout=PIPE).communicate()[0]

    # Field layout:
    # hostname, OS, load1, load2, load3, ???, used?, ???, report time. ???
    def try_once(first_try):
        """Check whether there's an available machine in the output.
        Try harder on the second try.
        """
        for line in output.split('\n'):
            if not line or line.startswith(BLACKLIST):
                continue
            line = line.split()
            if first_try and line[6] != "YES":
                continue
            return line[0]
        return None

    ret = try_once(True)
    if ret is None:
        ret = try_once(False)
    return ret

if __name__ == "__main__":
    host = getmach(sys.argv[1] if len(sys.argv) > 1 else None)

    if host is None:
        raise Exception("Couldn't find host!")

    os.execvp("ssh", ["ssh", "-l", "swolchok", host] + sys.argv[2:])