Add ancestry annotator

Source/common/BUILD

@@ -1,3 +1,4 @@
+load("@com_google_protobuf//bazel:proto_library.bzl", "proto_library")
 load("@rules_cc//cc:defs.bzl", "cc_proto_library")
 load("//:helper.bzl", "santa_unit_test")
 

Source/common/santa.proto

@@ -129,13 +129,13 @@ message FileDescriptor {
 // Process information
 message ProcessInfo {
   // Process ID of the process
-  optional ProcessID id = 1;
+  optional santa.pb.v1.ProcessID id = 1;
 
   // Process ID of the parent process
-  optional ProcessID parent_id = 2;
+  optional santa.pb.v1.ProcessID parent_id = 2;
 
   // Process ID of the process responsible for this one
-  optional ProcessID responsible_id = 3;
+  optional santa.pb.v1.ProcessID responsible_id = 3;
 
   // Original parent ID, remains stable in the event a process is reparented
   optional int32 original_parent_pid = 4;
@@ -181,10 +181,10 @@ message ProcessInfo {
 // Light variant of ProcessInfo message to help minimize on-disk/on-wire sizes
 message ProcessInfoLight {
   // Process ID of the process
-  optional ProcessID id = 1;
+  optional santa.pb.v1.ProcessID id = 1;
 
   // Process ID of the parent process
-  optional ProcessID parent_id = 2;
+  optional santa.pb.v1.ProcessID parent_id = 2;
 
   // Original parent ID, remains stable in the event a process is reparented
   optional int32 original_parent_pid = 3;

Source/santad/BUILD

@@ -802,6 +802,7 @@ objc_library(
         "//Source/common:Unit",
         "//Source/santad/ProcessTree:process_tree",
         "//Source/santad/ProcessTree/annotations:originator",
+        "//Source/santad/ProcessTree/annotations:ancestry",
         "@MOLXPCConnection",
     ],
 )
@@ -840,7 +841,6 @@ macos_bundle(
     ],
     entitlements = select({
         "//:adhoc_build": "com.google.santa.daemon.systemextension-adhoc.entitlements",
-        # Non-adhoc builds get their entitlements from the provisioning profile.
         "//conditions:default": None,
     }),
     infoplists = ["Info.plist"],
@@ -1451,6 +1451,7 @@ test_suite(
         "//Source/santad/Logs/EndpointSecurity/Writers/FSSpool:fsspool_test",
         "//Source/santad/ProcessTree:process_tree_test",
         "//Source/santad/ProcessTree/annotations:originator_test",
+        "//Source/santad/ProcessTree/annotations:ancestry_test",
     ],
     visibility = ["//:santa_package_group"],
 )

Source/santad/Logs/EndpointSecurity/Writers/FSSpool/BUILD

@@ -1,3 +1,4 @@
+load("@com_google_protobuf//bazel:proto_library.bzl", "proto_library")
 load("@rules_cc//cc:defs.bzl", "cc_proto_library")
 load("//:helper.bzl", "santa_unit_test")
 

Source/santad/ProcessTree/BUILD

@@ -1,3 +1,4 @@
+load("@com_google_protobuf//bazel:proto_library.bzl", "proto_library")
 load("@rules_cc//cc:defs.bzl", "cc_proto_library")
 load("//:helper.bzl", "santa_unit_test")
 

Source/santad/ProcessTree/annotations/BUILD

@@ -25,6 +25,19 @@ cc_library(
     ],
 )
 
+cc_library(
+    name = "ancestry",
+    srcs = ["ancestry.cc"],
+    hdrs = ["ancestry.h"],
+    deps = [
+        ":annotator",
+        "//Source/santad/ProcessTree:process",
+        "//Source/santad/ProcessTree:process_tree",
+        "//Source/santad/ProcessTree:process_tree_cc_proto",
+        "@com_google_absl//absl/container:flat_hash_map",
+    ],
+)
+
 santa_unit_test(
     name = "originator_test",
     srcs = ["originator_test.mm"],
@@ -35,3 +48,14 @@ santa_unit_test(
         "//Source/santad/ProcessTree:process_tree_test_helpers",
     ],
 )
+
+santa_unit_test(
+    name = "ancestry_test",
+    srcs = ["ancestry_test.mm"],
+    deps = [
+        ":ancestry",
+        "//Source/santad/ProcessTree:process",
+        "//Source/santad/ProcessTree:process_tree_cc_proto",
+        "//Source/santad/ProcessTree:process_tree_test_helpers",
+    ],
+)

Source/santad/ProcessTree/annotations/ancestry.cc

@@ -0,0 +1,89 @@
+/// Copyright 2023 Google LLC
+///
+/// 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
+///
+///     https://www.apache.org/licenses/LICENSE-2.0
+///
+/// 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.
+#include "Source/santad/ProcessTree/annotations/ancestry.h"
+
+#include "Source/santad/ProcessTree/process.h"
+#include "Source/santad/ProcessTree/process_tree.h"
+#include "Source/santad/ProcessTree/process_tree.pb.h"
+#include "absl/container/flat_hash_map.h"
+
+namespace ptpb = ::santa::pb::v1::process_tree;
+
+namespace santa::santad::process_tree {
+
+::santa::pb::v1::process_tree::Annotations::Ancestry
+AncestryAnnotator::getAncestry() const {
+  ::santa::pb::v1::process_tree::Annotations::Ancestry ancestry;
+  ancestry.CopyFrom(ancestry_);
+  return ancestry;
+}
+
+void AncestryAnnotator::AddEntryToAncestry(
+    ptpb::Annotations::Ancestry &ancestry, int pid, uint64_t secondary_id) {
+  ptpb::AncestryProcessID *ancestor = ancestry.add_ancestor();
+  ancestor->set_pid(pid);
+  ancestor->set_secondary_id(secondary_id);
+}
+
+void AncestryAnnotator::CopyAncestorsToAncestry(
+    ptpb::Annotations::Ancestry &ancestry,
+    std::vector<std::shared_ptr<const Process>> ancestors) {
+  // Add ancestors starting from the root process
+  for (auto it = ancestors.rbegin(); it != ancestors.rend(); it++) {
+    AddEntryToAncestry(ancestry, (*it)->pid_.pid, (*it)->creation_timestamp);
+  }
+}
+
+void AncestryAnnotator::AnnotateFork(ProcessTree &tree, const Process &parent,
+                                     const Process &child) {
+  ptpb::Annotations::Ancestry ancestry;
+  // If parent process has ancestry annotation, copy and add parent.
+  if (auto parent_annotation = tree.GetAnnotation<AncestryAnnotator>(parent)) {
+    ancestry.CopyFrom((*parent_annotation)->getAncestry());
+    AddEntryToAncestry(ancestry, parent.pid_.pid, parent.creation_timestamp);
+    // Otherwise, get all ancestors of the child and add them.
+  } else {
+    std::vector<std::shared_ptr<const Process>> ancestors =
+        tree.GetAncestors(child);
+    CopyAncestorsToAncestry(ancestry, ancestors);
+  }
+  tree.AnnotateProcess(child, std::make_shared<AncestryAnnotator>(ancestry));
+}
+
+void AncestryAnnotator::AnnotateExec(ProcessTree &tree,
+                                     const Process &orig_process,
+                                     const Process &new_process) {
+  ptpb::Annotations::Ancestry ancestry;
+  // If original process has ancestry annotation, copy entries.
+  if (auto orig_process_annotation =
+          tree.GetAnnotation<AncestryAnnotator>(orig_process)) {
+    ancestry.CopyFrom((*orig_process_annotation)->getAncestry());
+    // Otherwise, compute all ancestors of the new process and add them.
+  } else {
+    std::vector<std::shared_ptr<const Process>> ancestors =
+        tree.GetAncestors(new_process);
+    CopyAncestorsToAncestry(ancestry, ancestors);
+  }
+  tree.AnnotateProcess(new_process,
+                       std::make_shared<AncestryAnnotator>(ancestry));
+}
+
+std::optional<ptpb::Annotations> AncestryAnnotator::Proto() const {
+  auto annotation = ptpb::Annotations();
+  auto *ancestry_ptr = annotation.mutable_ancestry();
+  ancestry_ptr->CopyFrom(AncestryAnnotator::getAncestry());
+  return annotation;
+}
+
+}  // namespace santa::santad::process_tree

Source/santad/ProcessTree/annotations/ancestry.h

@@ -0,0 +1,53 @@
+/// Copyright 2023 Google LLC
+///
+/// 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
+///
+///     https://www.apache.org/licenses/LICENSE-2.0
+///
+/// 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.
+#ifndef SANTA__SANTAD_PROCESSTREE_ANNOTATIONS_ANCESTRY_H
+#define SANTA__SANTAD_PROCESSTREE_ANNOTATIONS_ANCESTRY_H
+
+#include <optional>
+
+#include "Source/santad/ProcessTree/annotations/annotator.h"
+#include "Source/santad/ProcessTree/process.h"
+#include "Source/santad/ProcessTree/process_tree.pb.h"
+
+namespace santa::santad::process_tree {
+
+class AncestryAnnotator : public Annotator {
+ public:
+  // clang-format off
+  AncestryAnnotator() {}
+ explicit AncestryAnnotator(
+            ::santa::pb::v1::process_tree::Annotations::Ancestry ancestry)
+            : ancestry_(ancestry) {};
+  // clang-format on
+  void AnnotateFork(ProcessTree &tree, const Process &parent,
+                    const Process &child) override;
+  void AnnotateExec(ProcessTree &tree, const Process &orig_process,
+                    const Process &new_process) override;
+  std::optional<::santa::pb::v1::process_tree::Annotations> Proto()
+      const override;
+  ::santa::pb::v1::process_tree::Annotations::Ancestry getAncestry() const;
+
+ private:
+  void AddEntryToAncestry(
+      ::santa::pb::v1::process_tree::Annotations::Ancestry &ancestry, int pid,
+      uint64_t secondary_id);
+  void CopyAncestorsToAncestry(
+      ::santa::pb::v1::process_tree::Annotations::Ancestry &ancestry,
+      std::vector<std::shared_ptr<const Process>> ancestors);
+  ::santa::pb::v1::process_tree::Annotations::Ancestry ancestry_;
+};
+
+}  // namespace santa::santad::process_tree
+
+#endif