Merge pull request #313 from guardicore/feature/210-monkey-mia-timeout

[DONE] Feature/210 monkey mia timeout

monkey/infection_monkey/control.py

@@ -20,6 +20,10 @@ requests.packages.urllib3.disable_warnings()
 LOG = logging.getLogger(__name__)
 DOWNLOAD_CHUNK = 1024
 
+# random number greater than 5,
+# to prevent the monkey from just waiting forever to try and connect to an island before going elsewhere.
+TIMEOUT_IN_SECONDS = 15
+
 
 class ControlClient(object):
     proxies = {}
@@ -73,7 +77,7 @@ class ControlClient(object):
                 requests.get("https://%s/api?action=is-up" % (server,),
                              verify=False,
                              proxies=ControlClient.proxies,
-                             timeout=TIMEOUT)
+                             timeout=TIMEOUT_IN_SECONDS)
                 WormConfiguration.current_server = current_server
                 break
 

monkey/infection_monkey/dropper.py

@@ -51,7 +51,6 @@ class MonkeyDrops(object):
         LOG.debug("Dropper is running with config:\n%s", pprint.pformat(self._config))
 
     def start(self):
-
         if self._config['destination_path'] is None:
             LOG.error("No destination path specified")
             return False

monkey/infection_monkey/main.py

@@ -98,6 +98,7 @@ def main():
             except OSError:
                 pass
         LOG_CONFIG['handlers']['file']['filename'] = log_path
+        # noinspection PyUnresolvedReferences
         LOG_CONFIG['root']['handlers'].append('file')
     else:
         del LOG_CONFIG['handlers']['file']

monkey/monkey_island/cc/auth.py

@@ -10,8 +10,8 @@ __author__ = 'itay.mizeretz'
 
 
 class User(object):
-    def __init__(self, id, username, secret):
-        self.id = id
+    def __init__(self, user_id, username, secret):
+        self.id = user_id
         self.username = username
         self.secret = secret
 

monkey/monkey_island/cc/environment/__init__.py

@@ -10,13 +10,29 @@ class Environment(object):
     __metaclass__ = abc.ABCMeta
 
     _ISLAND_PORT = 5000
-    _MONGO_URL = os.environ.get("MONKEY_MONGO_URL", "mongodb://localhost:27017/monkeyisland")
+    _MONGO_DB_NAME = "monkeyisland"
+    _MONGO_DB_HOST = "localhost"
+    _MONGO_DB_PORT = 27017
+    _MONGO_URL = os.environ.get("MONKEY_MONGO_URL", "mongodb://{0}:{1}/{2}".format(_MONGO_DB_HOST, _MONGO_DB_PORT, str(_MONGO_DB_NAME)))
     _DEBUG_SERVER = False
     _AUTH_EXPIRATION_TIME = timedelta(hours=1)
+
+    _testing = False
+
+    @property
+    def testing(self):
+        return self._testing
+
+    @testing.setter
+    def testing(self, value):
+        self._testing = value
+
     _MONKEY_VERSION = "1.6.3"
 
+
     def __init__(self):
         self.config = None
+        self._testing = False  # Assume env is not for unit testing.
 
     def set_config(self, config):
         self.config = config
@@ -56,3 +72,15 @@ class Environment(object):
     @abc.abstractmethod
     def get_auth_users(self):
         return
+
+    @property
+    def mongo_db_name(self):
+        return self._MONGO_DB_NAME
+
+    @property
+    def mongo_db_host(self):
+        return self._MONGO_DB_HOST
+
+    @property
+    def mongo_db_port(self):
+        return self._MONGO_DB_PORT

monkey/monkey_island/cc/environment/environment.py

@@ -2,7 +2,10 @@ import json
 import logging
 import os
 
+env = None
+
 from monkey_island.cc.environment import standard
+from monkey_island.cc.environment import testing
 from monkey_island.cc.environment import aws
 from monkey_island.cc.environment import password
 from monkey_island.cc.consts import MONKEY_ISLAND_ABS_PATH
@@ -14,11 +17,13 @@ logger = logging.getLogger(__name__)
 AWS = 'aws'
 STANDARD = 'standard'
 PASSWORD = 'password'
+TESTING = 'testing'
 
 ENV_DICT = {
     STANDARD: standard.StandardEnvironment,
     AWS: aws.AwsEnvironment,
     PASSWORD: password.PasswordEnvironment,
+    TESTING: testing.TestingEnvironment
 }
 
 

monkey/monkey_island/cc/environment/testing.py

@@ -0,0 +1,17 @@
+from monkey_island.cc.environment import Environment
+import monkey_island.cc.auth
+
+
+class TestingEnvironment(Environment):
+    def __init__(self):
+        super(TestingEnvironment, self).__init__()
+        self.testing = True
+
+    # SHA3-512 of '1234567890!@#$%^&*()_nothing_up_my_sleeve_1234567890!@#$%^&*()'
+    NO_AUTH_CREDS = '55e97c9dcfd22b8079189ddaeea9bce8125887e3237b800c6176c9afa80d2062' \
+                    '8d2c8d0b1538d2208c1444ac66535b764a3d902b35e751df3faec1e477ed3557'
+
+    def get_auth_users(self):
+        return [
+            monkey_island.cc.auth.User(1, self.NO_AUTH_CREDS, self.NO_AUTH_CREDS)
+        ]

monkey/monkey_island/cc/models/__init__.py

@@ -0,0 +1,19 @@
+from mongoengine import connect
+
+from monkey_island.cc.environment.environment import env
+
+# This section sets up the DB connection according to the environment.
+#   If testing, use mongomock which only emulates mongo. for more information, see
+#   http://docs.mongoengine.org/guide/mongomock.html .
+#   Otherwise, use an actual mongod instance with connection parameters supplied by env.
+if env.testing:
+    connect('mongoenginetest', host='mongomock://localhost')
+else:
+    connect(db=env.mongo_db_name, host=env.mongo_db_host, port=env.mongo_db_port)
+
+# Order of importing matters here, for registering the embedded and referenced documents before using them.
+from config import Config
+from creds import Creds
+from monkey_ttl import MonkeyTtl
+from pba_results import PbaResults
+from monkey import Monkey

monkey/monkey_island/cc/models/config.py

@@ -0,0 +1,11 @@
+from mongoengine import EmbeddedDocument
+
+
+class Config(EmbeddedDocument):
+    """
+    No need to define this schema here. It will change often and is already is defined in
+    monkey_island.cc.services.config_schema.
+    See https://mongoengine-odm.readthedocs.io/apireference.html#mongoengine.FieldDoesNotExist
+    """
+    meta = {'strict': False}
+    pass

monkey/monkey_island/cc/models/creds.py

@@ -0,0 +1,9 @@
+from mongoengine import EmbeddedDocument
+
+
+class Creds(EmbeddedDocument):
+    """
+    TODO get an example of this data, and make it strict
+    """
+    meta = {'strict': False}
+    pass

monkey/monkey_island/cc/models/monkey.py

@@ -0,0 +1,60 @@
+"""
+Define a Document Schema for the Monkey document.
+"""
+import mongoengine
+from mongoengine import Document, StringField, ListField, BooleanField, EmbeddedDocumentField, DateField, \
+    ReferenceField
+
+from monkey_island.cc.models.monkey_ttl import MonkeyTtl
+
+
+class Monkey(Document):
+    """
+    This class has 2 main section:
+        *   The schema section defines the DB fields in the document. This is the data of the object.
+        *   The logic section defines complex questions we can ask about a single document which are asked multiple
+            times, somewhat like an API.
+    """
+    # SCHEMA
+    guid = StringField(required=True)
+    config = EmbeddedDocumentField('Config')
+    creds = ListField(EmbeddedDocumentField('Creds'))
+    dead = BooleanField()
+    description = StringField()
+    hostname = StringField()
+    internet_access = BooleanField()
+    ip_addresses = ListField(StringField())
+    keepalive = DateField()
+    modifytime = DateField()
+    # TODO change this to an embedded document as well - RN it's an unnamed tuple which is confusing.
+    parent = ListField(ListField(StringField()))
+    config_error = BooleanField()
+    critical_services = ListField(StringField())
+    pba_results = ListField()
+    ttl_ref = ReferenceField(MonkeyTtl)
+
+    # LOGIC
+    @staticmethod
+    def get_single_monkey_by_id(db_id):
+        try:
+            return Monkey.objects(id=db_id)[0]
+        except IndexError:
+            raise MonkeyNotFoundError("id: {0}".format(str(db_id)))
+
+    def is_dead(self):
+        monkey_is_dead = False
+        if self.dead:
+            monkey_is_dead = True
+        else:
+            try:
+                if MonkeyTtl.objects(id=self.ttl_ref.id).count() == 0:
+                    # No TTLs - monkey has timed out. The monkey is MIA.
+                    monkey_is_dead = True
+            except (mongoengine.DoesNotExist, AttributeError):
+                # Trying to dereference unknown document - the monkey is MIA.
+                monkey_is_dead = True
+        return monkey_is_dead
+
+
+class MonkeyNotFoundError(Exception):
+    pass

monkey/monkey_island/cc/models/monkey_ttl.py

@@ -0,0 +1,40 @@
+from datetime import datetime, timedelta
+
+from mongoengine import Document, DateTimeField
+
+
+class MonkeyTtl(Document):
+    """
+    This model represents the monkey's TTL, and is referenced by the main Monkey document.
+    See https://docs.mongodb.com/manual/tutorial/expire-data/ and
+    https://stackoverflow.com/questions/55994379/mongodb-ttl-index-doesnt-delete-expired-documents/56021663#56021663
+    for more information about how TTL indexing works and why this class is set up the way it is.
+
+    If you wish to use this class, you can create it using the create_ttl_expire_in(seconds) function.
+    If you wish to create an instance of this class directly, see the inner implementation of
+    create_ttl_expire_in(seconds) to see how to do so.
+    """
+
+    @staticmethod
+    def create_ttl_expire_in(expiry_in_seconds):
+        """
+        Initializes a TTL object which will expire in expire_in_seconds seconds from when created.
+        Remember to call .save() on the object after creation.
+        :param expiry_in_seconds: How long should the TTL be in the DB, in seconds. Please take into consideration
+        that the cleanup thread of mongo might take extra time to delete the TTL from the DB.
+        """
+        # Using UTC to make the mongodb TTL feature work. See
+        # https://stackoverflow.com/questions/55994379/mongodb-ttl-index-doesnt-delete-expired-documents.
+        return MonkeyTtl(expire_at=datetime.utcnow() + timedelta(seconds=expiry_in_seconds))
+
+    meta = {
+        'indexes': [
+            {
+                'name': 'TTL_index',
+                'fields': ['expire_at'],
+                'expireAfterSeconds': 0
+            }
+        ]
+    }
+
+    expire_at = DateTimeField()

monkey/monkey_island/cc/models/pba_results.py

@@ -0,0 +1,9 @@
+from mongoengine import EmbeddedDocument, StringField, ListField
+
+
+class PbaResults(EmbeddedDocument):
+    ip = StringField()
+    hostname = StringField()
+    command = StringField()
+    name = StringField()
+    result = ListField()

monkey/monkey_island/cc/models/test_monkey.py

@@ -0,0 +1,54 @@
+import uuid
+from time import sleep
+from unittest import TestCase
+
+from monkey import Monkey
+from monkey_island.cc.models.monkey import MonkeyNotFoundError
+from monkey_ttl import MonkeyTtl
+
+
+class TestMonkey(TestCase):
+    """
+    Make sure to set server environment to `testing` in server.json! Otherwise this will mess up your mongo instance and
+    won't work.
+
+    Also, the working directory needs to be the working directory from which you usually run the island so the
+    server.json file is found and loaded.
+    """
+    def test_is_dead(self):
+        # Arrange
+        alive_monkey_ttl = MonkeyTtl.create_ttl_expire_in(30)
+        alive_monkey_ttl.save()
+        alive_monkey = Monkey(
+            guid=str(uuid.uuid4()),
+            dead=False,
+            ttl_ref=alive_monkey_ttl.id)
+        alive_monkey.save()
+
+        # MIA stands for Missing In Action
+        mia_monkey_ttl = MonkeyTtl.create_ttl_expire_in(30)
+        mia_monkey_ttl.save()
+        mia_monkey = Monkey(guid=str(uuid.uuid4()), dead=False, ttl_ref=mia_monkey_ttl)
+        mia_monkey.save()
+        # Emulate timeout - ttl is manually deleted here, since we're using mongomock and not a real mongo instance.
+        sleep(1)
+        mia_monkey_ttl.delete()
+
+        dead_monkey = Monkey(guid=str(uuid.uuid4()), dead=True)
+        dead_monkey.save()
+
+        # act + assert
+        self.assertTrue(dead_monkey.is_dead())
+        self.assertTrue(mia_monkey.is_dead())
+        self.assertFalse(alive_monkey.is_dead())
+
+    def test_get_single_monkey_by_id(self):
+        # Arrange
+        a_monkey = Monkey(guid=str(uuid.uuid4()))
+        a_monkey.save()
+
+        # Act + assert
+        # Find the existing one
+        self.assertIsNotNone(Monkey.get_single_monkey_by_id(a_monkey.id))
+        # Raise on non-existent monkey
+        self.assertRaises(MonkeyNotFoundError, Monkey.get_single_monkey_by_id, "abcdefabcdefabcdefabcdef")

monkey/monkey_island/cc/resources/local_run.py

@@ -7,6 +7,7 @@ from flask import request, jsonify, make_response
 import flask_restful
 
 from monkey_island.cc.environment.environment import env
+from monkey_island.cc.models import Monkey
 from monkey_island.cc.resources.monkey_download import get_monkey_executable
 from monkey_island.cc.services.node import NodeService
 from monkey_island.cc.utils import local_ip_addresses
@@ -57,7 +58,7 @@ class LocalRun(flask_restful.Resource):
         NodeService.update_dead_monkeys()
         island_monkey = NodeService.get_monkey_island_monkey()
         if island_monkey is not None:
-            is_monkey_running = not island_monkey["dead"]
+            is_monkey_running = not Monkey.get_single_monkey_by_id(island_monkey["_id"]).is_dead()
         else:
             is_monkey_running = False
 

monkey/monkey_island/cc/resources/monkey.py

@@ -2,18 +2,29 @@ import json
 from datetime import datetime
 
 import dateutil.parser
-from flask import request
 import flask_restful
+from flask import request
 
 from monkey_island.cc.database import mongo
+from monkey_island.cc.models.monkey_ttl import MonkeyTtl
 from monkey_island.cc.services.config import ConfigService
 from monkey_island.cc.services.node import NodeService
 
+MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS = 120
+
 __author__ = 'Barak'
 
 # TODO: separate logic from interface
 
 
+def create_monkey_ttl():
+    # The TTL data uses the new `models` module which depends on mongoengine.
+    current_ttl = MonkeyTtl.create_ttl_expire_in(MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS)
+    current_ttl.save()
+    ttlid = current_ttl.id
+    return ttlid
+
+
 class Monkey(flask_restful.Resource):
 
     # Used by monkey. can't secure.
@@ -47,6 +58,9 @@ class Monkey(flask_restful.Resource):
             tunnel_host_ip = monkey_json['tunnel'].split(":")[-2].replace("//", "")
             NodeService.set_monkey_tunnel(monkey["_id"], tunnel_host_ip)
 
+        ttlid = create_monkey_ttl()
+        update['$set']['ttl_ref'] = ttlid
+
         return mongo.db.monkey.update({"_id": monkey["_id"]}, update, upsert=False)
 
     # Used by monkey. can't secure.
@@ -88,7 +102,7 @@ class Monkey(flask_restful.Resource):
                 parent_to_add = (exploit_telem[0].get('monkey_guid'), exploit_telem[0].get('data').get('exploiter'))
             else:
                 parent_to_add = (parent, None)
-        elif (not parent or parent == monkey_json.get('guid')) and 'ip_addresses' in  monkey_json:
+        elif (not parent or parent == monkey_json.get('guid')) and 'ip_addresses' in monkey_json:
             exploit_telem = [x for x in
                              mongo.db.telemetry.find({'telem_type': {'$eq': 'exploit'}, 'data.result': {'$eq': True},
                                                       'data.machine.ip_addr': {'$in': monkey_json['ip_addresses']}})]
@@ -106,6 +120,8 @@ class Monkey(flask_restful.Resource):
             tunnel_host_ip = monkey_json['tunnel'].split(":")[-2].replace("//", "")
             monkey_json.pop('tunnel')
 
+        monkey_json['ttl_ref'] = create_monkey_ttl()
+
         mongo.db.monkey.update({"guid": monkey_json["guid"]},
                                {"$set": monkey_json},
                                upsert=True)

monkey/monkey_island/cc/services/node.py

@@ -4,9 +4,11 @@ from bson import ObjectId
 
 import monkey_island.cc.services.log
 from monkey_island.cc.database import mongo
+from monkey_island.cc.models import Monkey
 from monkey_island.cc.services.edge import EdgeService
 from monkey_island.cc.utils import local_ip_addresses
 import socket
+from monkey_island.cc import models
 
 __author__ = "itay.mizeretz"
 
@@ -66,7 +68,7 @@ class NodeService:
     def get_node_label(node):
         domain_name = ""
         if node["domain_name"]:
-            domain_name = " ("+node["domain_name"]+")"
+            domain_name = " (" + node["domain_name"] + ")"
         return node["os"]["version"] + " : " + node["ip_addresses"][0] + domain_name
 
     @staticmethod
@@ -104,7 +106,8 @@ class NodeService:
 
     @staticmethod
     def get_monkey_critical_services(monkey_id):
-        critical_services = mongo.db.monkey.find_one({'_id': monkey_id}, {'critical_services': 1}).get('critical_services', [])
+        critical_services = mongo.db.monkey.find_one({'_id': monkey_id}, {'critical_services': 1}).get(
+            'critical_services', [])
         return critical_services
 
     @staticmethod
@@ -123,7 +126,7 @@ class NodeService:
             monkey_type = "manual" if NodeService.get_monkey_manual_run(monkey) else "monkey"
 
         monkey_os = NodeService.get_monkey_os(monkey)
-        monkey_running = "" if monkey["dead"] else "_running"
+        monkey_running = "" if Monkey.get_single_monkey_by_id(monkey["_id"]).is_dead() else "_running"
         return "%s_%s%s" % (monkey_type, monkey_os, monkey_running)
 
     @staticmethod
@@ -135,13 +138,14 @@ class NodeService:
     @staticmethod
     def monkey_to_net_node(monkey, for_report=False):
         label = monkey['hostname'] if for_report else NodeService.get_monkey_label(monkey)
+        is_monkey_dead = Monkey.get_single_monkey_by_id(monkey["_id"]).is_dead()
         return \
             {
                 "id": monkey["_id"],
                 "label": label,
                 "group": NodeService.get_monkey_group(monkey),
                 "os": NodeService.get_monkey_os(monkey),
-                "dead": monkey["dead"],
+                "dead": is_monkey_dead,
                 "domain_name": "",
                 "pba_results": monkey["pba_results"] if "pba_results" in monkey else []
             }
@@ -293,7 +297,8 @@ class NodeService:
 
     @staticmethod
     def is_any_monkey_alive():
-        return mongo.db.monkey.find_one({'dead': False}) is not None
+        all_monkeys = models.Monkey.objects()
+        return any(not monkey.is_dead() for monkey in all_monkeys)
 
     @staticmethod
     def is_any_monkey_exists():

monkey/monkey_island/requirements.txt

@@ -22,4 +22,6 @@ awscli
 cffi
 virtualenv
 wheel
+mongoengine
+mongomock
 requests