refactored code
This commit is contained in:
parent
73429771d4
commit
414266eefc
6 changed files with 315 additions and 404 deletions
155
remove_host.py
155
remove_host.py
|
|
@ -1,82 +1,68 @@
|
|||
# ssh_manager/remove_host.py
|
||||
|
||||
import os
|
||||
import glob
|
||||
import subprocess
|
||||
import asyncio
|
||||
from collections import OrderedDict
|
||||
from .utils import (
|
||||
print_info,
|
||||
print_warning,
|
||||
print_error,
|
||||
safe_input
|
||||
)
|
||||
from .list_hosts import load_config_file, check_ssh_port
|
||||
|
||||
async def get_all_host_blocks(conf_dir):
|
||||
pattern = os.path.join(conf_dir, "*", "config")
|
||||
conf_files = sorted(glob.glob(pattern))
|
||||
|
||||
all_blocks = []
|
||||
for conf_file in conf_files:
|
||||
blocks = load_config_file(conf_file)
|
||||
all_blocks.extend(blocks)
|
||||
|
||||
return all_blocks
|
||||
from .list_hosts import build_host_list_table, load_config_file
|
||||
"""
|
||||
Remove host now reuses build_host_list_table to display the same columns:
|
||||
No. | Host | User | Port | HostName | IP Address | Conf Directory
|
||||
"""
|
||||
|
||||
async def remove_host(conf_dir):
|
||||
"""
|
||||
Remove an SSH host by:
|
||||
1) Listing all hosts
|
||||
1) Listing all hosts (same columns as main list)
|
||||
2) Letting user pick row number or label
|
||||
3) Attempting to remove the old pub key from remote authorized_keys
|
||||
4) Deleting the subdirectory in ~/.ssh/conf/<host_label>
|
||||
3) Removing old pub key from remote
|
||||
4) Deleting ~/.ssh/conf/<host_label>
|
||||
"""
|
||||
|
||||
print_info("Remove Host - Step 1: Show current hosts...\n")
|
||||
|
||||
all_blocks = await get_all_host_blocks(conf_dir)
|
||||
if not all_blocks:
|
||||
# Reuse the unified table from list_hosts
|
||||
headers, final_data = await build_host_list_table(conf_dir)
|
||||
|
||||
if not final_data:
|
||||
print_warning("No hosts found. Cannot remove anything.")
|
||||
return
|
||||
|
||||
# We'll display them in a table
|
||||
import socket
|
||||
# Print the same table
|
||||
from tabulate import tabulate
|
||||
print("\nSSH Conf Subdirectory Host List (Sorted by IP Ascending)")
|
||||
print(tabulate(final_data, headers=headers, tablefmt="grid"))
|
||||
|
||||
table_rows = []
|
||||
for idx, block in enumerate(all_blocks, start=1):
|
||||
host_label = block.get("Host", "N/A")
|
||||
hostname = block.get("HostName", "N/A")
|
||||
user = block.get("User", "N/A")
|
||||
port = int(block.get("Port", "22"))
|
||||
identity_file = block.get("IdentityFile", "N/A")
|
||||
# We have final_data rows => need to map row => block
|
||||
# So let's gather the raw blocks again to correlate.
|
||||
# We'll do a separate approach or we can parse final_data.
|
||||
# Easiest: Re-run load_config_file if needed or:
|
||||
blocks = []
|
||||
# The last gather call for build_host_list_table used load_config_file already
|
||||
# but it doesn't return the correlation. We'll replicate the logic quickly.
|
||||
|
||||
try:
|
||||
ip_address = socket.gethostbyname(hostname)
|
||||
port_open = await asyncio.wait_for(check_ssh_port(ip_address, port), timeout=1)
|
||||
except:
|
||||
ip_address = None
|
||||
port_open = False
|
||||
pattern = os.path.join(conf_dir, "*", "config")
|
||||
conf_files = sorted(os.listdir(conf_dir))
|
||||
|
||||
ip_disp = f"\033[0;32m{ip_address}\033[0m" if ip_address else "\033[0;31mN/A\033[0m"
|
||||
port_disp = f"\033[0;32m{port}\033[0m" if port_open else f"\033[0;31m{port}\033[0m"
|
||||
# Actually, let's do a small approach:
|
||||
from .list_hosts import gather_host_info, sort_by_ip
|
||||
all_blocks = []
|
||||
import glob
|
||||
|
||||
row = [
|
||||
idx,
|
||||
host_label,
|
||||
user,
|
||||
port_disp,
|
||||
hostname,
|
||||
ip_disp,
|
||||
identity_file
|
||||
]
|
||||
table_rows.append(row)
|
||||
for cfile in glob.glob(os.path.join(conf_dir, "*", "config")):
|
||||
blocks.extend(load_config_file(cfile))
|
||||
|
||||
headers = ["No.", "Host", "User", "Port", "HostName", "IP", "IdentityFile"]
|
||||
print("\nSSH Conf Subdirectory Host List")
|
||||
print(tabulate(table_rows, headers=headers, tablefmt="grid"))
|
||||
# gather the same big list
|
||||
results = await gather_host_info(blocks)
|
||||
sorted_rows = sort_by_ip(results)
|
||||
# sorted_rows is a list of 7-tuples:
|
||||
# (host_label, user, colored_port, hostname, colored_ip, conf_display, raw_ip)
|
||||
|
||||
# Prompt which host to remove
|
||||
choice = safe_input("Enter the row number or Host label to remove: ")
|
||||
if choice is None:
|
||||
return
|
||||
|
|
@ -85,43 +71,60 @@ async def remove_host(conf_dir):
|
|||
print_error("Invalid empty choice.")
|
||||
return
|
||||
|
||||
target_block = None
|
||||
target_tuple = None
|
||||
|
||||
# If digit => index in sorted_rows
|
||||
if choice.isdigit():
|
||||
idx = int(choice)
|
||||
if idx < 1 or idx > len(all_blocks):
|
||||
if idx < 1 or idx > len(sorted_rows):
|
||||
print_warning(f"Row number {idx} is invalid.")
|
||||
return
|
||||
target_block = all_blocks[idx - 1]
|
||||
target_tuple = sorted_rows[idx - 1]
|
||||
else:
|
||||
# They typed a label
|
||||
for b in all_blocks:
|
||||
if b.get("Host") == choice:
|
||||
target_block = b
|
||||
for t in sorted_rows:
|
||||
if t[0] == choice: # t[0] = host_label
|
||||
target_tuple = t
|
||||
break
|
||||
if not target_block:
|
||||
if not target_tuple:
|
||||
print_warning(f"No matching host label '{choice}' found.")
|
||||
return
|
||||
|
||||
host_label = target_block.get("Host", "")
|
||||
hostname = target_block.get("HostName", "")
|
||||
user = target_block.get("User", "root")
|
||||
port = int(target_block.get("Port", "22"))
|
||||
identity_file = target_block.get("IdentityFile", "")
|
||||
# target_tuple is (host_label, user, colored_port, hostname, colored_ip, conf_dir, raw_ip)
|
||||
host_label = target_tuple[0]
|
||||
hostname = target_tuple[3]
|
||||
user = target_tuple[1]
|
||||
# parse port from colored_port? We can do a quick approach. Alternatively, we re-run the load_config approach again.
|
||||
# We do a second approach: let's see if we can find the block in "blocks".
|
||||
# But let's parse the config block to get the real port, identity, etc.
|
||||
|
||||
# We find the matching block in "blocks" by host_label
|
||||
found_block = None
|
||||
for b in blocks:
|
||||
if b.get("Host") == host_label:
|
||||
found_block = b
|
||||
break
|
||||
|
||||
if not found_block:
|
||||
print_warning(f"No config block found for {host_label}.")
|
||||
return
|
||||
|
||||
port_str = found_block.get("Port", "22")
|
||||
port = int(port_str)
|
||||
identity_file = found_block.get("IdentityFile", "")
|
||||
|
||||
# now do removal logic
|
||||
if not host_label:
|
||||
print_warning("Target block has no Host label. Cannot remove.")
|
||||
return
|
||||
|
||||
# Check IdentityFile for old pub key
|
||||
if identity_file and identity_file != "N/A":
|
||||
expanded_key = os.path.expanduser(identity_file)
|
||||
pub_path = expanded_key + ".pub"
|
||||
pub_path = expanded_key + ".pub"
|
||||
old_pub_data = ""
|
||||
|
||||
if os.path.isfile(pub_path):
|
||||
try:
|
||||
with open(pub_path, "r") as f:
|
||||
# This is the EXACT line that might appear in authorized_keys
|
||||
old_pub_data = f.read().rstrip("\n")
|
||||
except Exception as e:
|
||||
print_warning(f"Could not read old pub key: {e}")
|
||||
|
|
@ -130,13 +133,13 @@ async def remove_host(conf_dir):
|
|||
print_info("Attempting to remove old key from remote authorized_keys...")
|
||||
await remove_old_key_remote(old_pub_data, user, hostname, port)
|
||||
|
||||
# Finally, remove the local subdirectory
|
||||
# remove local folder
|
||||
host_dir = os.path.join(conf_dir, host_label)
|
||||
import shutil
|
||||
if os.path.isdir(host_dir):
|
||||
confirm = safe_input(f"Are you sure you want to delete local folder {host_dir}? (y/n): ")
|
||||
if confirm and confirm.lower().startswith('y'):
|
||||
try:
|
||||
import shutil
|
||||
shutil.rmtree(host_dir)
|
||||
print_info(f"Removed local folder: {host_dir}")
|
||||
except Exception as e:
|
||||
|
|
@ -144,24 +147,17 @@ async def remove_host(conf_dir):
|
|||
else:
|
||||
print_warning("Local folder not removed.")
|
||||
else:
|
||||
print_warning(f"Local host folder {host_dir} not found. Nothing to remove.")
|
||||
print_warning(f"Local folder {host_dir} not found. Nothing to remove.")
|
||||
|
||||
async def remove_old_key_remote(old_pub_data, user, hostname, port):
|
||||
"""
|
||||
Remove the old key from remote authorized_keys if found, matching EXACT lines.
|
||||
We'll do:
|
||||
grep -Fxq for existence
|
||||
grep -vFx to remove it
|
||||
"""
|
||||
# 1) Check if old_pub_data is present in authorized_keys (exact line)
|
||||
# same logic as before
|
||||
check_cmd = [
|
||||
"ssh", "-o", "StrictHostKeyChecking=no",
|
||||
"-p", str(port),
|
||||
f"{user}@{hostname}",
|
||||
f"grep -Fxq \"{old_pub_data}\" ~/.ssh/authorized_keys"
|
||||
f'grep -Fxq "{old_pub_data}" ~/.ssh/authorized_keys'
|
||||
]
|
||||
found_key = False
|
||||
|
||||
try:
|
||||
subprocess.check_call(check_cmd)
|
||||
found_key = True
|
||||
|
|
@ -172,12 +168,11 @@ async def remove_old_key_remote(old_pub_data, user, hostname, port):
|
|||
print_warning("No old key found on remote authorized_keys.")
|
||||
return
|
||||
|
||||
# 2) Actually remove it by ignoring EXACT matches to that line
|
||||
remove_cmd = [
|
||||
"ssh", "-o", "StrictHostKeyChecking=no",
|
||||
"-p", str(port),
|
||||
f"{user}@{hostname}",
|
||||
f"grep -vFx \"{old_pub_data}\" ~/.ssh/authorized_keys > ~/.ssh/tmp && mv ~/.ssh/tmp ~/.ssh/authorized_keys"
|
||||
f'grep -vFx "{old_pub_data}" ~/.ssh/authorized_keys > ~/.ssh/tmp && mv ~/.ssh/tmp ~/.ssh/authorized_keys'
|
||||
]
|
||||
try:
|
||||
subprocess.check_call(remove_cmd)
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue