From 06baf19eeef19918ac76fd36c80be9fc8b21d230 Mon Sep 17 00:00:00 2001
From: Rauli Kaksonen <43605561+raulikak@users.noreply.github.com>
Date: Fri, 24 May 2024 14:50:47 +0300
Subject: [PATCH] Covert shell command adapter into endpoint tools, as tag
 address is so good

---
 tcsfw/shell_commands.py | 24 ++++++++++++++----------
 tcsfw/tools.py          |  5 +++--
 2 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/tcsfw/shell_commands.py b/tcsfw/shell_commands.py
index ffb53bc..def6430 100644
--- a/tcsfw/shell_commands.py
+++ b/tcsfw/shell_commands.py
@@ -3,22 +3,25 @@
 from io import BytesIO, TextIOWrapper
 import re
 from typing import Dict, Tuple
-from tcsfw.address import Addresses, EndpointAddress, HWAddresses, IPAddress
+from tcsfw.address import Addresses, AnyAddress, EndpointAddress, HWAddresses, IPAddress
 from tcsfw.components import OperatingSystem
 from tcsfw.event_interface import EventInterface, PropertyEvent
-from tcsfw.model import Addressable, IoTSystem, NetworkNode
+from tcsfw.model import IoTSystem
 from tcsfw.property import PropertyKey
-from tcsfw.tools import NetworkNodeTool
+from tcsfw.tools import EndpointTool
 from tcsfw.traffic import Evidence, EvidenceSource, IPFlow, Protocol, ServiceScan
 from tcsfw.verdict import Verdict
 
 
-class ShellCommandPs(NetworkNodeTool):
+class ShellCommandPs(EndpointTool):
     """Shell command 'ps' tool adapter"""
     def __init__(self, system: IoTSystem):
         super().__init__("shell-ps", ".txt", system)
 
-    def process_node(self, node: NetworkNode, data_file: BytesIO, interface: EventInterface, source: EvidenceSource):
+    def process_endpoint(self, endpoint: AnyAddress, stream: BytesIO, interface: EventInterface,
+                         source: EvidenceSource):
+        node = self.system.get_endpoint(endpoint)
+
         columns: Dict[str, int] = {}
         unexpected = {}
         os = OperatingSystem.get_os(node, add=self.load_baseline)
@@ -28,7 +31,7 @@ def process_node(self, node: NetworkNode, data_file: BytesIO, interface: EventIn
         for user, ps_list in os.process_map.items():
             regexp_map[user] = [re.compile(ps) for ps in ps_list]
 
-        with TextIOWrapper(data_file) as f:
+        with TextIOWrapper(stream) as f:
             while True:
                 line = f.readline().split(maxsplit=len(columns) -1 if columns else -1)
                 if not line:
@@ -82,7 +85,7 @@ def process_node(self, node: NetworkNode, data_file: BytesIO, interface: EventIn
                 interface.property_update(ev)
 
 
-class ShellCommandSs(NetworkNodeTool):
+class ShellCommandSs(EndpointTool):
     """Shell command 'ss' tool adapter"""
     def __init__(self, system: IoTSystem):
         super().__init__("shell-ss", ".txt", system)
@@ -96,16 +99,17 @@ def _parse_address(self, addr: str) -> Tuple[str, str, int]:
     LOCAL_ADDRESS = "Local_Address"
     PEER_ADDRESS = "Peer_Address"
 
-    def process_node(self, node: NetworkNode, data_file: BytesIO, interface: EventInterface, source: EvidenceSource):
+    def process_endpoint(self, endpoint: AnyAddress, stream: BytesIO, interface: EventInterface,
+                         source: EvidenceSource):
         columns: Dict[str, int] = {}
         local_ads = set()
         services = set()
         conns = set()
 
-        assert isinstance(node, Addressable)
+        node = self.system.get_endpoint(endpoint)
         tag = Addresses.get_tag(node.addresses)
 
-        with TextIOWrapper(data_file) as f:
+        with TextIOWrapper(stream) as f:
             while True:
                 line = f.readline()
                 if not line:
diff --git a/tcsfw/tools.py b/tcsfw/tools.py
index ee360a3..7b53a59 100644
--- a/tcsfw/tools.py
+++ b/tcsfw/tools.py
@@ -98,9 +98,10 @@ def create_file_name_map(self):
 
     def map_addressable(self, entity: Addressable):
         """Map addressable entity to file names"""
-        # First pass is DNS names, then IP addresses
+        # First pass is Tags, DNS names, then IP addresses
         addresses = entity.get_addresses()
-        ads_sorted = [a for a in addresses if isinstance(a.get_host(), DNSName)]
+        ads_sorted = [a for a in addresses if a.is_tag()]
+        ads_sorted.extend([a for a in addresses if isinstance(a.get_host(), DNSName)])
         ads_sorted.extend([a for a in addresses if isinstance(a.get_host(), IPAddress)])
         for a in ads_sorted:
             a_file_name = self.get_file_by_endpoint(a)