#!/usr/bin/env python3
#
# Copyright VyOS maintainers and contributors <maintainers@vyos.io>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License version 2 or later as
# published by the Free Software Foundation.
#
# 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 General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

import importlib.util
import re
import sys
import syslog
import signal
import select

from pyroute2 import IPRoute # pylint: disable = no-name-in-module
from pyroute2 import NetlinkError # pylint: disable = no-name-in-module
from pyroute2.netlink.rtnl import RTMGRP_LINK
from pyroute2.netlink.rtnl import RTMGRP_IPV4_IFADDR
from pyroute2.netlink.rtnl import RTMGRP_IPV6_IFADDR
from time import sleep
from typing import Optional

from vyos.configquery import op_mode_config_dict
from vyos.ifconfig import Section
from vyos.utils.boot import boot_configuration_complete
from vyos.utils.commit import commit_in_progress2
from vyos.utils.dict import dict_search
from vyos.utils.process import cmd
from vyos.utils.process import is_systemd_service_active
from vyos.utils.process import stop_systemd_unit

def _load_conf_mode_qos():
    spec = importlib.util.spec_from_file_location(
        'conf_mode_qos', '/usr/libexec/vyos/conf_mode/qos.py')
    mod = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(mod)
    return mod

_conf_mode_qos = _load_conf_mode_qos()

running = True

# compile regex once during startup for fast match
IFACE_RE = re.compile(r"^(?:eth|br|bond|wlan|pppoe|sstpc|wwan)")
_dynamic_qos_interfaces: set[str] = set()

# Per-interface previous operstate, used to suppress DHCP restarts on
# UP-to-UP re-notifications (e.g. post-migration gratuitous-ARP events).
_iface_prev_operstate: dict[str, str] = {}

def match_iface(ifname: str) -> bool:
    """ Helper function returning true if interface name is a match for further
    processing (e.g. restart of DHCP(v6) client)
    """
    return IFACE_RE.match(ifname) is not None

def _is_dynamic_qos_iface(ifname: str) -> bool:
    """ Helper function returning true if interface requires QoS re-apply
    after first address assignment.
    """
    return ifname.startswith(('pppoe', 'sstpc', 'wwan'))

def _handle_dynamic_qos_events(event: str, ifname: str,
                               operstate: Optional[str] = None) -> None:
    """ Re-apply QoS for dynamic interfaces after they receive an address. """
    if not _is_dynamic_qos_iface(ifname):
        return None

    match event:
        case 'RTM_NEWLINK' if operstate != 'DOWN':
            return None

        case 'RTM_NEWLINK' | 'RTM_DELLINK':
            if ifname in _dynamic_qos_interfaces:
                syslog.syslog(syslog.LOG_DEBUG,
                              f'Clearing QoS re-apply state on dynamic interface {ifname}...')
                _dynamic_qos_interfaces.discard(ifname)

        case 'RTM_NEWADDR':
            if ifname in _dynamic_qos_interfaces:
                syslog.syslog(syslog.LOG_DEBUG,
                              f'QoS already re-applied on dynamic interface {ifname}, skipping...')
                return None

            qos = _conf_mode_qos.get_config()
            if not qos or 'interface' not in qos or ifname not in qos['interface']:
                return None

            syslog.syslog(syslog.LOG_INFO,
                          f'Re-applying QoS on dynamic interface {ifname}...')
            _conf_mode_qos.apply_interface(qos, ifname)
            _dynamic_qos_interfaces.add(ifname)

def sigterm_handler(signo, frame):
    global running
    running = False
    sig = signal.Signals(signo)
    syslog.syslog(syslog.LOG_INFO, f'Received signal {sig.name} - shutting down...')

def _handle_dhcp_events(operstate: Optional[str], ifname: str) -> None:
    systemdV4_service = f'dhclient@{ifname}.service'
    systemdV6_service = f'dhcp6c@{ifname}.service'

    # Only handle explicit UP/DOWN state transitions; ignore other kernel states.
    if operstate not in ['UP', 'DOWN']:
        return None

    prev = _iface_prev_operstate.get(ifname)
    _iface_prev_operstate[ifname] = operstate

    if operstate == 'UP' and prev == 'UP':
        syslog.syslog(syslog.LOG_DEBUG,
                      f'Suppressing DHCP restart for {ifname}: already UP')
        return None

    if operstate == 'DOWN':
        # Interface moved state to down
        if is_systemd_service_active(systemdV4_service):
            syslog.syslog(syslog.LOG_DEBUG, f'Stopping {systemdV4_service}...')
            stop_systemd_unit(systemdV4_service, raise_on_failure=False)
        if is_systemd_service_active(systemdV6_service):
            syslog.syslog(syslog.LOG_DEBUG, f'Stopping {systemdV6_service}...')
            stop_systemd_unit(systemdV6_service, raise_on_failure=False)

    elif operstate == 'UP':
        v6_restart = False
        interface_path = Section.get_config_path(ifname, delimiter='.')

        config_dict = op_mode_config_dict(
            ['interfaces'], key_mangling=('-', '_'), get_first_key=True
        )

        if tmp := dict_search(f'{interface_path}.address', config_dict):
            # Always (re-)start the DHCP(v6) client service. If the DHCP(v6) client
            # is already running - which could happen if the interface is re-
            # configured in operational down state, it will have an exponential backoff
            # time increasing while not receiving a DHCP(v6) reply.
            #
            # To make the interface instantly available, and as for a DHCP(v6) lease
            # we will re-start the service and thus cancel the backoff time.
            if 'dhcp' in tmp:
                syslog.syslog(syslog.LOG_DEBUG, f'Restarting {systemdV4_service}...')
                cmd(f'systemctl restart {systemdV4_service}')
            if 'dhcpv6' in tmp:
                v6_restart = True

        if dict_search(f'{interface_path}.dhcpv6_options.pd', config_dict):
            v6_restart = True

        if v6_restart:
            syslog.syslog(syslog.LOG_DEBUG, f'Restarting {systemdV6_service}...')
            cmd(f'systemctl restart {systemdV6_service}')

    return None

def main():
    syslog.openlog(ident="vyos-netlinkd",
                   logoption=syslog.LOG_PID,
                   facility=syslog.LOG_DAEMON)
    syslog.syslog(syslog.LOG_INFO, "VyOS Netlink listener daemon started.")

    # Subscribe to link and address notifications only (not routes/rules/neigh/...).
    ipr = IPRoute()
    try:
        # newer pyroute2 versions support bind group in IPRoute() constructor
        ipr.bind(groups=RTMGRP_LINK | RTMGRP_IPV4_IFADDR | RTMGRP_IPV6_IFADDR)
        syslog.syslog(syslog.LOG_INFO,
            'IPRoute.bind() using link and address RTNL subscriptions')
    except TypeError:
        syslog.syslog(syslog.LOG_WARNING,
            'IPRoute.bind() has no groups= support; using default RTNL subscriptions',
        )
        ipr.bind()
    fd = ipr.fileno()

    global running
    while running:
        if not boot_configuration_complete():
            syslog.syslog(syslog.LOG_INFO, 'System bootup not yet finished...')
            sleep(5)
            continue

        try:
            # Wait for up to 1 second for a netlink message
            rlist, _, _ = select.select([fd], [], [], 1.0)
            if not rlist:
                # timeout - retry
                continue

            # Check if a config commit is in progress before processing any
            # messages. This avoids blocking per-message and reduces unnecessary
            # calls to commit_in_progress2()
            if commit_in_progress2():
                syslog.syslog(syslog.LOG_DEBUG,
                              'Config commit in progress, skipping netlink events')
                sleep(1)
                continue

            # Receive and process any messages
            for message in ipr.get():
                # Parse NETLINK message
                event = message['event']
                attrs = dict(message.get('attrs', []))
                ifname = attrs.get('IFLA_IFNAME', None) or attrs.get('IFA_LABEL', None)

                if not ifname and 'index' in message:
                    try:
                        links = ipr.get_links(message['index'])
                        if links:
                            ifname = dict(links[0].get('attrs', [])).get('IFLA_IFNAME', None)
                    except NetlinkError:
                        # Interface can disappear between the netlink event
                        # arriving and this get_links() lookup (e.g. pppoe
                        # teardown). Skip the message rather than breaking
                        # the entire ipr.get() batch with an ENODEV error.
                        continue

                # Bail out early - no interface name in the message
                if not ifname:
                    continue
                # Bail out early - not interested in interface type
                if not match_iface(ifname):
                    continue

                match event:
                    # Message received during interface creation or modification
                    # e.g. link up/down.
                    case 'RTM_NEWLINK':
                        mac = attrs.get('IFLA_ADDRESS', '<unknown>')
                        operstate = attrs.get('IFLA_OPERSTATE', None)
                        syslog.syslog(syslog.LOG_DEBUG,
                                      f'RTM_NEWLINK -> {ifname}, state={operstate}, mac={mac}')

                        _handle_dynamic_qos_events(event, ifname, operstate)
                        _handle_dhcp_events(operstate, ifname)

                    # Address added to an interface. For dynamic interfaces, this
                    # means the connection completed and QoS can be re-applied.
                    case 'RTM_NEWADDR':
                        addr = attrs.get('IFA_ADDRESS', '<unknown>')
                        prefixlen = message.get('prefixlen', '')
                        syslog.syslog(syslog.LOG_DEBUG,
                                      f'RTM_NEWADDR -> {ifname}, addr={addr}/{prefixlen}')
                        _handle_dynamic_qos_events(event, ifname)

                    # Deletion of a network link which has been previously added to the kernel
                    case 'RTM_DELLINK':
                        attrs = dict(message.get('attrs', []))
                        ifname = attrs.get('IFLA_IFNAME', None)
                        if ifname:
                            _iface_prev_operstate.pop(ifname, None)
                            _handle_dynamic_qos_events(event, ifname)
                    case _:
                        pass

        except NetlinkError as e:
            syslog.syslog(syslog.LOG_ERR, f'Netlink error: {e}')
        except Exception as e:
            syslog.syslog(syslog.LOG_ERR, f'Unhandled exception: {e}')
        except KeyboardInterrupt:
            break

    ipr.close()
    syslog.syslog(syslog.LOG_INFO, 'Netlink listener daemon stopped.')
    sys.exit(0)

if __name__ == "__main__":
    signal.signal(signal.SIGTERM, sigterm_handler)
    signal.signal(signal.SIGINT, sigterm_handler)
    main()