From 7d397da24538cdfaa90764d7993c8c8cd3d6ec2c Mon Sep 17 00:00:00 2001
From: VakarisZ <vakarisz@yahoo.com>
Date: Wed, 15 Apr 2020 15:42:37 +0300
Subject: [PATCH] Added telemetry collection

---
 monkey/monkey_island/cc/models/test_telem.py  | 18 ++++++++++++++++++
 monkey/monkey_island/cc/resources/monkey.py   |  3 +++
 .../cc/resources/test/utils/telem_store.py    | 19 +++++++++++++++++++
 3 files changed, 40 insertions(+)
 create mode 100644 monkey/monkey_island/cc/models/test_telem.py
 create mode 100644 monkey/monkey_island/cc/resources/test/utils/telem_store.py

diff --git a/monkey/monkey_island/cc/models/test_telem.py b/monkey/monkey_island/cc/models/test_telem.py
new file mode 100644
index 000000000..98bb23ebc
--- /dev/null
+++ b/monkey/monkey_island/cc/models/test_telem.py
@@ -0,0 +1,18 @@
+"""
+Define a Document Schema for the Monkey document.
+"""
+from mongoengine import Document, StringField
+
+
+class TestTelem(Document):
+    # SCHEMA
+    method = StringField(required=True)
+    endpoint = StringField(required=True)
+    content = StringField(required=True)
+
+    @staticmethod
+    def try_drop_collection():
+        try:
+            TestTelem.drop_collection()
+        except Exception:
+            pass
diff --git a/monkey/monkey_island/cc/resources/monkey.py b/monkey/monkey_island/cc/resources/monkey.py
index 3e3ef40c0..5d4be9645 100644
--- a/monkey/monkey_island/cc/resources/monkey.py
+++ b/monkey/monkey_island/cc/resources/monkey.py
@@ -3,6 +3,7 @@ from datetime import datetime
 
 import dateutil.parser
 import flask_restful
+from monkey_island.cc.resources.test.utils.telem_store import store_test_telem
 from flask import request
 
 from monkey_island.cc.consts import DEFAULT_MONKEY_TTL_EXPIRY_DURATION_IN_SECONDS
@@ -33,6 +34,7 @@ class Monkey(flask_restful.Resource):
         return {}
 
     # Used by monkey. can't secure.
+    @store_test_telem
     def patch(self, guid):
         monkey_json = json.loads(request.data)
         update = {"$set": {'modifytime': datetime.now()}}
@@ -56,6 +58,7 @@ class Monkey(flask_restful.Resource):
         return mongo.db.monkey.update({"_id": monkey["_id"]}, update, upsert=False)
 
     # Used by monkey. can't secure.
+    @store_test_telem
     def post(self, **kw):
         monkey_json = json.loads(request.data)
         monkey_json['creds'] = []
diff --git a/monkey/monkey_island/cc/resources/test/utils/telem_store.py b/monkey/monkey_island/cc/resources/test/utils/telem_store.py
new file mode 100644
index 000000000..55351a804
--- /dev/null
+++ b/monkey/monkey_island/cc/resources/test/utils/telem_store.py
@@ -0,0 +1,19 @@
+from functools import wraps
+
+from flask import request
+
+from monkey_island.cc.models.test_telem import TestTelem
+
+MONKEY_TELEM_COLLECTION_NAME = "monkey_telems_for_tests"
+
+
+def store_test_telem(f):
+    @wraps(f)
+    def decorated_function(*args, **kwargs):
+        method = request.method
+        content = request.data.decode()
+        endpoint = str(request.url_rule)
+        TestTelem(method=method, endpoint=endpoint, content=content).save()
+        return f(*args, **kwargs)
+
+    return decorated_function