forked from p15670423/monkey
Fixed ring bug with @staticmethod
This commit is contained in:
parent
3d332c7d66
commit
8046cddb8b
|
@ -97,25 +97,6 @@ class Monkey(Document):
|
||||||
os = "windows"
|
os = "windows"
|
||||||
return os
|
return os
|
||||||
|
|
||||||
@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]
|
|
||||||
if len(set(current_monkey.ip_addresses).intersection(local_ip_addresses())) > 0:
|
|
||||||
label = "MonkeyIsland - " + label
|
|
||||||
return label
|
|
||||||
|
|
||||||
@ring.lru()
|
|
||||||
@staticmethod
|
|
||||||
def get_hostname_by_id(object_id):
|
|
||||||
"""
|
|
||||||
:param object_id: the object ID of a Monkey in the database.
|
|
||||||
:return: The hostname of that machine.
|
|
||||||
:note: Use this and not monkey.hostname for performance - this is lru-cached.
|
|
||||||
"""
|
|
||||||
return Monkey.get_single_monkey_by_id(object_id).hostname
|
|
||||||
|
|
||||||
def set_hostname(self, hostname):
|
def set_hostname(self, hostname):
|
||||||
"""
|
"""
|
||||||
Sets a new hostname for a machine and clears the cache for getting it.
|
Sets a new hostname for a machine and clears the cache for getting it.
|
||||||
|
@ -123,8 +104,8 @@ class Monkey(Document):
|
||||||
"""
|
"""
|
||||||
self.hostname = hostname
|
self.hostname = hostname
|
||||||
self.save()
|
self.save()
|
||||||
Monkey.get_hostname_by_id.delete(self.id)
|
get_monkey_hostname_by_id.delete(self.id)
|
||||||
Monkey.get_label_by_id.delete(self.id)
|
get_monkey_label_by_id.delete(self.id)
|
||||||
|
|
||||||
def get_network_info(self):
|
def get_network_info(self):
|
||||||
"""
|
"""
|
||||||
|
@ -133,17 +114,6 @@ class Monkey(Document):
|
||||||
"""
|
"""
|
||||||
return {'ips': self.ip_addresses, 'hostname': self.hostname}
|
return {'ips': self.ip_addresses, 'hostname': self.hostname}
|
||||||
|
|
||||||
@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)
|
|
||||||
return True
|
|
||||||
except:
|
|
||||||
return False
|
|
||||||
|
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def get_tunneled_monkeys():
|
def get_tunneled_monkeys():
|
||||||
return Monkey.objects(tunnel__exists=True)
|
return Monkey.objects(tunnel__exists=True)
|
||||||
|
@ -153,5 +123,36 @@ class Monkey(Document):
|
||||||
self.save()
|
self.save()
|
||||||
|
|
||||||
|
|
||||||
|
# Can't make following methods static under Monkey class due to ring bug
|
||||||
|
@ring.lru(
|
||||||
|
expire=1 # data has TTL of 1 second. This is useful for rapid calls for report generation.
|
||||||
|
)
|
||||||
|
def is_monkey(object_id):
|
||||||
|
try:
|
||||||
|
_ = Monkey.get_single_monkey_by_id(object_id)
|
||||||
|
return True
|
||||||
|
except:
|
||||||
|
return False
|
||||||
|
|
||||||
|
|
||||||
|
@ring.lru()
|
||||||
|
def get_monkey_label_by_id(object_id):
|
||||||
|
current_monkey = Monkey.get_single_monkey_by_id(object_id)
|
||||||
|
label = get_monkey_hostname_by_id(object_id) + " : " + current_monkey.ip_addresses[0]
|
||||||
|
if len(set(current_monkey.ip_addresses).intersection(local_ip_addresses())) > 0:
|
||||||
|
label = "MonkeyIsland - " + label
|
||||||
|
return label
|
||||||
|
|
||||||
|
|
||||||
|
@ring.lru()
|
||||||
|
def get_monkey_hostname_by_id(object_id):
|
||||||
|
"""
|
||||||
|
:param object_id: the object ID of a Monkey in the database.
|
||||||
|
:return: The hostname of that machine.
|
||||||
|
:note: Use this and not monkey.hostname for performance - this is lru-cached.
|
||||||
|
"""
|
||||||
|
return Monkey.get_single_monkey_by_id(object_id).hostname
|
||||||
|
|
||||||
|
|
||||||
class MonkeyNotFoundError(Exception):
|
class MonkeyNotFoundError(Exception):
|
||||||
pass
|
pass
|
||||||
|
|
|
@ -2,7 +2,7 @@ import uuid
|
||||||
from time import sleep
|
from time import sleep
|
||||||
|
|
||||||
from .monkey import Monkey
|
from .monkey import Monkey
|
||||||
from monkey_island.cc.models.monkey import MonkeyNotFoundError
|
from monkey_island.cc.models.monkey import MonkeyNotFoundError, is_monkey, get_monkey_label_by_id
|
||||||
from monkey_island.cc.testing.IslandTestCase import IslandTestCase
|
from monkey_island.cc.testing.IslandTestCase import IslandTestCase
|
||||||
from .monkey_ttl import MonkeyTtl
|
from .monkey_ttl import MonkeyTtl
|
||||||
|
|
||||||
|
@ -125,26 +125,26 @@ class TestMonkey(IslandTestCase):
|
||||||
ip_addresses=[ip_example])
|
ip_addresses=[ip_example])
|
||||||
linux_monkey.save()
|
linux_monkey.save()
|
||||||
|
|
||||||
cache_info_before_query = Monkey.get_label_by_id.storage.backend.cache_info()
|
cache_info_before_query = get_monkey_label_by_id.storage.backend.cache_info()
|
||||||
self.assertEqual(cache_info_before_query.hits, 0)
|
self.assertEqual(cache_info_before_query.hits, 0)
|
||||||
|
|
||||||
# not cached
|
# not cached
|
||||||
label = Monkey.get_label_by_id(linux_monkey.id)
|
label = get_monkey_label_by_id(linux_monkey.id)
|
||||||
|
|
||||||
self.assertIsNotNone(label)
|
self.assertIsNotNone(label)
|
||||||
self.assertIn(hostname_example, label)
|
self.assertIn(hostname_example, label)
|
||||||
self.assertIn(ip_example, label)
|
self.assertIn(ip_example, label)
|
||||||
|
|
||||||
# should be cached
|
# should be cached
|
||||||
_ = Monkey.get_label_by_id(linux_monkey.id)
|
_ = get_monkey_label_by_id(linux_monkey.id)
|
||||||
cache_info_after_query = Monkey.get_label_by_id.storage.backend.cache_info()
|
cache_info_after_query = get_monkey_label_by_id.storage.backend.cache_info()
|
||||||
self.assertEqual(cache_info_after_query.hits, 1)
|
self.assertEqual(cache_info_after_query.hits, 1)
|
||||||
|
|
||||||
linux_monkey.set_hostname("Another hostname")
|
linux_monkey.set_hostname("Another hostname")
|
||||||
|
|
||||||
# should be a miss
|
# should be a miss
|
||||||
label = Monkey.get_label_by_id(linux_monkey.id)
|
label = get_monkey_label_by_id(linux_monkey.id)
|
||||||
cache_info_after_second_query = Monkey.get_label_by_id.storage.backend.cache_info()
|
cache_info_after_second_query = get_monkey_label_by_id.storage.backend.cache_info()
|
||||||
# still 1 hit only
|
# still 1 hit only
|
||||||
self.assertEqual(cache_info_after_second_query.hits, 1)
|
self.assertEqual(cache_info_after_second_query.hits, 1)
|
||||||
self.assertEqual(cache_info_after_second_query.misses, 2)
|
self.assertEqual(cache_info_after_second_query.misses, 2)
|
||||||
|
@ -156,18 +156,18 @@ class TestMonkey(IslandTestCase):
|
||||||
a_monkey = Monkey(guid=str(uuid.uuid4()))
|
a_monkey = Monkey(guid=str(uuid.uuid4()))
|
||||||
a_monkey.save()
|
a_monkey.save()
|
||||||
|
|
||||||
cache_info_before_query = Monkey.is_monkey.storage.backend.cache_info()
|
cache_info_before_query = is_monkey.storage.backend.cache_info()
|
||||||
self.assertEqual(cache_info_before_query.hits, 0)
|
self.assertEqual(cache_info_before_query.hits, 0)
|
||||||
|
|
||||||
# not cached
|
# not cached
|
||||||
self.assertTrue(Monkey.is_monkey(a_monkey.id))
|
self.assertTrue(is_monkey(a_monkey.id))
|
||||||
fake_id = "123456789012"
|
fake_id = "123456789012"
|
||||||
self.assertFalse(Monkey.is_monkey(fake_id))
|
self.assertFalse(is_monkey(fake_id))
|
||||||
|
|
||||||
# should be cached
|
# should be cached
|
||||||
self.assertTrue(Monkey.is_monkey(a_monkey.id))
|
self.assertTrue(is_monkey(a_monkey.id))
|
||||||
self.assertFalse(Monkey.is_monkey(fake_id))
|
self.assertFalse(is_monkey(fake_id))
|
||||||
|
|
||||||
cache_info_after_query = Monkey.is_monkey.storage.backend.cache_info()
|
cache_info_after_query = is_monkey.storage.backend.cache_info()
|
||||||
self.assertEqual(cache_info_after_query.hits, 2)
|
self.assertEqual(cache_info_after_query.hits, 2)
|
||||||
|
|
||||||
|
|
|
@ -2,7 +2,7 @@ from bson import ObjectId
|
||||||
|
|
||||||
from monkey_island.cc.database import mongo
|
from monkey_island.cc.database import mongo
|
||||||
import monkey_island.cc.services.node
|
import monkey_island.cc.services.node
|
||||||
from monkey_island.cc.models import Monkey
|
from monkey_island.cc.models.monkey import get_monkey_label_by_id, is_monkey
|
||||||
|
|
||||||
__author__ = "itay.mizeretz"
|
__author__ = "itay.mizeretz"
|
||||||
|
|
||||||
|
@ -145,13 +145,13 @@ class EdgeService:
|
||||||
from_id = edge["from"]
|
from_id = edge["from"]
|
||||||
to_id = edge["to"]
|
to_id = edge["to"]
|
||||||
|
|
||||||
from_label = Monkey.get_label_by_id(from_id)
|
from_label = get_monkey_label_by_id(from_id)
|
||||||
|
|
||||||
if to_id == ObjectId("000000000000000000000000"):
|
if to_id == ObjectId("000000000000000000000000"):
|
||||||
to_label = 'MonkeyIsland'
|
to_label = 'MonkeyIsland'
|
||||||
else:
|
else:
|
||||||
if Monkey.is_monkey(to_id):
|
if is_monkey(to_id):
|
||||||
to_label = Monkey.get_label_by_id(to_id)
|
to_label = get_monkey_label_by_id(to_id)
|
||||||
else:
|
else:
|
||||||
to_label = NodeService.get_node_label(NodeService.get_node_by_id(to_id))
|
to_label = NodeService.get_node_label(NodeService.get_node_by_id(to_id))
|
||||||
|
|
||||||
|
|
|
@ -4,7 +4,7 @@ from bson import ObjectId
|
||||||
|
|
||||||
import monkey_island.cc.services.log
|
import monkey_island.cc.services.log
|
||||||
from monkey_island.cc.database import mongo
|
from monkey_island.cc.database import mongo
|
||||||
from monkey_island.cc.models import Monkey
|
from monkey_island.cc.models.monkey import Monkey, get_monkey_hostname_by_id, get_monkey_label_by_id
|
||||||
from monkey_island.cc.services.edge import EdgeService
|
from monkey_island.cc.services.edge import EdgeService
|
||||||
from monkey_island.cc.utils import local_ip_addresses
|
from monkey_island.cc.utils import local_ip_addresses
|
||||||
import socket
|
import socket
|
||||||
|
@ -50,8 +50,8 @@ class NodeService:
|
||||||
|
|
||||||
for edge in edges:
|
for edge in edges:
|
||||||
from_node_id = edge["from"]
|
from_node_id = edge["from"]
|
||||||
from_node_label = Monkey.get_label_by_id(from_node_id)
|
from_node_label = get_monkey_label_by_id(from_node_id)
|
||||||
from_node_hostname = Monkey.get_hostname_by_id(from_node_id)
|
from_node_hostname = get_monkey_hostname_by_id(from_node_id)
|
||||||
|
|
||||||
accessible_from_nodes.append(from_node_label)
|
accessible_from_nodes.append(from_node_label)
|
||||||
accessible_from_nodes_hostnames.append(from_node_hostname)
|
accessible_from_nodes_hostnames.append(from_node_hostname)
|
||||||
|
@ -140,7 +140,7 @@ class NodeService:
|
||||||
@staticmethod
|
@staticmethod
|
||||||
def monkey_to_net_node(monkey, for_report=False):
|
def monkey_to_net_node(monkey, for_report=False):
|
||||||
monkey_id = monkey["_id"]
|
monkey_id = monkey["_id"]
|
||||||
label = Monkey.get_hostname_by_id(monkey_id) if for_report else Monkey.get_label_by_id(monkey_id)
|
label = get_monkey_hostname_by_id(monkey_id) if for_report else get_monkey_label_by_id(monkey_id)
|
||||||
monkey_group = NodeService.get_monkey_group(monkey)
|
monkey_group = NodeService.get_monkey_group(monkey)
|
||||||
return \
|
return \
|
||||||
{
|
{
|
||||||
|
|
Loading…
Reference in New Issue