#!/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()