polyfile-weave 0.5.5__py3-none-any.whl

polymerge/polymerge.py

@@ -0,0 +1,227 @@
+import copy
+from collections import Counter, defaultdict
+import heapq
+import math
+from statistics import stdev
+from typing import Dict, Set, Tuple
+
+from intervaltree import IntervalTree
+
+from polyfile import logger, __version__
+
+from . import polytracker
+from . import cfg
+
+log = logger.getStatusLogger("PolyMerge")
+
+
+def _function_labels(merged: dict, labeling: Dict[str, Set[Tuple[str]]], ancestry: Tuple[str] = ()):
+    if 'type' in merged:
+        name = merged['type']
+    else:
+        name = merged['name']
+    label: Tuple[str] = ancestry + (name,)
+    for f in merged.get('functions', ()):
+        labeling[f].add(label)
+    for s in merged.get('subEls', ()):
+        _function_labels(s, labeling, ancestry=label)
+
+
+def function_labels(merged_json_obj: dict) -> Dict[str, Set[Tuple[str]]]:
+    labels = defaultdict(set)
+    for merged in merged_json_obj['struc']:
+        _function_labels(merged, labels)
+    return labels
+
+
+class Hashable:
+    def __init__(self, value):
+        self.value = value
+        self._hash = None
+
+    @staticmethod
+    def deephash(obj):
+        if isinstance(obj, dict):
+            return hash(tuple(Hashable.deephash(o) for o in sorted(obj.items())))
+        elif isinstance(obj, list) or isinstance(obj, tuple):
+            return hash(tuple(Hashable.deephash(o) for o in obj))
+        else:
+            return hash(obj)
+
+    def __hash__(self):
+        if self._hash is None:
+            self._hash = Hashable.deephash(self.value)
+        return self._hash
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(value={self.value!r})"
+
+
+class IDHashable:
+    def __init__(self, value):
+        self.value = value
+
+    def __hash__(self):
+        return id(self.value)
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(value={self.value!r})"
+
+
+def build_intervals(elem: dict, tree: IntervalTree = None):
+    if tree is None:
+        tree = IntervalTree()
+    if 'size' in elem:
+        elem_size = elem['size']
+        if elem_size > 0:
+            tree[elem['offset']:elem['offset']+elem_size] = elem
+    if 'subEls' in elem:
+        for child in elem['subEls']:
+            build_intervals(child, tree)
+    return tree
+
+
+def shannon_entropy(data):
+    if not hasattr(data, '__len__'):
+        data = list(data)
+
+    if len(data) <= 1:
+        return 0
+
+    counts = Counter(data)
+
+    probabilities = [float(c) / len(data) for c in counts.values()]
+
+    return -sum(p * math.log(p, 2) for p in probabilities if p > 0.)
+
+
+def _build_type_graph(graph: cfg.DiGraph, elem: dict, parent: str = None):
+    node_name = elem['type']
+    graph.add_node(node_name)
+    if parent is not None:
+        graph.add_edge(parent, node_name)
+    if 'subEls' in elem:
+        for child in elem['subEls']:
+            _build_type_graph(graph, child, node_name)
+
+
+def polyfile_type_graph(polyfile_json_obj: dict) -> cfg.DiGraph:
+    graph = cfg.DiGraph()
+    roots = []
+    for match in polyfile_json_obj['struc']:
+        roots.append(match['type'])
+        _build_type_graph(graph, match)
+    graph.set_roots(roots)
+    return graph
+
+
+def _filter_function_matches(
+        type_dominators: cfg.DAG,
+        node: str,
+        matches: dict,
+        disallowed: Dict[str, int],
+        depth: int = 0):
+    if node in matches:
+        orig_matches = matches[node]
+        for disallowed_depth in range(depth, -1, -1):
+            matches[node] = [
+                func for func in orig_matches if func not in disallowed or disallowed[func] >= disallowed_depth
+            ]
+            if matches[node]:
+                break
+        for func in matches[node]:
+            if func not in disallowed:
+                disallowed[func] = depth
+    for child in type_dominators.successors(node):
+        _filter_function_matches(type_dominators, child, matches, disallowed, depth + 1)
+
+
+def filter_function_matches(type_dominators: cfg.DAG, matches: dict):
+    for root in type_dominators.roots:
+        _filter_function_matches(type_dominators, root, matches, {})
+
+
+def merge(polyfile_json_obj: dict, program_trace: polytracker.ProgramTrace) -> dict:
+    ret = copy.deepcopy(polyfile_json_obj)
+    if 'versions' in ret:
+        ret['versions']['polymerge'] = __version__
+    else:
+        ret['versions'] = {'polymerge': __version__}
+    intervals = None
+    for match in ret['struc']:
+        intervals = build_intervals(match, tree=intervals)
+    matches = defaultdict(set)
+    elems_by_function = defaultdict(set)
+    functions_by_type = defaultdict(set)
+    elems_by_type = defaultdict(set)
+    ret['versions']['polytracker'] = '.'.join(map(str, program_trace.polytracker_version))
+    # The following code assumes that taint was tracked from a single input file.
+    if log.isEnabledFor(logger.STATUS):
+        total_bytes = 0
+        for function_info in program_trace.functions.values():
+            for _, tainted_bytes in function_info.items():
+                total_bytes += len(tainted_bytes)
+        progress = 0
+    for function_name, function_info in program_trace.functions.items():
+        if log.isEnabledFor(logger.STATUS):
+            function_bytes = sum(len(tainted_bytes) for _, tainted_bytes in function_info.items())#.cmp_bytes.items())
+            function_progress = 0
+            function_percent = -1
+        for input_source, tainted_bytes in function_info.items():#cmp_bytes.items():
+            for offset in tainted_bytes:
+                if log.isEnabledFor(logger.STATUS):
+                    progress += 1
+                    function_progress += 1
+                    last_percent = function_percent
+                    function_percent = int((function_progress / function_bytes) * 100.0)
+                    if function_percent > last_percent:
+                        log.status(f"{(progress / total_bytes) * 100.0:.2f}% processing function {function_name}... ({function_percent}%)")
+                for interval in intervals[offset]:
+                    elem = IDHashable(interval.data)
+                    elems_by_function[function_name].add(elem)
+                    elem_type = elem.value['type']
+                    functions_by_type[elem_type].add(function_name)
+                    elems_by_type[elem_type].add(elem)
+                    matches[elem].add(function_name)
+    log.clear_status()
+    dominator_tree = program_trace.cfg.dominator_forest
+    ret['best_function_matches'] = {}
+    for elem_type, elems in elems_by_type.items():
+        # find the function that is most specialized in operating on elems of this type:
+        specialization = [
+            (shannon_entropy(elem.value['type'] for elem in elems_by_function[func]), func)
+            for func in functions_by_type[elem_type]
+        ]
+        if not specialization:
+            continue
+        elif len(specialization) == 1:
+            func_matches = [specialization[0][1]]
+        else:
+            std_dev = stdev(entropy for entropy, _ in specialization)
+            heapq.heapify(specialization)
+            best_value, best_match_func = heapq.heappop(specialization)
+            value_threshold = best_value + std_dev
+            func_matches = [best_match_func]
+            while specialization:
+                best_value, best_match_func = heapq.heappop(specialization)
+                if best_value > value_threshold:
+                    break
+                func_matches.append(best_match_func)
+        # now choose the functions that are roots in the vertex-induced subgraph of the CFG dominator tree:
+        ret['best_function_matches'][elem_type] = [
+            root.name
+            for root in cfg.roots(
+                dominator_tree.vertex_induced_subgraph(program_trace.functions[func] for func in func_matches)
+            )
+        ]
+    # finally, remove redundant functions in the best matches based upon the semantic hierarchy:
+    # first step in that is to build the dominator tree of the PolyFile hierarchy:
+    type_dominators: cfg.DAG = polyfile_type_graph(polyfile_json_obj).dominator_forest
+    # for each type in the 'best_function_matches' mapping, only include a function if that function is not
+    # included for any ancestor of the type in `type_dominators`
+    filter_function_matches(type_dominators, ret['best_function_matches'])
+    for elem, functions in matches.items():
+        elem.value['functions'] = list(functions)
+    return ret
+
+

polymerge/polytracker.py

@@ -0,0 +1,190 @@
+from collections import defaultdict
+from typing import Dict, Iterable, List, Set, Tuple
+
+from polyfile import logger
+
+from .cfg import CFG
+
+log = logger.getStatusLogger('PolyTracker')
+
+
+class FunctionInfo:
+    def __init__(self, name: str, cmp_bytes: Dict[str, List[int]], input_bytes: Dict[str, List[int]] = None, called_from: Iterable[str] = ()):
+        self.name = name
+        self.called_from = frozenset(called_from)
+        self.cmp_bytes = cmp_bytes
+        if input_bytes is None:
+            self.input_bytes = cmp_bytes
+        else:
+            self.input_bytes = input_bytes
+
+    @property
+    def taint_sources(self) -> Set[str]:
+        return self.input_bytes.keys()
+
+    def __getitem__(self, input_source_name):
+        return self.input_bytes[input_source_name]
+
+    def __iter__(self):
+        return self.taint_sources
+
+    def items(self):
+        return self.input_bytes.items()
+
+    def __hash__(self):
+        return hash(self.name)
+
+    def __str__(self):
+        return self.name
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(name={self.name!r}, cmp_bytes={self.cmp_bytes!r}, input_bytes={self.input_bytes!r}, called_from={self.called_from!r})"
+
+
+class ProgramTrace:
+    def __init__(self, polytracker_version: tuple, function_data: Iterable[FunctionInfo]):
+        self.polytracker_version = polytracker_version
+        self.functions: Dict[str, FunctionInfo] = {f.name: f for f in function_data}
+        self._cfg = None
+
+    @property
+    def cfg(self) -> CFG:
+        if self._cfg is not None:
+            return self._cfg
+        self._cfg = CFG(self)
+        self._cfg.add_nodes_from(self.functions.values())
+        for f in list(self.functions.values()):
+            for caller in f.called_from:
+                if caller not in self.functions:
+                    info = FunctionInfo(caller, {})
+                    self.functions[caller] = info
+                    self._cfg.add_node(info)
+                    self._cfg.add_edge(info, f)
+                else:
+                    self._cfg.add_edge(self.functions[caller], f)
+        return self._cfg
+
+    def __repr__(self):
+        return f"{self.__class__.__name__}(polytracker_version={self.polytracker_version!r}, function_data={list(self.functions.values())!r})"
+
+
+POLYTRACKER_JSON_FORMATS = []
+
+
+def normalize_version(*version) -> Tuple[str]:
+    version = tuple(str(v) for v in version)
+    version = tuple(version) + ('0',) * (3 - len(version))
+    version = tuple(version) + ('',) * (4 - len(version))
+    return version
+
+
+def polytracker_version(*version):
+    def wrapper(func):
+        POLYTRACKER_JSON_FORMATS.append((normalize_version(*version), func))
+        POLYTRACKER_JSON_FORMATS.sort(reverse=True)
+        return func
+    return wrapper
+
+
+def parse(polytracker_json_obj: dict) -> ProgramTrace:
+    if 'version' in polytracker_json_obj:
+        version = normalize_version(*polytracker_json_obj['version'].split('.'))
+        if len(version) > 4:
+            log.warn(f"Unexpectedly long PolyTracker version: {polytracker_json_obj['version']!r}")
+        for i, (known_version, parser) in enumerate(POLYTRACKER_JSON_FORMATS):
+            # POLYTRACKER_JSON_FORMATS is auto-sorted in decreasing order
+            if version >= known_version:
+                if i == 0 and version > known_version:
+                    log.warn(f"PolyTracker version {polytracker_json_obj['version']!r} is newer than the latest supported by PolyMerge ({'.'.join(known_version)})")
+                return parser(polytracker_json_obj)
+        raise ValueError(f"Unsupported PolyTracker version {polytracker_json_obj['version']!r}")
+    for function_name, function_data in polytracker_json_obj.items():
+        if isinstance(function_data, dict) and 'called_from' in function_data:
+            # this is the second version of the output format
+            return parse_format_v2(polytracker_json_obj)
+        else:
+            return parse_format_v1(polytracker_json_obj)
+
+
+@polytracker_version(0, 0, 1, '')
+def parse_format_v1(polytracker_json_obj: dict) -> ProgramTrace:
+    return ProgramTrace(
+        polytracker_version=(0, 0, 1),
+        function_data=[FunctionInfo(
+            function_name,
+            {None: taint_bytes}
+        ) for function_name, taint_bytes in polytracker_json_obj.items()
+        ]
+    )
+
+
+@polytracker_version(0, 0, 1, 'alpha2.1')
+def parse_format_v2(polytracker_json_obj: dict) -> ProgramTrace:
+    function_data = []
+    for function_name, data in polytracker_json_obj.items():
+        if 'input_bytes' not in data:
+            if 'cmp_bytes' in data:
+                input_bytes = data['cmp_bytes']
+            else:
+                input_bytes = {}
+        else:
+            input_bytes = data['input_bytes']
+        if 'cmp_bytes' in data:
+            cmp_bytes = data['cmp_bytes']
+        else:
+            cmp_bytes = input_bytes
+        if 'called_from' in data:
+            called_from = data['called_from']
+        else:
+            called_from = ()
+        function_data.append(FunctionInfo(
+            name=function_name,
+            cmp_bytes=cmp_bytes,
+            input_bytes=input_bytes,
+            called_from=called_from
+        ))
+    return ProgramTrace(
+        polytracker_version=(0, 0, 1, 'alpha2.1'),
+        function_data=function_data
+    )
+
+
+@polytracker_version(1, 0, 1)
+def parse_format_v3(polytracker_json_obj: dict) -> ProgramTrace:
+    version = polytracker_json_obj['version'].split('.')
+    function_data = []
+    tainted_functions = set()
+    for function_name, data in polytracker_json_obj['tainted_functions'].items():
+        if 'input_bytes' not in data:
+            if 'cmp_bytes' in data:
+                input_bytes = data['cmp_bytes']
+            else:
+                input_bytes = {}
+        else:
+            input_bytes = data['input_bytes']
+        if 'cmp_bytes' in data:
+            cmp_bytes = data['cmp_bytes']
+        else:
+            cmp_bytes = input_bytes
+        if function_name in polytracker_json_obj['runtime_cfg']:
+            called_from = frozenset(polytracker_json_obj['runtime_cfg'][function_name])
+        else:
+            called_from = frozenset()
+        function_data.append(FunctionInfo(
+            name=function_name,
+            cmp_bytes=cmp_bytes,
+            input_bytes=input_bytes,
+            called_from=called_from
+        ))
+        tainted_functions.add(function_name)
+    # Add any additional functions from the CFG that didn't operate on tainted bytes
+    for function_name in polytracker_json_obj['runtime_cfg'].keys() - tainted_functions:
+        function_data.append(FunctionInfo(
+            name=function_name,
+            cmp_bytes={},
+            called_from=polytracker_json_obj['runtime_cfg'][function_name]
+        ))
+    return ProgramTrace(
+        polytracker_version=version,
+        function_data=function_data
+    )