From 5fdeda7dc795002df3b5f2d8571926f8928b5818 Mon Sep 17 00:00:00 2001
From: sarthurdev <965089+sarthurdev@users.noreply.github.com>
Date: Tue, 8 Oct 2024 23:17:23 +0200
Subject: [PATCH] wlb: T4470: Support WLB op-mode commands
---
op-mode-definitions/load-balancing_wan.xml.in | 37 ++++++
src/op_mode/load-balancing_wan.py | 115 ++++++++++++++++++
2 files changed, 152 insertions(+)
create mode 100644 op-mode-definitions/load-balancing_wan.xml.in
create mode 100755 src/op_mode/load-balancing_wan.py
diff --git a/op-mode-definitions/load-balancing_wan.xml.in b/op-mode-definitions/load-balancing_wan.xml.in
new file mode 100644
index 0000000000..a40b1daab7
--- /dev/null
+++ b/op-mode-definitions/load-balancing_wan.xml.in
@@ -0,0 +1,37 @@
+
+
+
+
+
+
+ Restart Wide Area Network (WAN) load-balancing daemon
+
+ ${vyos_op_scripts_dir}/load-balancing_wan.py restart
+
+
+
+
+
+
+
+ Show Wide Area Network (WAN) load-balancing information
+
+ ${vyos_op_scripts_dir}/load-balancing_wan.py summary
+
+
+
+ Show Wide Area Network (WAN) load-balancing flow
+
+ ${vyos_op_scripts_dir}/load-balancing_wan.py connection
+
+
+
+ Show WAN load-balancing statistics
+
+ ${vyos_op_scripts_dir}/load-balancing_wan.py status
+
+
+
+
+
+
\ No newline at end of file
diff --git a/src/op_mode/load-balancing_wan.py b/src/op_mode/load-balancing_wan.py
new file mode 100755
index 0000000000..f498342c29
--- /dev/null
+++ b/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 .
+
+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 summary(raw: bool):
+ data = _get_raw_data()
+
+ if raw:
+ return data
+ else:
+ return _get_formatted_output(data)
+
+def 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 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)