main.py

#!/usr/bin/python3

"""main.py: Generate and submit CSV's to the Qualys API"""

# Built in libraries
import sys
import os
import string
import datetime
import argparse  # parse command line options
import logging  # log to a file for debugging instead of using print functions
import csv  # used to export list of IPs to CSV for each asset group
import glob  # autocomplete file names in user input
import socket  # resolve hostnames to IP
import xml.etree.ElementTree  # parse response from Qualys API
import getpass  # get the password from the user securely

# Libraries from pip
# pip install requests openpyxl pyyaml
import requests  # send POST and GET requests to Qualys API
import yaml
from openpyxl import load_workbook  # used to load the excel file

# Custom Libraries
from progress.bar import Bar # Modified version of https://pypi.org/project/progress/

if os.name != 'nt':  # readline doesn't exist/work on windows
    import readline  # autocomplete user input


def main(argv):
    """
    Main function

    :param argv: list of command line arguments
    :return: None
    """
    # define log levels
    log_levels = {1: logging.DEBUG,
                  2: logging.INFO,
                  3: logging.WARNING,
                  4: logging.ERROR,
                  5: logging.CRITICAL
                  }
    config_exists = os.path.isfile("config.yaml")
    config = None
    # if there are command line arguments parse them else run the main loop
    if len(argv) != 0:
        parser = argparse.ArgumentParser()
        subparsers = parser.add_subparsers(dest="command")
        parser.add_argument("-c", "--config", type=str, help="config file path")
        parser.add_argument("-l", "--log", type=int, help="logging level", choices=range(0, 5))
        parser.add_argument("--log_path", type=str, help="file to log to")
        subparsers.add_parser("import", help="import an excel file and export it to CSV")
        subparsers.add_parser("upload", help="upload asset group CSVs to the API")
        subparsers.add_parser("search", help="Search and match an IP or IPs from to asset group(s)")

        args = parser.parse_args()
        # load custom config if provided
        if args.config:
            config = load_config(args.config)
        elif config_exists:
            config = load_config("config.yaml")

        # setup log level
        if args.log:
            log_level = log_levels[args.log]
        elif config_exists:
            log_level = config['log_level']
        else:
            log_level = logging.INFO

        # setup log path
        if args.log_path:
            log_path = args.log_path
        elif config_exists:
            log_path = config['log_location']
        else:
            log_path = "logs/"

        # if log path doesn't exist create it
        if not os.path.exists(log_path):
            os.makedirs(log_path)

        logging.basicConfig(filename="%s/%s.log" % (log_path, datetime.date.today()), level=log_level,
                            format='%(asctime)s - %(levelname)s - %(message)s')

        # parse the sub-command and run the appropriate function
        if args.command == "import":
            import_excel(config)
        elif args.command == "upload":
            qualys_upload(config)
        elif args.command == "search":
            search(config)
    else:
        # load options from config if it exists otherwise use defaults
        if config_exists:
            config = load_config("config.yaml")
            log_path = config['log_location']
            log_level = log_levels[config['log_level']]
        else:
            log_path = os.path.dirname("logs/%s.log")
            log_level = logging.INFO

        # if log path doesn't exist create it
        if not os.path.exists(log_path):
            os.makedirs(log_path)

        logging.basicConfig(filename="%s/%s.log" % (log_path, datetime.date.today()), level=log_level,
                            format='%(asctime)s - %(levelname)s - %(message)s')
        if config_exists:
            loop(config)
        else:
            loop()


def loop(config=None):
    """
    Interface loop used if no command line arguments are used

    :return: None
    """
    # autocomplete file names https://stackoverflow.com/questions/6656819/filepath-autocompletion-using-users-input
    if os.name != 'nt':
        readline.set_completer_delims(' \t\n;')
        readline.parse_and_bind("tab: complete")
        readline.set_completer(complete)

    while True:
        clear()
        user_input = input(
            '''Welcome to AutoAssetGroup 1.0.2!
            
            \rWould you like to:

            \r1) Import an excel file and export to CSV
            \r2) Upload all CSV asset groups to Qualys
            \r3) Search and match an IP or IPs from an excel file to asset group(s) and
            \r   write the groups to the spreadsheet
            \r4) Resolve hostnames in a spreadsheet to IPs and write the IPs to a 
            \r   specified column
            \r0) Quit

            \r>>> '''
        )
        options = {
            '0': quit,
            '1': import_excel,
            '2': qualys_upload,
            '3': search,
            '4': auto_resolve
        }

        if user_input not in ('0', '1', '2', '3', '4'):
            print("Enter a valid option")
            input("Press enter to continue\n")
        elif user_input in ('1', '2', '3'):
            options[user_input](config)
        elif user_input in ('0', '4'):
            options[user_input]()


def quit():
    """
    Quit the program

    :return: None
    """
    clear()
    sys.exit()


def import_excel(config=None, profile=None):
    """
    Parse information from config or get information from user to pass to create_groups and generate_csv to
    import an excel spreadsheet and export to CSV

    :param config: config file
    :param profile: name of the import profile in the config
    :return:
    """
    # If a config is specified load from config else gather information from user
    if config:

        # build a list of profiles from the config
        options = []

        for key in config['import']:
            options.append(key)

        # if profile was not specified ask the user which one to use
        if not profile:
            print("Enter a profile from below:\n")

            for option in options:
                print("\t%s" % option)

            profile = input("\n>>> ")

        # if the profile exists load options from it and pass to create_groups and generate_csv
        if profile in options:
            profile_config = config['import'][profile]  # load the profile config
            if 'csv_directory' in profile_config and profile_config['csv_directory']:
                csv_directory = profile_config['csv_directory']
            else:
                csv_directory = config['csv_directory']
            
            if csv_directory == None:
                print("Error: csv_directory not specified in config, defaulting to 'csv/'")
                logging.error("csv_directory not specified in config, defaulting to 'csv/'")
                csv_directory = "csv/"

            group_dict = load_csv(csv_directory)  # load existing asset groups in directory

            # load options from the config file
            try:
                wb_path = profile_config['input']
            except KeyError:
                print("Error: input not specified in config for the import profile %s'" % profile)
                logging.error("input not specified in config for the import profile %s'" % profile)
                input("Press enter to continue.\n")
                return

            if 'header' in profile_config:
                header = profile_config['header']
            else:
                print("Header not specified in config, defaulting to 'no'")
                logging.warning("Header not specified in config, defaulting to 'no'")
                header = False

            if 'include' in profile_config:
                include = profile_config['include']
            else:
                include = None

            if 'ip_column' in profile_config and profile_config['ip_column']:
                ip_column = column_to_num(profile_config['ip_column'])
                hostname_column = None
            else:
                ip_column = None
                try:
                    hostname_column = column_to_num(profile_config['hostname_column'])
                except KeyError:
                    print("Error: ip_column nor hostname_column specified in config")
                    logging.error("ip_column nor hostname_column specified in config")
                    input("Press enter to continue.\n")
                    return

            try:
                group_column = column_to_num(profile_config['group_column'])
            except KeyError:
                print("Error: group_column not specified in config")
                logging.error("group_column not specified in config")
                input("Press enter to continue.\n")
                return

            # generate groups
            groups = create_groups(wb_path, header, include, ip_column, hostname_column, group_column, group_dict)
            generate_csv(groups, csv_directory)
            input("\nPress enter to continue\n")
            return
        else:
            print("Profile not found in config.")
            input("Press enter to continue.\n")
            return
    else:
        # collect information from the user since there is no config
        wb_path = input("Please enter the path to the excel file to import: ").strip("\" \'")
        csv_path = input("Please enter the path where the CSV files should be exported: ").strip("\" \'")
        header = string_to_bool(input("Is the first row of the excel file a header? (Y/n): ") or 'Y')
        include = input("Only include sheets(tabs) that include the following phrase in their name "
                        "(press enter to include all sheets): ")
        ips = string_to_bool(input("Are there IPs for each host already in the spreadsheet? (y/N): ") or 'N')
        if ips:
            ip_column = column_to_num(input("What column are the IPs in? (default: A): ") or 'A')
            hostname_column = None
        else:
            hostname_column = column_to_num(input("What column are the hostnames in? (default: A): ") or 'A')
            ip_column = None
        group_column = column_to_num(input("What column are the asset group names located? (default: B): ") or 'B')
        groups = create_groups(wb_path, header=header, include=include, ip_column=ip_column,
                               hostname_column=hostname_column, group_column=group_column)
        generate_csv(groups, csv_path)
    input("Press enter to continue.\n")


def qualys_upload(config=None, base_url=None, username=None, password=None, csv_path=None):
    """
    Upload Asset Groups to Qualys API from CSV format

    :param config: config file
    :param base_url: URL of the Qualys API
    :param username: Username to log into the API
    :param password: Password to log into the API
    :param csv_path: Path containing the CSV files to upload to the API
    :return: None
    """
    if config:
        csv_path = config['csv_directory']
        if 'username' in config['upload']:
            username = config['upload']['username']
        if 'url' in config['upload']:
            base_url = config['upload']['url']

    # fill in information not specified by the config
    if username is None:
        username = input("Enter Qualys API Username: ")
    if password is None:
        password = getpass.getpass("Enter Qualys API Password: ")
    if csv_path is None:
        csv_path = input("Please enter the path containing the CSV files you wish to upload: ").strip("\" \'")
    if base_url is None:
        base_url = input("Please enter the URL for the Qualys API: ")

    # https://community.qualys.com/docs/DOC-4551
    def upload():
        """
        Called when qualys_upload() is called

        :return: None
        """
        session = requests.Session()
        session.headers.update({'X-Requested-With': 'AutoAssetGroup'})
        login(session, username, password)

        for csv_file in os.listdir(csv_path):
            if csv_file.endswith(".csv"):  # only process csv files
                ag_title = csv_file.strip('.csv')  # strip extension to get group name
                with open("%s/%s" % (csv_path, csv_file), 'r') as file:  # all IPs are on one line
                    ips = file.readline()
                add_to_asset_groups(session, ips, ag_title)  # add the IPs to asset groups in Qualys
        logout(session)

    def login(session, username, password):
        """
        Login to the Qualys API

        :param session: requests session
        :param username: Qualys API username
        :param password: Qualys API password
        :return: None
        """
        payload = {
            'action': 'login',
            'username': '%s' % username,
            'password': '%s' % password
        }
        request = session.post('%s/api/2.0/fo/session/' % base_url, data=payload)
        response = xml.etree.ElementTree.fromstring(request.text)

        # print out the response of the API request
        for elem in response.findall('.//TEXT'):
            logging.info("API login: %s" % elem.text)
            print(elem.text)

    def logout(session):
        """
        Logout of the Qualys API

        :param session: requests session
        :return:
        """
        payload = {
            'action': 'logout'
        }
        request = session.post('%s/api/2.0/fo/session/' % base_url, data=payload)
        response = xml.etree.ElementTree.fromstring(request.text)

        # print out the response of the API request
        for elem in response.findall('.//TEXT'):
            logging.info("API logout: %s" % elem.text)
            print(elem.text)

    def add_to_asset_groups(session, ips, ag_title):
        """
        Add IP addresses to asset group using the API

        :param session: requests session
        :param ips: string containing IP addresses to add in CSV format
        :param ag_title: title of the asset group to add IP addresses to
        :return: None
        """
        payload = {
            'action': 'add',
            'ips': '%s' % ips,
            'enable_vm': '1',
            'ag_title': '%s' % ag_title
        }
        request = session.post('%s/api/2.0/fo/asset/ip/' % base_url, data=payload)
        response = xml.etree.ElementTree.fromstring(request.text)

        # print out the response of the API request
        for elem in response.findall('.//TEXT'):
            logging.info("API add to asset group %s: %s" % (ag_title, elem.text))
            print("Adding to asset group %s: %s" % (ag_title, elem.text))

    upload()
    input("Press enter to continue.\n")

def search(config=None):
    """
    Search CSV files in a directory to determine what asset group an IP address belongs to

    :param: config: configuration
    :return: None
    """

    if config:
        csv_path = config['csv_directory']
    else:
        csv_path = input("Please enter the path where Asset Group CSVs are stored to be searched: ").strip("\" \'")
    group_list = load_csv(csv_path)
    auto_search = (input("Would you like to automate a search using data from an excel file? (Y/n) ") or 'Y')
    if auto_search.upper() == 'N':
        print("To return to menu, enter 'exit'")
        while True:
            ip = input("Please enter an ip to search: ")
            if ip.upper() == 'EXIT':
                return
            found_in = []
            for group in Bar("Searching groups...").iter(group_list):
                if group_list[group].contains(ip):
                    found_in.append(group_list[group].name)
            if len(found_in) != 0:
                print("'%s' is in the group(s) %s" % (ip, found_in))
            else:
                print("'%s' was not found in any group" % ip)
    elif auto_search.upper() == 'Y':
        wb_path = input("Please enter the path to the excel file: ").strip("\" \'")
        header = (input("Is the first row of the excel file a header? (Y/n): ") or 'Y')
        if header.upper() == 'Y':
            header = True
        else:
            header = False
        ip_column = column_to_num(input("What column contains the IPs? (default: A): ") or 'A')
        group_column = column_to_num(input("What column to write group names to? (default: B): ") or 'B')
        try:
            wb = load_workbook(filename=wb_path)
        except FileNotFoundError:
            print("Error: %s does not exist!" % wb_path)
            input("Press enter to continue\n")
            return
        sheet = wb.active
        if header:  # Whether or not to include the top row(header)
            start_row = 2
        else:
            start_row = 1
        for row in Bar("Writing values...").iter(range(start_row, sheet.max_row + 1)):
            ip = sheet.cell(column=ip_column, row=row).value
            found_in = []
            for group in group_list:
                # if the asset group contains the ip, add it to the list of found groups
                if group_list[group].contains(ip):
                    found_in.append(group_list[group].name)
            if len(found_in) != 0:  # if group not found write "Not Found" otherwise write a list of asset groups
                sheet.cell(column=group_column, row=row, value="{0}".format(found_in))
            else:
                sheet.cell(column=group_column, row=row, value="{0}".format("Not Found"))
        wb.save(wb_path)
    else:
        print("Please enter 'Y' or 'N'")
        input("Press enter to continue\n")
        return
    input("Press enter to continue\n")

def auto_resolve():
    """
    Automatically resolve all hostnames in a column and output their IPs to another
    :return: None
    """

    wb_path = input("Please enter the path to the excel file: ").strip("\" \'")
    try:
        wb = load_workbook(filename=wb_path)
    except FileNotFoundError:
        print("Error: %s does not exist!" % wb_path)
        input("Press enter to continue\n")
        return
    header = (input("Is the first row of the excel file a header? (Y/n): ") or 'Y')
    if header.upper() == 'Y':
        header = True
    else:
        header = False
    hostname_column = column_to_num(input("What column contains the hostnames? (default: A): ") or 'A')
    ip_column = column_to_num(input("What column to write IPs to? (default: B): ") or 'B')
    sheet = wb.active
    if header:  # Whether or not to include the top row(header)
        start_row = 2
    else:
        start_row = 1
    for row in Bar("Resolving hostnames...").iter(range(start_row, sheet.max_row + 1)):
        hostname = sheet.cell(column=hostname_column, row=row).value
        if " " in hostname:
            logging.warning("'%s' contains a space and can't be resolved" % hostname)  # Log the problematic hostname
            ip = "0.0.0.0"  # Set the hostname to 0.0.0.0 which will not be added to the list of IP's
        else:
            ip = resolve(hostname)  # resolve the hostname to an IP using the sockets module
            sheet.cell(column=ip_column, row=row, value="{0}".format(ip))
        # search group_list for a group with the same name or create a new asset_group if it doesn't exist
        if ip != "0.0.0.0":  # Only add IPs of hostnames that could be resolved
            logging.debug("''%s' resolved to %s" % (hostname, ip))  # log the hostname and ip if log level is debug
        else:
            logging.info("'%s' could not be resolved" % hostname)
    wb.save(wb_path)
    input("Press enter to continue\n")


def load_csv(csv_path):
    """
    Load asset group CSVs from a directory

    :param csv_path: Path where CSV files are stored
    :return: Dict: Dictionary of AssetGroups
    """
    group_dict = {}
    for file in Bar("Loading existing groups...").iter(os.listdir(csv_path)):
        file = "%s/%s" % (csv_path, file)  # convert to absolute path
        base = os.path.basename(file)
        name = os.path.splitext(base)[0]  # Get the name of the asset group based on file name
        group_dict[name] = AssetGroup(name)
        with open(file, 'r') as f:
            reader = csv.reader(f)
            ips = list(reader)[0]
            for ip in ips:
                group_dict[name].add_ip(ip)
    return group_dict


def create_groups(wb_path, header=True, include=None, ip_column=None, hostname_column=1, group_column=2, group_dict={}):
    """
    Parse excel file into a list of AssetGroup objects

    :param wb_path: the path to the excel file to process
    :param header: whether or not the first row is the header and should be ignored (default: True)
    :param include: only include tabs/sheets with this string (default: "")
    :param ip_column: column where ips are located
    :param hostname_column: column where hostnames are located (default: 1)
    :param group_column: column where asset group names are located (default: 2)
    :param group_dict: dictionary of groups to use
    :return: Dictionary of AssetGroup objects
    """
    try:
        wb = load_workbook(filename=wb_path)
    except FileNotFoundError:
        print("Error: %s does not exist!" % wb_path)
        input("Press enter to continue\n")
        return
    sheet_list = []
    current_tab = 0  # current tab, used for progress bar
    total_tabs = 0  # total tabs to process, used for progress bar

    if header:  # Whether or not to include the top row(header)
        start_row = 2
    else:
        start_row = 1

    # build a list of sheets to process
    for sheet in wb.worksheets:
        if include:
            # make sure only tabs containing 'include' (usually "GUESTS") are processed
            if include in wb.sheetnames[wb.index(sheet)]:
                sheet_list.append(sheet)
                total_tabs += 1
        else:
            sheet_list.append(sheet)
            total_tabs += 1

    # process the sheets
    for sheet in sheet_list:
        current_tab += 1
        # Iterate through the list of tabs and display a progress bar
        for row in Bar('%s (Tab %d/%d)' % (wb.sheetnames[wb.index(sheet)],
                                           current_tab, total_tabs)) \
                .iter(range(start_row, sheet.max_row + 1)):
            group_name = sheet.cell(row, group_column).value  # Pull the group name from the row

            if (group_name is None) or (group_name == ""):  # If the cell is empty set the group name to Unknown
                group_name = "Unknown"

            if ip_column is None:
                hostname = sheet.cell(row, hostname_column).value  # Pull the VM name from the row
                # make sure vm names with spaces are excluded since they can not be resolved
                if " " in hostname:
                    logging.warning(
                        "'%s' contains a space and can't be resolved" % hostname)  # Log the problematic hostname
                    ip = "0.0.0.0"  # Set the hostname to 0.0.0.0 which will not be added to the list of IP's
                else:
                    ip = resolve(hostname)  # resolve the hostname to an IP using the sockets module
                # search group_list for a group with the same name or create a new asset_group if it doesn't exist
                if ip != "0.0.0.0":  # Only add IPs of hostnames that could be resolved
                    logging.debug(
                        "''%s' resolved to %s" % (hostname, ip))  # log the hostname and ip if log level is debug
                else:
                    logging.info("'%s' could not be resolved" % hostname)
            else:
                ip = sheet.cell(row, ip_column).value

            if group_name in group_dict:  # if the AssetGroup already exists add the IP to that group
                group_dict[group_name].add_ip(ip)
            else:  # otherwise create a new AssetGroup, add it to the dictionary, and add the IP to the AssetGroup
                group_dict[group_name] = AssetGroup(group_name)
                group_dict[group_name].add_ip(ip)

    return group_dict


def generate_csv(group_dict, save_path):
    """
    Generate a CSV for each AssetGroup in a list of AssetGroup objects

    :param group_dict: list of AssetGroup objects
    :param save_path: path where the csv files should be saved
    :return: None
    """
    if not os.path.exists(save_path):
        os.makedirs(save_path)
    for group in Bar("Saving groups to CSV...").iter(group_dict):
        filename = "%s/%s.csv" % (save_path, group_dict[group].name)
        with open(filename, "w+") as result:
            wr = csv.writer(result, quoting=csv.QUOTE_NONE)
            wr.writerow(group_dict[group].ips)


def load_config(config_path):
    """
    Load a config in YAML format and check that it is valid

    :param config_path: path of the config file to load
    :return: parse config as a dictionary
    """
    config_file = open(config_path, 'r')
    config = yaml.load(config_file)
    check_for = ['csv_directory', 'log_level', 'log_location', 'upload', 'import']
    for key in check_for:
        if key not in config:
            print("%s not specified in %s" % (key, config_path))
            input("Press enter to continue.\n")
            return
    if not os.path.exists(config['csv_directory']):
        os.makedirs(config['csv_directory'])
    if not os.path.exists(config['log_location']):
        os.makedirs(config['log_location'])
    if config['log_level'] not in [1, 2, 3, 4, 5]:
        print("Log level must be a value between 1 and 5")
        input("Press enter to continue.\n")
        return
    return config


class AssetGroup:
    """Class to represent an asset group"""

    def __init__(self, name):
        """
        Constructor for AssetGroup

        :param name: String: name of the AssetGroup
        """
        self._name = name
        self._ips = []
        self._ip_set = set()

    @property
    def name(self):
        """
        Name property

        :return: String: Name of the AssetGroup
        """
        return self._name

    @property
    def ips(self):
        """
        IPs contained within the AssetGroup

        :return: List: List of Strings of IP addresses in the asset group
        """
        return self._ips

    @name.setter
    def name(self, name):
        """
        Sets the name of the AssetGroup

        :param name: String: New name of the AssetGroup
        :return:
        """
        self._name = name
        return name

    def add_ip(self, ip):
        """
        Add an ip address to the ip list

        :param ip: ip address to add
        :return: None
        """
        if (ip not in self._ip_set) and (ip != '0.0.0.0'):
            self._ips.append(ip)
            self._ip_set.add(ip)

    def contains(self, ip):
        """
        Check if an IP is in an AssetGroup

        :param ip: IP to check
        :return: Boolean: Whether the AssetGroup contains the IP
        """
        if ip in self._ip_set:
            return True
        else:
            return False


# Helper functions


def string_to_bool(string):
    """
    Convert a string to a boolean

    :param string: string to convert
    :return: boolean
    """
    if string.upper() in ("TRUE", "T", "Y"):
        return True
    elif string.upper() in ("FALSE", "F", "N"):
        return False
    else:
        raise ValueError("String is not a valid boolean")


def resolve(hostname):
    """
    Return the IP address of a given hostname, returns '0.0.0.0' if the host cannot be resolved for any reason

    :param hostname: the hostname to be resolved
    :return: String: IP address of the host
    """
    try:
        return socket.gethostbyname(hostname)
    except socket.gaierror:
        return "0.0.0.0"


def clear():
    """
    Clear the console

    :return: None
    """
    if os.name == 'nt':
        os.system('cls')
    else:
        os.system('clear')


def column_to_num(column):
    num = 0
    for c in column:
        if c in string.ascii_letters:
            num = num * 26 + (ord(c.upper()) - ord('A')) + 1
    return num


def complete(text, state):
    """
    Used for tab completion
    """
    return (glob.glob(text + '*') + [None])[state]


if __name__ == "__main__":
    try:
        main(sys.argv[1:])
    except KeyboardInterrupt:
        print('\nInterrupted')  # Make ctrl + c look nicer
        sys.exit(0)