#!/opt/slapos/parts/python3/bin/python3
# Copyright (C) 2023-2024  Nexedi SA and Contributors.
#
# This program is free software: you can Use, Study, Modify and Redistribute
# it under the terms of the GNU General Public License version 3, or (at your
# option) any later version, as published by the Free Software Foundation.
#
# You can also Link and Combine this program with other software covered by
# the terms of any of the Free Software licenses or any of the Open Source
# Initiative approved licenses and Convey the resulting work. Corresponding
# source of such a combination shall include the source code for all other
# software used.
#
# This program is distributed WITHOUT ANY WARRANTY; without even the implied
# warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.
#
# See COPYING file for full licensing terms.
# See https://www.nexedi.com/licensing for rationale and options.
"""Program format-ims splits each slaptun into two, adds a network
namespace for IMS and configures a tun on that namespace
"""

# TODO Relying on format-ims should be replaced by slapos format.
#      See discussion at
#      https://lab.nexedi.com/nexedi/slapos/merge_requests/1471#note_194356
#      for details.

import os
import sys
try:
    from packaging.version import Version
except ModuleNotFoundError as e:
    from distutils.version import StrictVersion as Version

p = []
for module in ['netifaces', 'netaddr']:
    l = []
    for directory in os.listdir('/opt/slapos/eggs'):
        if directory.startswith(module):
            l.append(directory)
    p.append('/opt/slapos/eggs/' + sorted(l, key=lambda x:Version(x.split('%s-' % (module))[1].split('-')[0]))[-1])
sys.path[0:0] = p

os.environ['PATH'] = os.path.expandvars('/opt/slapos/parts/bison/bin:/opt/slapos/parts/bzip2/bin:/opt/slapos/parts/gettext/bin:/opt/slapos/parts/glib/bin:/opt/slapos/parts/libxml2/bin:/opt/slapos/parts/libxslt/bin:/opt/slapos/parts/m4/bin:/opt/slapos/parts/ncurses/bin:/opt/slapos/parts/openssl/bin:/opt/slapos/parts/pkgconfig/bin:/opt/slapos/parts/python3/bin:/opt/slapos/parts/readline/bin:/opt/slapos/parts/sqlite3/bin:/opt/slapos/parts/swig/bin:/opt/slapos/bin:/opt/slapos/parts/patch/bin:/opt/slapos/parts/socat/bin:$PATH')

import logging
import netifaces
import netaddr
from socket import AF_INET, AF_INET6
from math import log2, ceil

import sys
import subprocess
import json

# LinkDB represents snapshot of state of all network interfaces.
class LinkDB:
    def __init__(db):
        db.linkv = ip('link', 'show')

    # ifget returns information about interface with specified name.
    def ifget(db, ifname):
        for link in db.linkv:
            if link['ifname'] == ifname:
                return link
        raise KeyError('interface %r not found' % ifname)

def main():

    open(sys.argv[1], 'w+')
    run('sed', '-i', '50001,$ d', sys.argv[1])

    logger = logging.getLogger(__name__)
    logging.basicConfig(filename=sys.argv[1], level=logging.DEBUG)
    logging.basicConfig(format='%(asctime)s %(levelname)s %(message)s',
                        level=logging.DEBUG)

    for link in ip('link', 'list'):
        tun = link['ifname']
        if tun.startswith('slaptun'):
            try:
                int(tun.split('slaptun')[1])
                formatTun(tun)
            except (ValueError, IndexError) as e:
                pass

def formatTun(tun):

    # determine tun's network address and owner
    ldb = LinkDB()
    _ = ldb.ifget(tun)
    owner = _['linkinfo']['info_data']['user']
    try:
        net   = ifnet6(tun)
        net4  = ifnet4(tun)
    except RuntimeError as e:
        return False

    logging.info('%s: configure %s' % (tun, net))

    # do the split
    # with leaving first range for the original tun
    subtun_set = set()
    net4split = netsplit(net4, 2)
    for i, subnet in enumerate(netsplit(net, 2)):
        subnet4 = netaddr.IPNetwork("%s/%d" % (
                  net4split[i].ip + 1, net4split[i].prefixlen))
        if i == 0:
            logging.info('preserve         %s' % subnet)
            continue    # leave this range for original tun

        subtun = '%s-%d' % (tun, i)
        subtun_set.add(subtun)
        logging.info('-> %s   %s %s' % (subtun, subnet, subnet4))
        def note(msg):
            logging.info(' # %s: %s' % (subtun, msg))

        # create subtun
        try:
            link = ldb.ifget(subtun)
        except KeyError:
            run('ip', 'tuntap', 'add', 'dev', subtun, 'mode', 'tun', 'user', owner)
            link = ip('link', 'show', 'dev', subtun)[0]
        else:
            note('already exists')

        # set it up
        if 'UP' not in link['flags']:
            run('ip', 'link',   'set', subtun, 'up')
        else:
            note('already up')

        # add subnet address
        addrv = []
        for _ in ip('-6', 'addr', 'show', 'dev', subtun):
            addrv.extend(_['addr_info'])
        for addr in addrv:
            _ = netaddr.IPNetwork('%s/%s' % (addr['local'], addr['prefixlen']))
            if _ == subnet and addr['noprefixroute']:
                note('already has %s addr' % str(subnet))
                break
        else:
            run('ip', 'addr',   'add', str(subnet), 'dev', subtun, 'noprefixroute')

        # add subnet4 address
        addrv = []
        for _ in ip('-4', 'addr', 'show', 'dev', subtun):
            addrv.extend(_['addr_info'])
        for addr in addrv:
            _ = netaddr.IPNetwork('%s/%s' % (addr['local'], addr['prefixlen']))
            if _ == subnet4 and addr['noprefixroute']:
                note('already has %s addr' % str(subnet4))
                break
        else:
            run('ip', 'addr',   'add', str(subnet4), 'dev', subtun, 'noprefixroute')

        # add /128 route to subnet::1
        rtv = ip('-6', 'route', 'show', 'dev', subtun)
        for rt in rtv:
            if rt['dst'] == str(subnet[1])  and  'gateway' not in rt:
                note('already has %s route' % str(subnet[1]))
                break
        else:
            run('ip', 'route',  'add', str(subnet[1]), 'dev', subtun)

        # add route to subnet via subnet::1
        for rt in rtv:
            if rt['dst'] == str(subnet)  and  rt.get('gateway') == str(subnet[1]):
                note('already has %s route' % str(subnet))
                break
        else:
            run('ip', 'route',  'add', str(subnet), 'dev', subtun, 'via', str(subnet[1]))

        # configure IMS network namespace and tun
        if i == 1:
            tun_ims = '%s-ims' % (tun)
            netns_ims = '%s-netns' % (tun_ims)
            netns_ip4 = '%s/0' % (str(subnet4).split('/')[0])
            netns_ip6 = '%s/0' % (str(subnet).split('/')[0])

            try:
                link = ldb.ifget(tun_ims)
            except KeyError:
                run('ip', 'tuntap', 'add', 'dev', tun_ims, 'mode', 'tun', 'user', owner)
                link = ip('link', 'show', 'dev', tun_ims)[0]

            for _ in ip('netns', 'list'):
                if _['name'] == netns_ims:
                    break
            else:
                run('ip', 'netns', 'add', netns_ims)

            for _ in ip('netns', 'exec', netns_ims, 'ip', '-json',
                        '-details', 'link', 'list'):
                if _['ifname'] == tun_ims:
                    break
            else:
                run('ip', 'link', 'set', 'dev', tun_ims,
                    'name', tun_ims, 'netns', netns_ims)

            addrv = []
            for _ in ip('netns', 'exec', netns_ims, 'ip', '-json',
                        '-details', '-4', 'addr', 'show', 'dev', tun_ims):
                addrv.extend(_['addr_info'])
            for addr in addrv:
                _ = netaddr.IPNetwork('%s/%s' % (addr['local'], addr['prefixlen']))
                if str(_) == netns_ip4:
                    break
            else:
                run('ip', 'netns', 'exec', netns_ims, 'ip', 'addr',
                    'add', netns_ip4, 'dev', tun_ims)

            addrv = []
            for _ in ip('netns', 'exec', netns_ims, 'ip', '-json',
                        '-details', '-6', 'addr', 'show', 'dev', tun_ims):
                addrv.extend(_['addr_info'])
            for addr in addrv:
                _ = netaddr.IPNetwork('%s/%s' % (addr['local'], addr['prefixlen']))
                if str(_) == netns_ip6:
                    break
            else:
                run('ip', 'netns', 'exec', netns_ims, 'ip', 'addr', 'add', netns_ip6,
                    'dev', tun_ims)

            if ip('netns', 'exec', netns_ims, 'ip', '-json', '-details', 'link', 'show',
                   'dev', tun_ims)[0]['mtu'] != 1400:
                run('ip', 'netns', 'exec', netns_ims, 'ip', 'link', 'set', tun_ims,
                    'up', 'mtu', '1400')

            for _ in ip('netns', 'exec', netns_ims, 'ip', '-json', '-details', 'route',
                         'show', 'dev', tun_ims):
                if _['dst'] == 'default':
                    break
            else:
                    run('ip', 'netns', 'exec', netns_ims, 'ip', 'route', 'add',
                        'default', 'dev', tun_ims)


    # remove other existing children
    for ifname in netifaces.interfaces():
        if ifname.startswith('%s-' % tun) and (ifname not in subtun_set) and 'ims' not in ifname:
            logging.info('-> del %s' % ifname)
            run('ip', 'link', 'del', ifname)
    return True


# netsplit splits network into n subnetworks.
def netsplit(net, n): # -> []subnet
    # see how much prefix bits we need to take to be able to divide by n
    ptake = ceil(log2(n))
    return list( net.subnet(net.prefixlen + ptake) )[:n]

# ifnet6 returns IPv6 network address associated with given interface.
def ifnet6(ifname):
    addr  = None
    net   = None
    prefixlen = None
    for iaddr in netifaces.ifaddresses(ifname)[AF_INET6]:
        a = iaddr['addr']
        if '%' in a: # link-local
            a = a.split('%')[0]
            a = netaddr.IPAddress(a)
            assert a.is_link_local(), a
            continue

        if addr is not None:
            raise RuntimeError('%s: multiple addresses:  %s and %s' % (ifname, addr, a))

        addr = netaddr.IPAddress(a)
        netmask, plen = iaddr['netmask'].split('/')
        prefixlen = int(plen)
        net = netaddr.IPNetwork('%s/%d' % (a, prefixlen))

    if addr is None:
        raise RuntimeError('%s: no non link-local addresses' % ifname)

    # normalize network
    # ex 2401:5180:0:66:a7ff:ffff:ffff:ffff/71  ->  2401:5180:0:66:a600::/71
    net = net.cidr
    return net

# ifnet4 returns IPv4 network address associated with given interface.
def ifnet4(ifname):
    addr  = None
    net   = None
    prefixlen = None
    for iaddr in netifaces.ifaddresses(ifname)[AF_INET]:
        a = iaddr['addr']

        if addr is not None:
            raise RuntimeError('%s: multiple addresses:  %s and %s' % (ifname, addr, a))

        addr = netaddr.IPAddress(a)
        net = netaddr.IPNetwork('%s/%s' % (a, iaddr['netmask']))

    if addr is None:
        raise RuntimeError('%s: no addresses' % ifname)

    return net

# run executes `*argv` as action.
def run(*argv):
    logging.info(' # %s' % ' '.join(argv))
    subprocess.check_call(argv)

# ip returns decoded output of `ip -details *argv`
def ip(*argv):
    _ = subprocess.check_output(['ip', '-json', '-details'] + list(argv))
    return json.loads(_ or '{}')


if __name__ == '__main__':
    main()
