Repository: freifunkMUC/wgkex (language: python)

Issue:
Added small stupid json interface



Discussion / hints:

Fix patch:
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..113997b 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()