From f629a27405af41c735ce545fe6ba19cffe1894f5 Mon Sep 17 00:00:00 2001
From: =?utf8?q?Alejandro=20R=2E=20Sede=C3=B1o?= <asedeno@mit.edu>
Date: Sun, 22 Mar 2020 11:26:27 -0400
Subject: [PATCH 1/1] initial commit

---
 .gitignore           |   2 +
 requirements.txt     |   1 +
 wg_llv6.py           | 100 +++++++++++++++++++++++++++++++++++++++++++
 wg_pubkey_to_llv6.py |  43 +++++++++++++++++++
 4 files changed, 146 insertions(+)
 create mode 100644 .gitignore
 create mode 100644 requirements.txt
 create mode 100755 wg_llv6.py
 create mode 100755 wg_pubkey_to_llv6.py

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..63b4e5f
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,2 @@
+__pycache__
+ve
diff --git a/requirements.txt b/requirements.txt
new file mode 100644
index 0000000..f9e8e51
--- /dev/null
+++ b/requirements.txt
@@ -0,0 +1 @@
+pyroute2==0.5.10
diff --git a/wg_llv6.py b/wg_llv6.py
new file mode 100755
index 0000000..ac42e14
--- /dev/null
+++ b/wg_llv6.py
@@ -0,0 +1,100 @@
+#!/usr/bin/env python
+
+import base64
+import ipaddress
+import os
+import sys
+
+from pyroute2 import IPDB, WireGuard
+
+from wg_pubkey_to_llv6 import pubkey_to_lladdr
+
+
+def fix_addr(addr_str):
+    """This is awful. It fixes wrongly-formatted IPv6 addresses that come back from pyroute2.
+    All 16 bytes are returned, with a ':' between each.
+    e.g., fe:80:00:00:.../128 instead of fe80::.../128"""
+    if '.' in addr_str:
+        # IPv4 is rendered fine.
+        return addr_str
+
+    # Rebuild IPv6 address in proper form.
+    addr, mask = addr_str.split('/')
+    addr = str(ipaddress.IPv6Address(b''.join(bytes.fromhex(x) for x in addr.split(':'))))
+    return '/'.join([addr, mask])
+
+
+def peer_to_spec(peer):
+    """Partially reconstruct the peer definition. We only care about the
+    public key [identity] and the allowed ipaddresses [to be updated]."""
+    return {
+        'public_key': peer.WGPEER_A_PUBLIC_KEY['value'],
+        'allowed_ips': [fix_addr(x['addr']) for x in peer.WGPEER_A_ALLOWEDIPS['value']],
+    }
+
+
+def add_llv6_addresses_to_interface(ifname):
+    # Local state.
+    addresses = set()
+    collisions = set()
+    address_map = {}
+    wg = WireGuard()
+    wg_iface = wg.info(ifname)[0]
+
+    # Calculate interface IPv6 link-local address.
+    dev_pk = wg_iface.WGDEVICE_A_PUBLIC_KEY['value']
+    dev_lla = pubkey_to_lladdr(dev_pk)
+    addresses.add(dev_lla)
+    address_map[dev_pk] = dev_lla
+
+    # Calculate peer IPv6 link-local addresses.
+    peers = [peer_to_spec(peer) for peer in wg_iface.WGDEVICE_A_PEERS.value]
+    for peer in peers:
+        peer_pk = peer['public_key']
+        peer_lla = pubkey_to_lladdr(peer_pk)
+        if peer_lla in addresses:
+            collisions.add(peer_lla)
+        addresses.add(peer_lla)
+        address_map[peer_pk] = peer_lla
+
+    # Did we find collisions?
+    if collisions:
+        print('Peer link local address collisions:')
+        for pk, lla in addresses.items:
+            if lla in collision:
+                print(pk, lla)
+        if dev_lla in collisions:
+            raise Exception('Our link local address collides with a peer. Giving up.')
+        else:
+            print('Proceeding with link-local assignments, skipping colliding peers.')
+
+    # Add IPv6 link-local address and subnet to interface.
+    with IPDB() as ip:
+        iface = ip.interfaces[ifname]
+        iface.add_ip(str(dev_lla)+'/64')
+        iface.commit()
+
+    # Add IPv6 link-local addresses to peers without address collisions.
+    for peer in peers:
+        peer_pk = peer['public_key']
+        peer_lla = str(address_map[peer_pk])+'/128'
+        if peer_lla not in peer['allowed_ips'] and peer_lla not in collisions:
+            peer['allowed_ips'].append(peer_lla)
+            wg.set(ifname, peer=peer)
+
+
+def main():
+    ifname = os.getenv('IFACE')
+    if not ifname:
+        try:
+            ifname = sys.argv[1]
+        except IndexError:
+            pass
+    if not ifname:
+        print('No interface specified; quitting.')
+        return
+    add_llv6_addresses_to_interface(ifname)
+
+
+if __name__ == '__main__':
+    main()
diff --git a/wg_pubkey_to_llv6.py b/wg_pubkey_to_llv6.py
new file mode 100755
index 0000000..ae018fe
--- /dev/null
+++ b/wg_pubkey_to_llv6.py
@@ -0,0 +1,43 @@
+#!/usr/bin/env python
+
+import base64
+import ipaddress
+import sys
+
+try:
+    from hashlib import blake2s
+except ImportError:
+    try:
+        from pyblake2 import blake2s
+    except ImportError:
+        print('No native BLAKE2 support; please pip install pyblake2')
+        raise
+
+LL_NET_PART = int(ipaddress.IPv6Address('fe80::'))
+        
+def pubkey_to_lladdr(pubkey):
+    # Check / normalize public key to raw bytes.
+    if isinstance(pubkey, str):
+        pubkey = base64.b64decode(pubkey)
+    elif isinstance(pubkey, bytes):
+        if len(pubkey) != 32:
+            pubkey = base64.b64decode(pubkey)
+    else:
+        raise TypeError()
+    assert len(pubkey) == 32
+
+    # Hash and generate interface part
+    pk_hash = blake2s(pubkey, digest_size=8)
+    node = int.from_bytes(pk_hash.digest(), 'big')
+
+    # Combine with link-local net part.
+    addr = ipaddress.IPv6Address(LL_NET_PART + node)
+    return addr
+
+def main():
+    pubkey = sys.argv[1]
+    addr = pubkey_to_lladdr(pubkey)
+    print(str(addr))
+
+if __name__ == '__main__':
+    main()
-- 
2.45.1