yunohost/bin/yunomdns

#!/usr/bin/env python3
#
# Copyright (c) 2024 YunoHost Contributors
#
# This file is part of YunoHost (see https://yunohost.org)
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
#

"""
Pythonic declaration of mDNS .local domains for YunoHost
"""

import sys
from ipaddress import ip_address
from time import sleep
from typing import Dict, List

import ifaddr
import yaml
from zeroconf import ServiceBrowser, ServiceInfo, Zeroconf


def get_network_local_interfaces() -> Dict[str, Dict[str, List[str]]]:
    """
    Returns interfaces with their associated local IPs
    """

    interfaces = {
        adapter.name: {
            "ipv4": [
                ip.ip
                for ip in adapter.ips
                if ip.is_IPv4
                and ip_address(ip.ip).is_private
                and not ip_address(ip.ip).is_link_local
            ],
            "ipv6": [
                ip.ip[0]
                for ip in adapter.ips
                if ip.is_IPv6
                and ip_address(ip.ip[0]).is_private
                and not ip_address(ip.ip[0]).is_link_local
            ],
        }
        for adapter in ifaddr.get_adapters()
        if adapter.name != "lo"
    }
    return interfaces


# Listener class, to detect duplicates on the network
# Stores the list of servers in its list property
class Listener:
    def __init__(self) -> None:
        self.list = []

    def remove_service(self, zeroconf: Zeroconf, type: str, name: str) -> None:
        info = zeroconf.get_service_info(type, name)
        self.list.remove(info.server)

    def update_service(self, zeroconf: Zeroconf, type: str, name: str) -> None:
        pass

    def add_service(self, zeroconf: Zeroconf, type: str, name: str) -> None:
        info = zeroconf.get_service_info(type, name)
        self.list.append(info.server[:-1])


def main() -> bool:
    ###
    #  CONFIG
    ###

    with open("/etc/yunohost/mdns.yml", "r") as f:
        config = yaml.safe_load(f) or {}

    required_fields = ["domains"]
    missing_fields = [field for field in required_fields if field not in config]
    interfaces = get_network_local_interfaces()

    if missing_fields:
        print(f"The fields {missing_fields} are required in mdns.yml")
        return False

    if "interfaces" not in config:
        config["interfaces"] = [
            interface
            for interface, local_ips in interfaces.items()
            if local_ips["ipv4"]
        ]

    if "ban_interfaces" in config:
        config["interfaces"] = [
            interface
            for interface in config["interfaces"]
            if interface not in config["ban_interfaces"]
        ]

    # Let's discover currently published .local domains accross the network
    zc = Zeroconf()
    listener = Listener()
    browser = ServiceBrowser(zc, "_device-info._tcp.local.", listener)
    sleep(2)
    browser.cancel()
    zc.close()

    # Always attempt to publish yunohost.local
    if "yunohost.local" not in config["domains"]:
        config["domains"].append("yunohost.local")

    def find_domain_not_already_published(domain):

        # Try domain.local ... but if it's already published by another entity,
        # try domain-2.local, domain-3.local, ...

        i = 1
        domain_i = domain

        while domain_i in listener.list:
            print(f"Uh oh, {domain_i} already exists on the network...")

            i += 1
            domain_i = domain.replace(".local", f"-{i}.local")

        return domain_i

    config["domains"] = [
        find_domain_not_already_published(domain) for domain in config["domains"]
    ]

    zcs: Dict[Zeroconf, List[ServiceInfo]] = {}

    for interface in config["interfaces"]:

        if interface not in interfaces:
            print(
                f"Interface {interface} listed in config file is not present on system."
            )
            continue

        # Broadcast IPv4 and IPv6
        ips: List[str] = interfaces[interface]["ipv4"] + interfaces[interface]["ipv6"]

        # If at least one IP is listed
        if not ips:
            continue

        # Create a Zeroconf object, and store the ServiceInfos
        zc = Zeroconf(interfaces=ips)  # type: ignore
        zcs[zc] = []

        for d in config["domains"]:
            d_domain = d.replace(".local", "")
            if "." in d_domain:
                print(f"{d_domain}.local: subdomains are not supported.")
                continue
            # Create a ServiceInfo object for each .local domain
            zcs[zc].append(
                ServiceInfo(
                    type_="_device-info._tcp.local.",
                    name=f"{interface}: {d_domain}._device-info._tcp.local.",
                    parsed_addresses=ips,
                    port=80,
                    server=f"{d}.",
                )
            )
            print(f"Adding {d} with addresses {ips} on interface {interface}")

    # Run registration
    print("Registering...")
    for zc, infos in zcs.items():
        for info in infos:
            zc.register_service(
                info, allow_name_change=True, cooperating_responders=True
            )

    try:
        print("Registered. Press Ctrl+C or stop service to stop.")
        while True:
            sleep(1)
    except KeyboardInterrupt:
        pass
    finally:
        print("Unregistering...")
        for zc, infos in zcs.items():
            zc.unregister_all_services()
            zc.close()

    return True


if __name__ == "__main__":
    sys.exit(0 if main() else 1)