forked from extern/SSH-Snake
Fix issue of reverse-lookup-host.py and forward-lookup-host.py not obeying ignore_dest_user.
This commit is contained in:
parent
ff698f308d
commit
a40e8f82be
@ -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:
|
||||
|
@ -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:
|
||||
|
Loading…
Reference in New Issue
Block a user