#!/usr/bin/env python

#
# Copyright (c) 2013 Geir Skjotskift <geir@underworld.no>
#
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
#
# Dependencies : https://code.google.com/p/pefile/
#

import argparse
import sys
import pefile

EXIT_OK        = 0x0
EXIT_FAILED    = 0x1
EXIT_NO_PE     = 0x2
EXIT_EXCEPTION = 0x4
EXIT_USAGE     = 0x8

def get_pe(filename):
    """Create a pefile.PE object from filename

    Arguments:
        filename - String
    Returns:
        pefile.PE or None
    """

    try:
        return pefile.PE(filename)
    except (pefile.PEFormatError, OSError):
        return None
    except UnboundLocalError as err:
        sys.stderr.write("Internal Error: {0} when parsing {1}\n".format(str(err), filename))
        sys.exit(EXIT_EXCEPTION)

def check_imports(pe):
    """Check for missing import table

    Arguments:
        pe - pefile.PE object
    Returns:
        int - 1 if there are no import table, 0 if the import table is present
    """

    if hasattr(pe, "DIRECTORY_ENTRY_IMPORT"):
        return EXIT_OK
    return EXIT_FAILED


def check_exports(pe):
    """Check for missing export table

    Arguments:
        pe - pefile.PE object
    Returns:
        int - 1 if there are no export table, 0 if the export table is present
    """

    if hasattr(pe, "DIRECTORY_ENTRY_EXPORT"):
        return EXIT_OK
    return EXIT_FAILED


def check_warnings(pe):
    """Check for PE format parsing errors.

    Arguments:
        pe - pefile.PE object
    Returns:
        int - 1 if there are warnings, 0 if the parsing was OK
    """

    warnings = pe.get_warnings()
    if warnings:
        return EXIT_FAILED
    return EXIT_OK


def main():

    args, parser = parse_options()

    # If get_pe returns None, exit early with status 2 to
    # indicate that the supplied file is not a PE file.
    pe = get_pe(args.filename[0])
    if not pe:
        if args.verbose:
            sys.stderr.write("{0} - PE Format Error\n".format(args.filename[0]))
        sys.exit(EXIT_NO_PE)


    # the following switches are mutually exclusive, so testing for one after
    # another is OK.
    if args.imports:
        ret = check_imports(pe)
    elif args.exports:
        ret = check_exports(pe)
    elif args.warnings:
        ret = check_warnings(pe)
    else:
        parser.print_help()
        sys.exit(EXIT_USAGE)

    if args.verbose:
        if ret == EXIT_OK:
            sys.stderr.write("OK: {0}\n".format(args.filename[0]))
        elif ret == EXIT_FAILED:
            sys.stderr.write("FAILED: {0}\n".format(args.filename[0]))
        else:
            sys.stderr.write("UNKNOWN ({1}): {0}\n".format(args.filename[0], ret))
    sys.exit(ret)


def parse_options():
    """create optionparser and return options and filename

    returns - (namespace object for arguments, argumentparser object)
    """

    parser = argparse.ArgumentParser()
    parser.add_argument("-v", "--verbose",
                        help="Verbose output to stderr",
                        default=False, action="store_true")
    parser.add_argument("filename", metavar="FILE", type=str, nargs=1,
                        help="Name of the file to check")
    group = parser.add_mutually_exclusive_group()
    group.add_argument("-i", "--imports",
                       help="exit code 1 if no import table, or 2 if this is not a PE file",
                       dest="imports", default=False, action="store_true")
    group.add_argument("-e", "--exports",
                       help="exit code 1 if no export table, or 2 if this is not a PE file",
                       dest="exports", default=False, action="store_true")
    group.add_argument("-w", "--warnings",
                       help="exit code 1 there are PE parsing warnings, or 2 if this is not a PE file",
                       dest="warnings", default=False, action="store_true")
    return parser.parse_args(), parser


if __name__ == "__main__":
    main()