diff --git a/monkey/monkey_island/cc/environment/__init__.py b/monkey/monkey_island/cc/environment/__init__.py index 8a3f61b21..e30e0eebd 100644 --- a/monkey/monkey_island/cc/environment/__init__.py +++ b/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 diff --git a/monkey/monkey_island/cc/environment/environment.py b/monkey/monkey_island/cc/environment/environment.py index 3cd6bb587..6115e8dd9 100644 --- a/monkey/monkey_island/cc/environment/environment.py +++ b/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'] diff --git a/monkey/monkey_island/cc/environment/testing.py b/monkey/monkey_island/cc/environment/testing.py new file mode 100644 index 000000000..286e442dd --- /dev/null +++ b/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) + ] diff --git a/monkey/monkey_island/cc/models/__init__.py b/monkey/monkey_island/cc/models/__init__.py index c02d7f0c1..673b6454a 100644 --- a/monkey/monkey_island/cc/models/__init__.py +++ b/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 diff --git a/monkey/monkey_island/cc/models/config.py b/monkey/monkey_island/cc/models/config.py new file mode 100644 index 000000000..cfe128111 --- /dev/null +++ b/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 diff --git a/monkey/monkey_island/cc/models/creds.py b/monkey/monkey_island/cc/models/creds.py new file mode 100644 index 000000000..61322362e --- /dev/null +++ b/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 diff --git a/monkey/monkey_island/cc/models/monkey.py b/monkey/monkey_island/cc/models/monkey.py index 73f924847..f2a8ae5a4 100644 --- a/monkey/monkey_island/cc/models/monkey.py +++ b/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 diff --git a/monkey/monkey_island/cc/models/monkey_ttl.py b/monkey/monkey_island/cc/models/monkey_ttl.py new file mode 100644 index 000000000..e6d867628 --- /dev/null +++ b/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() diff --git a/monkey/monkey_island/cc/models/pba_results.py b/monkey/monkey_island/cc/models/pba_results.py new file mode 100644 index 000000000..d2cc48080 --- /dev/null +++ b/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() diff --git a/monkey/monkey_island/cc/models/test_monkey.py b/monkey/monkey_island/cc/models/test_monkey.py new file mode 100644 index 000000000..a79bf9bbb --- /dev/null +++ b/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()) + diff --git a/monkey/monkey_island/cc/resources/monkey.py b/monkey/monkey_island/cc/resources/monkey.py index ace843c85..9a028e71e 100644 --- a/monkey/monkey_island/cc/resources/monkey.py +++ b/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 diff --git a/monkey/monkey_island/requirements.txt b/monkey/monkey_island/requirements.txt index 953cad2cb..6b0b68031 100644 --- a/monkey/monkey_island/requirements.txt +++ b/monkey/monkey_island/requirements.txt @@ -23,3 +23,4 @@ cffi virtualenv wheel mongoengine +mongomock