Our main function accepts two parameters: the file path and output type. We first check the output type to ensure it's in the OUTPUT_OPTS list, just in case the function was called from other code that did not validate. If it is an unknown output format, we'll raise an error and exit the script:
051 def main(file_path, output_type):
052 """
053 The main function handles the main operations of the script
054 :param file_path: path to generate signatures for
055 :param output_type: type of output to provide
056 :return: None
057 """
058
059 # Check output formats
060 if output_type not in OUTPUT_OPTS:
061 logger.error(
062 "Unsupported output format '{}' selected. Please "
063 "use one of {}".format(
064 output_type, ", ".join(OUTPUT_OPTS)))
065 sys.exit(2)
We then start working with the file path, getting its absolute file path on line 67, and checking whether it's a directory on line 69. If so, we begin to iterate over the directory and subdirectories to find and process all files within. The code on lines 71 through 73 should look familiar from Chapter 5, Databases in Python. On line 74, we call the fuzz_file() function to generate our hash value, sigval. This sigval value is then provided, along with the filename and output format, to our output() function:
067 # Check provided file path
068 file_path = os.path.abspath(file_path)
069 if os.path.isdir(file_path):
070 # Process files in folders
071 for root, _, files in os.walk(file_path):
072 for f in files:
073 file_entry = os.path.join(root, f)
074 sigval = fuzz_file(file_entry)
075 output(sigval, file_entry, output_type)
The remainder of our main() function handles single file processing and error handling for invalid paths. If, as seen on lines 76 through 79, the path is a file, we'll process it the same as we did before, generating the hash with fuzz_file() and passing the values to our output() function. Lastly, on lines 80 through 84, we handle errors with accessing the specified file or folder path:
076 elif os.path.isfile(file_path):
077 # Process a single file
078 sigval = fuzz_file(file_path)
079 output(sigval, file_path, output_type)
080 else:
081 # Handle an error
082 logger.error("Error - path {} not found".format(
083 file_path))
084 sys.exit(1)