added additional functions

This commit is contained in:
Arctic 2025-03-07 06:58:49 -06:00
parent 92ffaa431b
commit 73429771d4

View file

@ -4,14 +4,11 @@ import os
import glob import glob
import socket import socket
import asyncio import asyncio
import ipaddress
from tabulate import tabulate from tabulate import tabulate
from collections import OrderedDict from collections import OrderedDict
from .utils import print_warning, print_error, Colors from .utils import print_warning, print_error, Colors
"""
This file is responsible for listing all SSH hosts discovered in ~/.ssh/conf/<label>/config.
We've modified it to show a "Conf Directory" column instead of "IdentityFile."
"""
async def check_ssh_port(ip_address, port): async def check_ssh_port(ip_address, port):
""" """
@ -45,9 +42,11 @@ def load_config_file(file_path):
for line in lines: for line in lines:
stripped_line = line.strip() stripped_line = line.strip()
# Skip empty lines and comments
if not stripped_line or stripped_line.startswith('#'): if not stripped_line or stripped_line.startswith('#'):
continue continue
# Start of a new Host block
if stripped_line.lower().startswith('host '): if stripped_line.lower().startswith('host '):
host_labels = stripped_line.split()[1:] host_labels = stripped_line.split()[1:]
for label in host_labels: for label in host_labels:
@ -68,13 +67,14 @@ def load_config_file(file_path):
async def check_host(host): async def check_host(host):
""" """
Given a host block, resolve IP, check SSH port, etc. Given a host block, resolve IP, check SSH port, etc.
Returns a row for tabulate containing: Returns a tuple of:
1) Host label 1) Host label
2) User 2) User
3) Port (color-coded if open) 3) Port (colored if open)
4) HostName 4) HostName
5) IP Address (color-coded if resolved) 5) IP Address (colored if resolved)
6) Conf Directory (green if has IdentityFile, else no color) 6) Conf Directory (green if has IdentityFile, else no color)
7) raw_ip (uncolored string for sorting)
""" """
host_label = host.get('Host', 'N/A') host_label = host.get('Host', 'N/A')
hostname = host.get('HostName', 'N/A') hostname = host.get('HostName', 'N/A')
@ -84,44 +84,44 @@ async def check_host(host):
# Resolve IP # Resolve IP
try: try:
ip_address = socket.gethostbyname(hostname) raw_ip = socket.gethostbyname(hostname) # uncolored
colored_ip = f"{Colors.GREEN}{ip_address}{Colors.RESET}" colored_ip = f"{Colors.GREEN}{raw_ip}{Colors.RESET}"
except socket.error: except socket.error:
ip_address = None raw_ip = "N/A"
colored_ip = f"{Colors.RED}N/A{Colors.RESET}" colored_ip = f"{Colors.RED}N/A{Colors.RESET}"
# Check port # Check port
if ip_address: if raw_ip != "N/A":
port_open = await check_ssh_port(ip_address, port) port_open = await check_ssh_port(raw_ip, port)
colored_port = ( colored_port = (
f"{Colors.GREEN}{port}{Colors.RESET}" f"{Colors.GREEN}{port}{Colors.RESET}" if port_open else f"{Colors.RED}{port}{Colors.RESET}"
if port_open
else f"{Colors.RED}{port}{Colors.RESET}"
) )
else: else:
colored_port = f"{Colors.RED}{port}{Colors.RESET}" colored_port = f"{Colors.RED}{port}{Colors.RESET}"
# Conf Directory = ~/.ssh/conf/<host_label> # Conf Directory = ~/.ssh/conf/<host_label>
conf_path = f"~/.ssh/conf/{host_label}" conf_path = f"~/.ssh/conf/{host_label}"
# If there's an IdentityFile, we color this path green # If there's an IdentityFile, color the conf path green
if identity_file != 'N/A': if identity_file != 'N/A':
conf_path_display = f"{Colors.GREEN}{conf_path}{Colors.RESET}" conf_path_display = f"{Colors.GREEN}{conf_path}{Colors.RESET}"
else: else:
conf_path_display = conf_path conf_path_display = conf_path
return [ # Return the data plus the uncolored IP for sorting
return (
host_label, host_label,
user, user,
colored_port, colored_port,
hostname, hostname,
colored_ip, colored_ip,
conf_path_display conf_path_display,
] raw_ip # for sorting
)
async def list_hosts(conf_dir): async def list_hosts(conf_dir):
""" """
List out all hosts found in ~/.ssh/conf/*/config. List out all hosts found in ~/.ssh/conf/*/config, sorted by IP in ascending order.
Shows columns: No., Host, User, Port, HostName, IP Address, Conf Directory Columns: No., Host, User, Port, HostName, IP Address, Conf Directory
""" """
pattern = os.path.join(conf_dir, "*", "config") pattern = os.path.join(conf_dir, "*", "config")
conf_files = sorted(glob.glob(pattern)) conf_files = sorted(glob.glob(pattern))
@ -131,7 +131,6 @@ async def list_hosts(conf_dir):
blocks = load_config_file(conf_file) blocks = load_config_file(conf_file)
all_host_blocks.extend(blocks) all_host_blocks.extend(blocks)
# Prepare table
headers = ["No.", "Host", "User", "Port", "HostName", "IP Address", "Conf Directory"] headers = ["No.", "Host", "User", "Port", "HostName", "IP Address", "Conf Directory"]
if not all_host_blocks: if not all_host_blocks:
print_warning("No hosts found. The server list is empty.") print_warning("No hosts found. The server list is empty.")
@ -139,9 +138,37 @@ async def list_hosts(conf_dir):
print(tabulate([], headers=headers, tablefmt="grid")) print(tabulate([], headers=headers, tablefmt="grid"))
return return
# Gather full data for each host
tasks = [check_host(h) for h in all_host_blocks] tasks = [check_host(h) for h in all_host_blocks]
results = await asyncio.gather(*tasks) results = await asyncio.gather(*tasks)
table = [[idx + 1] + row for idx, row in enumerate(results)] # We want to sort by IP ascending. results[i] is a tuple:
print("\nSSH Conf Subdirectory Host List") # (host_label, user, colored_port, hostname, colored_ip, conf_path, raw_ip)
print(tabulate(table, headers=headers, tablefmt="grid")) # We'll parse raw_ip as an ipaddress for sorting. "N/A" => sort to the end.
def parse_ip(ip_str):
try:
return ipaddress.ip_address(ip_str)
except ValueError:
return None
# Convert the results into a list of (ip_obj, original_tuple)
# so we can sort, then rebuild the final data.
sortable = []
for row in results:
raw_ip = row[-1] # last element
ip_obj = parse_ip(raw_ip)
# We'll sort None last by using a sort key that puts (True) after (False)
# e.g. (ip_obj is None, ip_obj)
sortable.append(((ip_obj is None, ip_obj), row))
# Sort by (is_none, ip_obj)
sortable.sort(key=lambda x: x[0])
# Rebuild the final display table, ignoring the raw_ip at the end
final_data = []
for idx, (_, row) in enumerate(sortable, start=1):
# row is (host_label, user, colored_port, hostname, colored_ip, conf_path, raw_ip)
final_data.append([idx] + list(row[:-1])) # omit raw_ip
print("\nSSH Conf Subdirectory Host List (Sorted by IP Ascending)")
print(tabulate(final_data, headers=headers, tablefmt="grid"))