From 50a6e17fabfd6c20eb0f5ab63b01e7f3377144b5 Mon Sep 17 00:00:00 2001
From: Nick Todd <nick.todd@t-s.qinetiq.com>
Date: Wed, 6 Dec 2023 16:42:28 +0000
Subject: [PATCH] 2041: Make NTP work with TCP transport layer

---
 src/primaite/game/agent/actions.py                     |  2 +-
 src/primaite/game/game.py                              |  4 ++++
 .../simulator/system/services/ntp/ntp_client.py        | 10 +++++-----
 .../simulator/system/services/ntp/ntp_server.py        |  3 ++-
 4 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/src/primaite/game/agent/actions.py b/src/primaite/game/agent/actions.py
index 8eed3ba4..893b12b4 100644
--- a/src/primaite/game/agent/actions.py
+++ b/src/primaite/game/agent/actions.py
@@ -601,7 +601,7 @@ class ActionManager:
         max_nics_per_node: int = 8,  # allows calculating shape
         max_acl_rules: int = 10,  # allows calculating shape
         protocols: List[str] = ["TCP", "UDP", "ICMP"],  # allow mapping index to protocol
-        ports: List[str] = ["HTTP", "DNS", "ARP", "FTP"],  # allow mapping index to port
+        ports: List[str] = ["HTTP", "DNS", "ARP", "FTP", "NTP"],  # allow mapping index to port
         ip_address_list: Optional[List[str]] = None,  # to allow us to map an index to an ip address.
         act_map: Optional[Dict[int, Dict]] = None,  # allows restricting set of possible actions
     ) -> None:
diff --git a/src/primaite/game/game.py b/src/primaite/game/game.py
index a36cbea9..edfe058d 100644
--- a/src/primaite/game/game.py
+++ b/src/primaite/game/game.py
@@ -25,6 +25,8 @@ from primaite.simulator.system.services.dns.dns_client import DNSClient
 from primaite.simulator.system.services.dns.dns_server import DNSServer
 from primaite.simulator.system.services.ftp.ftp_client import FTPClient
 from primaite.simulator.system.services.ftp.ftp_server import FTPServer
+from primaite.simulator.system.services.ntp.ntp_client import NTPClient
+from primaite.simulator.system.services.ntp.ntp_server import NTPServer
 from primaite.simulator.system.services.red_services.data_manipulation_bot import DataManipulationBot
 from primaite.simulator.system.services.web_server.web_server import WebServer
 
@@ -257,6 +259,8 @@ class PrimaiteGame:
                         "WebServer": WebServer,
                         "FTPClient": FTPClient,
                         "FTPServer": FTPServer,
+                        "NTPClient": NTPClient,
+                        "NTPServer": NTPServer,
                     }
                     if service_type in service_types_mapping:
                         _LOGGER.debug(f"installing {service_type} on node {new_node.hostname}")
diff --git a/src/primaite/simulator/system/services/ntp/ntp_client.py b/src/primaite/simulator/system/services/ntp/ntp_client.py
index c75e639d..f9cf29d4 100644
--- a/src/primaite/simulator/system/services/ntp/ntp_client.py
+++ b/src/primaite/simulator/system/services/ntp/ntp_client.py
@@ -23,7 +23,7 @@ class NTPClient(Service):
     def __init__(self, **kwargs):
         kwargs["name"] = "NTPClient"
         kwargs["port"] = Port.NTP
-        kwargs["protocol"] = IPProtocol.UDP
+        kwargs["protocol"] = IPProtocol.TCP
         super().__init__(**kwargs)
         self.start()
 
@@ -65,7 +65,7 @@ class NTPClient(Service):
         self,
         payload: NTPPacket,
         session_id: Optional[str] = None,
-        dest_ip_address: IPv4Address = ntp_server,
+        dest_ip_address: IPv4Address = None,
         dest_port: [Port] = Port.NTP,
         **kwargs,
     ) -> bool:
@@ -79,8 +79,6 @@ class NTPClient(Service):
         :return: True if successful, False otherwise.
         """
         self.ip_addr = payload.ntp_request.ntp_client
-        self.sys_log.info(f"{self.name}: Sending NTP request {payload.ntp_request.ntp_client}")
-
         return super().send(
             payload=payload,
             dest_ip_address=dest_ip_address,
@@ -101,6 +99,8 @@ class NTPClient(Service):
         :param session_id: The Session ID the payload is to originate from. Optional.
         :return: True if successful, False otherwise.
         """
+        self.sys_log.info(f"{self.name}: Receiving NTP request from {payload.ntp_request.ntp_client}")
+
         if not (isinstance(payload, NTPPacket) and payload.ntp_request.ntp_client):
             _LOGGER.debug(f"{payload} is not a NTPPacket")
             return False
@@ -116,7 +116,7 @@ class NTPClient(Service):
         """Send request to ntp_server."""
         ntp_request = NTPRequest(ntp_client=self.ip_addr)
         ntp_server_packet = NTPPacket(ntp_request=ntp_request)
-        self.send(payload=ntp_server_packet)
+        self.send(payload=ntp_server_packet, dest_ip_address=self.ntp_server)
 
     def apply_timestep(self, timestep: int) -> None:
         """
diff --git a/src/primaite/simulator/system/services/ntp/ntp_server.py b/src/primaite/simulator/system/services/ntp/ntp_server.py
index 6d76c1ed..400c397f 100644
--- a/src/primaite/simulator/system/services/ntp/ntp_server.py
+++ b/src/primaite/simulator/system/services/ntp/ntp_server.py
@@ -16,7 +16,7 @@ class NTPServer(Service):
     def __init__(self, **kwargs):
         kwargs["name"] = "NTPServer"
         kwargs["port"] = Port.NTP
-        kwargs["protocol"] = IPProtocol.UDP
+        kwargs["protocol"] = IPProtocol.TCP
         super().__init__(**kwargs)
         self.start()
 
@@ -60,6 +60,7 @@ class NTPServer(Service):
 
         :return: True if valid NTP request else False.
         """
+        self.sys_log.info(f"{self.name} received request from {payload.ntp_request.ntp_client}")
         if not (isinstance(payload, NTPPacket) and payload.ntp_request.ntp_client):
             _LOGGER.debug(f"{payload} is not a NTPPacket")
             return False