From 39dfbb741f53fbd43c25e43cd7c5dcbad153c29e Mon Sep 17 00:00:00 2001
From: Chris McCarthy <chris.mccarthy@methods.co.uk>
Date: Tue, 28 Nov 2023 00:21:41 +0000
Subject: [PATCH] #1859 - Made some fixes to resets. Still an issue with the
 Router reset.

---
 src/primaite/simulator/network/hardware/base.py       |  1 +
 .../simulator/network/hardware/nodes/router.py        |  4 ++++
 .../simulator/system/services/dns/dns_server.py       | 11 +++++++++++
 .../system/services/web_server/web_server.py          |  3 +--
 4 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/src/primaite/simulator/network/hardware/base.py b/src/primaite/simulator/network/hardware/base.py
index 2863dd22..09e2b12f 100644
--- a/src/primaite/simulator/network/hardware/base.py
+++ b/src/primaite/simulator/network/hardware/base.py
@@ -134,6 +134,7 @@ class NIC(SimComponent):
         if episode and self.pcap:
             self.pcap.current_episode = episode
             self.pcap.setup_logger()
+        self.enable()
 
     def describe_state(self) -> Dict:
         """
diff --git a/src/primaite/simulator/network/hardware/nodes/router.py b/src/primaite/simulator/network/hardware/nodes/router.py
index 8e03cfa3..1bf2ea2f 100644
--- a/src/primaite/simulator/network/hardware/nodes/router.py
+++ b/src/primaite/simulator/network/hardware/nodes/router.py
@@ -667,6 +667,10 @@ class Router(Node):
         """Reset the original state of the SimComponent."""
         self.arp.clear()
         self.acl.reset_component_for_episode(episode)
+        for i, nic in self.ethernet_ports.items():
+            nic.reset_component_for_episode(episode)
+            self.enable_port(i)
+
         super().reset_component_for_episode(episode)
 
     def _init_request_manager(self) -> RequestManager:
diff --git a/src/primaite/simulator/system/services/dns/dns_server.py b/src/primaite/simulator/system/services/dns/dns_server.py
index bbeaa62c..3b1f3bf6 100644
--- a/src/primaite/simulator/system/services/dns/dns_server.py
+++ b/src/primaite/simulator/system/services/dns/dns_server.py
@@ -28,10 +28,21 @@ class DNSServer(Service):
         super().__init__(**kwargs)
         self.start()
 
+    def set_original_state(self):
+        """Sets the original state."""
+        super().set_original_state()
+        vals_to_include = {"dns_table"}
+        self._original_state["dns_table_orig"] = self.model_dump(include=vals_to_include)["dns_table"]
+
     def reset_component_for_episode(self, episode: int):
         """Reset the original state of the SimComponent."""
+        print("dns reset")
+        print("DNSServer original state", self._original_state)
         self.dns_table.clear()
+        for key, value in self._original_state["dns_table_orig"].items():
+            self.dns_table[key] = value
         super().reset_component_for_episode(episode)
+        self.show()
 
     def describe_state(self) -> Dict:
         """
diff --git a/src/primaite/simulator/system/services/web_server/web_server.py b/src/primaite/simulator/system/services/web_server/web_server.py
index 754aa22f..56f47195 100644
--- a/src/primaite/simulator/system/services/web_server/web_server.py
+++ b/src/primaite/simulator/system/services/web_server/web_server.py
@@ -31,7 +31,6 @@ class WebServer(Service):
 
     @last_response_status_code.setter
     def last_response_status_code(self, val: Any):
-        print(f"val: {val}, type: {type(val)}")
         self._last_response_status_code = val
 
     def describe_state(self) -> Dict:
@@ -47,6 +46,7 @@ class WebServer(Service):
         state["last_response_status_code"] = (
             self.last_response_status_code.value if isinstance(self.last_response_status_code, HttpStatusCode) else None
         )
+        print(state)
         return state
 
     def __init__(self, **kwargs):
@@ -99,7 +99,6 @@ class WebServer(Service):
 
         # return true if response is OK
         self.last_response_status_code = response.status_code
-
         return response.status_code == HttpStatusCode.OK
 
     def _handle_get_request(self, payload: HttpRequestPacket) -> HttpResponsePacket: