behavior tree rewrite

modules/processing/behavior.py

@@ -6,7 +6,9 @@
 import logging
 import os
 import struct
+from copy import deepcopy
 from contextlib import suppress
+from collections import OrderedDict
 
 from lib.cuckoo.common.abstracts import Processing
 from lib.cuckoo.common.compressor import CuckooBsonCompressor
@@ -1001,8 +1003,25 @@ class ProcessTree:
     key = "processtree"
 
     def __init__(self):
+
         self.processes = []
         self.tree = []
+        """
+        {
+            "pid_graph": {
+                <parent1>: [<child11>, <child12>, ...],
+                <parent2>: [<child21>, <child22>...],
+                ...
+            },
+            pid_map: {}
+        }
+        """
+        self.graph_tree = {
+            "pid_graph": OrderedDict(),
+            "pid_map": {},
+        }
+
+        # self.graph_tree["pid_map"][10664]
 
     def add_node(self, node, tree):
         """Add a node to a process tree.
@@ -1016,7 +1035,9 @@ def add_node(self, node, tree):
             # If the current process has the same ID of the parent process of
             # the provided one, append it the children.
             if process["pid"] == node["parent_id"]:
+                self.graph_tree["pid_graph"].setdefault(str(process["pid"]), list()).append(str(node["pid"]))
                 process["children"].append(node)
+                # self.graph_tree["pid_map"][node["pid"]] = node
                 ret = True
                 break
             # Otherwise try with the children of the current process.
@@ -1053,25 +1074,30 @@ def run(self):
             for process_again in self.processes:
                 if process_again == process:
                     continue
-                # If we find a parent for the first process, we mark it as
-                # as a child.
+
+                str_pid = str(process["pid"])
+                # If we find a parent for the first process, we mark it as a child.
                 if process_again["pid"] == process["parent_id"]:
+                    self.graph_tree["pid_graph"][str_pid] = list()
+                    self.graph_tree["pid_map"][str_pid] = deepcopy(process)
                     has_parent = True
                     break
 
             # If the process has a parent, add it to the children list.
             if has_parent:
                 children.append(process)
+                self.graph_tree["pid_graph"][str_pid] = list()
+                self.graph_tree["pid_map"][str_pid] = deepcopy(process)
             # Otherwise it's an orphan and we add it to the tree root.
             else:
                 self.tree.append(process)
-
+                self.graph_tree["pid_graph"][str_pid] = list()
+                self.graph_tree["pid_map"][str_pid] = deepcopy(process)
         # Now we loop over the remaining child processes.
         for process in children:
             if not self.add_node(process, self.tree):
                 self.tree.append(process)
-
-        return self.tree
+        return self.graph_tree
 
 
 class EncryptedBuffers:

web/analysis/templatetags/generic_tags.py

@@ -12,20 +12,25 @@ def endswith(value, thestr):
 
 @register.filter("proctreetolist")
 def proctreetolist(tree):
-    outlist = []
     if not tree:
-        return outlist
-    stack = deque(tree)
+        return []
+    graph = {}
+    outlist = []
+    stack = deque(tree.get("pid_graph", {}).keys())
     while stack:
-        node = stack.popleft()
+        pid_id = stack.popleft()
         is_special = False
-        if "startchildren" in node or "endchildren" in node:
+        if "startchildren" in pid_id or "endchildren" in pid_id:
             is_special = True
-            outlist.append(node)
+            outlist.append(pid_id)
         else:
+            node = tree["pid_map"][pid_id]
+            str_pid = str(node["pid"])
             newnode = {}
-            newnode["pid"] = node["pid"]
+            newnode["pid"] = str_pid
             newnode["name"] = node["name"]
+            newnode["parent_id"] = str(node["parent_id"])
+
             if "module_path" in node:
                 newnode["module_path"] = node["module_path"]
             if "environ" in node and "CommandLine" in node["environ"]:
@@ -49,10 +54,13 @@ def proctreetolist(tree):
                     cmdline = cmdline[:200] + " ...(truncated)"
                 newnode["commandline"] = convert_to_printable(cmdline)
             outlist.append(newnode)
+            graph[str_pid] = newnode
+
         if is_special:
             continue
-        if node["children"]:
+        for child in tree["pid_graph"][pid_id]:
             stack.appendleft({"endchildren": 1})
-            stack.extendleft(reversed(node["children"]))
+            stack.extendleft([child])
             stack.appendleft({"startchildren": 1})
-    return outlist
+
+        return outlist