diff --git a/grr/server/grr_response_server/databases/db.py b/grr/server/grr_response_server/databases/db.py index a3764adfd1..e24949eabf 100644 --- a/grr/server/grr_response_server/databases/db.py +++ b/grr/server/grr_response_server/databases/db.py @@ -2061,6 +2061,7 @@ def ReadAllFlowObjects( max_create_time: Optional[rdfvalue.RDFDatetime] = None, include_child_flows: bool = True, not_created_by: Optional[Iterable[str]] = None, + created_by: Optional[Iterable[str]] = None, ) -> list[flows_pb2.Flow]: """Returns all flow objects. @@ -2072,6 +2073,7 @@ def ReadAllFlowObjects( include_child_flows: include child flows in the results. If False, only parent flows are returned. Must be `True` if the parent flow is given. not_created_by: exclude flows created by any of the users in this list. + created_by: only include flows created by any of the users in this list. Returns: A list of Flow objects. @@ -4334,6 +4336,7 @@ def ReadAllFlowObjects( max_create_time: Optional[rdfvalue.RDFDatetime] = None, include_child_flows: bool = True, not_created_by: Optional[Iterable[str]] = None, + created_by: Optional[Iterable[str]] = None, ) -> list[flows_pb2.Flow]: if client_id is not None: precondition.ValidateClientId(client_id) @@ -4348,6 +4351,9 @@ def ReadAllFlowObjects( if not_created_by is not None: precondition.AssertIterableType(not_created_by, str) + if created_by is not None: + precondition.AssertIterableType(created_by, str) + return self.delegate.ReadAllFlowObjects( client_id=client_id, parent_flow_id=parent_flow_id, @@ -4355,6 +4361,7 @@ def ReadAllFlowObjects( max_create_time=max_create_time, include_child_flows=include_child_flows, not_created_by=not_created_by, + created_by=created_by, ) def ReadChildFlowObjects( diff --git a/grr/server/grr_response_server/databases/db_flows_test.py b/grr/server/grr_response_server/databases/db_flows_test.py index 9b9f6e64bc..dbefc11db8 100644 --- a/grr/server/grr_response_server/databases/db_flows_test.py +++ b/grr/server/grr_response_server/databases/db_flows_test.py @@ -443,6 +443,34 @@ def testReadAllFlowObjectsWithNotCreatedBy(self): flows = self.db.ReadAllFlowObjects(not_created_by=frozenset(["baz", "foo"])) self.assertCountEqual([f.flow_id for f in flows], ["000A0002"]) + def testReadAllFlowObjectsWithCreatedBy(self): + client_id_1 = "C.1111111111111111" + self.db.WriteClientMetadata(client_id_1) + + self.db.WriteFlowObject( + flows_pb2.Flow(client_id=client_id_1, flow_id="000A0001", creator="foo") + ) + self.db.WriteFlowObject( + flows_pb2.Flow(client_id=client_id_1, flow_id="000A0002", creator="bar") + ) + self.db.WriteFlowObject( + flows_pb2.Flow(client_id=client_id_1, flow_id="000A0003", creator="foo") + ) + + flows = self.db.ReadAllFlowObjects(created_by=frozenset(["foo"])) + self.assertCountEqual([f.flow_id for f in flows], ["000A0001", "000A0003"]) + + flows = self.db.ReadAllFlowObjects(created_by=frozenset(["bar"])) + self.assertCountEqual([f.flow_id for f in flows], ["000A0002"]) + + flows = self.db.ReadAllFlowObjects(created_by=frozenset(["foo", "bar"])) + self.assertCountEqual( + [f.flow_id for f in flows], ["000A0001", "000A0002", "000A0003"] + ) + + flows = self.db.ReadAllFlowObjects(created_by=frozenset(["nonexistent"])) + self.assertEmpty(flows) + def testReadAllFlowObjectsWithAllConditions(self): client_id_1 = "C.1111111111111111" client_id_2 = "C.2222222222222222" diff --git a/grr/server/grr_response_server/databases/mem_flows.py b/grr/server/grr_response_server/databases/mem_flows.py index bfb7cf1998..8d604daa7f 100644 --- a/grr/server/grr_response_server/databases/mem_flows.py +++ b/grr/server/grr_response_server/databases/mem_flows.py @@ -239,6 +239,7 @@ def ReadAllFlowObjects( max_create_time: Optional[rdfvalue.RDFDatetime] = None, include_child_flows: bool = True, not_created_by: Optional[Iterable[str]] = None, + created_by: Optional[Iterable[str]] = None, ) -> list[flows_pb2.Flow]: """Returns all flow objects.""" res = [] @@ -261,6 +262,8 @@ def ReadAllFlowObjects( continue if not_created_by is not None and flow.creator in list(not_created_by): continue + if created_by is not None and flow.creator not in list(created_by): + continue res.append(flow) return res diff --git a/grr/server/grr_response_server/databases/mysql_flows.py b/grr/server/grr_response_server/databases/mysql_flows.py index a77b4f2f82..945a3f3706 100644 --- a/grr/server/grr_response_server/databases/mysql_flows.py +++ b/grr/server/grr_response_server/databases/mysql_flows.py @@ -407,6 +407,7 @@ def ReadAllFlowObjects( max_create_time: Optional[rdfvalue.RDFDatetime] = None, include_child_flows: bool = True, not_created_by: Optional[Iterable[str]] = None, + created_by: Optional[Iterable[str]] = None, cursor: Optional[cursors.Cursor] = None, ) -> list[flows_pb2.Flow]: """Returns all flow objects.""" @@ -441,6 +442,10 @@ def ReadAllFlowObjects( # The cursor implementation knows how to convert lists and ordinary sets. args.append(list(not_created_by)) + if created_by is not None: + conditions.append("creator IN %s") + args.append(list(created_by)) + query = f"SELECT {self.FLOW_DB_FIELDS} FROM flows" if conditions: query += " WHERE " + " AND ".join(conditions)