"""
Support for os, importlib, argparse and other modules which written for avoiding repetitive code.
"""
import os
import importlib
import argparse
from modules import constants, run_command, supported_environment_functions, output_format_functions, docker_functions

MENU_MESSAGE = '''The 'Am I Exploitable?' CVEs database:
Remote Code Execution (RCE):
Ghost - CVE-2015-0235
GhostCat - CVE-2020-1938
CVE-2021-3711
CVE-2021-41773
CVE-2021-42013
Spring4Shell - CVE-2022-22965
Log4Shell - CVE-2021-44228, CVE-2021-45046, CVE-2021-4104, CVE-2021-45105, CVE-2021-44832, CVE-2021-42013
ShellShock - CVE-2014-6271, CVE-2014-6277, CVE-2014-6278, CVE-2014-7169, CVE-2014-7186, CVE-2014-7187

Privilege Escalation (PLE):
Heartbleed - CVE-2014-0160
Dirty_COW - CVE-2016-5195
Huge_Dirty_COW - CVE-2017-1000405
PWNKIT - CVE-2021-4034
Dirty_Cred - CVE-2021-4154, CVE-2022-5288
Dirty_Pipe - CVE-2022-0847
CVE-2022-25636
NimbusPWN - CVE-2022-29799, CVE-2022-29800
Meltdown - CVE-2017-5754
Spectre - CVE-2017-5715, CVE-2017-5753, CVE-2017-5754

Run options:
all - runs checks for all the CVEs in the database
CVE-YYYY-XXXX - run specific vulnerability check by inserting its CVE id 
name - run specific vulnerability check by inserting its name (for example - Log4Shell)
'''
ALL = 'all'


def execution_according_os_type(vulnerability_identifier, running_os_type, vulnerability_validation, description, graph,
                                report_format, debug, container_name):
    """This function run the cve file according to the os type and calls the output format."""
    if vulnerability_identifier in constants.WINDOWS_VULNERABILITIES and \
            vulnerability_identifier in constants.LINUX_VULNERABILITIES:
        state = vulnerability_validation.main(description, graph, running_os_type, debug, container_name)
    else:
        state = vulnerability_validation.main(description, graph, debug, container_name)
    if report_format:
        output_format_functions.format_type(container_name, report_format, state)


def run_cve_check(vulnerability_identifier, description, graph, report_format, debug, container_name):
    """This function run the cve file that matches the entered vulnerability name."""
    vulnerability_path = f'cves.{vulnerability_identifier}'
    vulnerability_validation = importlib.import_module(vulnerability_path)
    if vulnerability_identifier in constants.ALL_VULNERABILITIES:
        running_os_type = supported_environment_functions.check_supported_environment(vulnerability_identifier, debug,
                                                                                      container_name)
        if running_os_type == constants.LINUX or running_os_type == constants.WINDOWS:
            execution_according_os_type(vulnerability_identifier, running_os_type, vulnerability_validation,
                                        description, graph, report_format, debug, container_name)


def run(vulnerability_identifier, description, graph, report_format, debug, container_name):
    """This function checks if the cve_id that received has a file with its name."""
    cve_dir_path = f"{os.getcwd()}/cves"
    cve_validation_files = os.listdir(cve_dir_path)
    cves_files = [f.split('.')[0] for f in cve_validation_files]
    if vulnerability_identifier in cves_files:
        run_cve_check(vulnerability_identifier, description, graph, report_format, debug, container_name)
    else:
        print(constants.FULL_EXPLANATION_MESSAGE.format('Vulnerability name does not match the CVEs files'))


def fix_cve_format(vulnerability_identifier):
    """This function fixes the cve format so all cases will be included."""
    vulnerability_identifier = vulnerability_identifier.lower()
    if vulnerability_identifier.startswith('cve') and '-' in vulnerability_identifier:
        vulnerability_identifier = vulnerability_identifier.replace('-', '_')
    return vulnerability_identifier


def checks_cve_id_parameter(vulnerability_identifier, description, graph, report_format, debug, container_name):
    """This function run the next function according to the cve_id parameter."""
    fixed_cve = fix_cve_format(vulnerability_identifier)
    if fixed_cve == ALL:
        for vulnerability in constants.ALL_VULNERABILITIES:
            run(vulnerability, description, graph, report_format, debug, container_name)
    elif fixed_cve in constants.ALL_VULNERABILITIES:
        run(fixed_cve, description, graph, report_format, debug, container_name)
    elif fixed_cve in constants.DUPLICATE_VULNERABILITIES_NAMES:
        run(constants.DUPLICATE_VULNERABILITIES_NAMES[fixed_cve], description, graph, report_format, debug, container_name)
    elif fixed_cve in constants.VULNERABILITIES_WITH_MULTIPLE_CVES.keys():
        for cve in constants.VULNERABILITIES_WITH_MULTIPLE_CVES[fixed_cve]:
            run(cve, description, graph, report_format, debug, container_name)
    elif vulnerability_identifier is not None:
        print(constants.FULL_EXPLANATION_MESSAGE.format(MENU_MESSAGE))
        print(constants.FULL_EXPLANATION_MESSAGE.format('Enter one of the running options in order to be scanned'))
    else:
        print(constants.FULL_EXPLANATION_MESSAGE.format('The vulnerability name does not exists in the database'))


def container_mode(debug, container_name, vulnerability_identifier, description, graph, format):
    """This function handles the containers mode."""
    running_containers = docker_functions.get_running_containers(debug)
    if len(running_containers) > 0:
        if container_name:
            for container in container_name:
                if container in running_containers:
                    print(f'\nScanning vulnerabilities on {container} container')
                    checks_cve_id_parameter(vulnerability_identifier, description, graph, format, debug, container)
                else:
                    print(constants.FULL_EXPLANATION_MESSAGE.format(f'The {container} container you mentioned is not '
                                                                    f'valid name or not a running container'))
        else:
            for running_container in running_containers:
                print(f'\nScanning vulnerabilities on {running_container} container')
                checks_cve_id_parameter(vulnerability_identifier, description, graph, format, debug, running_container)
    return running_containers


def set_boolean_parameter(parameter_value):
    """This function sets the boolean arguments."""
    if str(parameter_value).lower() == 'true':
        return True
    return False


def check_parameter_value(parameter_value, parameter_name):
    """This function checks if the boolean parameters are valid."""
    if not (str(parameter_value).lower() == 'true' or str(parameter_value).lower() == 'false'):
        print(constants.FULL_EXPLANATION_MESSAGE.format(f'The {parameter_name} value can be set to True or False only'))
        return False
    return True


def arguments():
    """This function sets the arguments."""
    parser = argparse.ArgumentParser(description="'AM I Exploitable?' is a service that let's you validate "
                                                 "whether or not your system is susceptible to a given CVE")
    parser.add_argument('-v', '--vulnerability_identifier', type=str, default='', help='Enter vulnerability identifier '
                                                                                 'according to the following format:\n'
                                                                                 '`CVE-<YEAR>-<NUMBER>`\nVulnerability '
                                                                                 'name (if exists in db)\n`all` - scan '
                                                                                 'your system for all the '
                                                                                 'vulnerabilities in the cves directory'
                                                                                 '\nIf the argument is not set, a menu '
                                                                                 'message will appear presenting the '
                                                                                 'currently supported vulnerabilities'
                                                                                 '\n')
    parser.add_argument('--description', type=str, default=True, help='A description of the vulnerability '
                                                                       '(True by default)')
    parser.add_argument('-g', '--graph', type=str, default=False, help='Graph which presents the security checks of the '
                                                                  'vulnerability')

    parser.add_argument('-f', '--format', type=str, default='', help='Specify the report output formatter (json, csv, text)')
    parser.add_argument('--debug', type=str, default=False, help='An option to debug the program and see errors')
    parser.add_argument('-c', '--is_container', type=str, default=False, help='Scan only running containers on the host')
    parser.add_argument('-n', '--container_name', type=str, default=False, nargs="+", help='Scan only running '
                                                                                           'containers names that '
                                                                                           'mentioned.\nNeeds to be '
                                                                                           'seperated by commas only')
    return parser.parse_args()


def main():
    """This is the main function."""
    print("Welcome to Rezilion's 'Am I Exploitable?' Service")
    args = arguments()
    args_dict = {'Description': args.description, 'Graph': args.graph, 'Debug': args.debug, 'Container': args.is_container}
    for parameter_name, parameter_value in args_dict.items():
        result = check_parameter_value(parameter_value, parameter_name)
        if not result:
            return False
    args.description = set_boolean_parameter(args.description)
    args.graph = set_boolean_parameter(args.graph)
    args.debug = set_boolean_parameter(args.debug)
    args.container = set_boolean_parameter(args.is_container)
    if supported_environment_functions.get_os(args.debug, args.container_name) == constants.WINDOWS and args.is_container:
        print(constants.NOT_SUPPORTED_MESSAGE.format('Windows containers'))
    elif args.is_container or args.container_name:
        running_containers = container_mode(args.debug, args.container_name, args.vulnerability_identifier,
                                            args.description, args.graph, args.format)
        if len(running_containers) == 0:
            print(constants.FULL_EXPLANATION_MESSAGE.format('Docker containers were not found, unsupported value'))
    else:
        container_name = ''
        checks_cve_id_parameter(args.vulnerability_identifier, args.description, args.graph, args.format, args.debug,
                                container_name)
    return True


if __name__ == '__main__':
    main()
