1+ // / Copyright 2023 Google LLC
2+ // /
3+ // / Licensed under the Apache License, Version 2.0 (the "License");
4+ // / you may not use this file except in compliance with the License.
5+ // / You may obtain a copy of the License at
6+ // /
7+ // / https://www.apache.org/licenses/LICENSE-2.0
8+ // /
9+ // / Unless required by applicable law or agreed to in writing, software
10+ // / distributed under the License is distributed on an "AS IS" BASIS,
11+ // / WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12+ // / See the License for the specific language governing permissions and
13+ // / limitations under the License.
14+ #include " Source/santad/ProcessTree/annotations/ancestry.h"
15+
16+ #include " Source/santad/ProcessTree/process.h"
17+ #include " Source/santad/ProcessTree/process_tree.h"
18+ #include " Source/santad/ProcessTree/process_tree.pb.h"
19+ #include " absl/container/flat_hash_map.h"
20+
21+ namespace ptpb = ::santa::pb::v1::process_tree;
22+
23+ namespace santa ::santad::process_tree {
24+
25+ ::santa::pb::v1::process_tree::Annotations::Ancestry
26+ AncestryAnnotator::getAncestry () const {
27+ ::santa::pb::v1::process_tree::Annotations::Ancestry ancestry;
28+ ancestry.CopyFrom (ancestry_);
29+ return ancestry;
30+ }
31+
32+ void AncestryAnnotator::AddEntryToAncestry (
33+ ptpb::Annotations::Ancestry &ancestry, int pid, uint64_t secondary_id) {
34+ ptpb::AncestryProcessID *ancestor = ancestry.add_ancestor ();
35+ ancestor->set_pid (pid);
36+ ancestor->set_secondary_id (secondary_id);
37+ }
38+
39+ void AncestryAnnotator::CopyAncestorsToAncestry (
40+ ptpb::Annotations::Ancestry &ancestry,
41+ std::vector<std::shared_ptr<const Process>> ancestors) {
42+ // Add ancestors starting from the root process
43+ for (auto it = ancestors.rbegin (); it != ancestors.rend (); it++) {
44+ AddEntryToAncestry (ancestry, (*it)->pid_ .pid , (*it)->creation_timestamp );
45+ }
46+ }
47+
48+ void AncestryAnnotator::AnnotateFork (ProcessTree &tree, const Process &parent,
49+ const Process &child) {
50+ ptpb::Annotations::Ancestry ancestry;
51+ // If parent process has ancestry annotation, copy and add parent.
52+ if (
auto parent_annotation = tree.
GetAnnotation <AncestryAnnotator>(parent)) {
53+ ancestry.CopyFrom ((*parent_annotation)->getAncestry ());
54+ AddEntryToAncestry (ancestry, parent.pid_ .pid , parent.creation_timestamp );
55+ // Otherwise, get all ancestors of the child and add them.
56+ } else {
57+ std::vector<std::shared_ptr<const Process>> ancestors =
58+ tree.GetAncestors (child);
59+ CopyAncestorsToAncestry (ancestry, ancestors);
60+ }
61+ tree.
AnnotateProcess (child, std::make_shared<AncestryAnnotator>(ancestry));
62+ }
63+
64+ void AncestryAnnotator::AnnotateExec (ProcessTree &tree,
65+ const Process &orig_process,
66+ const Process &new_process) {
67+ ptpb::Annotations::Ancestry ancestry;
68+ // If original process has ancestry annotation, copy entries.
69+ if (auto orig_process_annotation =
70+ tree.
GetAnnotation <AncestryAnnotator>(orig_process)) {
71+ ancestry.CopyFrom ((*orig_process_annotation)->getAncestry ());
72+ // Otherwise, compute all ancestors of the new process and add them.
73+ } else {
74+ std::vector<std::shared_ptr<const Process>> ancestors =
75+ tree.GetAncestors (new_process);
76+ CopyAncestorsToAncestry (ancestry, ancestors);
77+ }
78+ tree.AnnotateProcess (new_process,
79+ std::make_shared<AncestryAnnotator>(ancestry));
80+ }
81+
82+ std::optional<ptpb::Annotations> AncestryAnnotator::Proto () const {
83+ auto annotation = ptpb::Annotations ();
84+ auto *ancestry_ptr = annotation.mutable_ancestry ();
85+ ancestry_ptr->CopyFrom (AncestryAnnotator::getAncestry ());
86+ return annotation;
87+ }
88+
89+ } // namespace santa::santad::process_tree
0 commit comments