From 5d22e2418173d2780bfe14e080ed14aff08d1905 Mon Sep 17 00:00:00 2001
From: Vincent Wiemann <vincent.wiemann@ironai.com>
Date: Sun, 31 Jan 2021 01:35:44 +0100
Subject: [PATCH] ffh.mesh_wireguard: add statistics export

---
 .../files/bin/ffh_wg_stats.py                 | 91 +++++++++++++++++++
 roles/ffh.mesh_wireguard/tasks/main.yml       |  3 +
 roles/ffh.mesh_wireguard/tasks/stats.yml      | 25 +++++
 .../templates/telegraf_mesh_wg_stats.conf.j2  |  3 +
 4 files changed, 122 insertions(+)
 create mode 100644 roles/ffh.mesh_wireguard/files/bin/ffh_wg_stats.py
 create mode 100644 roles/ffh.mesh_wireguard/tasks/stats.yml
 create mode 100644 roles/ffh.mesh_wireguard/templates/telegraf_mesh_wg_stats.conf.j2

diff --git a/roles/ffh.mesh_wireguard/files/bin/ffh_wg_stats.py b/roles/ffh.mesh_wireguard/files/bin/ffh_wg_stats.py
new file mode 100644
index 00000000..243abad8
--- /dev/null
+++ b/roles/ffh.mesh_wireguard/files/bin/ffh_wg_stats.py
@@ -0,0 +1,91 @@
+#!/usr/bin/env python3
+from datetime import datetime, timedelta
+from json import dumps as json_dumps
+from pyroute2 import WireGuard
+import argparse, os
+
+TIMEOUT = timedelta(minutes=3)
+
+def get_stats(format, interface):
+    clients = WireGuard().info(interface)[0].WGDEVICE_A_PEERS.value
+
+    peers = []
+    established_peers_count = 0
+    for client in clients:
+        public_key = client.WGPEER_A_PUBLIC_KEY["value"].decode("utf-8")
+
+        # ignore dummy key
+        if public_key == (43 * "0" + "="):
+            continue
+
+        established = False
+        latest_handshake = client.WGPEER_A_LAST_HANDSHAKE_TIME["tv_sec"]
+        if datetime.now() - datetime.fromtimestamp(latest_handshake) < TIMEOUT:
+            established = True
+            established_peers_count += 1
+
+        peers.append({
+            "public_key" : public_key,
+            "latest_handshake" : latest_handshake,
+            "is_established" : established,
+            "rx_bytes" : int(client.WGPEER_A_RX_BYTES['value']),
+            "tx_bytes" : int(client.WGPEER_A_TX_BYTES['value']),
+        })
+
+    return {
+        "interface" : interface,
+        "peers" : peers,
+        "established_peers_count" : established_peers_count
+        }
+
+def print_influx_format(stats):
+    timestamp = str(round(datetime.now().timestamp())) + "0" * 9
+
+    print("wireguard_device" +
+          ",name=" + stats['interface'] +
+          ",type=linux_kernel " +
+          "established_peers=" + str(stats['established_peers_count']) +
+          "i " + timestamp)
+    for peer in stats['peers']:
+        print("wireguard_peer,device=" + stats['interface'] +
+               ",public_key=" + peer['public_key'] +
+               " last_handshake=" + str(peer['latest_handshake']) +
+               "i,rx_bytes=" + str(peer['rx_bytes']) +
+               "i,tx_bytes=" + str(peer['tx_bytes']) +
+               "i " + timestamp)
+
+def check_iface_type(iface, type):
+    if not os.path.exists(f'/sys/class/net/{iface}'):
+        print(f'Iface {iface} does not exist! Exiting...')
+        exit(1)
+
+    with open(f'/sys/class/net/{iface}/uevent', 'r') as f:
+        for line in f.readlines():
+            l = line.replace('\n', '').split('=')
+            if l[0] == 'DEVTYPE' and l[1] != type:
+                print(f'Iface {iface} is wrong type! Should be {type}, but is {l[1]}. Exiting...')
+                exit(1)
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(description='Dump statistics of WireGuard interfaces')
+    parser.add_argument('-f', '--format', metavar='FORMAT', type=str,
+                    help='allowed formats are: "json", "influx"')
+    parser.add_argument('-i', '--interface', metavar='IFACE', type=str, nargs='+',
+                    help='the WireGuard interfaces', default=[])
+    args = parser.parse_args()
+
+    if len(args.interface) == 0:
+        print('Error: You need to specify at least one WireGuard interface.')
+        exit(1)
+
+    iface_stats = []
+    for i in range(len(args.interface)):
+        check_iface_type(args.interface[i], 'wireguard')
+        iface_stats.append(get_stats(args.format, args.interface[i]))
+
+    if args.format == "influx":
+        for stats in iface_stats:
+            print_influx_format(stats);
+    else:
+        print(json_dumps(iface_stats))
+
diff --git a/roles/ffh.mesh_wireguard/tasks/main.yml b/roles/ffh.mesh_wireguard/tasks/main.yml
index 98b3c1e7..2b64d46b 100644
--- a/roles/ffh.mesh_wireguard/tasks/main.yml
+++ b/roles/ffh.mesh_wireguard/tasks/main.yml
@@ -31,3 +31,6 @@
 
 - name: Establish netlink worker
   include_tasks: netlink.yml
+
+- name: Setup statistics export
+  include_tasks: stats.yml
diff --git a/roles/ffh.mesh_wireguard/tasks/stats.yml b/roles/ffh.mesh_wireguard/tasks/stats.yml
new file mode 100644
index 00000000..b62f9c23
--- /dev/null
+++ b/roles/ffh.mesh_wireguard/tasks/stats.yml
@@ -0,0 +1,25 @@
+---
+
+- name: Install ffh_wg_stats.py files
+  copy:
+    src: "bin/ffh_wg_stats.py"
+    dest: "/usr/bin/"
+    owner: "root"
+    mode: 0755
+
+- name: Check telegraf path
+  stat:
+    path: /etc/telegraf/telegraf.d/
+  register: telegraf_path
+
+- name: Ensure telegraf path exists
+  shell: mkdir -p /etc/telegraf/telegraf.d/
+  when: telegraf_path is defined and telegraf_path.stat.exists==false
+
+# We do not care whether telegraf is really installed. So it could be installed
+# later (or also never).
+- name: Set /etc/telegraf/telegraf.d/mesh_wg_stats.conf
+  notify: Maybe restart telegraf # handler from ffh.routingnode
+  template:
+    src: telegraf_mesh_wg_stats.conf.j2
+    dest: /etc/telegraf/telegraf.d/mesh_wg_stats.conf
diff --git a/roles/ffh.mesh_wireguard/templates/telegraf_mesh_wg_stats.conf.j2 b/roles/ffh.mesh_wireguard/templates/telegraf_mesh_wg_stats.conf.j2
new file mode 100644
index 00000000..3ff78535
--- /dev/null
+++ b/roles/ffh.mesh_wireguard/templates/telegraf_mesh_wg_stats.conf.j2
@@ -0,0 +1,3 @@
+[[inputs.exec]]
+  commands = ["/usr/bin/ffh_wg_stats.py -f influx -i{% for domain in domains | default( [] ) %} wg-{{ domain.id }}{% endfor %}"]
+  data_format = "influx"