#!/usr/bin/env python
"""This program generates test data and invokes diff on to test syslimits."""

__author__    = 'Chuck Swiger <chuck@pkix.net>'
__copyright__ = 'Copyright (c) 2005 Charles Swiger'
__license__   = 'Python Software License.'
__version__   = '$Id: difftest.py,v 1.5 2005/04/13 03:00:04 chuck Exp $'

# add program version info
versionNumber = '1.1'
versionInfo = '$Revision: 1.5 $'

import os,sys,getopt,time

config = {}

def usage(msg=None, exitCode=64):
    """Displays program usage message and help, then exits."""
    if msg: sys.stderr.write('\nERROR: %s\n\n' % msg)
    sys.stderr.write('''Usage: %s [options] 

-c, --chunksize   Number of lines added per change.  0 discards lines instead.
-d, --debug       Enable debugging.
-f, --factor      Multiplier growing the test file size between runs.  (> 1)
-h, --help        Display this help message.
-i, --input       Source file to read from.
-l, --location    Location where the test files should be kept.  To be safe,
                  the program will exit if it that path already exists.
-m, --maxsize     The maximum size in MB of test files the program may create.
-r, --ratio       The ratio of changed lines to unchanged lines.
-s, --startsize   The initial size in MB of the test files to create.
-v, --verbose     Enable verbosity.
-V, --version     Display the version of %s and exit

This program attempts to invoke the diff program on successively larger test
files, until the program encounters an exception such as running out of disk
space, or exceeding a process limit threshold like MAXDATASIZE.

The program will repeat lines from the specified input file, or stdin.
''' % (sys.argv[0], sys.argv[0]))
    sys.exit(exitCode)


def version(exitCode=0):
    """Displays program version info, then exits."""
    sys.stderr.write('%s version %s %s\n' % \
                     (sys.argv[0], versionNumber, versionInfo))
    # don't clean up here if just displaying program version
    sys.exit(exitCode)


# option parsing code
def countOptOccurrences(oL, shortName, longName):
    """Counts the number of times an option or long-alias appears."""
    count = 0
    x = [shortName, longName]
    for o in oL:
        if o[0] in x: count = count + 1
    return count

def checkOpt(oD, shortName, longName, defaultValue=None, setValue=1):
    """Scans oD for options with name shortName or longName.  If found,
    returns the value.  If not found, returns defaultValue.  If found, but
    the value is '', returns the setValue."""

    if oD.has_key(shortName):
        possibleValue = oD[shortName]
    elif oD.has_key(longName):
        possibleValue = oD[longName]
    else:
        return defaultValue

    if possibleValue == '':
        return setValue
    else:
        return possibleValue

# parse the command-line options into the config dictionary and use appropriate
# default values as needed
def parseOptions():
    """This parses the command-line options sent to the program with getopt."""
    optDict = {}
    try:
        optionList, args = getopt.getopt(sys.argv[1:], 'dhvVc:i:l:m:p:s:f:r:',
        ['help', 'verbose', 'debug', 'version', 'chunksize=', 'location=',
         'input=', 'ratio=', 'startsize=', 'factor='])
    except:
        print sys.exc_info()
        usage()

    for k,v in optionList: optDict[k] = v
    #print 'optDict = %s' % optDict

    if checkOpt(optDict, '-h', '--help', 0):
        usage(None, exitCode=0)
    if checkOpt(optDict, '-V', '--version', 0):
	version()
    config['debug'] = checkOpt(optDict, '-d', '--debug', None, setValue=countOptOccurrences(optionList, '-d', '--debug'))
    config['verbose'] = checkOpt(optDict, '-v', '--verbose', None, setValue=countOptOccurrences(optionList, '-v', '--verbose'))

    config['chunksize'] = checkOpt(optDict, '-c', '--chunksize', '1')
    config['chunksize'] = int(eval(config['chunksize']))
    config['input'] = checkOpt(optDict, '-i', '--input', '/etc/services')
    config['loc'] = checkOpt(optDict, '-l', '--location', '/tmp/dt')
    config['filea_path'] = config['loc'] + '/a'
    config['fileb_path'] = config['loc'] + '/b'
    config['diff_path'] = config['loc'] + '/diff'
    config['ratio'] = checkOpt(optDict, '-r', '--ratio', '100')
    config['ratio'] = int(eval(config['ratio']))
    
    # you'd think we'd all agree, but you can change these here if you'd like
    config['kb'] = 1024
    config['mb'] = config['kb'] ** 2
    
    config['maxsize'] = checkOpt(optDict, '-m', '--maxsize', '10000')
    config['maxsize'] = config['mb'] * eval(config['maxsize'])
    config['startsize'] = checkOpt(optDict, '-s', '--startsize', '10')
    config['startsize'] = config['mb'] * eval(config['startsize'])
    config['factor'] = checkOpt(optDict, '-f', '--factor', '1.5')
    config['factor'] = eval(config['factor'])
    # bufsize=1 in python means line-buffering
    config['bufsize'] = 10240

def readInput():
    """This reads sample data to be used for the test files, splits them into
    an array of strings, line by line, and keepts track of the lengths."""
    try:
        if config['input'] == '-':
            fd = open(sys.stdin, 'r')
        else:
            fd = open(config['input'], 'r')

        config['linelist'] = []
        config['ll_len'] = []
        for line in fd.readlines():
            config['linelist'].append(line)
            config['ll_len'].append(len(line))
            
        config['inputlc'] = len(config['linelist']) - 1
    except:
        print 'WARNING: unable to open', config['input']
        print sys.exc_info()[0], sys.exc_info()[1]
        sys.exit(66)


def setupTestDir():
    """Sets up the test directory."""
    if os.path.exists(config['loc']):
        print 'WARNING: test location already exists:', config['loc']
        sys.exit(73)
        
    try:
        os.makedirs(config['loc'])
    except:
        print 'WARNING: unable to set up testbed location ', config['loc']
        print sys.exc_info()[0], sys.exc_info()[1]
        sys.exit(73)


def enlargeTestFiles(size, ratio=10):
    """Make the test files larger by size bytes, using a given ratio of lines
    unchanged to lines changed."""
    if config['debug']:
        print 'size_delta=%d, ratio=%d' % (size, ratio)

    fda = open(config['filea_path'], 'a', config['bufsize'])
    fdb = open(config['fileb_path'], 'a', config['bufsize'])
        
    inputlc = config['inputlc']
    lc = 0

    # while there is still more bytes left to append to the test files...
    while size > 0:
        line = config['linelist'][lc]
        size -= config['ll_len'][lc]
        fda.write(line)
        if (lc % ratio) == 0:
            for otherlc in range(0, config['chunksize']):
                fdb.write(config['linelist'][otherlc])
        else:
            fdb.write(line)
            
        if lc >= inputlc:
            lc = 0
        else:
            lc += 1

    # it's safer to flush and close these, then to depend on the OS to get all
    # data written out before we run the diff subprocess...
    fda.flush()
    fda.close()
    fdb.flush()
    fdb.close()

    
def test():
    """This runs a test of diff against files, possibly with different ratios
    of changed data."""
    
    cmd = 'diff -u %s %s > %s' % \
          (config['filea_path'], config['fileb_path'], config['diff_path'])
    if config['debug']:
        print 'diff cmd is:', cmd
    
    ratio = config['ratio']
    if ratio == 0:
        rl = range(5, 100, 5)
    else:
        rl = [ ratio ]

    for ratio in rl:
        size = config['startsize']
        oldsize = 0
        if config['verbose'] or config['debug']:
            print 'INFO: beginning diff trial run with ratio =', ratio
            
        while size < config['maxsize']:
            delta = size - oldsize
            if config['verbose'] or config['debug']:
                print 'filea_size=%d (aka %0.3f MB)' % (size, float(size) / config['mb'])
            
            enlargeTestFiles((size - oldsize), ratio=ratio)
            try:
                os.remove(config['diff_path'])
            except: pass

            now = time.time()
            try:
                diff_exitval = os.system(cmd)
            except:
                print 'got an exception'
                print sys.exc_info()[0], sys.exc_info()[1]
                break
            runtime = time.time() - now
            # 1 lets you grab the files now if you wanna look at em
            if 0 and config['debug']:
                time.sleep(60)

            if diff_exitval == 0 or diff_exitval == 256:
                dobreak = False
                if config['debug']:
                    print 'DEBUG: diff exitted normally:', diff_exitval
            else:
                dobreak = True
                if diff_exitval < 256:
                    print 'NOTICE: diff exitted with signal', diff_exitval
                else:
                    print 'NOTICE: diff exitted with errno', diff_exitval >> 8                
            oldsize = os.stat(config['filea_path']).st_size
            diff_size = os.stat(config['diff_path']).st_size
            if diff_size > config['mb'] * 9:
                print 'time=%0.3f filea_size=%dMB diff_size=%dMB' % (runtime, oldsize / config['mb'], diff_size / config['mb'])
            else:
                print 'time=%0.3f filea_size=%dMB diff_size=%dKB' % (runtime, oldsize / config['mb'], diff_size / config['kb'])

            size *= config['factor']
            if dobreak:
                break

        # this is after we've finished a run and the while loop has stopped
        try:
            os.remove(config['filea_path'])
            os.remove(config['fileb_path'])
        except:
            print 'WARNING: failed to remove test files between runs.'
            print sys.exc_info()
            
def terminate(exitCode=0):
    """Clean up the test directory and exit OK or with specified exit code."""
    if not config['debug']:
        try:
            os.remove(config['filea_path'])
        except: pass
        try:
            os.remove(config['fileb_path'])
        except: pass
        try:
            os.remove(config['diff_path'])
        except: pass
        try:
            os.rmdir(config['loc'])
        except:
            # we have no idea which of the above files may be present at any
            # one time, but if the test directory doesn't go away, something
            # went wrong...
            print 'FATAL: exception trying to clean up test directory'
            print sys.exc_info()[0], sys.exc_info()[1]
    else:
        print 'INFO: not cleaning up test directory:', config['loc']
            
    sys.exit(exitCode)

    
def main():
    parseOptions()
    readInput()
    setupTestDir()
    test()
    
if __name__ == '__main__':
    try:
        main()
    except KeyboardInterrupt:
        print "  ...Control-C seen, quitting program."
        terminate()