test_driver.py

import copy
import csv
import datetime
import json
import os
import queue
import shutil
import sys
import tempfile
import threading
import time
import timeit
from collections import OrderedDict
from typing import Union

import psutil
import summarize_json


class TestDriver:
    def __init__(self, tests:list[str], num_containers:int, summarization_reproduce_failure_config:dict):
        #Create the queue and enque all of the tests
        self.testing_queue = queue.Queue()
        for test in tests:
            self.testing_queue.put(test)
        
        self.total_number_of_tests = self.testing_queue.qsize()
        #Creates a lock that will be used to synchronize access to this object
        self.lock = threading.Lock()
        self.start_time = timeit.default_timer()
        self.failures = []
        self.successes = []
        self.errors = []
        self.container_ready_time = None
        
        #No containers have failed
        self.container_failure = False

        #Just make a random folder to store attack data that we donwload
        self.attack_data_root_folder = tempfile.mkdtemp(prefix="attack_data_", dir=os.getcwd())
        print("Attack data for this run will be stored at: [%s]"%(self.attack_data_root_folder))
        
        #Not used right now, but we will keep it around for a bit in case we want to use it again
        self.start_barrier = threading.Barrier(num_containers)

        #The config that will be used for writing out the error config reproduction fiel
        self.summarization_reproduce_failure_config = copy.deepcopy(summarization_reproduce_failure_config)


        #According to the docs:
        # Warning the first time this function is called with interval = 0.0 or None it will return a meaningless 0.0 value which you are supposed to ignore.
        # We call this exactly once here to prime for future calls and throw away the result
        cpu_info = psutil.cpu_times_percent(percpu=False)
            

    def checkContainerFailure(self)->bool:
        
        self.lock.acquire()
        
        try:
            result = self.container_failure
        finally:
            self.lock.release()
        
        
        
        return result
        

    def containerFailure(self)->None:
        self.lock.acquire()
        try:
            self.container_failure = True
        finally:
            self.lock.release()
        
    def checkIfTestsRemain(self):
        failure = self.checkContainerFailure()
        if failure:
            #Just return None, don't continue testing if a container crashed
            #Indicate there are no tests remaining
            return False
        
        try:
            #This call isn't reliable according to documentation, but can save us some time.
            #Err on the side of caution
            return not self.testing_queue.empty()
        except Exception as e:
            print("Error determinging if testing queue was empty.  Return False and try to get something.",file=sys.stderr)
            return True
        

    def getTest(self)-> Union[str,None]:
        
        failure = self.checkContainerFailure()
        
        

        if failure:
            #Just return None, don't continue testing if a container crashed
            return None

        try:
            return self.testing_queue.get(block=False)
        except Exception as e:
            return None
        
    def addSuccess(self, result:dict, duration_string:str)->None:
        print("Test PASSED: [%s --> %s] in %s"%(result['detection_name'], result['detection_file'], duration_string))
        self.lock.acquire()
        try:
            self.successes.append(result)
        finally:
            self.lock.release()
        

    def addFailure(self, result:dict, duration_string:str)->None:
        print("Test FAILED: [%s --> %s] in %s"%(result['detection_name'], result['detection_file'], duration_string))
        self.lock.acquire()
        try:
            self.failures.append(result)
        finally:
            self.lock.release()

    def addError(self, result:dict, duration_string:str)->None:
        #Make sure that even errors have all of the required fields.
        for required_field in ['search_string', 'diskUsage','runDuration', 'detection_name', 'scanCount', 'detection_error', 'detection_file']:
            if required_field not in result:
                result[required_field] = ""
        if  'error' not in result:
            result['error'] = True
        if  'success' not in result:
            result['success'] = False
        print("Test ERROR: [%s --> %s] in %s"%(result['detection_name'], result['detection_file'], duration_string))
        self.lock.acquire()
        try:
            self.errors.append(result)
        finally:
            self.lock.release()
    

    def outputResultsCSV(self, field_names:list[str], output_filename:str, data:list[dict], baseline:OrderedDict)->bool:
        success = True

        print("Generating %s..."%(output_filename), end='')
        self.lock.acquire()
        
        try:                        
            with open(output_filename, 'w') as csvfile:
                header_writer = csv.writer(csvfile, quoting=csv.QUOTE_ALL)
                for key in baseline:
                    #Very basic support for pretty pritning dicts. Doesn't handle more than 1 nested dict
                    if type(baseline[key]) is OrderedDict:
                        header_writer.writerow([key, "-"])
                        for nestedkey in baseline[key]:
                            header_writer.writerow([nestedkey, baseline[key][nestedkey]])
                    #Basic support for 1 layer nested list. Doesn't handle more than 1.
                    elif type(baseline[key]) is list and len(baseline[key])>0:
                        header_writer.writerow([key, baseline[key][0]])
                        for i in range(1,len(baseline[key])):
                            header_writer.writerow(['-', baseline[key][i]])
                        
                    else:
                        header_writer.writerow([key, baseline[key]])
                header_writer.writerow(['',''])
                csv_writer = csv.DictWriter(csvfile, fieldnames=field_names)
                csv_writer.writeheader()
                for row in data:
                    csv_writer.writerow(row)
            print("Done with [%d] detections"%(len(data)))

        except Exception as e:
            print("Failure writing to CSV file for [%s]:"%(output_filename, str(e)))
            success = False

        finally:
            self.lock.release()

        return success

    def outputResultsJSON(self, field_names:list[str], output_filename:str, data:list[dict], baseline:OrderedDict)->bool:
        success = True
        try:
            with open(output_filename, "w") as jsonFile:
                json.dump({'baseline': baseline, 'results':data}, jsonFile, indent="   ")
        except Exception as e:
            print("There was an error generating [%s]: [%s]"%(output_filename, str(e)))
            success = False
        return success
        
    def outputResultsFile(self, field_names:list[str], output_filename:str, data:list[dict], baseline:OrderedDict, output_json:bool=True, output_csv:bool=True)->bool:
        success = True
        if output_csv:
            success |= self.outputResultsCSV(field_names, output_filename + ".csv", data, baseline)
        if output_json:
            success |= self.outputResultsJSON(field_names, output_filename + ".json", data, baseline)
        return success
        

    def outputResultsFiles(self, baseline:OrderedDict, fields:list[str]=['detection_name', 'detection_file','runDuration','diskUsage', 'search_string', 'error', 'success', 'scanCount', 'detection_error'])->bool:
        results_directory = "test_results"
        try:
            shutil.rmtree(results_directory,ignore_errors=True)
            os.mkdir(results_directory)
        except Exception as e:
            print("There was an error removing the results directory [%s]: [%s].\n\t We will try to continue output anyway."%(results_directory, str(e)))


        res = self.outputResultsFile(fields,os.path.join(results_directory, "success"), self.successes, baseline)
        res |= self.outputResultsFile(fields, os.path.join(results_directory, "failure"), self.failures, baseline)
        res |= self.outputResultsFile(fields, os.path.join(results_directory, "error"), self.errors, baseline)
        combined_data = self.successes + self.failures + self.errors
        res |= self.outputResultsFile(fields, os.path.join(results_directory, "combined"), combined_data, baseline)
        
        try:
            success, test_count,pass_count,fail_count,error_count = \
                summarize_json.outputResultsJSON("summary.json", combined_data, 
                                                 baseline, output_folder=results_directory, 
                                                 summarization_reproduce_failure_config=self.summarization_reproduce_failure_config)
            summarize_json.print_summary(test_count, pass_count, fail_count, error_count)
            res |= success
        except Exception as e:
            print("Failure writing the summary file: [%s]"%str(e),file=sys.stderr)
            res = False

        return res

    def finish(self, baseline:OrderedDict):
        self.cleanup()
        success = True
        if self.outputResultsFiles(baseline) == False:
            print("There was an error generating one or more of the output files. "\
                  "Check the logs for details.",file=sys.stderr)
            success = False
        

        if self.checkContainerFailure():
            print("One or more containers crashed or the test was HALTED early, so testing did not complete successfully. We wrote out all the results that we could")
            return False
        else:
            return success

        

    def cleanup(self):
        self.lock.acquire()
        try:
            print("Removing all attack data that was downloaded during this test at: [%s]"%(self.attack_data_root_folder))
            shutil.rmtree(self.attack_data_root_folder)
            print("Successfully removed all attack data")
        finally:
            self.lock.release()

    def get_system_stats(self)->str:
        
        bytes_per_GB = 1024 * 1024 * 1024
        cpu_info = psutil.cpu_times_percent(percpu=False)
        memory_info = psutil.virtual_memory()
        disk_usage_info = psutil.disk_usage('/')

        #macOS is really weird about disk usage.... so to get free space we use TOTAL-FREE = USED instead of just USED
        corrected_used_space = disk_usage_info.total - disk_usage_info.free

        cpu_info_string =        "Total CPU Usage   : %d%% (%d CPUs)"%(100 - cpu_info.idle, psutil.cpu_count(logical=False))
        memory_info_string =     "Total Memory Usage: %0.1fGB USED / %0.1fGB TOTAL"%((memory_info.total - memory_info.available) / bytes_per_GB, memory_info.total / bytes_per_GB)
        disk_usage_info_string = "Total Disk Usage  : %0.1fGB USED / %0.1fGB TOTAL"%(corrected_used_space / bytes_per_GB, disk_usage_info.total / bytes_per_GB)
        
        return "System Information:\n\t%s\n\t%s\n\t%s"%(cpu_info_string, memory_info_string, disk_usage_info_string)


    def summarize(self,testing_currently_active:bool=False)->bool:
        self.lock.acquire()
        try:
        
            #Get a summary of some system stats
            system_stats=self.get_system_stats()
            
            current_time = timeit.default_timer()
            
            

            if not testing_currently_active:
                #Testing has not started yet. We are setting up containers
                print("***********PROGRESS UPDATE***********\n"\
                      "\tWaiting for container setup: %s\n\t%s\n"%(datetime.timedelta(seconds=current_time - self.start_time),system_stats))
            else:
                
                if self.container_ready_time is None:
                    #This is the first status update since container setup has completed.  Get the current time.
                    #This makes our remaining time estimates better since that estimate should not involve
                    #the container setup time  
                    print("SETTING THE CONTAINER READY TIME!")
                    
                    self.container_ready_time = current_time

                numberOfCompletedTests = len(self.successes) + len(self.failures) + len(self.errors)
                remaining_tests = self.testing_queue.qsize()         
                testsCurrentlyRunning = self.total_number_of_tests - remaining_tests - numberOfCompletedTests
                total_execution_time_seconds = round(current_time - self.start_time)

                test_execution_time_seconds = current_time - self.container_ready_time
                
                
                if numberOfCompletedTests == 0 or test_execution_time_seconds == 0:
                    estimated_seconds_to_finish_all_tests = "UNKNOWN"
                    estimated_completion_time_string = "UNKNOWN"
                    average_time_per_test_string = "UNKNOWN"
                else:
                    average_time_per_test = test_execution_time_seconds / numberOfCompletedTests
                    average_time_per_test_string = datetime.timedelta(seconds=round(test_execution_time_seconds/numberOfCompletedTests))
                    #divide testsCurrentlyRunning by 2.0 because, on average, each running test will be 50% completed
                    estimated_seconds_to_finish_all_tests = round(average_time_per_test * (remaining_tests + testsCurrentlyRunning/2.0))
                    estimated_completion_time_string = datetime.timedelta(seconds=estimated_seconds_to_finish_all_tests)
                    
                

            
                print(f"***********PROGRESS UPDATE***********\n"\
                f"\tElapsed Time               : {datetime.timedelta(seconds=total_execution_time_seconds)}\n"\
                f"\tTest Execution Time        : {datetime.timedelta(seconds=round(test_execution_time_seconds))}\n"\
                f"\tEstimated Remaining Time   : {estimated_completion_time_string}\n"\
                f"\tTests to run               : {remaining_tests}\n"\
                f"\tAverage Time Per Test      : {average_time_per_test_string}\n",
                f"\tTests currently running    : {testsCurrentlyRunning}\n"\
                f"\tTests completed            : {numberOfCompletedTests}\n"\
                f"\t\tSuccess : {len(self.successes)}\n"\
                f"\t\tFailure : {len(self.failures)}\n"\
                f"\t\tError   : {len(self.errors)}\n"\
                f"\t{system_stats}\n")

        except Exception as e:
            print("Error in printing execution summary: [%s]"%(str(e)))
        finally:
            self.lock.release()
            
        
        #Return true while there are tests remaining
        completed_tests = len(self.successes) + len(self.failures) + len(self.errors)
        remaining_tests = self.total_number_of_tests - completed_tests
        return remaining_tests > 0
                
        
        
    def addResult(self, result:dict, duration_string:str)->None:
        try:
            if result['detection_result']['error'] is True:
                self.addError(result['detection_result'], duration_string = duration_string)
            elif result['detection_result']['success'] is False:
                #This is actually a failure of the detection, not an error. Naming is confusiong
                self.addFailure(result['detection_result'], duration_string = duration_string)
            elif result['detection_result']['success'] is True:
                self.addSuccess(result['detection_result'], duration_string = duration_string)
        except Exception as e:
            #Neither a success or a failure, so add the object to the failures queue
            print('"There was an error adding the result: [%s]'%(str(e)))
            self.addError({'detection_file':"Unknown File", "detection_error":str(result)})