Generate WireGuard IP Addresses from public key, compatible with wg-ip bash script

wg-ip.py

#!/usr/bin/env python3
# need at least python3.6+ for blake2
from base64 import b64decode
from hashlib import sha256, blake2s
from ipaddress import ip_address, ip_network


# https://github.com/chmduquesne/wg-ip
def gen_ip(pubkey, subnet=ip_network('fe80::/64')):
    """Generate wg-ip compatible addresses from WireGuard public key.
    Uses sha256 over base64 encoded newline terminated publickey string.
    """
    prefix_bytes = subnet.network_address.packed
    mask_bytes = subnet.netmask.packed
    suffix_bytes = sha256(pubkey.encode('ascii')+b'\n').digest()
    address = b''
    for prefix, suffix, mask in zip(prefix_bytes, suffix_bytes, mask_bytes):
        address += ((prefix & mask) | (suffix & (mask^255))).to_bytes(1, byteorder='big')
    return ip_address(address)


# https://gist.github.com/reidrankin/3a39210ce437680f5cf1ac549fd1f1ff
def gen_lla(pubkey, subnet=ip_network('fe80::/10')):
    """Calculates cryptographically-bound IPv6 Link-Local Addresses from WireGuard public keys.
    Uses blake2s hash over public key as bytes.
    """
    prefix_bytes = subnet.network_address.packed
    mask_bytes = subnet.netmask.packed
    suffix_bytes = blake2s(pubkey).digest()
    address = b''
    for prefix, suffix, mask in zip(prefix_bytes, suffix_bytes, mask_bytes):
        address += ((prefix & mask) | (suffix & (mask^255))).to_bytes(1, byteorder='big')
    return ip_address(address)


if __name__ == '__main__':
    print("# sha256 based address calculation")
    pubkey = "foo"
    subnet = ip_network("fd1a:6126:2887::/48")
    exp_result = ip_address("fd1a:6126:2887:f9b1:d61e:21e7:96d7:8dcc")
    real_result = gen_ip(pubkey, subnet)
    print(exp_result, real_result, exp_result==real_result)

    print("# blake2s based address calculation")
    pubkey = b"\x00"*32  # 32byte public key fo all zeroes
    exp_result = ip_address("fe8b:5ea9:9e65:3bc2:b593:db41:30d1:0a4e")
    real_result = gen_lla(pubkey)
    print(exp_result, real_result, exp_result==real_result)

    pubkey = b64decode("hTvlXzX5ZoTg6BxbWHYWSZo6pkLGPdQwVoXZYqrMMgs=")
    exp_result = ip_address("fe80:3fe9:4d26:adc4:1953:7b20:314a:3167")
    real_result = gen_lla(pubkey)
    print(exp_result, real_result, exp_result==real_result)