diff --git a/kea_listener.py b/kea_listener.py index f6d699e..a22e1e8 100644 --- a/kea_listener.py +++ b/kea_listener.py @@ -38,6 +38,31 @@ def on_connect(client, userdata, flags, rc, properties=None): client.subscribe(userdata.mqtt_topic) +def extract_lease_ip_info(msg_dict) -> dict[str, str | None]: + ip_address = None + hostname = None + mac_address = None + + # "QUERY4_TYPE": "DHCPREQUEST", + + if msg_dict["HOOK_TYPE"] in ["RENEW"] and msg_dict["QUERY4_TYPE"] == "DHCPREQUEST": + ip_address = msg_dict.get("LEASE4_ADDRESS") + hostname = msg_dict.get("LEASE4_HOSTNAME") + mac_address = msg_dict.get("LEASE4_HWADDR") + elif msg_dict["HOOK_TYPE"] in ["RELEASE", "EXPIRE", "RECOVER"]: + ip_address = msg_dict.get("LEASE4_ADDRESS") + + if not ip_address and not hostname and not mac_address: + return None + + return { + "HOOK_TYPE": msg_dict["HOOK_TYPE"], + "ip_address": ip_address, + "hostname": hostname, + "mac_address": mac_address, + } + + def on_message(client, userdata, msg): msg_dict = json.loads(msg.payload.decode()) msg_dict["id"] = str(uuid.uuid1(node=uuid.getnode())) @@ -45,48 +70,38 @@ def on_message(client, userdata, msg): logging.debug( "HOOK: [" + msg_dict["id"] + "] " + msg.topic + " " + msg.payload.decode() ) - if msg_dict["HOOK_TYPE"] == "COMMIT": - ip_address = msg_dict.get("LEASES4_AT0_ADDRESS") - hostname = msg_dict.get("LEASES4_AT0_HOSTNAME") - mac_address = msg_dict.get("LEASES4_AT0_HWADDR") - if not ip_address or not hostname or not mac_address: - logging.error( - "Missing required fields in COMMIT message: " - + str(msg_dict) - + " skipping processing" - ) - return + + lease_ip_info = extract_lease_ip_info(msg_dict) + if lease_ip_info and lease_ip_info["HOOK_TYPE"] in ["RENEW"]: if logging.getLogger().isEnabledFor(logging.DEBUG): logging.debug( "Processing COMMIT for IP: " - + ip_address + + lease_ip_info["ip_address"] + " Hostname: " - + hostname + + lease_ip_info["hostname"] + " MAC: " - + mac_address + + lease_ip_info["mac_address"] ) try: ip_info = userdata.ipam_util.update_ip_information( - ip_address=ip_address, hostname=hostname, mac_address=mac_address + ip_address=lease_ip_info["ip_address"], + hostname=lease_ip_info["hostname"], + mac_address=lease_ip_info["mac_address"], ) logging.info("Updated IP Info: " + str(ip_info)) userdata.increment_count() except Exception as e: logging.error("Error updating IP information: " + str(e)) - if msg_dict["HOOK_TYPE"] == "RELEASE": - ip_address = msg_dict.get("LEASE4_ADDRESS") - if not ip_address: - logging.error( - "Missing required field 'LEASE4_ADDRESS' in RELEASE message: " - + str(msg_dict) - + " skipping processing" - ) - return - + if lease_ip_info and lease_ip_info["HOOK_TYPE"] in ["RELEASE", "EXPIRE", "RECOVER"]: try: - userdata.ipam_util.clear_ip_information(ip_address=ip_address) + userdata.ipam_util.clear_ip_information( + ip_address=lease_ip_info["ip_address"] + ) userdata.increment_count() - logging.info("Processed RELEASE for IP: " + ip_address) + logging.info( + "Processed RELEASE/EXPIRE/RECOVER for IP: " + + lease_ip_info["ip_address"] + ) except Exception as e: logging.error("Error clearing IP information: " + str(e))