FortinetCSVParser.py

'''
Title: FortinetCSVParser
Description: This script is a tool that allow the user to parse ".csv" data files generated by the Fortinet solution without effort.
Author: BrAmaral

Disclaimer: This code is not professional, but rather a humble project to help my team optimize their work. Feel free to use it or contribute!
'''

import argparse as ap   # Necessary import for argument parsing.
import csv              # Necessary import for csv file manipulation
import logging as log   # Necessary import for logging operations


# This method is responsible to obtain and parse user arguments given to the program
def parse_args():
    
    # This block creates an parser object instance responsible to collect the user arguments
    parser = ap.ArgumentParser()        # This line creates a parser instance
    parser.add_argument('-i', '--input', help='The CSV file path that will be parsed', required=True)       # This line adds an argument for file input.
    parser.add_argument('-n', '--name', help='The name for your parsed CSV file. Default: Planilha', default='Planilha')      # This line adds an argument for file name. 
    parser.add_argument('-s', '--size', help='The file size limit(in mb) of the parsed CSV. Default: 65.0  *** THIS VALUE SHOULD BE A FLOAT ***', default=65.0)     # This line adds an argument for file size.
    return parser.parse_args()



# This method is responsible to search the title for empty cells. If the method doesnt find the title, the placeholder "EMPTY_CELL" is used.
def findEmptyCell(locator, title, counter):
    while title == "":          # While "title" is empty, do ...
        try:
            nextLine = next(locator)      # Goes to the next line of the csv file
            if nextLine[counter] != "":         # If, on the next line, the position where we found the empty cell is not empty, then ...
                return nextLine[counter]        # Return the title name
        except StopIteration:                   # Exception treatment. The code returns "StopIteration" when it reaches the end of the csv
            return "EMPTY_CELL="                # Using "Empty_Cell=" as a placeholder for treatment later in the code



# This method is responsible to configure the script logs
def logConfig():
    log.basicConfig(format='%(asctime)s %(levelname)s: %(message)s', level=log.INFO)    # This line configures the log format and level



# This method is responsible to delete the empty columns --- DEPRECATED ---
def removeEmptyColumns(titles_lst):
    while "Empty_Cell" in titles_lst:       # This line goes through the list searching for the empty cells      
        titles_lst.remove("Empty_Cell")     # This line delete them once they are found
    return titles_lst


# This method is responsible to parse the titles and store them in a List
def parseTitles(csv_in_memory, parsed_args):
    csv_line = next(csv_in_memory)  # Loads the first CSV line from the csv_in_memory object into csv_line

    titles_lst = [] # This line creates a List called "titles_lst"
    counter = 0     # This line creates a integer variable called "counter" and will be used to go through the titles in the line
    
    for title in csv_line:  # This line creates a loop for each title within the csv line
        if title != "":                                 # If the line is not empty then...
            split_position = title.find("=")            # Find the char "=" and return his position
            titles_lst.append(title[:split_position])   # Separate the title from the begining until the "=" char, not including it, and then store it in the list
            counter += 1                                # Adds one to the counter to go to the next title within the line
        else:                                           # If it is empty (or something else - which is not predicted and thus bad), then...
            with open(parsed_args.input, newline='\n') as csvfile:
                locator = csv.reader(csvfile, delimiter=',', quotechar='"')
                title = findEmptyCell(locator, title, counter)    # Search for the title on the other lines
            split_position = title.find("=")                        # Find the char "=" and return his position
            titles_lst.append(title[:split_position])               # Collect the title string from the begining to the "=" sign, not including it, and then store it in the title list
            counter += 1                                            # Adds one to the counter to go to the next title within the line#
    return titles_lst


# This method is responsible to parse the data in the original CSV for writing on the new one
def parseData(line):
    data_line = []

    for data in line:                              # This line creates a loop for each title within the csv line
        if data != "":                                 # If the line is not empty then...
            split_position = data.find("=")            # Find the char "=" and return his position
            data_line.append(data[split_position+1:])    # This line collects the data after the "=" and store it in the list.
        else:
            data_line.append(data)                      # Else adds the empty data to the list when found
    return data_line


# This method is responsible to create the new CSV file and populate it with the parsed data from the original CSV
def createParsedCSV(titles_lst, parsed_args):
    count = 1   # This counter is used to count how many files are being created (for file naming purposes)
    with open(str(parsed_args.name) + str(count) + ".csv", "w", newline="\n") as parsed_csv:    # Creating a new CSV file to receive the parsed data.
        writer = csv.writer(parsed_csv)                             # Instantiating the writer object
        log.info("Writing titles on " + parsed_args.name + str(count) + "...")
        writer.writerow(titles_lst)                                 # Writing the titles in the first line of the csv file

        log.info("Writing data on " + str(parsed_args.name) + str(count) + "...")
        with open(parsed_args.input, newline='\n') as csvfile:        # This line opens the original csv again to refresh the next() iterator
            csv_in_memory = csv.reader(csvfile, delimiter=',', quotechar='"')   # Loads the csv in memory
            for line in csv_in_memory:      # and for each line, parses it and write it to the new CSV
                file_size = parsed_csv.tell() / 1024 /1024  # This line calculates the file size in mb
                if file_size < float(parsed_args.size):
                    writer.writerow(parseData(line))    # Writes the line on the CSV file 
                else:
                    writer.writerow(parseData(line))    # Write the line he is at before creating a new CSV file
                    count += 1
                    log.info("File memory exceded, creating new CSV file called " + parsed_args.name + str(count) + ".")
                    createNewParsedCSV(titles_lst, csv_in_memory, count, parsed_args)    # Creates a new CSV file to continue writing data 


# This method is responsible to create a new CSV if the file size limit given by the user is reached
def createNewParsedCSV(titles_lst, csv_in_memory, count, parsed_args):
    with open(str(parsed_args.name) + str(count) + ".csv", "w", newline="\n") as parsed_csv:
        writer = csv.writer(parsed_csv)
        log.info("Writing titles on "+ str(parsed_args.name) + str(count) + "...")
        writer.writerow(titles_lst)

        log.info("Writing data on " + str(parsed_args.name) + str(count) + "...")
        for line in csv_in_memory:      # This line continue reading the previous csv loaded in memory
                file_size = parsed_csv.tell() / 1024 /1024  # This line calculates the file size in mb
                if file_size < float(parsed_args.size):
                    writer.writerow(parseData(line))    # Writes the line on the CSV file
                else:
                    writer.writerow(parseData(line))    # Write the line he is at before creating a new CSV file
                    count += 1
                    log.info("File memory exceded, creating new CSV file called " + str(parsed_args.name) + str(count) + ".")
                    createNewParsedCSV(titles_lst, csv_in_memory, count, parsed_args)    # Uses recursion to continue the process until it is finished



# Main code
def main():
    logConfig()                     # Calling the method for the log configurations
    parsed_args = parse_args()      # This line creates an object with the parsed arguments given by the user. In this case, the only variable inside this object is one called "input" which will contain the csv path to be manipulated
    
    if parsed_args.input.find(".csv") != -1:            # This line adds a conditional to validade if the file specified is a CSV
        with open(parsed_args.input, newline='\n') as csvfile: 
            log.info("Loading CSV file...")
            csv_in_memory = csv.reader(csvfile, delimiter=',', quotechar='"')   # Loads the CSV file into the object "csv_in_memory"
            log.info("CSV file loaded!")

            log.info("Parsing titles...")
            titles_lst = parseTitles(csv_in_memory, parsed_args)     # Calling the method "parseTitles" to parse the titles and store them in a List
            log.info("Titles parsed!")
            
        # --- DEPRECATED ---    
        #log.info("Removing empty columns (if applicable)...")
        #titles_lst = removeEmptyColumns(titles_lst)       # This line call the method responsible to remove the empty columns

        log.info("Creating new CSV file...")
        createParsedCSV(titles_lst, parsed_args) # This line calls the method to generate the parsed CSV 
        log.info("Done! You may now open the files.") # --- CODE EXECUTION ENDS HERE ---

    else:       #If the file specified is not a CSV the system logs and closes
        log.error("The file indicated either doesn't exist or is not a '.csv'. Exiting...")



# Default python "main" 
if __name__ == '__main__':
    main()      # --- CODE EXECUTION STARTS HERE ---