Merge pull request #924 from NordSecurity/tag-each-ping-in-nat-lab

Tag each ping in nat lab with random payload

nat-lab/tests/utils/connection/connection.py

@@ -17,7 +17,7 @@ def __init__(self, target_os: TargetOS) -> None:
         self._target_os = target_os
 
     @abstractmethod
-    def create_process(self, command: List[str]) -> "Process":
+    def create_process(self, command: List[str], kill_id=None) -> "Process":
         pass
 
     @property

nat-lab/tests/utils/connection/docker_connection.py

@@ -33,8 +33,8 @@ def container_name(self) -> str:
     def target_name(self) -> str:
         return self.container_name()
 
-    def create_process(self, command: List[str]) -> "Process":
-        process = DockerProcess(self._container, command)
+    def create_process(self, command: List[str], kill_id=None) -> "Process":
+        process = DockerProcess(self._container, command, kill_id)
         print(
             datetime.now(),
             "Executing",

nat-lab/tests/utils/connection/ssh_connection.py

@@ -16,7 +16,7 @@ def __init__(self, connection: asyncssh.SSHClientConnection, target_os: TargetOS
         self._connection = connection
         self._target_os = target_os
 
-    def create_process(self, command: List[str]) -> "Process":
+    def create_process(self, command: List[str], kill_id=None) -> "Process":
         print(datetime.now(), "Executing", command, "on", self.target_os)
         if self._target_os == TargetOS.Windows:
             escape_argument = cmd_exe_escape.escape_argument

nat-lab/tests/utils/ping.py

@@ -1,5 +1,6 @@
 import asyncio
 import re
+import secrets
 from contextlib import asynccontextmanager
 from ipaddress import ip_address
 from typing import AsyncIterator, Optional
@@ -33,18 +34,29 @@ def __init__(self, connection: Connection, ip: str) -> None:
         self._ip = ip
         self._connection = connection
         self._ip_proto = testing.unpack_optional(get_ip_address_type(ip))
+        kill_id = secrets.token_hex(8).upper()
 
         if connection.target_os == TargetOS.Windows:
             self._process = connection.create_process(
                 ["ping", ("-4" if self._ip_proto == IPProto.IPv4 else "-6"), "-t", ip]
             )
         elif connection.target_os == TargetOS.Mac:
-            self._process = connection.create_process(
-                [("ping" if self._ip_proto == IPProto.IPv4 else "ping6"), ip]
-            )
+            self._process = connection.create_process([
+                ("ping" if self._ip_proto == IPProto.IPv4 else "ping6"),
+                "-p",
+                kill_id,
+                ip,
+            ])
         else:
             self._process = connection.create_process(
-                ["ping", ("-4" if self._ip_proto == IPProto.IPv4 else "-6"), ip]
+                [
+                    "ping",
+                    ("-4" if self._ip_proto == IPProto.IPv4 else "-6"),
+                    "-p",
+                    kill_id,
+                    ip,
+                ],
+                kill_id,
             )
         self._next_ping_event = asyncio.Event()
 

nat-lab/tests/utils/process/docker_process.py

@@ -1,6 +1,5 @@
 import asyncio
-import random
-import string
+import secrets
 import subprocess
 import sys
 from .process import Process, ProcessExecError, StreamCallback
@@ -26,7 +25,9 @@ class DockerProcess(Process):
         str  # Private ID added to find the process easily when it needs to be killed
     )
 
-    def __init__(self, container: DockerContainer, command: List[str]) -> None:
+    def __init__(
+        self, container: DockerContainer, command: List[str], kill_id=None
+    ) -> None:
         self._container = container
         self._command = command
         self._stdout = ""
@@ -35,9 +36,7 @@ def __init__(self, container: DockerContainer, command: List[str]) -> None:
         self._is_done = asyncio.Event()
         self._stream = None
         self._execute = None
-        self._kill_id = "".join(
-            random.choices(string.ascii_uppercase + string.digits, k=12)
-        )
+        self._kill_id = kill_id if kill_id else secrets.token_hex(8).upper()
 
     async def execute(
         self,