heartbeat/hbd/dns.py

"""DNS update helper and pure asyncio worker for heartbeat daemon."""
from __future__ import annotations
import subprocess
from subprocess import Popen, PIPE, STDOUT
from typing import Optional
import asyncio


def create_nsupdate_payload(hostname: str, newip: str, dyndomain: str, dnsttl: str = "5") -> str:
    D = {"domain": dyndomain, "fqdn": f"{hostname}.dy.{dyndomain}", "dnsttl": dnsttl, "newip": newip, "ts": __import__("time").strftime("%Y-%m-%d.%H:%M:%S", __import__("time").gmtime())}
    if ":" in newip:
        nsup = (
            """update delete %(fqdn)s AAAA
update add %(fqdn)s %(dnsttl)s AAAA %(newip)s
update delete %(fqdn)s TXT
update add %(fqdn)s %(dnsttl)s TXT "Created: %(ts)s"
send
answer

""" % D
        )
    else:
        nsup = (
            """update delete %(fqdn)s A
update add %(fqdn)s %(dnsttl)s A %(newip)s
update delete %(fqdn)s TXT
update add %(fqdn)s %(dnsttl)s TXT "Created: %(ts)s"
send
answer

""" % D
        )
    return nsup


def nsupdate(hostname: str, newip: str, dyndomain: str, nsupdate_bin: str = "/usr/local/bin/nsupdate", rndc_key: str = "/etc/dhcpc/rndc-key") -> Optional[str]:
    """Perform DNS update via nsupdate command.

    Returns None on success, else returns combined stdout/stderr as a string.
    """
    nsup = create_nsupdate_payload(hostname, newip, dyndomain)
    cmd = [nsupdate_bin, "-k", rndc_key, "-v"]
    try:
        p = Popen(cmd, shell=False, bufsize=0, stdin=PIPE, stdout=PIPE, stderr=STDOUT)
    except OSError as e:
        return f"nsupdate: execution failed: {e}"
    except Exception as e:
        return f"nsupdate: some error occured: {e}"

    (output, err) = p.communicate(nsup.encode())
    out = output.decode() if output else ""
    if out.find("status: NOERROR") >= 0:
        return None
    return out


async def dns_update_worker(hbdclass, cfg: dict, async_queue=None, log: Optional[callable] = None, pushmsg: Optional[callable] = None, loop: Optional[asyncio.AbstractEventLoop] = None):
    """Pure async DNS worker that processes updates from asyncio.Queue.

    Exits when it receives a None sentinel.
    """
    if loop is None:
        loop = asyncio.get_running_loop()

    dnsq = async_queue
    if not dnsq:
        if log:
            try:
                await loop.run_in_executor(None, log, None, "dns_update_worker: no queue available")
            except Exception:
                pass
        return

    while True:
        try:
            item = await dnsq.get()
        except Exception as e:
            if log:
                try:
                    await loop.run_in_executor(None, log, None, f"dns_update_worker: error getting item: {e}")
                except Exception:
                    pass
            break

        if item is None:
            break

        try:
            name, addr = item
        except Exception:
            try:
                dnsq.task_done()
            except Exception:
                pass
            continue

        m = f"changed address to {addr}"
        for dyndomain in cfg.get("dyndomains", []):
            err = await loop.run_in_executor(None, nsupdate, name, addr, dyndomain, cfg.get("nsupdate_bin", "/usr/local/bin/nsupdate"), cfg.get("rndc_key", "/etc/dhcpc/rndc-key"))
            if err:
                m += f", DNS update failed: {err}"
                if pushmsg:
                    try:
                        await loop.run_in_executor(None, pushmsg, "error: nsupdate failed", f"{name}.dy.{dyndomain}: {m}")
                    except Exception:
                        pass
            else:
                m += ", DNS updated."

        try:
            dnsq.task_done()
        except Exception:
            pass

        if log:
            try:
                await loop.run_in_executor(None, log, name, m)
            except Exception:
                pass

    if log:
        try:
            await loop.run_in_executor(None, log, None, "dns_update_worker exiting")
        except Exception:
            pass


def start_dns_worker(hbdclass, cfg: dict, log: Optional[callable] = None, pushmsg: Optional[callable] = None, loop: Optional[asyncio.AbstractEventLoop] = None):
    """Start the async DNS worker and return the Task.

    Replaces Host.dnsQ with an asyncio.Queue wrapped in a thread-safe bridge
    so legacy synchronous put() calls from UDP handlers still work.
    """
    if loop is None:
        loop = asyncio.get_event_loop()

    # Create asyncio.Queue and wrap in a bridge for thread-safe puts
    async_q = asyncio.Queue()

    class _QueueBridge:
        """Thread-safe wrapper around asyncio.Queue for synchronous callers."""
        def __init__(self, loop, aq):
            self._loop = loop
            self._aq = aq

        def put(self, item):
            """Thread-safe put that schedules onto event loop."""
            try:
                # Try to detect if we're in the event loop thread
                asyncio.get_running_loop()
                # We're in event loop context, use put_nowait directly
                self._aq.put_nowait(item)
            except RuntimeError:
                # We're in a different thread, schedule safely
                try:
                    self._loop.call_soon_threadsafe(self._aq.put_nowait, item)
                except Exception:
                    pass

        def task_done(self):
            """Delegate task_done to asyncio.Queue."""
            try:
                self._aq.task_done()
            except Exception:
                pass

    bridge = _QueueBridge(loop, async_q)
    hbdclass.Host.dnsQ = bridge

    task = loop.create_task(dns_update_worker(hbdclass, cfg, async_queue=async_q, log=log, pushmsg=pushmsg, loop=loop))
    return task