#!/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 os
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.file import read_file
from vyos.utils.dict import dict_search
from vyos.utils.process import cmdl
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), and
# to detect edges missed while netlink events were drained during a commit.
_iface_prev_operstate: dict[str, str] = {}

# True while the main loop is draining netlink events because a config
# commit holds the lock. When it clears we reconcile DHCP clients against
# live operstate so a transition that only happened inside the drain window
# is not lost (T9086).
_in_commit_skip = False

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 _read_operstate(ifname: str) -> Optional[str]:
    """Read live kernel operstate from sysfs (uppercase), or None if unavailable.

    Values match IFLA_OPERSTATE names used by RTM_NEWLINK (UP, DOWN, ...).
    See Documentation/ABI/testing/sysfs-class-net.
    """
    fname = f'/sys/class/net/{ifname}/operstate'
    return read_file(fname, defaultonfailure='').strip().upper() or 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}...')
                cmdl(['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}...')
            cmdl(['systemctl', 'restart', systemdV6_service])

    return None

def _reconcile_dhcp_operstate() -> None:
    """After a commit drain window, sync DHCP clients to live operstate (T9086).

    Events discarded while commit_in_progress2() was true may have included a
    real UP/DOWN edge. Compare each interface's live sysfs operstate against
    the last one we acted on and run the normal DHCP handler only where they
    disagree.

    Interfaces never seen before are seeded into the tracker without acting -
    conf_mode owns DHCP bring-up across commits for those.
    """
    try:
        ifnames = os.listdir('/sys/class/net')
    except OSError as e:
        syslog.syslog(syslog.LOG_ERR, f'Failed to list interfaces for reconcile: {e}')
        return

    for ifname in ifnames:
        if not match_iface(ifname):
            continue

        operstate = _read_operstate(ifname)
        if operstate not in ['UP', 'DOWN']:
            continue

        prev = _iface_prev_operstate.get(ifname)
        if prev is None:
            _iface_prev_operstate[ifname] = operstate
            continue
        if prev == operstate:
            continue

        syslog.syslog(syslog.LOG_DEBUG,
                      f'Reconcile {ifname}: prev={prev} current={operstate}')
        _handle_dhcp_events(operstate, ifname)

def main():
    global _in_commit_skip

    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()

    # Establish a baseline so the first UP re-notification is not mistaken
    # for a real interface transition.
    for link in ipr.get_links():
        attrs = dict(link.get('attrs', []))
        ifname = attrs.get('IFLA_IFNAME')
        operstate = attrs.get('IFLA_OPERSTATE')
        if ifname and match_iface(ifname) and operstate in ['UP', 'DOWN']:
            _iface_prev_operstate[ifname] = operstate

    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. The timeout also
            # lets us notice a commit ending even with no further netlink
            # traffic, so we can reconcile (T9086).
            rlist, _, _ = select.select([fd], [], [], 1.0)

            if commit_in_progress2():
                # Drain the socket instead of leaving messages queued, so a
                # stale event (e.g. a DOWN from a disable that was already
                # re-enabled) can't be processed as fresh once the commit
                # ends (T9086).
                if not _in_commit_skip:
                    syslog.syslog(syslog.LOG_DEBUG,
                                  'Config commit in progress, draining netlink events without acting')
                    _in_commit_skip = True
                if rlist:
                    try:
                        ipr.get()
                    except NetlinkError as e:
                        syslog.syslog(syslog.LOG_ERR,
                                      f'Netlink error while draining during commit: {e}')
                continue

            if _in_commit_skip:
                _in_commit_skip = False
                syslog.syslog(syslog.LOG_INFO,
                              'Config commit finished, reconciling DHCP client state with operstate')
                _reconcile_dhcp_operstate()

            if not rlist:
                # timeout - retry
                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()