Extracted models to files, created TestingEnv, and added unit testing for Monkey.is_dead

The init of models checks the env and sets up the DB connection.

monkey/monkey_island/cc/environment/__init__.py

@@ -16,9 +16,19 @@ class Environment(object):
     _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
 
     def __init__(self):
         self.config = None
+        self._testing = False  # Assume env is not for unit testing.
 
     def set_config(self, config):
         self.config = config

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
 }
 
 
@@ -32,6 +37,7 @@ def load_env_from_file():
     config_json = load_server_configuration_from_file()
     return config_json['server_config']
 
+
 try:
     config_json = load_server_configuration_from_file()
     __env_type = config_json['server_config']

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

@@ -1 +1,16 @@
+from mongoengine import connect
+
+# set up the DB connection.
+from monkey_island.cc.environment.environment import 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 or importing matters, 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

@@ -3,51 +3,9 @@ Define a Document Schema for the Monkey document.
 """
 import mongoengine
 from mongoengine import Document, StringField, ListField, BooleanField, EmbeddedDocumentField, DateField, \
-    EmbeddedDocument, connect, ReferenceField, DateTimeField
+    ReferenceField
 
-from monkey_island.cc.environment.environment import env
-
-connect(db=env.mongo_db_name, host=env.mongo_db_host, port=env.mongo_db_port)
-
-
-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
-
-
-class Creds(EmbeddedDocument):
-    """
-    TODO get an example of this data, and make it strict
-    """
-    meta = {'strict': False}
-    pass
-
-
-class PbaResults(EmbeddedDocument):
-    ip = StringField()
-    hostname = StringField()
-    command = StringField()
-    name = StringField()
-    result = ListField()
-
-
-class MonkeyTtl(Document):
-    meta = {
-        'indexes': [
-            {
-                'name': 'TTL_index',
-                'fields': ['expire_at'],
-                'expireAfterSeconds': 0
-            }
-        ]
-    }
-
-    expire_at = DateTimeField()
+from monkey_island.cc.models.monkey_ttl import MonkeyTtl
 
 
 class Monkey(Document):
@@ -81,9 +39,9 @@ class Monkey(Document):
         else:
             try:
                 if MonkeyTtl.objects(id=self.ttl_ref.id).count() == 0:
-                    # No TTLs - monkey has timed out. The monkey is MIA
+                    # No TTLs - monkey has timed out. The monkey is MIA.
                     monkey_is_dead = True
             except mongoengine.DoesNotExist:
-                # Trying to dereference unknown document
+                # Trying to dereference unknown document - the monkey is MIA.
                 monkey_is_dead = True
         return monkey_is_dead

monkey/monkey_island/cc/models/monkey_ttl.py

@@ -0,0 +1,15 @@
+from mongoengine import Document, DateTimeField
+
+
+class MonkeyTtl(Document):
+    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,37 @@
+import uuid
+from datetime import timedelta, datetime
+from time import sleep
+from unittest import TestCase
+
+# noinspection PyUnresolvedReferences
+import mongomock
+
+from monkey import Monkey
+from monkey_ttl import MonkeyTtl
+
+
+class TestMonkey(TestCase):
+    def test_is_dead(self):
+        alive_monkey_ttl = MonkeyTtl(expire_at=datetime.now() + timedelta(seconds=30))
+        alive_monkey_ttl.save()
+        alive_monkey = Monkey(
+            guid=str(uuid.uuid4()),
+            dead=False,
+            ttl_ref=alive_monkey_ttl.id)
+        alive_monkey.save()
+
+        mia_monkey_ttl = MonkeyTtl(expire_at=datetime.now() + timedelta(seconds=30))
+        mia_monkey_ttl.save()
+        mia_monkey = Monkey(guid=str(uuid.uuid4()), dead=False, ttl_ref=mia_monkey_ttl)
+        mia_monkey.save()
+        # Emulate timeout
+        sleep(1)
+        mia_monkey_ttl.delete()
+
+        dead_monkey = Monkey(guid=str(uuid.uuid4()), dead=True)
+        dead_monkey.save()
+
+        self.assertTrue(dead_monkey.is_dead())
+        self.assertTrue(mia_monkey.is_dead())
+        self.assertFalse(alive_monkey.is_dead())
+

monkey/monkey_island/cc/resources/monkey.py

@@ -5,7 +5,7 @@ import dateutil.parser
 import flask_restful
 from flask import request
 
-from monkey_island.cc import models
+from monkey_island.cc.models.monkey_ttl import MonkeyTtl
 from monkey_island.cc.database import mongo
 from monkey_island.cc.services.config import ConfigService
 from monkey_island.cc.services.node import NodeService
@@ -48,7 +48,8 @@ class Monkey(flask_restful.Resource):
             tunnel_host_ip = monkey_json['tunnel'].split(":")[-2].replace("//", "")
             NodeService.set_monkey_tunnel(monkey["_id"], tunnel_host_ip)
 
-        current_ttl = models.monkey.MonkeyTtl(expire_at=datetime.now() + timedelta(seconds=30))
+        # The TTL data uses the new `models` module which depends on mongoengine.
+        current_ttl = MonkeyTtl(expire_at=datetime.now() + timedelta(seconds=30))
         current_ttl.save()
 
         update['$set']['ttl_ref'] = current_ttl.id

monkey/monkey_island/requirements.txt

@@ -23,3 +23,4 @@ cffi
 virtualenv
 wheel
 mongoengine
+mongomock