From 0d8070080ae8b199ffe3ac3ada90e21c5895c7a2 Mon Sep 17 00:00:00 2001
From: Ilija Lazoroski <ilija.la@live.com>
Date: Mon, 22 Nov 2021 14:27:39 +0100
Subject: [PATCH] Agent: Implement ControlChannel should_agent_stop

---
 monkey/infection_monkey/control_channel.py   | 32 ++++++++++++++++++++
 monkey/infection_monkey/i_control_channel.py | 27 +++++++++++++++++
 2 files changed, 59 insertions(+)
 create mode 100644 monkey/infection_monkey/control_channel.py
 create mode 100644 monkey/infection_monkey/i_control_channel.py

diff --git a/monkey/infection_monkey/control_channel.py b/monkey/infection_monkey/control_channel.py
new file mode 100644
index 000000000..639991e71
--- /dev/null
+++ b/monkey/infection_monkey/control_channel.py
@@ -0,0 +1,32 @@
+import json
+import logging
+from abc import ABC
+
+import requests
+
+from common.common_consts.timeouts import SHORT_REQUEST_TIMEOUT
+from infection_monkey.config import GUID, WormConfiguration
+from monkey.infection_monkey.i_control_channel import IControlChannel
+
+requests.packages.urllib3.disable_warnings()
+
+logger = logging.getLogger(__name__)
+
+
+class ControlChannel(IControlChannel, ABC):
+    def should_agent_stop(self) -> bool:
+        server = WormConfiguration.current_server
+        if not server:
+            return
+
+        try:
+            response = requests.get(  # noqa: DUO123
+                f"{server}/api/monkey_control/{GUID}",
+                verify=False,
+                timeout=SHORT_REQUEST_TIMEOUT,
+            )
+
+            response = json.loads(response.content.decode())
+            return response["stop_agent"]
+        except Exception as e:
+            logger.error(f"Error happened while trying to connect to server. {e}")
diff --git a/monkey/infection_monkey/i_control_channel.py b/monkey/infection_monkey/i_control_channel.py
new file mode 100644
index 000000000..eb1a4d5b2
--- /dev/null
+++ b/monkey/infection_monkey/i_control_channel.py
@@ -0,0 +1,27 @@
+import abc
+
+
+class IControlChannel(metaclass=abc.ABCMeta):
+    @abc.abstractmethod
+    def should_agent_stop(self) -> bool:
+        """
+        Checks if the agent should stop
+        return: True if the agent should stop, False otherwise
+        rtype: bool
+        """
+
+    @abc.abstractmethod
+    def get_config(self) -> dict:
+        """
+        :return: A dictionary containing Agent Configuration
+        :rtype: dict
+        """
+        pass
+
+    @abc.abstractmethod
+    def get_credentials_for_propagation(self) -> dict:
+        """
+        :return: A dictionary containing propagation credentials data
+        :rtype: dict
+        """
+        pass