Add wireguard python example

Ref. eng/recordflux/RecordFlux#1565

Makefile

@@ -512,6 +512,7 @@ $(VSIX):
 
 clean:
 	rm -rf $(BUILD_DIR) .coverage .coverage.* .hypothesis .mypy_cache .pytest_cache .ruff_cache doc/language_reference/build doc/user_guide/build
+	$(MAKE) -C examples/apps/wireguard clean
 	$(MAKE) -C examples/apps/ping clean
 	$(MAKE) -C examples/apps/dhcp_client clean
 	$(MAKE) -C examples/apps/spdm_responder clean

examples/apps/wireguard/Makefile

@@ -0,0 +1,19 @@
+include ../../../Makefile.common
+
+SPECS = $(wildcard specs/*.rflx)
+CWD := $(dir $(abspath $(firstword $(MAKEFILE_LIST))))
+RUST_BINDING = tai64_bindings/target/debug/libtai64_bindings.so
+
+.PHONY: test generate clean
+
+$(RUST_BINDING):
+	$(MAKE) -C tai64_bindings build
+
+test: $(RUST_BINDING)
+	$(POETRY) run ../../../devutils/linux/run $(CWD) timeout 180 tests/wireguard_test.sh $$(dirname $$(command -v python3))
+
+generate:
+
+clean:
+	$(RM) -r build
+	$(MAKE) -C tai64_bindings clean

examples/apps/wireguard/README.md

@@ -0,0 +1,36 @@
+# WireGuard
+
+This example shows how to use `PyRFLX` to implement a subset version of [WireGuard](https://www.wireguard.com/papers/wireguard.pdf).
+The example app creates an encrypted tunnel between `10.1.0.1` (PyRFLX) and `10.2.0.1` (Linux WireGuard).
+The tunnel is tested using the `ping` utility.
+Note that `IP` packets and `ICMP` packets are also parsed by RecordFlux.
+
+## Rust
+
+This example also demonstrates how to integrate Rust code with `PyRFLX`.
+The [tai64](https://docs.rs/tai64/latest/tai64/) crate is used to parse and create [TAI64N timestamps](https://cr.yp.to/libtai/tai64.html#tai64n).
+
+## How to use?
+
+### Build and install the tai64 binding
+
+Make sure you're working from a Python virtual environment (e.g `poetry shell`).
+
+```command
+$ make -C tai64_bindings build
+```
+
+To make sure that the binding is installed the following command should run without errors.
+```command
+$ python -c 'import tai64_bindings'
+```
+
+### Setup WireGuard
+
+To set up a WireGuard interface, you can look at [tests/wireguard_test.sh](./tests/wireguard_test.sh).
+It follows the usual WireGuard interface creation process, including the generation of keys, allowed IPs, routes, etc.
+
+You can now run the `wireguard.py` script and ping the `PyRFLX` side of the tunnel using ping (assuming that the configured python side of the tunnel is `10.2.0.2`):
+```command
+$ ping -c 5 10.2.0.1
+```

examples/apps/wireguard/const_values.py

@@ -0,0 +1,4 @@
+CONSTRUCTION = b"Noise_IKpsk2_25519_ChaChaPoly_BLAKE2s"
+IDENTIFIER = b"WireGuard v1 zx2c4 [email protected]"
+LABEL_MAC1 = b"mac1----"
+WIREGUARD_PID_FILE = "wireguard_py.pid"

examples/apps/wireguard/crypto.py

@@ -0,0 +1,77 @@
+"""
+Wireguard crypto functions.
+
+Cryptography functions as described in https://www.wireguard.com/papers/wireguard.pdf
+in section 5.4.
+"""
+
+from __future__ import annotations
+
+import hashlib
+from hmac import HMAC
+
+from cryptography.hazmat.primitives.asymmetric.x25519 import X25519PrivateKey, X25519PublicKey
+from cryptography.hazmat.primitives.ciphers.aead import ChaCha20Poly1305
+from nacl.public import PrivateKey
+
+
+def blk2s_hash(*args: bytes) -> bytes:
+    h = hashlib.blake2s(digest_size=32)
+    for a in args:
+        h.update(a)
+    return h.digest()
+
+
+def dh(private_key: bytes, public_key: bytes) -> bytes:
+    x25519_private_key = X25519PrivateKey.from_private_bytes(private_key)
+    peer_public_key = X25519PublicKey.from_public_bytes(public_key)
+    return x25519_private_key.exchange(peer_public_key)
+
+
+def dh_generate() -> tuple[bytes, bytes]:
+    random_key = PrivateKey.generate()
+    return bytes(random_key), bytes(random_key.public_key)
+
+
+def _create_chacha20poly1305_nonce(counter: int) -> bytes:
+    # From the wireguard spec: nonce are composed of 32 bits of zeros followed by
+    # the 64-bits little-endian value of counter
+    return (b"\x00" * 4) + counter.to_bytes(length=8, byteorder="little")
+
+
+def chacha20poly1305_encrypt(
+    key: bytes,
+    counter: int,
+    plain_text: bytes,
+    auth_text: bytes,
+) -> bytes:
+    chacha = ChaCha20Poly1305(key)
+    nonce = _create_chacha20poly1305_nonce(counter)
+    return chacha.encrypt(nonce, plain_text, auth_text)
+
+
+def chacha20poly1305_decrypt(
+    key: bytes,
+    counter: int,
+    plain_text: bytes,
+    auth_text: bytes,
+) -> bytes:
+    chacha = ChaCha20Poly1305(key)
+    nonce = _create_chacha20poly1305_nonce(counter)
+    return chacha.decrypt(nonce, plain_text, auth_text)
+
+
+def mac(key: bytes, value: bytes) -> bytes:
+    h = hashlib.blake2s(key=key, digest_size=16)
+    h.update(value)
+    return h.digest()
+
+
+def hmac(key: bytes, *values: bytes) -> bytes:
+    h = HMAC(key, digestmod=hashlib.blake2s)
+    h.digest_size = 32
+
+    for value in values:
+        h.update(value)
+
+    return h.digest()

examples/apps/wireguard/handlers.py

@@ -0,0 +1,151 @@
+from __future__ import annotations
+
+import dataclasses
+import math
+import os
+from pathlib import Path
+from typing import TYPE_CHECKING, Callable
+
+from rflx.model import NeverVerify
+from rflx.pyrflx import MessageValue, PyRFLX
+
+if TYPE_CHECKING:
+    from .const_values import WIREGUARD_PID_FILE
+    from .crypto import blk2s_hash, chacha20poly1305_decrypt, chacha20poly1305_encrypt, dh, hmac
+    from .ping import handle_ping
+else:
+    from const_values import WIREGUARD_PID_FILE
+    from crypto import blk2s_hash, chacha20poly1305_decrypt, chacha20poly1305_encrypt, dh, hmac
+    from ping import handle_ping
+
+PYRFLX = PyRFLX.from_specs(["specs/wireguard.rflx"], NeverVerify())
+WIREGUARD = PYRFLX.package("Wireguard")
+
+
+@dataclasses.dataclass()
+class Connection:
+    static_public_key: bytes
+    static_private_key: bytes
+    session_id: bytes | None = None
+    receiver_id: bytes | None = None
+    initiator_chain: bytes | None = None
+    initiator_hash: bytes | None = None
+    ephemeral_public: bytes | None = None
+    ephemeral_private: bytes | None = None
+    receiver_public_key: bytes | None = None
+    receving_key: bytes | None = None
+    sending_key: bytes | None = None
+    sending_index: int = 0
+    packet_counter: int = 0
+
+
+HandlerType = Callable[[Connection, MessageValue], bytes | None]
+
+
+def _create_transport_data_message(conn: Connection, clear_data: bytes) -> bytes:
+    packet = WIREGUARD.new_message("Handshake")
+    packet.set("Message_Type", "Transport_Data_Message")
+    packet.set("Reserved", 0)
+    assert conn.receiver_id is not None
+    packet.set("Receiver", conn.receiver_id)
+    conn.packet_counter += 1
+    packet.set("Counter", conn.packet_counter.to_bytes(8, "little"))
+    zero_padding = b"\x00" * (16 * math.ceil(len(clear_data) / 16) - len(clear_data))
+    raw_data = clear_data + zero_padding
+    assert conn.sending_key is not None
+    packet.set(
+        "Packet",
+        chacha20poly1305_encrypt(conn.sending_key, conn.packet_counter, raw_data, b""),
+    )
+    return packet.bytestring
+
+
+def _announce_readiness() -> None:
+    """
+    Announce readiness.
+
+    Write the process pid to `WIREGUARD_PID_FILE` to indicate to the outside world that
+    the handshake is complete and we're ready to handle data packet.
+    """
+    Path(WIREGUARD_PID_FILE).write_text(str(os.getpid()))
+
+
+def handle_handshake_init(_conn: Connection, _packet: MessageValue) -> bytes | None:
+    raise NotImplementedError
+
+
+def handle_handshake_response(conn: Connection, packet: MessageValue) -> bytes | None:
+    # We don't support optional preshared key.
+    preshared_key = b"\x00" * 32
+
+    uncrypted_ephemeral = packet.get("Ephemeral")
+    assert conn.initiator_hash is not None
+    assert conn.initiator_chain is not None
+    assert conn.ephemeral_private is not None
+    assert isinstance(uncrypted_ephemeral, bytes)
+
+    conn.initiator_hash = blk2s_hash(conn.initiator_hash, uncrypted_ephemeral)
+    temp = hmac(conn.initiator_chain, uncrypted_ephemeral)
+    chaining_key = hmac(temp, b"\x01")
+    ephemeral_shared = dh(conn.ephemeral_private, uncrypted_ephemeral)
+    temp = hmac(chaining_key, ephemeral_shared)
+    chaining_key = hmac(temp, b"\x01")
+    temp = hmac(chaining_key, dh(conn.static_private_key, uncrypted_ephemeral))
+    chaining_key = hmac(temp, b"\x01")
+    temp = hmac(chaining_key, preshared_key)
+    chaining_key = hmac(temp, b"\x01")
+    temp2 = hmac(temp, chaining_key, b"\x02")
+    key = hmac(temp, temp2, b"\x03")
+    conn.initiator_hash = blk2s_hash(conn.initiator_hash, temp2)
+    empty_field_bytes = packet.get("Empty")
+    assert isinstance(empty_field_bytes, bytes)
+    chacha20poly1305_decrypt(key, 0, empty_field_bytes, conn.initiator_hash)
+
+    ck_hash = hmac(chaining_key, b"")
+    sending_key = hmac(ck_hash, b"\x01")
+    receiving_key = hmac(ck_hash, sending_key, b"\x02")
+
+    conn.receving_key = receiving_key
+    conn.sending_key = sending_key
+    conn.initiator_chain = chaining_key
+    sender = packet.get("Sender")
+    assert isinstance(sender, bytes)
+    conn.receiver_id = sender
+
+    # Tell the world that we're ready
+    _announce_readiness()
+
+    # Send a keep-alive packet
+    return _create_transport_data_message(conn, b"")
+
+
+def handle_transport_data_message(conn: Connection, packet: MessageValue) -> bytes | None:
+    # We don't check receiver index here. We assume only a single peer.
+    counter_bytes = packet.get("Counter")
+    assert isinstance(counter_bytes, bytes)
+    counter = int.from_bytes(counter_bytes, "little")
+
+    encrypted_data = packet.get("Packet")
+    assert isinstance(encrypted_data, bytes)
+    assert conn.receving_key is not None
+    decrypted_data = chacha20poly1305_decrypt(conn.receving_key, counter, encrypted_data, b"")
+
+    # Keep-alive packet, do nothing
+    if decrypted_data == b"":
+        return None
+
+    ping_reply = handle_ping(decrypted_data)
+
+    return _create_transport_data_message(conn, ping_reply) if ping_reply is not None else None
+
+
+def handle_cookie(*_: object) -> bytes | None:
+    raise NotImplementedError
+
+
+PACKET_HANDLERS: dict[str, HandlerType] = {
+    "Wireguard::Handshake_Init": handle_handshake_init,
+    "Wireguard::Handshake_Response": handle_handshake_response,
+    "Wireguard::Cookie": handle_cookie,
+    "Wireguard::Transport_Data_Message": handle_transport_data_message,
+}