def Run(self):
ref = rdf_hunts.FlowLikeObjectReference(
object_type="HUNT_REFERENCE",
hunt_reference=rdf_objects.HuntReference(hunt_id="H:332211"))
if data_store.RelationalDBReadEnabled("hunts"):
with test_lib.FakeTime(42):
hunt_id = self.CreateHunt(description="the hunt",
original_object=ref)
hunt_obj = data_store.REL_DB.ReadHuntObject(hunt_id)
hunt_obj.client_resources_stats.user_cpu_stats.sum = 5000
hunt_obj.client_resources_stats.network_bytes_sent_stats.sum = 1000000
data_store.REL_DB.WriteHuntObject(hunt_obj)
else:
with test_lib.FakeTime(42):
with self.CreateHunt(description="the hunt",
original_object=ref) as hunt_obj:
hunt_id = hunt_obj.urn.Basename()
hunt_stats = hunt_obj.context.usage_stats
hunt_stats.user_cpu_stats.sum = 5000
hunt_stats.network_bytes_sent_stats.sum = 1000000
self.Check("GetHunt",
args=hunt_plugin.ApiGetHuntArgs(hunt_id=hunt_id),
replace={hunt_id: "H:123456"})
def StopHunt(hunt_id, reason=None):
"""Stops a hunt with a given id."""
def UpdateFn(h):
if h.hunt_state not in [h.HuntState.STARTED, h.HuntState.PAUSED]:
raise OnlyStartedOrPausedHuntCanBeStoppedError(h)
h.hunt_state = h.HuntState.STOPPED
if reason is not None:
h.hunt_state_comment = reason
return h
# If the hunt was not started or paused, the exception from UpdateFn is
# guaranteed to be propagated by UpdateHuntObject implementation.
hunt_obj = data_store.REL_DB.UpdateHuntObject(hunt_id, UpdateFn)
data_store.REL_DB.RemoveForemanRule(hunt_id=hunt_obj.hunt_id)
flows = data_store.REL_DB.ReadHuntFlows(hunt_obj.hunt_id, 0, sys.maxsize)
data_store.REL_DB.UpdateFlows(
[(f.client_id, f.flow_id) for f in flows],
pending_termination=rdf_flow_objects.PendingFlowTermination(
reason="Parent hunt stopped."))
if (reason is not None
and hunt_obj.creator not in aff4_users.GRRUser.SYSTEM_USERS):
notification.Notify(
hunt_obj.creator,
rdf_objects.UserNotification.Type.TYPE_HUNT_STOPPED, reason,
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.HUNT,
hunt=rdf_objects.HuntReference(hunt_id=hunt_obj.hunt_id)))
return hunt_obj
def testHuntNotificationIsParsedCorrectly(self):
n = self.InitFromObj_(
rdf_objects.UserNotification.Type.TYPE_HUNT_STOPPED,
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.HUNT,
hunt=rdf_objects.HuntReference(hunt_id="H:123456")))
self.assertEqual(n.reference.type, "HUNT")
self.assertEqual(n.reference.hunt.hunt_id, "H:123456")
def Run(self):
ref = rdf_hunts.FlowLikeObjectReference(
object_type="HUNT_REFERENCE",
hunt_reference=rdf_objects.HuntReference(hunt_id="H:332211"))
# TODO(user): make hunt stats non-zero when AFF4 is gone to
# improve test coverage.
with test_lib.FakeTime(42):
hunt_id = self.CreateHunt(description="the hunt",
original_object=ref)
self.Check("GetHunt",
args=hunt_plugin.ApiGetHuntArgs(hunt_id=hunt_id),
replace={hunt_id: "H:123456"})
def Run(self):
with test_lib.FakeTime(42):
ref = rdf_hunts.FlowLikeObjectReference(
object_type="HUNT_REFERENCE",
hunt_reference=rdf_objects.HuntReference(hunt_id="H:332211"))
with self.CreateHunt(description="the hunt",
original_object=ref) as hunt_obj:
hunt_urn = hunt_obj.urn
hunt_stats = hunt_obj.context.usage_stats
hunt_stats.user_cpu_stats.sum = 5000
hunt_stats.network_bytes_sent_stats.sum = 1000000
self.Check(
"GetHunt",
args=hunt_plugin.ApiGetHuntArgs(hunt_id=hunt_urn.Basename()),
replace={hunt_urn.Basename(): "H:123456"})
def Run(self):
ref = rdf_hunts.FlowLikeObjectReference(
object_type="HUNT_REFERENCE",
hunt_reference=rdf_objects.HuntReference(hunt_id="H:332211"))
if data_store.RelationalDBReadEnabled("hunts"):
# TODO(user): make hunt stats non-zero when AFF4 is gone to
# improve test coverage.
with test_lib.FakeTime(42):
hunt_id = self.CreateHunt(description="the hunt",
original_object=ref)
else:
with test_lib.FakeTime(42):
with self.CreateHunt(description="the hunt",
original_object=ref) as hunt_obj:
hunt_id = hunt_obj.urn.Basename()
self.Check("GetHunt",
args=hunt_plugin.ApiGetHuntArgs(hunt_id=hunt_id),
replace={hunt_id: "H:123456"})
def StopHunt(hunt_id, reason=None):
"""Stops a hunt with a given id."""
hunt_obj = data_store.REL_DB.ReadHuntObject(hunt_id)
if hunt_obj.hunt_state not in [
hunt_obj.HuntState.STARTED, hunt_obj.HuntState.PAUSED
]:
raise OnlyStartedOrPausedHuntCanBeStoppedError(hunt_obj)
data_store.REL_DB.UpdateHuntObject(
hunt_id, hunt_state=hunt_obj.HuntState.STOPPED, hunt_state_comment=reason)
data_store.REL_DB.RemoveForemanRule(hunt_id=hunt_obj.hunt_id)
if (reason is not None and
hunt_obj.creator not in access_control.SYSTEM_USERS):
notification.Notify(
hunt_obj.creator, rdf_objects.UserNotification.Type.TYPE_HUNT_STOPPED,
reason,
rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.HUNT,
hunt=rdf_objects.HuntReference(hunt_id=hunt_obj.hunt_id)))
return data_store.REL_DB.ReadHuntObject(hunt_id)
def ObjectReference(self):
return rdf_objects.ObjectReference(
reference_type=rdf_objects.ObjectReference.Type.HUNT,
hunt=rdf_objects.HuntReference(hunt_id=str(self.hunt_id)))
def FromHuntId(cls, hunt_id):
res = FlowLikeObjectReference()
res.object_type = "HUNT_REFERENCE"
res.hunt_reference = rdf_objects.HuntReference(hunt_id=hunt_id)
return res