Fix issue of reverse-lookup-host.py and forward-lookup-host.py not obeying ignore_dest_user.

tools/forward-lookup-host.py

@@ -9,24 +9,27 @@ import heapq
 from collections import defaultdict
 import argparse
 
-def indirect_get_connected_nodes(graph, interesting_host):
+def indirect_get_connected_nodes(graph, interesting_host, ignore_dest_user):
-    backward_connected_nodes = set()
+    connected_nodes = set()
     sentinel = '__SENTINEL__'
 
-    backward_heap = [(interesting_host, sentinel)]
+    heap = [(interesting_host, sentinel)]
-    while backward_heap:
+    while heap:
-        current_node, parent_node = heapq.heappop(backward_heap)
+        current_node, parent_node = heapq.heappop(heap)
 
-        if current_node not in backward_connected_nodes:
+        if current_node not in connected_nodes:
-            backward_connected_nodes.add(current_node)
+            connected_nodes.add(current_node)
             if parent_node is not sentinel:
-                heapq.heappush(backward_heap, (parent_node, sentinel))
+                heapq.heappush(heap, (parent_node, sentinel))
             if current_node in graph:
                 for connection in graph[current_node]:
-                    node = connection[4]  # Assuming the fifth element in the tuple is the destination host
+                    if ignore_dest_user:
-                    heapq.heappush(backward_heap, (node, current_node))
+                        node = connection[4] # dest_host
+                    else:
+                        node = f"{connection[3]}@{connection[4]}" #dest_host@dest_user
+                    heapq.heappush(heap, (node, current_node))
 
-    return backward_connected_nodes
+    return connected_nodes
 
 def build_lookup_table(input_lines, ignore_dest_user):
     graph = defaultdict(set)
@@ -94,7 +97,7 @@ if __name__ == "__main__":
     if interesting_host in lookup_table:
         print(f"{interesting_host} is able to connect {mode} to:\n")
         if mode == "indirectly":
-            result = indirect_get_connected_nodes(lookup_table, interesting_host)
+            result = indirect_get_connected_nodes(lookup_table, interesting_host, ignore_dest_user)
             for dest in result:
                 print(dest)
         else:

tools/reverse-lookup-host.py

@@ -9,7 +9,7 @@ import heapq
 from collections import defaultdict
 import argparse
 
-def indirect_get_connected_nodes(graph, interesting_host):
+def indirect_get_connected_nodes(graph, interesting_host, ignore_dest_user):
     backward_connected_nodes = set()
     sentinel = '__SENTINEL__'
 
@@ -23,7 +23,10 @@ def indirect_get_connected_nodes(graph, interesting_host):
                 heapq.heappush(backward_heap, (parent_node, sentinel))
             if current_node in graph:
                 for connection in graph[current_node]:
-                    node = connection[1]  # Assuming the second element in the tuple is the source host (so find it as a dest)
+                    if ignore_dest_user:
+                        node = connection[1] # host
+                    else:
+                        node = f"{connection[0]}@{connection[1]}" # user@host
                     heapq.heappush(backward_heap, (node, current_node))
 
     return backward_connected_nodes
@@ -57,7 +60,7 @@ def build_lookup_table(input_lines, ignore_dest_user):
 
             line_to_add = (user, host, path, dest_user, dest_host)
             if ignore_dest_user:
-                graph[dest_host].append(line_to_add)
+                graph[dest_host].add(line_to_add)
             else:
                 graph[f"{dest_user}@{dest_host}"].add(line_to_add)
 
@@ -96,7 +99,7 @@ if __name__ == "__main__":
                 user, host, path, dest_user, dest_host = entry
                 print(f"{user}@{host}{path} -> {dest_user}@{dest_host}")
         else:
-            result = indirect_get_connected_nodes(reverse_lookup_table, interesting_host)
+            result = indirect_get_connected_nodes(reverse_lookup_table, interesting_host, ignore_dest_user)
             for entry in result:
                 print(entry)
     else: