wlb: T4470: Support WLB op-mode commands

2 files changed, 122 insertions, 0 deletions

diff --git a/src/op_mode/load-balancing_wan.py b/src/op_mode/load-balancing_wan.py
new file mode 100755
index 000000000..9fa473802
--- /dev/null
+++ b/src/op_mode/load-balancing_wan.py
@@ -0,0 +1,117 @@
+#!/usr/bin/env python3
+#
+# Copyright (C) 2024 VyOS maintainers and contributors
+#
+# 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 json
+import re
+import sys
+
+from datetime import datetime
+
+from vyos.config import Config
+from vyos.utils.process import cmd
+
+import vyos.opmode
+
+wlb_status_file = '/run/wlb_status.json'
+
+status_format = '''Interface: {ifname}
+Status: {status}
+Last Status Change: {last_change}
+Last Interface Success: {last_success}
+Last Interface Failure: {last_failure}
+Interface Failures: {failures}
+'''
+
+def _verify(func):
+    """Decorator checks if WLB config exists"""
+    from functools import wraps
+
+    @wraps(func)
+    def _wrapper(*args, **kwargs):
+        config = Config()
+        if not config.exists(['load-balancing', 'wan']):
+            unconf_message = 'WAN load-balancing is not configured'
+            raise vyos.opmode.UnconfiguredSubsystem(unconf_message)
+        return func(*args, **kwargs)
+    return _wrapper
+
+def _get_raw_data():
+    with open(wlb_status_file, 'r') as f:
+        data = json.loads(f.read())
+        if not data:
+            return {}
+        return data
+
+def _get_formatted_output(raw_data):
+    for ifname, if_data in raw_data.items():
+        latest_change = if_data['last_success'] if if_data['last_success'] > if_data['last_failure'] else if_data['last_failure']
+
+        change_dt = datetime.fromtimestamp(latest_change) if latest_change > 0 else None
+        success_dt = datetime.fromtimestamp(if_data['last_success']) if if_data['last_success'] > 0 else None
+        failure_dt = datetime.fromtimestamp(if_data['last_failure']) if if_data['last_failure'] > 0 else None
+        now = datetime.utcnow()
+
+        fmt_data = {
+            'ifname': ifname,
+            'status': "active" if if_data['state'] else "failed",
+            'last_change': change_dt.strftime("%Y-%m-%d %H:%M:%S") if change_dt else 'N/A',
+            'last_success': str(now - success_dt) if success_dt else 'N/A',
+            'last_failure': str(now - failure_dt) if failure_dt else 'N/A',
+            'failures': if_data['failure_count']
+        }
+        print(status_format.format(**fmt_data))
+
+@_verify
+def show_summary(raw: bool):
+    data = _get_raw_data()
+
+    if raw:
+        return data
+    else:
+        return _get_formatted_output(data)
+
+@_verify
+def show_connection(raw: bool):
+    res = cmd('sudo conntrack -L -n')
+    lines = res.split("\n")
+    filtered_lines = [line for line in lines if re.search(r' mark=[1-9]', line)]
+
+    if raw:
+        return filtered_lines
+
+    for line in lines:
+        print(line)
+
+@_verify
+def show_status(raw: bool):
+    res = cmd('sudo nft list chain ip vyos_wanloadbalance wlb_mangle_prerouting')
+    lines = res.split("\n")
+    filtered_lines = [line.replace("\t", "") for line in lines[3:-2] if 'meta mark set' not in line]
+
+    if raw:
+        return filtered_lines
+
+    for line in filtered_lines:
+        print(line)
+
+if __name__ == "__main__":
+    try:
+        res = vyos.opmode.run(sys.modules[__name__])
+        if res:
+            print(res)
+    except (ValueError, vyos.opmode.Error) as e:
+        print(e)
+        sys.exit(1)
diff --git a/src/op_mode/restart.py b/src/op_mode/restart.py
index 3b0031f34..efa835485 100755
--- a/src/op_mode/restart.py
+++ b/src/op_mode/restart.py
@@ -53,6 +53,10 @@ service_map = {
         'systemd_service': 'strongswan',
         'path': ['vpn', 'ipsec'],
     },
+    'load-balancing_wan': {
+        'systemd_service': 'vyos-wan-load-balance',
+        'path': ['load-balancing', 'wan'],
+    },
     'mdns_repeater': {
         'systemd_service': 'avahi-daemon',
         'path': ['service', 'mdns', 'repeater'],
@@ -86,6 +90,7 @@ services = typing.Literal[
     'haproxy',
     'igmp_proxy',
     'ipsec',
+    'load-balancing_wan',
     'mdns_repeater',
     'router_advert',
     'snmp',