#!/usr/bin/python3

"""Automatically hunt for an environment variable difference that causes a program to fail in one case and succeed in another."""

import os
import sys
import copy
import getopt


def usage(retval):
    """Output a usage message."""
    if retval == 0:
        write = sys.stdout.write
    else:
        write = sys.stderr.write

    write(sys.argv[0]+' -s envsnapshotfile\n')
    write(sys.argv[0]+' -w workingenv -b brokenenv -c cmpprog [-r]\n')
    write('\n-r makes '+sys.argv[0]+' work from the good env to\n')
    write('the bad, which is the reverse of how it is usually done\n')
    write('\nPlease see http://stromberg.dnsalias.org/~strombrg/env-search.html\n\n')
    write('To use this program:\n')
    write('1) Go to the account that is working and take an\n')
    write('environment variable snapshot like:\n')
    write(sys.argv[0]+' -s /tmp/working\n\n')
    write('2) Go to the account that is not working, and take another\n')
    write('environment variable snapshot like:\n')
    write(sys.argv[0]+' -s /tmp/broken\n\n')
    write('3) Write a small shell script that returns 0 (true) if your program is working,\n')
    write('and returns nonzero (false) if your program is not working\n.')
    write('\n')
    write('4) Run '+sys.argv[0]+' -w workingenv -b brokenenv -c compareprog\n')
    write('and hopefully this program will discover the difference that makes the difference. :)\n')
    write('\n')
    write('Note: this program largely assumes that there is a Single env var difference\n')
    write('that is causing the problem!  That will not be the\n')
    write('case 100% of the time though.\n')

    sys.exit(retval)


if len(sys.argv) == 1:
    usage(0)

snapshotfilename = ''
workingfilename = ''
brokenfilename = ''
compareprog = ''
reverse = 0
shell = ''
(opts, args) = getopt.getopt(sys.argv[1:], 's:w:b:c:rSh')
for opt in opts:
    if opt[0] == '-s':
        if workingfilename or brokenfilename or compareprog:
            sys.stderr.write('-s cannot be used at the same time as -w, -b or -c')
            usage(1)
        snapshotfilename = opt[1]
    elif opt[0] == '-w':
        if snapshotfilename:
            sys.stderr.write('-w cannot be used at the same time as -s')
            usage(1)
        workingfilename = opt[1]
    elif opt[0] == '-b':
        if snapshotfilename:
            sys.stderr.write('-b cannot be used at the same time as -s')
            usage(1)
        brokenfilename = opt[1]
    elif opt[0] == '-c':
        if snapshotfilename:
            sys.stderr.write('-c cannot be used at the same time as -s')
            usage(1)
        compareprog = opt[1]
    elif opt[0] == '-r':
        reverse = 1
    elif opt[0] == '-S':
        shell = os.environ['SHELL']
    elif opt[0] == '-h':
        usage(0)

if snapshotfilename:
    keys = list(os.environ)
    keys.sort()
    file_ = open(snapshotfilename, 'w')

    # write the number of variables
    file_.write(str(len(keys))+'\n')

    # write each variable's name's, the variable's name, the variable's
    # value's length and the variable's value
    for varname in keys:
        file_.write(str(len(varname))+'\n'+varname+'\n')
        file_.write(str(len(os.environ[varname]))+'\n'+os.environ[varname]+'\n')
    file_.close()
    sys.exit(0)


def readenv(filename):
    """Read a environment file. This is a sort of proto-JSON, I guess - but JSON is a lot nicer :)."""
    file_ = open(filename, 'r')
    dict = {}

    # get the number of variables
    # numvars=string.atoi(string.strip(file_.readline()))
    line = file_.readline().strip()
    try:
        numvars = int(line)
    except ValueError:
        sys.stderr.write('Could not get number of variables from %s\n' % line)
        sys.exit(1)

    sys.stderr.write('Will read %d variables\n' % numvars)

    for varno in range(numvars):
        # get the variable's name
        line = file_.readline().strip()
        try:
            varnamelen = int(line)
        except ValueError:
            sys.stderr.write('Could not get length of variable name from %s\n' % line)
            sys.exit(1)
        # varnamelen=string.atoi(string.strip(file_.readline()))
        varname = file_.read(varnamelen)
        sys.stderr.write('Got variable %d varname %s\n' % (varno, varname))
        # toast the newline
        _ = file_.read(1)

        # get the variable's value
        line = file_.readline().strip()
        try:
            varvaluelen = int(line)
        except ValueError:
            sys.stderr.write('Could not get length of variable value from %s\n' % line)
            sys.exit(1)
#        varvaluelen = string.atoi(string.strip(file_.readline()))
        varvalue = file_.read(varvaluelen)
        sys.stderr.write('Got variable %d varvalue %s\n' % (varno, varvalue))
        # toast the newline
        _ = file_.read(1)

        # add it to the dictionary
        dict[varname] = varvalue

    file_.close()

    return dict


if shell:
    workingdict = readenv(workingfilename)
    retval = os.execve(shell, [shell], workingdict)
    sys.stderr.write('execve failed\n')
    sys.exit(1)


def getstatus(prog, env):
    """Run a subprogram and return its exit code."""
    pid = os.fork()
    if pid == 0:
        # requires that the prog have no arguments.  We're only defining
        # argv[0]
        retval = os.execve(prog, [prog], env)
        return retval
    else:
        retval = os.waitpid(pid, 0)
        # sys.stderr.write("Got "+str(retval)+" from compareprog\n")
        return retval[1]/256


def verify(compareprog, dict, status):
    """Confirm that this single variable appears to be the salient one."""
    if getstatus(compareprog, dict) == status:
        sys.stderr.write('Good, this appears to be a single variable problem.\n')
    else:
        sys.stderr.write('...Sorry, this appears to be a multiple variable problem.\n')
        sys.stderr.write('You may be able to get to the bottom of it\n')
        sys.stderr.write('by running env-search more than once.  Another\n')
        sys.stderr.write('possibility is that there is some other form of variable impacting\n')
        sys.stderr.write('the outcome.\n')


def fixed_recreated(reverse):
    """Compute a function for output verbiage."""
    if reverse:
        return 'recreated'
    else:
        return 'fixed'


def main():
    """Coordinate the search."""
    if brokenfilename and workingfilename and compareprog:
        brokendict = readenv(brokenfilename)
        workingdict = readenv(workingfilename)

        brokenkeys = list(brokendict)
        workingkeys = list(workingdict)

        brokenkeys.sort()
        workingkeys.sort()

        numbrokenkeys = len(brokenkeys)
        numworkingkeys = len(workingkeys)

        # first a do a quick sanity check: make sure that running the program
        # with the "broken env" yields a false return, and running the
        # program with the "working env" yields a true return
        if getstatus(compareprog, brokendict) == 1:
            sys.stderr.write('Good, brokendict does yield a false result from compareprog\n')
        else:
            sys.stderr.write('Not good, brokendict does not yield a false result from compareprog\n')
            sys.stderr.write('This probably either means that your compareprog\n')
            sys.stderr.write('is not working well, or that this is not an env var problem\n')
            sys.exit(1)

        if getstatus(compareprog, workingdict) == 0:
            sys.stderr.write('Good, workingdict does yield a true result from compareprog\n')
        else:
            sys.stderr.write('Not good, workingdict does not yield a true result from compareprog\n')
            sys.stderr.write('This probably either means that your compareprog\n')
            sys.stderr.write('is not working well, or that this is not an env var problem\n')
            sys.exit(1)

        # now generate the lists of degrees of similarity
        sharednameonly = []
        sharednameandvalue = []
        missingfromworking = []
        missingfrombroken = []
        brokenkeyno = 0
        workingkeyno = 0
        # this algorithm depends on the key lists being sorted.  It's
        # we're categorizing the variables into 4 groups:
        # 1) identical names and values
        # 2) identical names but different values
        # 3) Present in working but not in broken
        # 3) Present in broken but not in working
        while brokenkeyno < numbrokenkeys-1 or workingkeyno < numworkingkeys-1:
            # compare name
            string = 'Comparing brokenkeys[%d] with workingkeys[%d], len(brokenkeys) == %d, len(workingkeys) == %d,' + \
                'brokenkeys value is %s, workingkeys value is %s\n'
            sys.stderr.write(string % (
                brokenkeyno,
                workingkeyno,
                len(brokenkeys),
                len(workingkeys),
                brokenkeys[brokenkeyno],
                workingkeys[workingkeyno],
                ))
            if brokenkeys[brokenkeyno] == workingkeys[workingkeyno]:
                # compare value
                if brokendict[brokenkeys[brokenkeyno]] == workingdict[workingkeys[workingkeyno]]:
                    # names and values are the same
                    sharednameandvalue.append(brokenkeys[brokenkeyno])
                else:
                    # names are the same, but not the values
                    sharednameonly.append(brokenkeys[brokenkeyno])
                sys.stderr.write('Want to increment both from broken: %d and working: %d\n' % (brokenkeyno, workingkeyno))
                if brokenkeyno < numbrokenkeys-1:
                    sys.stderr.write('Incrementing brokenkeyno\n')
                    brokenkeyno += 1
                else:
                    sys.stderr.write('Hit end of brokenkeyno\n')
                sys.stderr.write('workingkeyno is %d, numworkingkeys-1 is %d\n' % (workingkeyno, numworkingkeys-1))
                if workingkeyno < numworkingkeys-1:
                    sys.stderr.write('Incrementing workingkeyno\n')
                    workingkeyno += 1
                else:
                    sys.stderr.write('Hit end of workingkeyno\n')
            else:
                # names are not the same
                if brokenkeys[brokenkeyno] < workingkeys[workingkeyno]:
                    missingfromworking.append(brokenkeys[brokenkeyno])
                    if brokenkeyno < numbrokenkeys-1:
                        sys.stderr.write('Want to increment only brokenkeyno\n')
                        brokenkeyno += 1
                elif brokenkeys[brokenkeyno] > workingkeys[workingkeyno]:
                    missingfrombroken.append(workingkeys[workingkeyno])
                    if workingkeyno < numworkingkeys-1:
                        sys.stderr.write('Want to increment only workingkeyno\n')
                        workingkeyno += 1
                else:
                    sys.stderr.write(sys.argv[0]+' internal error 1\n')
                    sys.exit(0)
        print('same in both: %s' % sharednameandvalue)
        print('same name in both but different values: %s' % sharednameonly)
        print('missing from working: %s' % missingfromworking)
        print('missing from broken %s' % missingfrombroken)

        # clearly, we don't need to do anything with variables that are the
        # same in name and value in both the working and broken environments.
        # So we cover the other three cases, but not this one.

        # try setting variables that have the same value in working and
        # broken environments, but do not have the same values, to one at a
        # time, -have- the same values
        if not reverse:
            # work from the bad account to the good account
            testdict = copy.deepcopy(brokendict)
            for key in sharednameonly:
                testdict[key] = workingdict[key]
                if getstatus(compareprog, testdict) == 0:
                    sys.stderr.write('Good, modifying the value of:\n')
                    sys.stderr.write(key+'\n')
                    sys.stderr.write('...to...\n')
                    sys.stderr.write(testdict[key]+'\n')
                    sys.stderr.write('...appears to have %s the problem.\n' % (fixed_recreated(reverse)))
                    singledict = copy.deepcopy(brokendict)
                    singledict[key] = workingdict[key]
                    verify(compareprog, singledict, 0)
                    sys.exit(0)
                else:
                    sys.stderr.write('var '+key+' did not appear to help, continuing...\n')

            # try adding variables to broken env that existed only in working env
            for key in missingfrombroken:
                testdict[key] = workingdict[key]
                if getstatus(compareprog, testdict) == 0:
                    sys.stderr.write('Good, adding the value of:\n')
                    sys.stderr.write(key+'\n')
                    sys.stderr.write('...and defining it to...\n')
                    sys.stderr.write(testdict[key]+'\n')
                    sys.stderr.write('...appears to have %s the problem\n' % (fixed_recreated(reverse)))
                    singledict = copy.deepcopy(brokendict)
                    singledict[key] = workingdict[key]
                    verify(compareprog, singledict, 0)
                    sys.exit(0)
                else:
                    sys.stderr.write('var '+key+' did not appear to help, continuing...\n')

            # try removing variables from broken env that don't exist in working env
            for key in missingfromworking:
                del(testdict[key])
                if getstatus(compareprog, testdict) == 0:
                    sys.stderr.write('Good, removing the variable:\n')
                    sys.stderr.write(key+'\n')
                    sys.stderr.write('...appears to have %s the problem\n' % (fixed_recreated(reverse)))
                    singledict = copy.deepcopy(brokendict)
                    del(singledict[key])
                    verify(compareprog, singledict, 0)
                    sys.exit(0)
                else:
                    sys.stderr.write('var '+key+' did not appear to help, continuing...\n')
        else:
            testdict = copy.deepcopy(workingdict)
            for key in sharednameonly:
                testdict[key] = brokendict[key]
                if getstatus(compareprog, testdict) == 1:
                    sys.stderr.write('Good, modifying the value of:\n')
                    sys.stderr.write(key+'\n')
                    sys.stderr.write('...to...\n')
                    sys.stderr.write(testdict[key]+'\n')
                    sys.stderr.write('...appears to have %s the problem.\n' % (fixed_recreated(reverse)))
                    singledict = copy.deepcopy(workingdict)
                    singledict[key] = brokendict[key]
                    verify(compareprog, singledict, 1)
                    sys.exit(0)
                else:
                    sys.stderr.write('var '+key+' did not appear to help, continuing...\n')

            # try adding variables to broken env that existed only in working env
            for key in missingfromworking:
                testdict[key] = brokendict[key]
                if getstatus(compareprog, testdict) == 1:
                    sys.stderr.write('Good, adding the value of:\n')
                    sys.stderr.write(key+'\n')
                    sys.stderr.write('...and defining it to...\n')
                    sys.stderr.write(testdict[key]+'\n')
                    sys.stderr.write('...appears to have %s the problem\n' % (fixed_recreated(reverse)))
                    singledict = copy.deepcopy(workingdict)
                    singledict[key] = brokendict[key]
                    verify(compareprog, singledict, 1)
                    sys.exit(0)
                else:
                    sys.stderr.write('var '+key+' did not appear to help, continuing...\n')

            # try removing variables from broken env that don't exist in working env
            for key in missingfrombroken:
                del(testdict[key])
                if getstatus(compareprog, testdict) == 1:
                    sys.stderr.write('Good, removing the variable:\n')
                    sys.stderr.write(key+'\n')
                    sys.stderr.write('...appears to have %s the problem\n' % (fixed_recreated(reverse)))
                    singledict = copy.deepcopy(workingdict)
                    del(singledict[key])
                    verify(compareprog, singledict, 1)
                    sys.exit(0)
                else:
                    sys.stderr.write('var '+key+' did not appear to help, continuing...\n')

        sys.stderr.write('Sorry, this does not appear to be a single environment\n')
        sys.stderr.write('variable problem.  It may be a multiple environment\n')
        sys.stderr.write('variable problem, or it may not be an environment variable problem\n')
        sys.exit(1)

    else:
        sys.stderr.write('You must specify -b brokenfile, -w workingfile, and -c compareprog\n')
        usage(1)


main()