3
\$\begingroup\$

I wrote this Python script as a utility script for me. It's something that could be achieved in a simple shell script or alias, but it wasn't 'verbose' enough for me in that format, or customizable.

This script, basically, is just designed to take and process whether a specified DNS TSIG keyfile (necessary for 'signed' requests to DNS servers for different types of requests that need authentication/authorization to be processed) exists, and then calls out to dig, specifying the relevant -k filepath (either a filename, in which case it checks current directory, a full file path, or if no -k is provided then a hard-coded default path).

Basically, looking for feedback on how I could improve this. Any and all feedback is appreciated. (Just don't crucify me for making a more complicated Python script to do something a simple shell alias or script could do, please.)


dig+key.py:

#!/usr/bin/env python

import argparse
import os
import subprocess

DEV_DEBUG_MODE = False

CURRENT_WORKING_DIR = os.getcwd() + "/"

DEFAULT_TSIG_MD5_KEYFILE = "/path/to/tsig/keyfile"


# noinspection PyClassHasNoInit
class _Errors:
    class KeyfilePathIssue(IOError):
        """Error that triggers only when an issue occurred with regards to the specified keyfile"""

        def __init__(self, custommsg):
            if custommsg:
                self.msg = custommsg
            else:
                self.msg = self.message

        def __str__(self):
            return self.msg

    class NoDigBinaryFound(IOError):
        """Error that triggers when we couldn't find the 'dig' binary"""

        def __init__(self, custommsg):
            if custommsg:
                self.msg = custommsg
            else:
                self.msg = self.message

        def __str__(self):
            return self.msg


def _check_if_dig_exists():
    try:
        for p in os.environ["PATH"].split(os.pathsep):
            if DEV_DEBUG_MODE:
                print "'dig' in %s: %s" % (p, os.path.exists(os.path.join(p, 'dig')))

            if os.path.exists(os.path.join(p, 'dig')):
                return

        raise _Errors.NoDigBinaryFound("Could not find dig on the system, check if $PATH is set correctly to "
                                       "include the directory containing dig executable.")
    except _Errors.NoDigBinaryFound as digerr:
        print "Attempting to run failed: %s" % str(digerr)
    except Exception as e:
        print "An unknown error occurred: %s" % str(e)


def _parse_arguments():
    # Parse Arguments
    parser = argparse.ArgumentParser(
        description="Use 'dig' with a TSIG key, with a default key unless otherwise specified.",
        add_help=False
    )
    optargs = parser.add_argument_group("Possible Arguments")
    optargs.add_argument('-h', required=False, action='store_true', default=False,
                         help="Show this help message and then exit.")
    optargs.add_argument('-k', metavar='keyfile', required=False, default=DEFAULT_TSIG_MD5_KEYFILE,
                         help='Specify the path to the keyfile for use with the signature. '
                              'Both {name}.key and {name}.private must exist to use this method. '
                              'If this is not specified, the default hardcoded location will be '
                              'used.')
    optargs.add_argument('args', metavar='STRING', nargs="*",
                         help="Any other 'dig' arguments or query string parameters.")

    parsed_arguments = parser.parse_args()

    if parsed_arguments.h:
        parser.print_help()
        exit()

    if DEV_DEBUG_MODE:
        print "Parsed Arguments: ", parsed_arguments

    return parsed_arguments


if __name__ == "__main__":
    _check_if_dig_exists()

    args = _parse_arguments()

    dig_args = ['dig', '-k']
    try:
        if args.k is not None:
            keypath = args.k
            if not str(keypath).__contains__('/'):
                keypath = CURRENT_WORKING_DIR + keypath

            if DEV_DEBUG_MODE:
                print "KEYPATH: ", keypath

            if not (os.path.exists(keypath + '.key') and os.path.exists(keypath + '.private')):
                raise _Errors.KeyfilePathIssue("Specified keyfile pair does not exist, "
                                               "check specified path, and ensure .key and "
                                               ".private files exist.")

            dig_args.append(keypath)

        if args.args:
            for arg in args.args:
                dig_args.append(arg)

    except _Errors.KeyfilePathIssue as err:
        print "An error occurred: %s" % err.__str__()
        exit()
    except Exception as err:
        print "An unknown error occurred: %s" % err.__str__()
        exit()

    if DEV_DEBUG_MODE:
        print dig_args

    # noinspection PyBroadException
    try:
        subprocess.call(dig_args)
    except (KeyboardInterrupt, RuntimeError):
        exit()
    except Exception as e:
        print "An unknown error occurred: " % str(e)
        exit()
\$\endgroup\$
2
\$\begingroup\$

To quote this answer:

exit is a helper for the interactive shell - sys.exit is intended for use in programs.

You should switch your exit() calls to sys.exit().

\$\endgroup\$
  • 5
    \$\begingroup\$ Better yet, structure the code so that the program ends naturally, and you don't have to call sys.exit() at all (unless you want to exit with a non-zero status). \$\endgroup\$ – 200_success Oct 23 '16 at 4:41
2
\$\begingroup\$

Excpetions already support a custom error message out-of-the-box.

You can just do:

class KeyfilePathIssue(IOError):
    """Error that triggers only when an issue occurred with regards to the specified keyfile"""
    pass

And use it like you already do:

raise KeyfilePathIssue("Specified keyfile pair does not exist, "
                       "check specified path, and ensure .key and "
                       ".private files exist.")

There also seems to be no need for the _Error class holding the custom exceptions, it is obvious from their use and name that they are exceptions.

For a function named _check_if_dig_exists I would expect it to return True or False and only raise an exception if something unexpected happens. I would also rename it to dig_exists.

def dig_exists():
    try:
        for p in os.environ["PATH"].split(os.pathsep):
            if DEV_DEBUG_MODE:
                print "'dig' in %s: %s" % (p, os.path.exists(os.path.join(p, 'dig')))

            if os.path.exists(os.path.join(p, 'dig')):
                return True
        return False
    except Exception as e:
        print "An unknown error occurred: %s" % str(e)
        return False


if __name__ == "__main__":
    if not dig_exists():
        raise NoDigBinaryFound("Could not find dig on the system, check if $PATH is set correctly to include the directory containing dig executable.")
    ...

The main code has also become big enough to put it into a main function and just call that in the if __name__ == "__main__" block.

\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.