#!/usr/bin/python
# -*- coding: utf-8 -*-
__author__ = 'Austin Taylor'


import argparse
import logging
import os
import sys

from vulnwhisp.base.config import vwConfig
from vulnwhisp.test.mock import mockAPI
from vulnwhisp.vulnwhisp import vulnWhisperer


def isFileValid(parser, arg):
    if not os.path.exists(arg):
        parser.error("The file %s does not exist!" % arg)
    else:
        return arg


def main():

    parser = argparse.ArgumentParser(description=""" VulnWhisperer is designed to create actionable data from\
     your vulnerability scans through aggregation of historical scans.""")
    parser.add_argument('-c', '--config', dest='config', required=False, default='frameworks.ini',
                        help='Path of config file', type=lambda x: isFileValid(parser, x.strip()))
    parser.add_argument('-s', '--section', dest='section', required=False,
                        help='Section in config')
    parser.add_argument('-f', '--filter', dest='scan_filter', required=False,
                        help='Retrieve scans matching this regex pattern')
    parser.add_argument('--days', dest='days', type=int, required=False,
                        help='Retrieve scans from this many days ago to now')
    parser.add_argument('-l', '--list', dest='list_scans', required=False, action="store_true",
                        help='List available scans')
    parser.add_argument('--source', dest='source', required=False,
                        help='JIRA required only! Source scanner to report')
    parser.add_argument('-n', '--scanname', dest='scanname', required=False,
                        help='JIRA required only! Scan name from scan to report')
    parser.add_argument('-v', '--verbose', dest='verbose', action='store_true',
                        help='Prints status out to screen (defaults to True)')
    parser.add_argument('-F', '--fancy', action='store_true',
                        help='Enable colourful logging output')
    parser.add_argument('-d', '--debug', action='store_true',
                        help='Enable debugging messages')
    parser.add_argument('--mock', action='store_true',
                        help='Enable mocked API responses')
    parser.add_argument('--mock_dir', dest='mock_dir', required=False, default=None,
                        help='Path of test directory')
    args = parser.parse_args()

    # First setup logging
    logging.basicConfig(
        stream=sys.stdout,
        # format only applies when not using -F flag for colouring
        format='%(levelname)s:%(name)s:%(funcName)s:%(message)s',
        level=logging.DEBUG if args.debug else logging.INFO if args.verbose else logging.WARNING
    )
    logger = logging.getLogger()
    # we set up the logger to log as well to file
    fh = logging.FileHandler('vulnwhisperer.log')
    fh.setLevel(logging.DEBUG if args.debug else logging.INFO if args.verbose else logging.WARNING)
    fh.setFormatter(logging.Formatter("%(asctime)s %(levelname)s %(name)s - %(funcName)s:%(message)s", "%Y-%m-%d %H:%M:%S"))
    logger.addHandler(fh)

    if args.fancy:
        import coloredlogs
        coloredlogs.install(level='DEBUG' if args.debug else 'INFO')

    if args.mock:
        mock_api = mockAPI(args.mock_dir, args.verbose)
        mock_api.mock_endpoints()

    exit_code = 0

    try:
        if args.config and not args.section:
            # this remains a print since we are in the main binary
            print(
                "WARNING: No section was specified, vulnwhisperer will scrape enabled modules from config file. \
                \nPlease specify a section using -s. \
                \nExample vuln_whisperer -c config.ini -s nessus"
            )
            logger.info('No section was specified, vulnwhisperer will scrape enabled modules from the config file.')

            config = vwConfig(config_in=args.config)
            enabled_sections = config.get_sections_with_attribute('enabled')

            for section in enabled_sections:
                vw = vulnWhisperer(config=args.config,
                                   profile=section,
                                   verbose=args.verbose,
                                   debug=args.debug,
                                   source=args.source,
                                   scan_filter=args.scan_filter,
                                   days=args.days,
                                   scanname=args.scanname,
                                   list_scans=args.list_scans)
                exit_code += vw.whisper_vulnerabilities()
        else:
            logger.info('Running vulnwhisperer for section {}'.format(args.section))
            vw = vulnWhisperer(config=args.config,
                               profile=args.section,
                               verbose=args.verbose,
                               debug=args.debug,
                               source=args.source,
                               scan_filter=args.scan_filter,
                               days=args.days,
                               scanname=args.scanname,
                               list_scans=args.list_scans)
            exit_code += vw.whisper_vulnerabilities()

        close_logging_handlers(logger)
        sys.exit(exit_code)

    except Exception as e:
        logger.error('{}'.format(str(e)))
        print('ERROR: {error}'.format(error=e))
        # TODO: fix this to NOT be exit 2 unless in error
        close_logging_handlers(logger)
        sys.exit(2)

    close_logging_handlers(logger)

def close_logging_handlers(logger):
    for handler in logger.handlers:
        handler.close()
        logger.removeFilter(handler)

if __name__ == '__main__':
    main()