Merge 948b7c3 into 3b2d02f

Source/santad/BUILD

@@ -779,6 +779,7 @@ objc_library(
         "//Source/common:SNTXPCUnprivilegedControlInterface",
         "//Source/common:Unit",
         "//Source/santad/ProcessTree:process_tree",
+        "//Source/santad/ProcessTree/annotations:originator",
         "@MOLXPCConnection",
     ],
 )
@@ -1420,6 +1421,8 @@ test_suite(
         ":SantadTest",
         ":WatchItemsTest",
         "//Source/santad/Logs/EndpointSecurity/Writers/FSSpool:fsspool_test",
+        "//Source/santad/ProcessTree:process_tree_test",
+        "//Source/santad/ProcessTree/annotations:originator_test",
     ],
     visibility = ["//:santa_package_group"],
 )

Source/santad/ProcessTree/BUILD

@@ -76,6 +76,7 @@ objc_library(
         ":process_tree",
         "@com_google_absl//absl/synchronization",
     ],
+    visibility = ["//Source/santad/ProcessTree:__subpackages__"],
 )
 
 santa_unit_test(

Source/santad/ProcessTree/annotations/BUILD

@@ -1,3 +1,5 @@
+load("//:helper.bzl", "santa_unit_test")
+
 package(
     default_visibility = ["//:santa_package_group"],
 )
@@ -9,3 +11,27 @@ cc_library(
         "//Source/santad/ProcessTree:process_tree_cc_proto",
     ],
 )
+
+cc_library(
+    name = "originator",
+    hdrs = ["originator.h"],
+    srcs = ["originator.cc"],
+    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"],
+    deps = [
+        ":originator",
+        "//Source/santad/ProcessTree:process",
+        "//Source/santad/ProcessTree:process_tree_test_helpers",
+        "//Source/santad/ProcessTree:process_tree_cc_proto",
+    ],
+)

Source/santad/ProcessTree/annotations/originator.cc

@@ -0,0 +1,67 @@
+/// 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/originator.h"
+
+#include <memory>
+#include <optional>
+#include <string>
+#include <utility>
+
+#include "absl/container/flat_hash_map.h"
+#include "Source/santad/ProcessTree/process.h"
+#include "Source/santad/ProcessTree/process_tree.h"
+#include "Source/santad/ProcessTree/process_tree.pb.h"
+
+namespace ptpb = ::santa::pb::v1::process_tree;
+
+namespace santa::santad::process_tree {
+
+void OriginatorAnnotator::AnnotateFork(ProcessTree &tree, const Process &parent,
+                                       const Process &child) {
+  // "Base case". Propagate existing annotations down to descendants.
+  if (auto annotation = tree.GetAnnotation<OriginatorAnnotator>(parent)) {
+    tree.AnnotateProcess(child, std::move(*annotation));
+  }
+}
+
+void OriginatorAnnotator::AnnotateExec(ProcessTree &tree,
+                                       const Process &orig_process,
+                                       const Process &new_process) {
+  static const absl::flat_hash_map<std::string, ptpb::Annotations::Originator>
+      originator_programs = {
+          {"/usr/bin/login",
+           ptpb::Annotations::Originator::Annotations_Originator_LOGIN},
+          {"/usr/sbin/cron",
+           ptpb::Annotations::Originator::Annotations_Originator_CRON},
+      };
+
+  if (auto annotation = tree.GetAnnotation<OriginatorAnnotator>(orig_process)) {
+    tree.AnnotateProcess(new_process, std::move(*annotation));
+    return;
+  }
+
+  if (auto it = originator_programs.find(new_process.program_->executable);
+      it != originator_programs.end()) {
+    tree.AnnotateProcess(new_process,
+                         std::make_shared<OriginatorAnnotator>(it->second));
+  }
+}
+
+std::optional<ptpb::Annotations> OriginatorAnnotator::Proto() const {
+  auto annotation = ptpb::Annotations();
+  annotation.set_originator(originator_);
+  return annotation;
+}
+
+}  // namespace santa::santad::process_tree

Source/santad/ProcessTree/annotations/originator.h

@@ -0,0 +1,48 @@
+/// 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_ORIGINATOR_H
+#define SANTA__SANTAD_PROCESSTREE_ANNOTATIONS_ORIGINATOR_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 OriginatorAnnotator : public Annotator {
+ public:
+  OriginatorAnnotator()
+      : originator_(::santa::pb::v1::process_tree::Annotations::Originator::
+                        Annotations_Originator_UNSPECIFIED) {};
+  explicit OriginatorAnnotator(
+      ::santa::pb::v1::process_tree::Annotations::Originator originator)
+      : originator_(originator) {};
+
+  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;
+
+ private:
+  ::santa::pb::v1::process_tree::Annotations::Originator originator_;
+};
+
+}  // namespace santa::santad::process_tree
+
+#endif

Source/santad/ProcessTree/annotations/originator_test.mm

@@ -0,0 +1,78 @@
+/// 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.
+#import <Foundation/Foundation.h>
+#import <XCTest/XCTest.h>
+
+#include "Source/santad/ProcessTree/annotations/originator.h"
+#include "Source/santad/ProcessTree/process.h"
+#include "Source/santad/ProcessTree/process_tree.pb.h"
+#include "Source/santad/ProcessTree/process_tree_test_helpers.h"
+
+using namespace santa::santad::process_tree;
+namespace ptpb = ::santa::pb::v1::process_tree;
+
+@interface OriginatorAnnotatorTest : XCTestCase
+@property std::shared_ptr<ProcessTreeTestPeer> tree;
+@property std::shared_ptr<const Process> initProc;
+@end
+
+@implementation OriginatorAnnotatorTest
+
+- (void)setUp {
+  std::vector<std::unique_ptr<Annotator>> annotators;
+  annotators.emplace_back(std::make_unique<OriginatorAnnotator>());
+  self.tree = std::make_shared<ProcessTreeTestPeer>(std::move(annotators));
+  self.initProc = self.tree->InsertInit();
+}
+
+- (void)testAnnotation {
+  uint64_t event_id = 1;
+  const struct Cred cred = {.uid = 0, .gid = 0};
+
+  // PID 1.1: fork() -> PID 2.2
+  const struct Pid login_pid = {.pid = 2, .pidversion = 2};
+  self.tree->HandleFork(event_id++, *self.initProc, login_pid);
+
+  // PID 2.2: exec("/usr/bin/login") -> PID 2.3
+  const struct Pid login_exec_pid = {.pid = 2, .pidversion = 3};
+  const struct Program login_prog = {.executable = "/usr/bin/login", .arguments = {}};
+  auto login = *self.tree->Get(login_pid);
+  self.tree->HandleExec(event_id++, *login, login_exec_pid, login_prog, cred);
+
+  // Ensure we have an annotation on login itself...
+  login = *self.tree->Get(login_exec_pid);
+  auto annotation_opt = self.tree->GetAnnotation<OriginatorAnnotator>(*login);
+  XCTAssertTrue(annotation_opt.has_value());
+  auto proto_opt = (*annotation_opt)->Proto();
+  XCTAssertTrue(proto_opt.has_value());
+  XCTAssertEqual(proto_opt->originator(),
+                 ptpb::Annotations::Originator::Annotations_Originator_LOGIN);
+
+  // PID 2.3: fork() -> PID 3.3
+  const struct Pid shell_pid = {.pid = 3, .pidversion = 3};
+  self.tree->HandleFork(event_id++, *login, shell_pid);
+  // PID 3.3: exec("/bin/zsh") -> PID 3.4
+  const struct Pid shell_exec_pid = {.pid = 3, .pidversion = 4};
+  const struct Program shell_prog = {.executable = "/bin/zsh", .arguments = {}};
+  auto shell = *self.tree->Get(shell_pid);
+  self.tree->HandleExec(event_id++, *shell, shell_exec_pid, shell_prog, cred);
+
+  // ... and also ensure we have the same annotation on the descendant zsh.
+  shell = *self.tree->Get(shell_exec_pid);
+  auto descendant_annotation_opt = self.tree->GetAnnotation<OriginatorAnnotator>(*shell);
+  XCTAssertTrue(descendant_annotation_opt.has_value());
+  XCTAssertEqual(*descendant_annotation_opt, *annotation_opt);
+}
+
+@end

Source/santad/ProcessTree/process_tree.proto

@@ -3,4 +3,11 @@ syntax = "proto3";
 package santa.pb.v1.process_tree;
 
 message Annotations {
+  enum Originator {
+    UNSPECIFIED = 0;
+    LOGIN = 1;
+    CRON = 2;
+  }
+
+  Originator originator = 1;
 }

Source/santad/SantadDeps.mm

@@ -24,6 +24,7 @@
 #import "Source/santad/DataLayer/SNTRuleTable.h"
 #include "Source/santad/DataLayer/WatchItems.h"
 #include "Source/santad/EventProviders/EndpointSecurity/EndpointSecurityAPI.h"
+#include "Source/santad/ProcessTree/annotations/originator.h"
 #include "Source/santad/ProcessTree/process_tree.h"
 #import "Source/santad/SNTDatabaseController.h"
 #include "Source/santad/SNTDecisionCache.h"
@@ -159,8 +160,11 @@
   std::vector<std::unique_ptr<process_tree::Annotator>> annotators;
 
   for (NSString *annotation in [configurator enabledProcessAnnotations]) {
-    // TODO(nickmg): add annotation name switch
-    (void)annotation;
+    if ([[annotation lowercaseString] isEqualToString:@"originator"]) {
+      annotators.emplace_back(std::make_unique<process_tree::OriginatorAnnotator>());
+    } else {
+      LOGW(@"Unrecognized process annotation %@", annotation);
+    }
   }
 
   auto tree_status = process_tree::CreateTree(std::move(annotators));