cape/helpers.py: add a find_process() function for quick-fetching processes from the cape report

capa/features/extractors/cape/helpers.py

@@ -0,0 +1,28 @@
+# Copyright (C) 2020 Mandiant, Inc. All Rights Reserved.
+# Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at: [package root]/LICENSE.txt
+# Unless required by applicable law or agreed to in writing, software distributed under the License
+#  is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and limitations under the License.
+from typing import Any, Dict, List
+
+from capa.features.extractors.base_extractor import ProcessHandle
+
+
+def find_process(processes: List[Dict[str, Any]], ph: ProcessHandle) -> Dict[str, Any]:
+    """
+    find a specific process identified by a process handler.
+
+    args:
+      processes: a list of processes extracted by CAPE
+      ph: handle of the sought process
+
+    return:
+      a CAPE-defined dictionary for the sought process' information
+    """
+
+    for process in processes:
+        if ph.pid == process["process_id"] and ph.inner["ppid"] == process["parent_id"]:
+            return process
+    return {}

capa/features/extractors/cape/process.py

@@ -24,9 +24,8 @@ def get_threads(behavior: Dict, ph: ProcessHandle) -> Iterator[ThreadHandle]:
     get a thread's child processes
     """
 
-    for process in behavior["processes"]:
-        if ph.pid == process["process_id"] and ph.inner["ppid"] == process["parent_id"]:
-            threads: List = process["threads"]
+    process = capa.features.extractors.cape.helpers.find_process(behavior["processes"], ph)
+    threads: List = process["threads"]
 
     for thread in threads:
         yield ThreadHandle(int(thread), inner={})
@@ -37,9 +36,8 @@ def extract_environ_strings(behavior: Dict, ph: ProcessHandle) -> Iterator[Tuple
     extract strings from a process' provided environment variables.
     """
 
-    for process in behavior["processes"]:
-        if ph.pid == process["process_id"] and ph.inner["ppid"] == process["parent_id"]:
-            environ: Dict[str, str] = process["environ"]
+    process = capa.features.extractors.cape.helpers.find_process(behavior["processes"], ph)
+    environ: Dict[str, str] = process["environ"]
 
     if not environ:
         return

capa/features/extractors/cape/thread.py

@@ -9,6 +9,7 @@
 import logging
 from typing import Any, Dict, List, Tuple, Iterator
 
+import capa.features.extractors.cape.helpers
 from capa.features.insn import API, Number
 from capa.features.common import String, Feature
 from capa.features.address import Address, AbsoluteVirtualAddress
@@ -31,9 +32,8 @@ def extract_call_features(behavior: Dict, ph: ProcessHandle, th: ThreadHandle) -
       Feature, address; where Feature is either: API, Number, or String.
     """
 
-    for process in behavior["processes"]:
-        if ph.pid == process["process_id"] and ph.inner["ppid"] == process["parent_id"]:
-            calls: List[Dict] = process["calls"]
+    process = capa.features.extractors.cape.helpers.find_process(behavior["processes"], ph)
+    calls: List[Dict[str, Any]] = process["calls"]
 
     tid = str(th.tid)
     for call in calls: