vulnerability.py

class Vulnerability:
    def __init__(self, name, sources, sanitizers, sinks, implicit, output_file) -> None:
        self.name = name
        self.sources = sources
        self.sanitizers = sanitizers
        self.sinks = sinks
        self.implicit = implicit
        self.output_file = output_file
        self.output = []

    def is_implicit(self) -> bool:
        if self.implicit == "yes":
            return True
        else:
            return False
        
    def get_sources(self) -> list:
        return self.sources
    
    def get_sinks(self) -> list:
        return self.sinks
    
    def get_sanitizers(self) -> list:
        return self.sanitizers

    def get_vultype(self, name) -> str:
        if name in self.sources:
            return "source"
        elif name in self.sinks:
            return "sink"
        elif name in self.sanitizers:
            return "sanitizer"
        else:
            return ""
    
    def add_instance(self, source: str, sink: str, unsanitized: bool, sanitized: list) -> None:
        if unsanitized:
            unsanitized_flows = "yes"
        else:
            unsanitized_flows = "no"

        vuln = {
                "vulnerability": self.name,
                "source": source,
                "sink": sink,
                "unsanitized flows": unsanitized_flows,
                "sanitized flows": sanitized
        }
        
        same_vuln = self.get_same_vuln(source, sink)


        if vuln in self.output: # if we already have the exact same vuln
            pass
        elif same_vuln != {}: # if we already have a vuln with same source/sink pair but its not the same
            self.output.remove(same_vuln) # remove same from vuln list, we'll add the merged after
            merged_vuln = self.merge_vulnerabilities(same_vuln, vuln)
            self.output.append(merged_vuln)
        else: # new vuln
            """
            if its a new vuln and it has sanitized flows we have to add [] to it
            otherwise vuln will have -> "sanitized flows: [
                '...',
                '...'
            ]
            """
            self.output.append(vuln)
    
    def get_same_vuln(self, source: str, sink: str) -> dict:
        for vuln in self.output:
            if vuln['source'] == source and vuln['sink'] == sink:
                return vuln
        return {}

    def merge_vulnerabilities(self, vuln1, vuln2) -> dict:
        """
        Given two output vulnerabilities with the same source and sink merge them into a single vuln
        Essentialy, merge the sanitized flows
        """

        final_vulnerability = vuln1['vulnerability']
        final_source = vuln1['source']
        final_sink = vuln1['sink']
        v1_unsan_flows = vuln1['unsanitized flows']
        v2_unsan_flows = vuln2['unsanitized flows']
        final_san_flows = []

        if v1_unsan_flows == "yes" or v2_unsan_flows == "yes":
            final_unsan_flows = "yes"
        else:
            final_unsan_flows = "no"

        for sanitizer in vuln1['sanitized flows'] + vuln2['sanitized flows']:
            if sanitizer not in final_san_flows:
                final_san_flows.append(sanitizer)
        
        return {
            "vulnerability": final_vulnerability,
            "source": final_source,
            "sink": final_sink,
            "unsanitized flows": final_unsan_flows,
            "sanitized flows": final_san_flows
        }

    def __str__(self) -> str:
        return 'Vulnerability {}: sources -> {}, sanitizers -> {}, sinks -> {}, implicit flows -> {}'.format(self.name, self.sources, self.sanitizers, self.sinks, self.implicit)