p34709852
/
monkey
forked from p15670423/monkey
1
0
Fork 1
Code Issues Pull Requests Projects Releases Wiki Activity

Fixed ring bugs

This commit is contained in:
VakarisZ 2019-10-15 11:12:49 +03:00
parent c474a23339
commit a4ebfb0c3d
2 changed files with 20 additions and 11 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
monkey/monkey_island/cc/models/monkey.py
View File

@ -45,20 +45,29 @@ class Monkey(Document):
aws_instance_id = StringField(required=False) # This field only exists when the monkey is running on an AWS
# instance. See https://github.com/guardicore/monkey/issues/426.
@staticmethod
def __ring_key__():
"""
Cash key representation
https://ring-cache.readthedocs.io/en/stable/quickstart.html#method-classmethod-staticmethod-property
:return:
"""
return Monkey.guid
# LOGIC
@staticmethod
def get_single_monkey_by_id(db_id):
try:
return Monkey.objects.get(id=db_id)
except DoesNotExist as ex:
raise MonkeyNotFoundError("info: {0} | id: {1}".format(ex.message, str(db_id)))
raise MonkeyNotFoundError("info: {0} | id: {1}".format(ex, str(db_id)))
@staticmethod
def get_single_monkey_by_guid(monkey_guid):
try:
return Monkey.objects.get(guid=monkey_guid)
except DoesNotExist as ex:
raise MonkeyNotFoundError("info: {0} | guid: {1}".format(ex.message, str(monkey_guid)))
raise MonkeyNotFoundError("info: {0} | guid: {1}".format(ex, str(monkey_guid)))
@staticmethod
def get_latest_modifytime():
@ -88,8 +97,8 @@ class Monkey(Document):
os = "windows"
return os
@staticmethod
@ring.lru()
@staticmethod
def get_label_by_id(object_id):
current_monkey = Monkey.get_single_monkey_by_id(object_id)
label = Monkey.get_hostname_by_id(object_id) + " : " + current_monkey.ip_addresses[0]
@ -97,8 +106,8 @@ class Monkey(Document):
label = "MonkeyIsland - " + label
return label
@staticmethod
@ring.lru()
@staticmethod
def get_hostname_by_id(object_id):
"""
:param object_id: the object ID of a Monkey in the database.
@ -124,10 +133,10 @@ class Monkey(Document):
"""
return {'ips': self.ip_addresses, 'hostname': self.hostname}
@staticmethod
@ring.lru(
expire=1 # data has TTL of 1 second. This is useful for rapid calls for report generation.
)
@staticmethod
def is_monkey(object_id):
try:
_ = Monkey.get_single_monkey_by_id(object_id)

12
monkey/monkey_island/cc/models/test_monkey.py
View File

@ -126,7 +126,7 @@ class TestMonkey(IslandTestCase):
linux_monkey.save()
cache_info_before_query = Monkey.get_label_by_id.storage.backend.cache_info()
self.assertEquals(cache_info_before_query.hits, 0)
self.assertEqual(cache_info_before_query.hits, 0)
# not cached
label = Monkey.get_label_by_id(linux_monkey.id)
@ -138,7 +138,7 @@ class TestMonkey(IslandTestCase):
# should be cached
_ = Monkey.get_label_by_id(linux_monkey.id)
cache_info_after_query = Monkey.get_label_by_id.storage.backend.cache_info()
self.assertEquals(cache_info_after_query.hits, 1)
self.assertEqual(cache_info_after_query.hits, 1)
linux_monkey.set_hostname("Another hostname")
@ -146,8 +146,8 @@ class TestMonkey(IslandTestCase):
label = Monkey.get_label_by_id(linux_monkey.id)
cache_info_after_second_query = Monkey.get_label_by_id.storage.backend.cache_info()
# still 1 hit only
self.assertEquals(cache_info_after_second_query.hits, 1)
self.assertEquals(cache_info_after_second_query.misses, 2)
self.assertEqual(cache_info_after_second_query.hits, 1)
self.assertEqual(cache_info_after_second_query.misses, 2)
def test_is_monkey(self):
self.fail_if_not_testing_env()
@ -157,7 +157,7 @@ class TestMonkey(IslandTestCase):
a_monkey.save()
cache_info_before_query = Monkey.is_monkey.storage.backend.cache_info()
self.assertEquals(cache_info_before_query.hits, 0)
self.assertEqual(cache_info_before_query.hits, 0)
# not cached
self.assertTrue(Monkey.is_monkey(a_monkey.id))
@ -169,5 +169,5 @@ class TestMonkey(IslandTestCase):
self.assertFalse(Monkey.is_monkey(fake_id))
cache_info_after_query = Monkey.is_monkey.storage.backend.cache_info()
self.assertEquals(cache_info_after_query.hits, 2)
self.assertEqual(cache_info_after_query.hits, 2)