diff --git a/tools/forward-lookup-host.py b/tools/forward-lookup-host.py index 66ea963..a076881 100755 --- a/tools/forward-lookup-host.py +++ b/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): - backward_connected_nodes = set() +def indirect_get_connected_nodes(graph, interesting_host, ignore_dest_user): + connected_nodes = set() sentinel = '__SENTINEL__' - backward_heap = [(interesting_host, sentinel)] - while backward_heap: - current_node, parent_node = heapq.heappop(backward_heap) + heap = [(interesting_host, sentinel)] + while heap: + current_node, parent_node = heapq.heappop(heap) - if current_node not in backward_connected_nodes: - backward_connected_nodes.add(current_node) + if current_node not in connected_nodes: + 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 - heapq.heappush(backward_heap, (node, current_node)) + if ignore_dest_user: + 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: diff --git a/tools/reverse-lookup-host.py b/tools/reverse-lookup-host.py index cb6f2b4..d4702a1 100755 --- a/tools/reverse-lookup-host.py +++ b/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: