wlb: T4470: Support WLB op-mode commands

Author: sarthurdev

op-mode-definitions/load-balancing_wan.xml.in

@@ -0,0 +1,37 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<interfaceDefinition>
+  <node name="restart">
+    <children>
+      <node name="wan-load-balance">
+        <properties>
+          <help>Restart Wide Area Network (WAN) load-balancing daemon</help>
+        </properties>
+        <command>sudo systemctl restart vyos-wan-load-balance.service</command>
+      </node>
+    </children>
+  </node>
+  <node name="show">
+    <children>
+      <node name="wan-load-balance">
+        <properties>
+          <help>Show Wide Area Network (WAN) load-balancing information</help>
+        </properties>
+        <command>${vyos_op_scripts_dir}/load-balancing_wan.py show_summary</command>
+        <children>
+          <node name="connection">
+            <properties>
+              <help>Show Wide Area Network (WAN) load-balancing flow</help>
+            </properties>
+            <command>${vyos_op_scripts_dir}/load-balancing_wan.py show_connection</command>
+          </node>
+          <node name="status">
+            <properties>
+              <help>Show WAN load-balancing statistics</help>
+            </properties>
+            <command>${vyos_op_scripts_dir}/load-balancing_wan.py show_status</command>
+          </node>
+        </children>
+      </node>
+    </children>
+  </node>
+</interfaceDefinition>

src/op_mode/load-balancing_wan.py

@@ -0,0 +1,115 @@
+#!/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)
+
+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)
+
+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)