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

This commit is contained in:
Joshua Rogers 2024-01-05 20:30:58 +07:00
parent ff698f308d
commit a40e8f82be
2 changed files with 22 additions and 16 deletions

View File

@ -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:

View File

@ -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: