iPhone Presence across VLANs with Home Assistant, Synology, and MQTT

import subprocess
import shlex
import socket
import paramiko
import paho.mqtt.client as mqtt
import time 
import re
import logging
from datetime import datetime 

# Uncomment this for more detailed debugging info 
# logging.basicConfig(level=logging.DEBUG)

# MQTT broker configuration
mqtt_broker = "192.168.1.42"  # Replace with your MQTT broker address
mqtt_port = 1883
mqtt_username = "mqtt_user"
mqtt_password = "my mqtt password"
mqtt_timeout = 2  # Wait this many seconds for the subscription callback. It should be immediate unless fresh installation
pre_ping_topic = "arp_presence/config/pre_ping"
connected_topic = "arp_presence/connected"
updated_topic = "arp_presence/updated"

# SSH server configuration
ssh_server_ip = "192.168.1.1"  # Replace with the router ssh serer IP address
ssh_username = "router_user"
ssh_password = "my router password"

# Pre-ping configuration-- who gets pinged and at what ports 
pre_ping_addresses = []
pre_ping_port = 5353

# Function to connect via SSH and retrieve connected IP addresses from the server ARP table 
def retrieve_connected_ips(mqtt_client):
    try:
        # Connect to the SSH server
        logging.info(f"Connecting to SSH server...")
        ssh_client = paramiko.SSHClient()
        ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
        ssh_client.connect(ssh_server_ip, username=ssh_username, password=ssh_password)

        # Fetch the ARP table
        logging.info(f"Fetching the ARP table...")
        stdin, stdout, stderr = ssh_client.exec_command("arp -a")
        arp_table = stdout.read()
        logging.debug(f"ARP table: {arp_table}")

        # Regex to capture both the IP and the MAC address. It's only connected if both are there. 
        mac_pattern = re.compile(r"\((?P<ip>(?:[0-9]{1,3}\.){3}[0-9]{1,3})\) at (?P<mac>[0-9a-fA-F]{2}(?:[:-][0-9a-fA-F]{2}){5})")
        connected_ips = [match.group('ip') for match in mac_pattern.finditer(arp_table.decode())]

        logging.debug(f"Connected IP addresses:")
        for ip in connected_ips:
            logging.debug(ip)

        if mqtt_client.is_connected():
            logging.info(f"Publishing connected IP addresses to MQTT...")
            mqtt_client.publish(f"{connected_topic}", ' '.join(connected_ips), retain=True)
            mqtt_client.publish(f"{updated_topic}", datetime.now().strftime("%Y-%m-%d %H:%M:%S"), retain=True)

        else:
            logging.warning(f"MQTT client not connected. Unable to publish IP addresses.")

    except Exception as e:
        logging.error(f"An unexpected error occurred: {e}")

    finally:
        # Close the SSH connection
        ssh_client.close()
        logging.info(f"SSH connection closed.")

# Callback function for when subscribed messages are received 
def on_message(client, userdata, msg):
    # The only configuration topic we care about 
    if ( msg.topic == pre_ping_topic ): 
        global pre_ping_addresses
        pre_ping_addresses = msg.payload.decode().split()
        logging.debug(f"Received IP addresses to pre-ping: {pre_ping_addresses}")
    else:
        logging.debug(f"Received an unexpected and unsupported toopic: {msg.topic}")

# Create MQTT client and set the message callback 
mqtt_client = mqtt.Client()
mqtt_client.on_log = lambda client, userdata, level, buf: logging.debug(f"MQTT Log: {buf}")
mqtt_client.on_message = on_message

# Connect to MQTT broker and subscribe to the appropriate topics
mqtt_client.username_pw_set(username=mqtt_username, password=mqtt_password)
mqtt_client.connect(mqtt_broker, mqtt_port, 60)
mqtt_client.subscribe(pre_ping_topic)

# Start the MQTT loop to listen for 10 seconds or until both the pre-ping and the tracked addresses are received 
mqtt_client.loop_start()
start_time = time.time()

try:
    while not pre_ping_addresses and time.time() - start_time < mqtt_timeout:
        pass

    # Only do the pinging if there are pre-ping addresses to ping 		
    if not pre_ping_addresses:
        logging.warning(f"Nothing to ping. The topic '{pre_ping_topic}' is missing or never responded.")
    else:
        # Send UDP messages to pre-ping IP addresses. iPhones should show up in the ARP table after this. 
        logging.info(f"Sending UDP messages...")
        for ip in pre_ping_addresses:
            with socket.socket(socket.AF_INET, socket.SOCK_DGRAM) as udp_socket:
                udp_socket.settimeout(1)
                udp_socket.sendto(b"ping", (ip, pre_ping_port ))
                #logging.debug(f"Sent 'ping' message to {ip}")

    # Retrieve connected MAC addresses over SSH
    retrieve_connected_ips(mqtt_client)

finally:
    # Disconnect from MQTT broker
    mqtt_client.disconnect()
    logging.info(f"Disconnected from MQTT broker.")