forked from p15670423/monkey
Added models using mongoengine and started using them in the code, and added TTL field
TTL doesn't get expired for some reason, trying to solve in https://stackoverflow.com/questions/55994379/mongodb-ttl-index-doesnt-delete-expired-documents
This commit is contained in:
parent
fd2e0887ff
commit
1018906602
|
@ -0,0 +1 @@
|
|||
from monkey import Monkey
|
|
@ -0,0 +1,89 @@
|
|||
"""
|
||||
Define a Document Schema for the Monkey document.
|
||||
"""
|
||||
import mongoengine
|
||||
from mongoengine import Document, StringField, ListField, BooleanField, EmbeddedDocumentField, DateField, \
|
||||
EmbeddedDocument, connect, ReferenceField, DateTimeField
|
||||
|
||||
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
|
||||
"""
|
||||
meta = {'strict': False}
|
||||
pass
|
||||
|
||||
|
||||
class PbaResults(EmbeddedDocument):
|
||||
ip = StringField()
|
||||
hostname = StringField()
|
||||
command = StringField()
|
||||
name = StringField()
|
||||
result = ListField()
|
||||
|
||||
|
||||
class Ttl(Document):
|
||||
meta = {
|
||||
'indexes': [
|
||||
{
|
||||
'name': 'TTL_index',
|
||||
'fields': ['expire_at'],
|
||||
'expireAfterSeconds': 0
|
||||
}
|
||||
]
|
||||
}
|
||||
|
||||
expire_at = DateTimeField()
|
||||
|
||||
|
||||
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.
|
||||
"""
|
||||
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(Ttl)
|
||||
|
||||
def is_dead(self):
|
||||
monkey_is_dead = False
|
||||
if self.dead:
|
||||
monkey_is_dead = True
|
||||
else:
|
||||
try:
|
||||
if Ttl.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:
|
||||
# Trying to dereference unknown document
|
||||
monkey_is_dead = True
|
||||
return monkey_is_dead
|
|
@ -1,13 +1,13 @@
|
|||
import json
|
||||
from datetime import datetime
|
||||
from datetime import datetime, timedelta
|
||||
|
||||
import dateutil.parser
|
||||
import flask_restful
|
||||
from flask import request
|
||||
|
||||
from monkey_island.cc import models
|
||||
from monkey_island.cc.database import mongo
|
||||
from monkey_island.cc.services.config import ConfigService
|
||||
from monkey_island.cc.services.monkey_timeout import start_timer_decorator
|
||||
from monkey_island.cc.services.node import NodeService
|
||||
|
||||
__author__ = 'Barak'
|
||||
|
@ -18,7 +18,6 @@ __author__ = 'Barak'
|
|||
class Monkey(flask_restful.Resource):
|
||||
|
||||
# Used by monkey. can't secure.
|
||||
@start_timer_decorator
|
||||
def get(self, guid=None, **kw):
|
||||
NodeService.update_dead_monkeys() # refresh monkeys status
|
||||
if not guid:
|
||||
|
@ -49,6 +48,11 @@ 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.Ttl(expire_at=datetime.now() + timedelta(seconds=30))
|
||||
current_ttl.save()
|
||||
|
||||
update['$set']['ttl_ref'] = current_ttl.id
|
||||
|
||||
return mongo.db.monkey.update({"_id": monkey["_id"]}, update, upsert=False)
|
||||
|
||||
# Used by monkey. can't secure.
|
||||
|
|
|
@ -7,6 +7,7 @@ from monkey_island.cc.database import mongo
|
|||
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"
|
||||
|
||||
|
@ -141,7 +142,7 @@ class NodeService:
|
|||
"label": label,
|
||||
"group": NodeService.get_monkey_group(monkey),
|
||||
"os": NodeService.get_monkey_os(monkey),
|
||||
"dead": monkey["dead"],
|
||||
"dead": models.Monkey.objects(id=monkey["_id"])[0].is_dead(),
|
||||
"domain_name": "",
|
||||
"pba_results": monkey["pba_results"] if "pba_results" in monkey else []
|
||||
}
|
||||
|
@ -293,7 +294,7 @@ class NodeService:
|
|||
|
||||
@staticmethod
|
||||
def is_any_monkey_alive():
|
||||
return mongo.db.monkey.find_one({'dead': False}) is not None
|
||||
return models.Monkey.objects(dead=False).count() > 0
|
||||
|
||||
@staticmethod
|
||||
def is_any_monkey_exists():
|
||||
|
|
|
@ -22,3 +22,4 @@ bson
|
|||
cffi
|
||||
virtualenv
|
||||
wheel
|
||||
mongoengine
|
||||
|
|
Loading…
Reference in New Issue