Repository: freifunkMUC/wgkex (language: python)

Issue:
Added small stupid json interface



Discussion / hints:

Fix patch:
diff --git a/Docs/architecture.dot b/Docs/architecture.dot
new file mode 100644
index 0000000..ab453e7
--- /dev/null
+++ b/Docs/architecture.dot
@@ -0,0 +1,50 @@
+# http://www.graphviz.org/content/cluster
+
+digraph G {
+  graph [fontname = "Handlee"];
+  node [fontname = "Handlee"];
+  edge [fontname = "Handlee"];
+
+  bgcolor=transparent;
+
+  subgraph wgkex_broker {
+    style=filled;
+    color=lightgrey;
+    style="dotted"
+    node [style=filled,color="#d8b365"];
+    wgkex_broker [label="*WGKex Broker*"];
+    fontsize = 20;
+  }
+
+  subgraph wgkex_mqtt {
+    node [style=filled,color="#f5f5f5"];
+    style="dotted"
+    mqtt [label="*Mosquitto*"];
+    fontsize = 20;
+    color=blue
+  }
+
+  subgraph wgkex_worker {
+    node [style=filled,color="#5ab4ac"];
+    style="dotted"
+    color=blue;
+    label = "*WGKEX Worker*";
+    fontsize = 20;
+    color=blue
+    mqtt -> wgkex_worker -> netlink -> wireguard -> vxlan [style=invis];
+    wgkex_worker [label="*WGKex Worker*"];
+    netlink [label="netlink\n(pyroute2)"];
+    wireguard [label="wireguard\n(pyroute2)"];
+    vxlan [label="vxlan FDB\n(pyroute2)"];
+  }
+  client -> wgkex_broker [label="RESTFul API"];
+  wgkex_broker -> mqtt [label="publish"];
+  wgkex_worker -> mqtt [label="Subscribe"];
+  wgkex_worker -> netlink [label="Route\nInjection"];
+  wgkex_worker -> wireguard [label="Peer\nCreation"];
+  wgkex_worker -> vxlan [label="VxLAN FDB\nEntry"];
+
+  client [shape=Mdiamond];
+  {rank=same wgkex_broker mqtt}
+  {rank=same netlink wireguard vxlan}
+}
diff --git a/Docs/architecture.png b/Docs/architecture.png
new file mode 100644
index 0000000..f39fe52
Binary files /dev/null and b/Docs/architecture.png differ
diff --git a/README.md b/README.md
index 2823678..d1a0a98 100644
--- a/README.md
+++ b/README.md
@@ -1,7 +1,65 @@
+- [WireGuard Key Exchange](#wireguard-key-exchange)
+  * [Overview](#overview)
+    + [Frontend broker](#frontend-broker)
+      - [GET /wg-public-key/<path:key>](#get--wg-public-key--path-key-)
+      - [POST /api/v1/wg/key/exchange](#post--api-v1-wg-key-exchange)
+    + [Backend worker](#backend-worker)
+  * [Installation](#installation)
+  * [Configuration](#configuration)
+  * [Running the broker](#running-the-broker)
+  * [Client usage](#client-usage)
+  * [Contact](#contact)
+
+
 # WireGuard Key Exchange
 
-WireGuard Key Exchange is a tool consisting of two parts: a frontend (broker) and a backend (worker). The frontend (broker) is where the client can push (register) its key before connecting. The backend (worker) is injecting those keys into a WireGuard instance.
-This tool is intended to facilitate running BATMAN over VXLAN over WireGuard as a means to create encrypted high-performance mesh links.
+wgkex is a WireGuard key exchange and management tool designed and run by FFMUC.
+
+## Overview
+
+WireGuard Key Exchange is a tool consisting of two parts: a frontend (broker) and a backend (worker). These components 
+communicate to each other via MQTT - a messaging bus.
+
+![](Docs/architecture.png)
+
+### Frontend broker
+
+The frontend broker is where the client can push (register) its key before connecting. These keys are then pushed into
+an MQTT bus for all workers to consume.
+
+The frontend broker exposes the following API endpoints for use:
+
+```
+/wg-public-key/<path:key>
+/api/v1/wg/key/exchange
+```
+
+#### GET /wg-public-key/<path:key>
+
+The key in the URI (ie: part after `/wg-public-key/`) is validated and written to disk.
+
+#### POST /api/v1/wg/key/exchange
+
+JSON POST'd to this endpoint should be in this format:
+
+```json
+{
+  "domain": "CONFIGURED_DOMAIN", 
+  "public_key": "PUBLIC_KEY"
+}
+```
+
+The broker will validate the domain and public key, and if valid, will push the key onto the MQTT bus.
+
+### Backend worker
+
+The backend (worker) waits for new keys to appear on the MQTT message bus. Once a new key appears, the worker performs
+validation task on the key, then injects those keys into a WireGuard instance(While also updating the VxLAN FDB). 
+
+This tool is intended to facilitate running BATMAN over VXLAN over WireGuard as a means to create encrypted 
+high-performance mesh links.
+
+For further information, please see this [presentation on the architecture](https://www.slideshare.net/AnnikaWickert/ffmuc-goes-wild-infrastructure-recap-2020-rc3)
 
 ## Installation
 
@@ -9,19 +67,32 @@ This tool is intended to facilitate running BATMAN over VXLAN over WireGuard as
 
 ## Configuration
 
-* TBA
+* Configuration file
+
+The `wgkex` configuration file defaults to `/etc/wgkex.yaml` ([Sample configuration file](wgkex.yaml.example)), however
+can also be overwritten by setting the environment variable `WGKEX_CONFIG_FILE`.
 
 ## Running the broker
 
 * The broker web frontend can be started directly from a Git checkout:
+
 ```
 # defaults to /etc/wgkex.yaml if not set
 export WGKEX_CONFIG_FILE=/opt/wgkex/wgkex.yaml
 poetry run wgkex-broker
 ```
 
+* The broker can also be built and run via [bazel](https://bazel.build):
+
+```shell
+bazel build //wgkex/broker:app
+# Artifact will now be placed into ./bazel-bin/wgkex/broker/app
+./bazel-bin/wgkex/broker/app
+```
+
 ## Client usage
 
+The client can be used via CLI:
 ```
 $ wget -q  -O- --post-data='{"domain": "ffmuc_welt","public_key": "o52Ge+Rpj4CUSitVag9mS7pSXUesNM0ESnvj/wwehkg="}'   --header='Content-Type:application/json'   'http://127.0.0.1:5000/api/v1/wg/key/exchange'
 {
@@ -29,6 +100,15 @@ $ wget -q  -O- --post-data='{"domain": "ffmuc_welt","public_key": "o52Ge+Rpj4CUS
 }
 ```
 
+Or via python:
+```python
+import requests
+key_data = {"domain": "ffmuc_welt","public_key": "o52Ge+Rpj4CUSitVag9mS7pSXUesNM0ESnvj/wwehkg="}
+broker_url = "http://127.0.0.1:5000"
+push_key = requests.get(f'{broker_url}/api/v1/wg/key/exchange', json=key_data)
+print(f'Key push was: {push_key.json().get("Message")]}')
+```
+
 ## Contact
 
-#wgkex - IRCNet
+[wgkex - IRCNet](ircs://irc.ircnet.net:6697/wgkex)
\ No newline at end of file
diff --git a/WORKSPACE b/WORKSPACE
index 9712e76..4aafce8 100644
--- a/WORKSPACE
+++ b/WORKSPACE
@@ -1,14 +1,13 @@
-load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
-
 # Load python specific bazel rules.
+load("@bazel_tools//tools/build_defs/repo:http.bzl", "http_archive")
 http_archive(
     name = "rules_python",
     url = "https://github.com/bazelbuild/rules_python/releases/download/0.1.0/rules_python-0.1.0.tar.gz",
     sha256 = "b6d46438523a3ec0f3cead544190ee13223a52f6a6765a29eae7b7cc24cc83a0",
 )
 
+# PIP support.
 load("@rules_python//python:pip.bzl", "pip_install")
-
 pip_install(
    requirements = "//:requirements.txt",
 )
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index aae862d..6eb6ea0 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -3,7 +3,6 @@ flask-mqtt
 pyroute2
 PyYAML
 Flask
-voluptuous
 
 # Common
 ipaddress
diff --git a/wgkex/broker/BUILD b/wgkex/broker/BUILD
new file mode 100644
index 0000000..7ee73ef
--- /dev/null
+++ b/wgkex/broker/BUILD
@@ -0,0 +1,23 @@
+load("@rules_python//python:defs.bzl", "py_binary", "py_test")
+load("@pip//:requirements.bzl", "requirement")
+
+py_binary(
+    name="app",
+    srcs=["app.py"],
+    data=["templates/index.html"],
+    visibility=["//visibility:public"],
+    deps=[
+        requirement("flask"),
+        requirement("flask-mqtt"),
+        "//wgkex/config:config",
+    ],
+)
+
+py_test(
+    name="app_test",
+    srcs=["app_test.py"],
+    deps=[
+        ":app",
+        requirement("mock"),
+    ],
+)
diff --git a/wgkex/broker/app.py b/wgkex/broker/app.py
index 949c084..879eb06 100644
--- a/wgkex/broker/app.py
+++ b/wgkex/broker/app.py
@@ -1,92 +1,182 @@
 #!/usr/bin/env python3
 import re
-
-from flask import Flask, abort, jsonify, render_template, request
+import dataclasses
+from typing import Tuple, Any
+
+from flask import Flask
+from flask import abort
+from flask import jsonify
+from flask import render_template
+from flask import request
+from flask.app import Flask as Flask_app
 from flask_mqtt import Mqtt
-from voluptuous import Invalid, MultipleInvalid, Required, Schema
-
-from wgkex.config import load_config
+import paho.mqtt.client as mqtt_client
 
-app = Flask(__name__)
-config = load_config()
+from wgkex.config import config
 
-app.config["MQTT_BROKER_URL"] = config.get("mqtt", {}).get("broker_url")
-app.config["MQTT_BROKER_PORT"] = config.get("mqtt", {}).get("broker_port")
-app.config["MQTT_USERNAME"] = config.get("mqtt", {}).get("username")
-app.config["MQTT_PASSWORD"] = config.get("mqtt", {}).get("password")
-app.config["MQTT_KEEPALIVE"] = config.get("mqtt", {}).get("keepalive")
-app.config["MQTT_TLS_ENABLED"] = config.get("mqtt", {}).get("tls")
-
-mqtt = Mqtt(app)
 
 WG_PUBKEY_PATTERN = re.compile(r"^[A-Za-z0-9+/]{42}[AEIMQUYcgkosw480]=$")
+app = None
+mqtt = None
+
+
+@dataclasses.dataclass
+class KeyExchange:
+    """A key exchange message.
+
+    Attributes:
+        public_key: The public key for this exchange.
+        domain: The domain for this exchange.
+    """
+
+    public_key: str
+    domain: str
+
+    @classmethod
+    def from_dict(cls, msg: dict) -> "KeyExchange":
+        """Creates a new KeyExchange message from dict.
+
+        Arguments:
+            msg: The message to convert.
+        Returns:
+            A KeyExchange object.
+        """
+        public_key = is_valid_wg_pubkey(msg.get("public_key"))
+        domain = is_valid_domain(msg.get("domain"))
+        return cls(public_key=public_key, domain=domain)
+
+
+def main():
+    global app, mqtt
+    app = _fetch_app_config()
+    mqtt = Mqtt(app)
+
+    @app.route("/", methods=["GET"])
+    def index() -> None:
+        """Returns main page"""
+        return render_template("index.html")
+
+    # Keep to be compatible
+    @app.route("/wg-public-key/<path:key>", methods=["GET"])
+    def receive_public_key(key: str) -> Tuple[str, int]:
+        """Validates a new key and saves to disk.
+
+        Arguments:
+            key: The public key.
+        Returns:
+            Status message.
+        """
+        if not is_valid_wg_pubkey(key):
+            return jsonify({"Message": "Invalid Key"}), 400
+        # TODO(ruairi): Move to env and static variable.
+        with open("/var/lib/wgkex/public.keys", "a") as pubkeys:
+            pubkeys.write("%s\n" % key)
+        return jsonify({"Message": "OK"}), 200
+
+    @app.route("/api/v1/wg/key/exchange", methods=["POST"])
+    def wg_key_exchange() -> Tuple[str, int]:
+        """Retrieves a new key and validates.
+
+        Returns:
+            Status message.
+        """
+        try:
+            data = KeyExchange.from_dict(request.get_json(force=True))
+        except TypeError as ex:
+            return abort(400, jsonify({"error": {"message": str(ex)}}))
+
+        key = data.public_key
+        domain = data.domain
+        # in case we want to decide here later we want to publish it only to dedicated gateways
+        gateway = "all"
+        print(f"wg_key_exchange: Domain: {domain}, Key:{key}")
+
+        pubkey_cfg_file = config.fetch_from_config("pubkeys_file")
+        if pubkey_cfg_file:
+            with open(pubkey_cfg_file, "a") as pubkeys:
+                pubkeys.write(f"{key} {domain}\n")
+        # mqtt.publish(f"wireguard/{domain}/{gateway}", key)
+        return jsonify({"Message": "OK"}), 200
+
+    app.run()
+
+    @mqtt.on_connect()
+    def handle_mqtt_connect(
+        client: Mqtt.Client, userdata: bytes, flags: Any, rc: Any
+    ) -> None:
+        """Prints status of connect message."""
+        # TODO(ruairi): Clarify current usage of this function.
+        print(
+            "MQTT connected to {}:{}".format(
+                app.config["MQTT_BROKER_URL"], app.config["MQTT_BROKER_PORT"]
+            )
+        )
+        # mqtt.subscribe("wireguard/#")
 
+    @mqtt.on_message()
+    def handle_mqtt_message(
+        client: Mqtt.Client, userdata: bytes, message: mqtt_client.MQTTMessage
+    ) -> None:
+        """Prints message contents."""
+        # TODO(ruairi): Clarify current usage of this function.
+        print(f"MQTT message received on {message.topic}: {message.payload.decode()}")
 
-def is_valid_wg_pubkey(pubkey):
-    if WG_PUBKEY_PATTERN.match(pubkey) is None:
-        raise Invalid("Not a valid Wireguard public key")
-    return pubkey
-
-
-def is_valid_domain(domain):
-    if domain not in config.get("domains"):
-        raise Invalid("Not a valid domain")
-    return domain
-
-
-WG_KEY_EXCHANGE_SCHEMA_V1 = Schema(
-    {Required("public_key"): is_valid_wg_pubkey, Required("domain"): is_valid_domain}
-)
 
+def is_valid_wg_pubkey(pubkey: str) -> str:
+    """Verifies if key is a valid WireGuard public key or not.
 
-@app.route("/", methods=["GET"])
-def index():
-    # return templates/index.html
-    return render_template("index.html")
+    Arguments:
+        pubkey: The key to verify.
 
+    Raises:
+        ValueError: If the Wireguard Key is invalid.
 
-# Keep to be compatible
-@app.route("/wg-public-key/<path:key>", methods=["GET"])
-def receive_public_key(key):
-    if not is_valid_wg_pubkey(key):
-        return jsonify({"Message": "Invalid Key"}), 400
-    with open("/var/lib/wgkex/public.keys", "a") as pubkeys:
-        pubkeys.write("%s\n" % key)
-    return jsonify({"Message": "OK"}), 200
+    Returns:
+        The public key.
+    """
+    # TODO(ruairi): Refactor to return bool.
+    if WG_PUBKEY_PATTERN.match(pubkey) is None:
+        raise ValueError(f"Not a valid Wireguard public key: {pubkey}.")
+    return pubkey
 
 
-@app.route("/api/v1/wg/key/exchange", methods=["POST"])
-def wg_key_exchange():
-    try:
-        data = WG_KEY_EXCHANGE_SCHEMA_V1(request.get_json(force=True))
-    except MultipleInvalid as ex:
-        return abort(400, jsonify({"error": {"message": str(ex)}}))
+def is_valid_domain(domain: str) -> str:
+    """Verifies if the domain is configured.
 
-    key = data["public_key"]
-    domain = data["domain"]
-    # in case we want to decide here later we want to publish it only to dedicated gateways
-    gateway = "all"
-    print(f"wg_key_exchange: Domain: {domain}, Key:{key}")
+    Arguments:
+        domain: The domain to verify.
 
-    if config.get("pubkeys_file") != "":
-        with open(config["pubkeys_file"], "a") as pubkeys:
-            pubkeys.write(f"{key} {domain}\n")
+    Raises:
+        ValueError: If the domain is not configured.
 
-    mqtt.publish(f"wireguard/{domain}/{gateway}", key)
+    Returns:
+        The domain.
+    """
+    # TODO(ruairi): Refactor to return bool.
+    if domain not in config.fetch_from_config("domains"):
+        raise ValueError(
+            f'Domains {domain} not in configured domains({config.fetch_from_config("domains")}) a valid domain'
+        )
+    return domain
 
-    return jsonify({"Message": "OK"}), 200
 
+def _fetch_app_config() -> Flask_app:
+    """Creates the Flask app from configuration.
 
-@mqtt.on_connect()
-def handle_mqtt_connect(client, userdata, flags, rc):
-    print(
-        "MQTT connected to {}:{}".format(
-            app.config["MQTT_BROKER_URL"], app.config["MQTT_BROKER_PORT"]
-        )
-    )
-    # mqtt.subscribe("wireguard/#")
+    Returns:
+        A created Flask app.
+    """
+    app = Flask(__name__)
+    # TODO(ruairi): Refactor load_config to return Dataclass.
+    mqtt_cfg = config.Config.from_dict(config.load_config()).mqtt
+    app.config["MQTT_BROKER_URL"] = mqtt_cfg.broker_url
+    app.config["MQTT_BROKER_PORT"] = mqtt_cfg.broker_port
+    app.config["MQTT_USERNAME"] = mqtt_cfg.username
+    app.config["MQTT_PASSWORD"] = mqtt_cfg.password
+    app.config["MQTT_KEEPALIVE"] = mqtt_cfg.keepalive
+    app.config["MQTT_TLS_ENABLED"] = mqtt_cfg.tls
+    return app
 
 
-@mqtt.on_message()
-def handle_mqtt_message(client, userdata, message):
-    print(f"MQTT message received on {message.topic}: {message.payload.decode()}")
+if __name__ == "__main__":
+    main()


Test patch:
diff --git a/wgkex/broker/app_test.py b/wgkex/broker/app_test.py
new file mode 100644
index 0000000..0dfc991
--- /dev/null
+++ b/wgkex/broker/app_test.py
@@ -0,0 +1,91 @@
+import unittest
+import mock
+import app
+import sys
+
+_VALID_CFG = "domains:\n- a\n- b\nmqtt:\n  broker_port: 1883\n  broker_url: mqtt://broker\n  keepalive: 5\n  password: pass\n  tls: true\n  username: user\npubkeys_file: some/file\n"
+_INVALID_CFG = "asdasdasdasd"
+
+
+class TestApp(unittest.TestCase):
+
+    # TODO(ruairi): Add test for Flask.
+    # def setUp(self) -> None:
+    #     mock_open = mock.mock_open(read_data=_VALID_CFG)
+    #     with mock.patch("builtins.open", mock_open):
+    #         app_cfg = app.app.test_client()
+    #         app.main()
+    #     self.app_cfg = app_cfg
+
+    def test_app_load_success(self):
+        """Tests _fetch_app_config success."""
+        mock_open = mock.mock_open(read_data=_VALID_CFG)
+        with mock.patch("builtins.open", mock_open):
+            cfg = app._fetch_app_config()
+            self.assertIsNotNone(cfg)
+
+    @mock.patch.object(sys, "exit", autospec=True)
+    def test_app_load_fails_bad_config(self, exit_mock):
+        """Tests _fetch_app_config fails with bad configuration."""
+        mock_open = mock.mock_open(read_data=_INVALID_CFG)
+        with mock.patch("builtins.open", mock_open):
+            with self.assertRaises(TypeError):
+                app._fetch_app_config()
+                exit_mock.assert_called_with(2)
+
+    def test_is_valid_wg_pubkey_success(self):
+        """Tests is_valid_wg_pubkey success."""
+        key = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAE="
+        self.assertEqual(key, app.is_valid_wg_pubkey(key))
+
+    def test_is_valid_wg_pubkey_fails_bad_key(self):
+        """Tests is_valid_wg_pubkey fails on bad key."""
+        key = "not_a_key"
+        with self.assertRaises(ValueError):
+            app.is_valid_wg_pubkey(key)
+
+    def test_is_valid_domain_success(self):
+        """Tests is_valid_domain success."""
+        domain = "a"
+        mock_open = mock.mock_open(read_data=_VALID_CFG)
+        with mock.patch("builtins.open", mock_open):
+            self.assertEqual(domain, app.is_valid_domain(domain))
+
+    def test_is_valid_domain_fails_domain_not_configured(self):
+        """Tests is_valid_domain fails on bad domain."""
+        domain = "not_Configured"
+        mock_open = mock.mock_open(read_data=_VALID_CFG)
+        with mock.patch("builtins.open", mock_open):
+            with self.assertRaises(ValueError):
+                app.is_valid_domain(domain)
+
+    def test_KeyExchange_success(self):
+        """Tests creating KeyExchange successfully."""
+        key = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAE="
+        msg = dict(public_key=key, domain="a")
+        expected = app.KeyExchange(public_key=key, domain="a")
+        mock_open = mock.mock_open(read_data=_VALID_CFG)
+        with mock.patch("builtins.open", mock_open):
+            self.assertEqual(expected, app.KeyExchange.from_dict(msg))
+
+    def test_KeyExchange_fails_bad_key(self):
+        """Tests creating KeyExchange fails due to bad key."""
+        key = "asd"
+        msg = dict(public_key=key, domain="a")
+        mock_open = mock.mock_open(read_data=_VALID_CFG)
+        with mock.patch("builtins.open", mock_open):
+            with self.assertRaises(ValueError):
+                app.KeyExchange.from_dict(msg)
+
+    def test_KeyExchange_fails_bad_domain(self):
+        """Tests creating KeyExchange fails due to bad key."""
+        key = "AAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAAE="
+        msg = dict(public_key=key, domain="unconfigured_domain")
+        mock_open = mock.mock_open(read_data=_VALID_CFG)
+        with mock.patch("builtins.open", mock_open):
+            with self.assertRaises(ValueError):
+                app.KeyExchange.from_dict(msg)
+
+
+if __name__ == "__main__":
+    unittest.main()