import os
import subprocess
import sys
import re
import socket
import time
# --- Configuration ---
HOSTS_FILE = "/etc/hosts"
# Define the set of ALL known machines that need mDNS resolution in your network
# The script will try to find and manage entries for all of these *except* the current machine.
# Keep these lowercase as we will convert discovered hostnames to lowercase for comparison.
ALL_NETWORK_MACHINES = ["hplaptop", "babyhp"] # Add any other fixed machines here (e.g., "myfileserver")
# Regex patterns for parsing avahi-browse -atr output
HOSTNAME_PATTERN = re.compile(r'^\s*hostname = \[(.*?\.local)\]$')
ADDRESS_PATTERN = re.compile(r'^\s*address = \[(\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3})\]$') # Only captures IPv4
# --- Functions ---
def get_mdns_discovered_hosts():
"""
Uses avahi-browse -atr to discover .local hostnames and their IPs.
Returns a dictionary: { "hostname.local": "IP_Address" }
"""
discovered_hosts = {}
print("--- DEBUG: Starting mDNS discovery using avahi-browse -atr...")
sys.stdout.flush()
try:
cmd = ['avahi-browse', '-atr']
result = subprocess.run(
cmd,
capture_output=True,
text=True,
check=True,
timeout=10
)
print("--- DEBUG: avahi-browse command stdout ---")
print(result.stdout)
print("--- DEBUG: avahi-browse command stderr ---")
print(result.stderr)
print("------------------------------------------")
sys.stdout.flush()
current_hostname_full = None
for line in result.stdout.splitlines():
hostname_match = HOSTNAME_PATTERN.match(line)
if hostname_match:
current_hostname_full = hostname_match.group(1)
continue
address_match = ADDRESS_PATTERN.match(line)
if address_match:
ip_address = address_match.group(1)
if current_hostname_full:
discovered_hosts[current_hostname_full] = ip_address
print(f"--- DEBUG: Found mDNS entry: {current_hostname_full} -> {ip_address}")
sys.stdout.flush()
current_hostname_full = None
print("--- DEBUG: mDNS discovery completed.")
sys.stdout.flush()
except subprocess.CalledProcessError as e:
print(f"--- ERROR: avahi-browse failed with exit code {e.returncode} ---", file=sys.stderr)
print(f"STDOUT: {e.stdout}", file=sys.stderr)
print(f"STDERR: {e.stderr}", file=sys.stderr)
print("Please ensure avahi-browse is installed and avahi-daemon is running on advertising hosts.", file=sys.stderr)
sys.stderr.flush()
except subprocess.TimeoutExpired:
print("--- ERROR: avahi-browse timed out after 10 seconds. It might not be running or discovery is very slow.", file=sys.stderr)
sys.stderr.flush()
except FileNotFoundError:
print("--- ERROR: avahi-browse command not found. Is avahi-utils installed?", file=sys.stderr)
sys.stderr.flush()
except Exception as e:
print(f"--- ERROR: An unexpected error occurred while Browse mDNS: {e}", file=sys.stderr)
sys.stderr.flush()
return discovered_hosts
def get_own_lan_ip():
"""
Retrieves the current machine's primary IPv4 address on the LAN.
"""
print("--- DEBUG: Getting own LAN IP...")
sys.stdout.flush()
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
try:
s.connect(('8.8.8.8', 1))
IP = s.getsockname()[0]
print(f"--- DEBUG: Own LAN IP found: {IP}")
sys.stdout.flush()
except Exception as e:
IP = '127.0.0.1'
print(f"--- DEBUG: Could not determine own LAN IP, falling back to {IP}. Error: {e}")
sys.stdout.flush()
finally:
s.close()
return IP
def generate_new_hosts_content(current_machine_hostname, own_lan_ip, discovered_hosts_data):
"""
Generates the desired content for the new /etc/hosts file.
discovered_hosts_data is expected to be a dict of { "short_hostname": "IP_Address" }
"""
print("--- DEBUG: Generating new /etc/hosts content...")
sys.stdout.flush()
hosts_content = [
"127.0.0.1\tlocalhost",
"127.0.1.1\t" + current_machine_hostname,
f"{own_lan_ip}\t{current_machine_hostname}",
]
added_hostnames = {current_machine_hostname} # Keep track of already added hostnames (e.g., current machine)
for short_name, ip_address in discovered_hosts_data.items():
# Ensure we don't add the current machine's own entry twice if discovered via mDNS
if short_name.lower() != current_machine_hostname.lower() and short_name not in added_hostnames: # Added .lower() for robustness
hosts_content.append(f"{ip_address}\t{short_name}")
added_hostnames.add(short_name)
print(f"--- DEBUG: Added/Updated '{short_name}' to content with IP {ip_address}")
sys.stdout.flush()
print("--- DEBUG: Finished generating content.")
sys.stdout.flush()
return "\n".join(hosts_content) + "\n"
def recreate_hosts_file(content, hosts_file_path):
"""Deletes existing /etc/hosts and writes new content."""
try:
print(f"--- DEBUG: Deleting existing {hosts_file_path}...")
sys.stdout.flush()
rm_result = subprocess.run(['sudo', 'rm', '-f', hosts_file_path], capture_output=True, text=True, check=True)
print(f"--- DEBUG: rm stdout: {rm_result.stdout.strip()}")
print(f"--- DEBUG: rm stderr: {rm_result.stderr.strip()}")
print(f"Successfully deleted {hosts_file_path}.")
sys.stdout.flush()
print(f"--- DEBUG: Recreating {hosts_file_path} with new content...")
sys.stdout.flush()
tee_process = subprocess.run(
['sudo', 'tee', hosts_file_path],
input=content.encode('utf-8'),
capture_output=True,
check=True
)
print(f"--- DEBUG: tee stdout: {tee_process.stdout.strip()}")
print(f"--- DEBUG: tee stderr: {tee_process.stderr.strip()}")
print(f"Successfully recreated {hosts_file_path}.")
sys.stdout.flush()
except subprocess.CalledProcessError as e:
print(f"--- ERROR: Recreating {hosts_file_path} failed with exit code {e.returncode} ---", file=sys.stderr)
print(f"STDOUT: {e.stdout}", file=sys.stderr)
print(f"STDERR: {e.stderr}", file=sys.stderr)
sys.stderr.flush()
sys.exit(1)
except Exception as e:
print(f"--- ERROR: An unexpected error occurred during hosts file recreation: {e}", file=sys.stderr)
sys.stderr.flush()
sys.exit(1)
# --- Main Execution ---
if __name__ == "__main__":
print("--- SCRIPT EXECUTION START ---")
sys.stdout.flush()
if os.geteuid() != 0:
print("This script needs to be run with sudo to modify /etc/hosts.", file=sys.stderr)
sys.stderr.flush()
sys.exit(1)
print("Starting mDNS-based /etc/hosts recreation for IPv4 hosts...")
sys.stdout.flush()
# Dynamically determine the current hostname of THIS machine (and convert to lowercase for comparison)
CURRENT_MACHINE_HOSTNAME = socket.gethostname().split('.')[0]
print(f"--- DEBUG: Current machine's short hostname: {CURRENT_MACHINE_HOSTNAME}")
sys.stdout.flush()
print("--- DEBUG: Checking for avahi-utils installation...")
sys.stdout.flush()
try:
subprocess.run(['dpkg', '-s', 'avahi-utils'], capture_output=True, check=True, text=True)
print("--- DEBUG: avahi-utils found.")
sys.stdout.flush()
except subprocess.CalledProcessError:
print("--- ERROR: avahi-utils not found. Please install it: sudo apt install avahi-utils", file=sys.stderr)
sys.stderr.flush()
sys.exit(1)
# Get this machine's own LAN IP
own_ip = get_own_lan_ip()
print(f"This machine ({CURRENT_MACHINE_HOSTNAME})'s LAN IP: {own_ip}")
sys.stdout.flush()
# Get all discovered hosts from mDNS
discovered_mdns_hosts_full = get_mdns_discovered_hosts()
# Process discovered hosts into a map of { "short_hostname": "ip_address" }
filtered_and_shortened_discovered_hosts = {}
print("--- DEBUG: Processing mDNS discovered hosts for target machines...")
sys.stdout.flush()
# Get current machine hostname in lowercase for robust comparison
current_machine_short_lower = CURRENT_MACHINE_HOSTNAME.lower()
for full_name_mdns, ip in discovered_mdns_hosts_full.items():
# Convert the discovered hostname to lowercase for robust comparison
short_name_lower = full_name_mdns.replace(".local", "").lower()
# Check if this discovered host is one of our target machines (case-insensitively)
# AND if it's not the current machine itself.
if short_name_lower in ALL_NETWORK_MACHINES and short_name_lower != current_machine_short_lower:
# We add the original full_name (e.g., BabyHP.local) to the filtered list's keys
# and the short_name (e.g., babyhp) as the key for the hosts content generation.
# This requires a slight change in how filtered_and_shortened_discovered_hosts is used.
# Let's simplify: store the original short_name (case as seen from Avahi) if desired,
# but use its lowercase form for comparison.
# The generate_new_hosts_content function expects short_hostname as key
# so let's use the original case as given by Avahi for the value in the hosts file.
# So, we need to save the original short_name from full_name_mdns, not the lowercase one.
original_short_name = full_name_mdns.replace(".local", "") # Keep original case for the hosts file entry
filtered_and_shortened_discovered_hosts[original_short_name] = ip
print(f" ADDING TARGET host to hosts file list: '{original_short_name}' at {ip}")
sys.stdout.flush()
else:
print(f" Ignoring non-target/self mDNS host: '{short_name_lower}' at {ip}") # Use lowercase for debug output here for consistency
sys.stdout.flush()
# Generate the new content for /etc/hosts
new_hosts_content = generate_new_hosts_content(
CURRENT_MACHINE_HOSTNAME,
own_ip,
filtered_and_shortened_discovered_hosts
)
# Recreate the /etc/hosts file with the generated content
recreate_hosts_file(new_hosts_content, HOSTS_FILE)
print("--- SCRIPT EXECUTION FINISHED ---")
sys.stdout.flush()