#!/usr/bin/python
#
# For the present
# 1) What to do when you run a command twice?
#    a) List of pids?  That's what I'm trying to achieve with
#       "extended names" like 'sln-gsn.eps' - Seems complicated
#    b) Kill running process - Seems lame
#    c) Refuse to run second process - Seems best
# 2) Handling sln:
#    a) Need eithter commands to take arguments or other commands that
#       write to ports data structure...  Then I need locks?
#    b) Need more than one of each command type or extended command
#       names (then definitely need to distinguish between "commands"
#       dict and actual running commands.
# 3) Start/end with regexps/globbing
#    a) When does it happen wrt aliases?
#    b) Where do the matches come from?  Running programs or existing commands?
#    c) How do you identify regexps vs. commands that will end up as aliases?
#       Esp when you are killing programs: if you match aliases against running
#       programs, you won't get anything even though they'll expand into
#       something later.
#
# 
# Killing programs: Can you match against _both_ running programs and
# existing commands, and then uniquify them?
# 
# For the future
# 1) Use dependencies instead of priorities?
# 2) Make it possible for commands to take arguments?
# 3) Killing remote programs: Problem is can't get PID reliably.
#    Can either make remote commands depend on local things: ie, they
#    die when the connection is torn down, or maybe RPC or something
#    would work.
#
#
# Looking like this is a little mini proc filesystem with remote processes...
# Maybe I should just get more proficient with ps and kill...

import getopt,sys,os,pickle,time,re

VERSION = 0.1

def snarf(cmd):
    """Run a shell gommand and grab the output as a string."""
    f=os.popen(cmd)
    result = f.readline()[:-1] # strip final newline
    f.close()
    return result

def ip_address(): return snarf('ip-address')
def todays_date(): return snarf('date +%y-%m-%d')

def mail_socket_command():
    """Return command to do my socket-listening mail forwarding thing"""
    str = 'socket -l -d ~/.html -f %s-percentd.html %s 2025' % \
          (todays_date(), ip_address())
    return str.replace("percentd", "%d")

##################################################
# Info about running processes.  Dictionary where the keys are process
# names (strings), values are process id's (ints)
pids = {}

# Commands and priorities for executing them.  Dictionary of where the
# keys are process names (strings) and the values are dictionaries of
# command info.
#
# Command info dictionaries contain keys:
# 'command': command line to run the program (string)
# 'priority': (int) that determines the order in which to run the
#             tasks.  Lower numbers => run sooner.
# 'sleep_before': (int) seconds to sleep _before_ running this task
#                 (to allow some earlier thing to finish).
# 'sleep_after': (int) seconds to sleep _after_ running this task before
#                running any other tasks (because those tasks depend
#                on this one).

def cmd(command,priority=0,sleep_after=0, sleep_before=0,
        shutdown=None, sleep_before_shutdown=2):
    """Produce "Command info dictionary" with some defaults"""
    return dict(command=command, priority=priority,
                sleep_after=sleep_after, sleep_before=sleep_before,
                shutdown=shutdown, sleep_before_shutdown=3)

# Suggested sshfs options for performance according to some blog post 
# -C -oping_diskarb,noappledouble,noreadahead,

# Worry about double compression here in sshfs and port forwarding
# Also, with VPN don't really need to worry about firewalls
sshfsOptions = 'compression=yes,auto_cache,jail_symlinks,follow_symlinks,transform_symlinks'

commands = dict(
    ##############################
    # Aliases
    #
    work=['tigressfs', ],
    
    ##############################
    # Port forwarding
    #
    # Slime 
    slime_greymonster =cmd('ssh -n -N -L4005:192.168.1.1:4005 greymonster'),
    slime_jen         =cmd('ssh -n -N -L4005:127.0.0.1:4005 192.168.0.198'),

    # Making IDL license server available
    # Keep around as an example
    # idl =cmd('ssh -n -N -L:1700:mambo:1700 -L:32790:mambo:32790 ssh'),

    # IPython
    # Keep around as an example
    # ipy_upsand = cmd('ssh -n -N -L10105:upsand:10105 -L10113:upsand:10113 andromeda'),

    ##############################
    # Remote programs
    #
    # Emacs: make sure that emacs gets a full login environment
    aramis_emacs = cmd("ssh aramis -f bash -l -c emacs"),

    # Xload: doesn't care as much about the environment
    dionysus_xload=cmd('ssh -f dionysus xload -geometry 70x80+877+665'),

    ####
    # Make clio available from Observatory
    #
    # This is not quite ideal since the whole reason I wrote this was
    # so that I could reliably kill ssh processes that had allocated
    # ports.  So this script needs to start the ssh process and then
    # kill it.  Or it needs to monitor the ssh process itself and
    # restart it.  Or I can just run the persistent forward command
    # myself, since I'm not getting much out of this anyway.  Maybe I
    # could print the PID for killing it?
    #
    clio_door = cmd('persistent-remote-forward 22 lerma90 9022'), 
    
    ##############################
    # Sshfs filesystems    
    #
    dionysusfs = cmd('mkdir /Volumes/dionysus;' + 
                     'sshfs dionysus:/home/novak /Volumes/dionysus ' + 
                     '-o volname=Dionysus,%s' % sshfsOptions,
                     sleep_before=2),
    sesamehomefs = cmd('mkdir /Volumes/sesame-home;' + 
                       'sshfs sesame1:/home/gnovak /Volumes/sesame-home ' + 
                       '-o volname="Sesame Home",%s' % sshfsOptions,
                       shutdown='umount /Volumes/sesame-home', sleep_before=2),
    sesameworkfs = cmd('mkdir /Volumes/sesame;' + 
                       'sshfs sesame1:/scratch/lustre/gnovak /Volumes/sesame ' + 
                       '-o volname="Sesame",%s' % sshfsOptions,
                       shutdown='umount /Volumes/sesame', sleep_before=2),
    sesamefs = ['sesameworkfs', 'sesamehomefs'],
    tigresshomefs = cmd('mkdir /Volumes/tigress-home;' + 
                       'sshfs tigressdata:/home/gnovak /Volumes/tigress-home ' + 
                       '-o volname="Tigress Home",%s' % sshfsOptions,
                       shutdown='umount /Volumes/tigress-home', sleep_before=2),
    tigressworkfs = cmd('mkdir /Volumes/tigress;' + 
                       'sshfs tigressdata:/tigress-hsm/gnovak /Volumes/tigress ' + 
                       '-o volname="Tigress",%s' % sshfsOptions,
                       shutdown='umount /Volumes/tigress', sleep_before=2),
    tigressfs = ['tigressworkfs', 'tigresshomefs'],
    columbiahomefs = cmd('echo mkdir /Volumes/columbia-home;' + 
                       'echo sshfs cfe2:/u/gnovak /Volumes/columbia-home ' + 
                       '-o volname=\"Columbia-Home\",%s' % sshfsOptions,
                       shutdown='umount /Volumes/columbia-home'),
    columbiaworkfs = cmd('echo mkdir /Volumes/columbia-work;' + 
                       'echo sshfs cfe2:/nobackup2d/gnovak /Volumes/columbia-work ' + 
                       '-o volname=\"Columbia-Work\",%s' % sshfsOptions,
                       shutdown='umount /Volumes/columbia-work'),
    columbiafs = ['columbiaworkfs', 'columbiahomefs'],
    loufs = cmd('echo mkdir /Volumes/lou;' + 
                'echo sshfs lou2:/u/gnovak /Volumes/lou ' + 
                '-o volname="Lou",%s' % sshfsOptions,
                shutdown='umount /Volumes/lou', sleep_before=2),
    nasapleiadeshomefs = cmd('echo mkdir /Volumes/nasa-pleiades-home;' + 
                       'echo sshfs nasa-pleiades:/home/gnovak /Volumes/nasa-pleiades-home ' + 
                       '-o volname="Nasa-Pleiades Home",%s' % sshfsOptions,
                       shutdown='umount /Volumes/nasa-pleiades-home', sleep_before=2),
    nasapleiadesworkfs = cmd('echo mkdir /Volumes/nasa-pleiades-work;' + 
                       'echo sshfs nasa-pleiades:/nobackup2d/gnovak /Volumes/nasa-pleiades-work ' + 
                       '-o volname="Nasa-Pleiades Work",%s' % sshfsOptions,
                       shutdown='umount /Volumes/nasa-pleiades-work', sleep_before=2),
    nasapleiadesfs = ['nasapleiadesworkfs', 'nasapleiadeshomefs'],
    ucscpleiadeshomefs = cmd('mkdir /Volumes/ucsc-pleiades-home;' + 
                       'sshfs pleiades:/home/novak /Volumes/ucsc-pleiades-home ' + 
                       '-o volname="Ucsc Pleiades Home",%s' % sshfsOptions,
                       shutdown='umount /Volumes/ucsc-pleiades-home', sleep_before=2),
    ucscpleiadesibrixfs = cmd('mkdir /Volumes/ucsc-pleiades-ibrixfs;' + 
                  'sshfs pleiades:/ibrixfs/novak /Volumes/ucsc-pleiades-ibrixfs ' + 
                       '-o volname="Ucsc Pleiades Ibrix",%s' % sshfsOptions,
                       shutdown='umount /Volumes/ucsc-pleiades-ibrixfs',
                              sleep_before=2),
    ucscpleiadesifs = cmd('mkdir /Volumes/ucsc-pleiades-ifs;' + 
                       'sshfs pleiades:/ifs/novak /Volumes/ucsc-pleiades-ifs ' + 
                       '-o volname="Ucsc Pleiades Ifs",%s' % sshfsOptions,
                       shutdown='umount /Volumes/ucsc-pleiades-ifs', sleep_before=2),
    ucscpleiadesfs = ['ucscpleiadesibrixfs', 'ucscpleiadesifs', 'ucscpleiadeshomefs'],         
    )

commands['all'] = commands.keys()
##################################################

def prioritize(name1, name2):
    """Sort named commands by priority"""
    if commands[name1]['priority'] == commands[name2]['priority']:
        return 0
    if commands[name1]['priority'] < commands[name2]['priority']:
        return -1
    if commands[name1]['priority'] > commands[name2]['priority']:
        return 1

def unprioritize(name1, name2):
    """Sort named commands by reverse priority"""
    # give= commands not in the command dict lowest priority
    def get_priority(name):        
        if commands.has_key(name):
            return commands[name]['priority']
        return 99

    p1 = get_priority(name1)
    p2 = get_priority(name2)
    
    if p1 == p2:
        return 0
    if p1 < p2:
        return 1
    if p1 > p2:
        return -1

def is_not_alias(name):
    """Return true if name is an alias for other commands rather
    than a command itself.  Used to filter out the aliases after
    they've been expaneded."""
    if is_alias(name):
        return False
    return True

def is_alias(name):
    """Return true if name is an alias for other commands rather
    than a command itself.  Used to filter out the aliases after
    they've been expaneded."""
    if commands.has_key(name) and type(commands[name]) is type([]):
        return True
    return False

def expand_regexp(regexp, candidates, include_aliases=True):
    """Build a list of elements of list candidates matching regexp.
    If include_aliases is True, then automatically return [regexp] if
    regexp exactly matches an alias for other commands."""
    
    result = []

    if include_aliases and regexp in filter(is_alias, commands.keys()):
        result.append(regexp)
    else:
        for cc in candidates:
            if re.compile('^' + regexp + '$').match(cc):
                result.append(cc)
                
    return result

def expand_regexps(regexps, *args, **kw):
    """Given a list of regexps, call expand_regexp on each of them.
    Concatenate all the results together and return a single flat list
    of matches, filtering out duplicates."""    
    # if no regexps, return empty list
    if len(regexps) == 0: return []
    return uniqueify(reduce(lambda x,y: x+y,
                            map(lambda re: expand_regexp(re, *args, **kw),
                                regexps)))
    #result = []
    #for nn in names:
    #    result += expand_regexp(nn, *args, **kw)
    #return result

def expand_aliases(names):
    """Expand commands that are really aliases for other commands"""
    #expand all aliases
    # If a name isn't in the command dict, it could be a running process
    # we're going to kill.  So just pass it through at this stage
    for name in names:
        if commands.has_key(name) and type(commands[name]) is type([]):
            names += commands[name]        
    #Now filter out the (expanded) aliases
    return uniqueify(filter(is_not_alias, names))
    
def uniqueify(names):
    """Given a list, return a new list with duplicate entries
    removed (order is not preserved)!"""
    if len(names) == 0: return []

    names = list(names)    # work on a copy
    names.sort()
    result = [names[0]]
    for name in names:
        if name != result[-1]: result.append(name)            
    return result

def run_processes(names):
    """Run the commands associated with a list of names"""
    global pids
    # Check to make sure they're all there.
    for name in names:
        if not commands.has_key(name):
            print "No command for ", name
            print "Possibilities:"
            print commands.keys()
            sys.exit(1)
            
    # if no args, list availalbe processes
    if len(names) == 0:
        print "Possibilities:"
        print commands.keys()
    else:
        names = list(names)
        names.sort(prioritize) # sort a copy
        for name in names:
            if pids.has_key(name):
                print "%s already running as pid %d" % (name,pids[name])
            else:
                time.sleep(commands[name]['sleep_before'])
                pid = os.fork()
                if pid != 0: # parent
                    pids[name] = pid
                    print "Running %s, pid=%d: %s" % (name, pid,
                                                      commands[name]['command'])
                else: # child
                    os.execl('/bin/sh','junk', '-c', commands[name]['command'])
                    time.sleep(commands[name]['sleep_after'])
                
def kill_processes(names):
    """Kill the processes associated with a list of names"""
    # Check to make sure they're all there.
    for name in names:
        if not pids.has_key(name):
            print "Couldn't find pid for: ", name
            print "Possibilities:"
            print pids.keys()
            sys.exit(1)

    if len(names):
        names = list(names)
        names.sort(unprioritize) # sort a copy
        for name in names:            
            if commands[name]['shutdown']:
                time.sleep(commands[name]['sleep_before_shutdown'])
                os.system(commands[name]['shutdown'])                
            print "Killing %s, pid %d" % (name, pids[name])
            os.system("kill %d" % pids[name])
            del pids[name]
    else:
        # if no args, list availalbe processes
        print "Running Processes:"
        print pids

def test_processes(names):
    """Kill the processes associated with a list of names"""
    if len(names) != 1:
        print "Error!"
        return
    if pids.has_key(names[0]):
        sys.exit(0)
    sys.exit(1)
    
def can(obj, file, *args, **kw):
    """Store an object in a file.  This is a slightly more compact
    form of the build-in 'pickle.'  It allows you to just give a
    filename and let this function open and close the file"""

    if type(file) is type('string'): f=open(file,'w')
    else: f=file

    pickle.dump(obj, f, *args, **kw)    

    if type(file) is type('string'): f.close()

def uncan(file, *args, **kw):
    """Read an object from a file.  This is a slightly more compact
    form of the build-in 'pickle.'  It allows you to just give a
    filename and let this function open and close the file"""
    # If filename, should this read until all exhausted?
    if type(file) is type('string'): f=open(file)
    else: f=file    

    obj = pickle.load(f, *args, **kw)    

    if type(file) is type('string'): f.close()

    return obj

def usage():
    print "Usage: ports [-k] [-h] [-v] [names]"
    print "       -h display this message"
    print "       -v display version"
    print "       -k to kill processes"
    print "       no swithces => run processes."
    print "       run 'ports' to see a list of possibilities"
    print "       run 'ports -k' to see a list of running processes"    

def main():
    """Where the action is"""
    global pids

    try: opts, args = getopt.getopt(sys.argv[1:],"hktv")
    except getopt.GetoptError:
        usage()
        sys.exit(1)
        
    # Figure out what to do
    run_flag = True
    test_flag = False
    
    for o,a in opts:
        if o == '-k':
            run_flag = False
        if o == '-t':
            run_flag = False
            test_flag = True
        if o == '-v':
            print "Ports version ", VERSION
            sys.exit(0)
        if o == '-h':
            usage()
            sys.exit(1)

    # Try to load running pid table.
    # If none, create an empty one
    try: pids = uncan('/tmp/ports-pids')
    except IOError: pids = {}

    # When running processes, expand regexps based on available commands
    # When killing them, expand based on running commands.
    if run_flag or test_flag:
        names = expand_regexps(args, commands.keys())
    else: names = expand_regexps(args, pids.keys())
    names = expand_aliases(names)
    
    if run_flag: run_processes(names)
    elif test_flag: test_processes(names)
    else: kill_processes(names)

    # write pid table
    can(pids, '/tmp/ports-pids')

main()
