Covert shell command adapter into endpoint tools, as tag address is s…

…o good

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:

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)