#!/usr/bin/env python
# waitfor -- polls until specified network resource is available or event has occured
"""Usage: waitfor <options> <tests>
This utility will wait until a url is available, until a port is being
listened to, until an amount of time has passed or until a shell command
succeeds. It's very useful when you want to coordinate the startup or
shutdown of services. You specify the resources you're waiting for as 
urls such as:
      <http url>             retrieve web page
      <ftp url>              retrieve ftp file
      <file url>             verify existance of file
      cmd:<command>          verify the specified command succeeds
      port:[<host>:]<port>   check for TCP connection
      time:<delay>           interval passed (in seconds)
Prefix any test with '!' to wait for it to fail instead.

Options:
    -?, --help             display this option help
    -p, --poll <interval>  specify poll interval
    -w, --wait <delay>     specify maximum wait time
    -v, --verbose          work noisily (diagnostic output)"""

# Copyright (C) 2004 Denis Hennessy <denis@hennessynet.com>

# v0.5 - 2004-09-20

# This program is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 2
# of the License, or (at your option) any later version.

# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.

# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA  02111-1307, USA.

import getopt
import os
import socket
import sys
import time
import urllib2

class CommandAssertion:
    '''Test whether a command executes successfully'''
    def __init__(self, cmd):
        self.spec = cmd
        self.cmd = cmd[4:]
        
    def __str__(self):
        return self.spec
    
    def test(self):
        status = os.system(self.cmd)
        if status == 0:
            return True
        else:
            return False
    
class NegateAssertion:
    '''Return the negation of the source assertion.'''
    def __init__(self, source):
        self.assertion = source
        
    def __str__(self):
        return 'Not [' + self.assertion.__str__() + ']'
    
    def test(self):
        return not self.assertion.test()
    
class SocketAssertion:
    '''Tests whether a connection can be established to a TCP host/port'''
    def __init__(self, spec):
        self.spec = spec
        i = self.spec.find(':', 5)
        if i != -1:
            self.host = self.spec[5:i]
            self.port = int(self.spec[i+1:])
        else:
            self.host = 'localhost'
            self.port = int(self.spec[5:])

    def __str__(self):
        return self.spec

    def test(self):
        try:
            s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            s.connect((self.host, self.port))
            s.close()
        except:
            return False
        else:
            return True

class TimeAssertion:
    '''Waits until a certain interval has passed, or until a certain time is reached.'''
    def __init__(self, spec):
        proto, timespec = urllib2.splittype(spec)
        try:
            interval = float(timespec)
        except ValueError:
            fatal('Unable to parse time interval -', spec)
        self.endtime = time.time() + interval

    def __str__(self):
        return 'wait until ' + time.asctime(time.localtime(self.endtime))

    def test(self):
        if time.time() > self.endtime:
            return True
        else:
            return False

class UrlAssertion:
    '''Tests whether a specific URL (http, ftp or file) works'''
    def __init__(self, url):
        self.url = url

    def __str__(self):
        return self.url

    def test(self):
        try:
            fp = urllib2.urlopen(self.url).read()
        except:
            return False
        else:
            return True

def fatal(msg, *args):
    '''Display an error message and exit.'''
    print '%s:' % os.path.basename(sys.argv[0]), msg,
    for a in args:
        print a,
    print
    print 'Try %s -? for help.' % os.path.basename(sys.argv[0])
    sys.exit(1)

verboseflag = False    
def verbose(msg, *args):
    '''Display an diagnostic message if verboseflag is True.'''
    if verboseflag:
        print msg,
        for a in args:
            print a,
        print

def main():
    polltime = 1.0
    waittime = 0.0

    # Parse the options first

    try:
        opts, args = getopt.getopt(sys.argv[1:], "?p:rvw:", ["help", "poll=", "retest", "wait="])
    except getopt.error, msg:
        print msg
        print __doc__
        sys.exit(1)

    for o, v in opts:
        if o in ("-?", "--help"):
            print __doc__
            sys.exit(1)
        if o in ("-v", "--verbose"):
            global verboseflag
            verboseflag = True
        if o in ("-p", "--poll"):
            try:
                polltime = float(v)
            except ValueError:
                fatal('Unable to parse polltime -', v)
        if o in ("-w", "--wait"):
            try:
                waittime = time.time() + float(v)
            except ValueError:
                fatal('Unable to parse waittime -', v)

    if not args:
        print 'waitfor: No tests specified.'
        print __doc__
        sys.exit(1)
        
    # What remains in args is a list of assertions, in the form [!]<proto>:<details>. Parse
    # them into a list of objects which can test each assertion.

    asserts = list()
    for a in args:
        negate = False
        if a[0] == '!':
            negate = True
            a = a[1:]
        proto, spec = urllib2.splittype(a)
        if proto == None:
            print 'Badly formed assertion -', a
            sys.exit(1)
        if proto == 'http' or proto == 'https' or proto == 'ftp' or proto == 'file':
            test = UrlAssertion(a)
        elif proto == 'cmd':
            test = CommandAssertion(a)
        elif proto == 'port':
            test = SocketAssertion(a)
        elif proto == 'time':
            test = TimeAssertion(a)
        else:
            print 'waitfor: Unrecognised assertion type -', proto
            sys.exit(1)
        if negate:
            asserts.append(NegateAssertion(test))
        else:
            asserts.append(test)

    verbose('Checking every', polltime, 'seconds')

    try:
        for assertion in asserts:
            while True:
                verbose('Testing:', assertion)
                if assertion.test():
                    break
                if waittime != 0 and time.time() > waittime:
                    verbose('Timeout exceeded')
                    sys.exit(1)
                time.sleep(polltime)
    except KeyboardInterrupt:
        sys.exit(1)

if __name__ == '__main__':
    main()
    sys.exit(0)
