#!/usr/bin/env python3
# SPDX-License-Identifier: LGPL-2.1+
# ~~~
#   Description: Tests for keepalived - Simple and robust facilities for loadbalancing and high-availability to Linux system
#
#   Author: Susant Sahani <susant@redhat.com>
#   Copyright (c) 2018 Red Hat, Inc.
# ~~~

import errno
import os
import sys
import time
import unittest
import subprocess
import signal
import shutil
import psutil
import socket
from pyroute2 import NetNS
from pyroute2 import IPDB

KEEPALIVED_MASTER_CONFIG_FILE='/var/run/keepalived-ci/keepalived-master.conf'
KEEPALIVED_MASTER_PID_FILE='/var/run/keepalived-ci/test-keepalived-master.pid'
KEEPALIVED_MASTER_LOG_FILE='/var/run/keepalived-ci/test-keepalived-master.log'

KEEPALIVED_SLAVE_CONFIG_FILE='/var/run/keepalived-ci/keepalived-slave.conf'
KEEPALIVED_SLAVE_PID_FILE='/var/run/keepalived-ci/test-keepalived-slave.pid'
KEEPALIVED_SLAVE_LOG_FILE='/var/run/keepalived-ci/test-keepalived-slave.log'

#                                            |-------namespace-----|
# MASTER<-vethtest---------------------------|-----vethpeer->SLAVE |
#--------------------------------------------|---------------------|
# Floated IP: 192.168.111.52

def setUpModule():
    """Initialize the environment, and perform sanity checks on it."""

    if shutil.which('keepalived') is None:
        raise OSError(errno.ENOENT, 'keepalived not found')

    if shutil.which('nc') is None:
        raise OSError(errno.ENOENT, 'nc not found')

    if shutil.which('socat') is None:
        raise OSError(errno.ENOENT, 'socat not found')

    if subprocess.call(['systemctl', 'is-active', '--quiet',
                        'keepalived.service']) == 0:
        raise unittest.SkipTest('keepalived is already active')

def tearDownModule():
        pass

class GenericUtilities():
    """Provide a set of utility functions start stop daemons. write config files etc """

    def StartKeepAlivedMasterHost(self):
        """ start keepalived in the host master mode """
        subprocess.check_output(['systemctl', 'start', 'keepalived-master.service'])
        output=subprocess.check_output(['systemctl', 'status', 'keepalived-master.service']).rstrip().decode('utf-8')
        print(output)

    def StopKeepAlivedMasterHost(self):
        """ stop keepalived in the host master mode """
        subprocess.check_output(['systemctl', 'stop', 'keepalived-master.service'])

    def StartKeepAlivedSlaveNameSpace(self):
        """start keepalived in the namespace slave mode """
        subprocess.check_output(['ip', 'netns', 'exec', 'test-keepalived-ns', 'systemctl', 'start', 'keepalived-slave.service'])
        output=subprocess.check_output(['ip', 'netns', 'exec', 'test-keepalived-ns', 'systemctl', 'status', 'keepalived-slave.service']).rstrip().decode('utf-8')
        print(output)

    def StopKeepAlivedSlaveNameSpace(self):
        """stop keepalived in the namespace slave mode """
        subprocess.check_output(['ip', 'netns', 'exec', 'test-keepalived-ns', 'systemctl', 'stop', 'keepalived-slave.service'])

    def StartNetCatHost(self):
        """Start netcat Master """
        subprocess.check_output(['systemctl', 'start', 'nc-keepalived-test.service'])

    def StopNetCatHost(self):
        """Stop netcat Master """
        subprocess.check_output(['systemctl', 'stop', 'nc-keepalived-test.service'])

    def StartSocatNameSpace(self):
        """Start socat backup """
        subprocess.check_output(['ip', 'netns', 'exec', 'test-keepalived-ns', 'systemctl', 'start', 'socat-keepalived-test.service'])

    def StopSocatNameSpace(self):
        """Stop socat backup """
        subprocess.check_output(['ip', 'netns', 'exec', 'test-keepalived-ns', 'systemctl', 'stop', 'socat-keepalived-test.service'])

    def SetupNetwork(self):
        """Setup veth interface and namespace """

        # start the main network settings database:
        ipdb_main = IPDB()
        # start the same for a netns:
        ipdb_test = IPDB(nl=NetNS('test-keepalived-ns'))

        # create VETH
        ipdb_main.create(ifname='vethtest', kind='veth', peer='vethpeer').commit()

        # move peer VETH into the netns
        with ipdb_main.interfaces.vethpeer as veth:
            veth.net_ns_fd = 'test-keepalived-ns'

        with ipdb_main.interfaces.vethtest as veth:
            veth.add_ip('192.168.111.50/24')
            veth.up()

        with ipdb_test.interfaces.vethpeer as veth:
            veth.add_ip('192.168.111.51/24')
            veth.up()

    def TearDownNetwork(self):
        netns = NetNS('test-keepalived-ns')
        netns.close()
        netns.remove()

class KeepalivedTests(unittest.TestCase, GenericUtilities):

    def setUp(self):
        self.SetupNetwork()

    def tearDown(self):
        self.StopKeepAlivedMasterHost()
        self.StopKeepAlivedSlaveNameSpace()

        self.TearDownNetwork()

        os.remove(KEEPALIVED_MASTER_LOG_FILE)
        os.remove(KEEPALIVED_SLAVE_LOG_FILE)

    def test_keepalived_failover(self):
        """ Test Floating IP address is asigned to backup incase of failover (High Availability (HA)) """

        self.StartNetCatHost()
        self.StartSocatNameSpace()

        time.sleep(5)
        self.StartKeepAlivedMasterHost()
        self.StartKeepAlivedSlaveNameSpace()
        time.sleep(15)

        print('\n----------------------------------------------------')
        print('Floated IP 192.168.111.52 should be at master vethtest')
        print('----------------------------------------------------')

        output=subprocess.check_output(['ip','address', 'show', 'vethtest']).rstrip().decode('utf-8')
        print(output)
        self.assertRegex(output, "192.168.111.50")
        self.assertRegex(output, "192.168.111.52")

        print('----------------------------------------------------')
        output=subprocess.check_output(['ip', 'netns', 'exec', 'test-keepalived-ns', 'ip','address', 'show', 'vethpeer']).rstrip().decode('utf-8')
        print(output)
        self.assertRegex(output, "192.168.111.51")

        # start failover
        print('---- *****----------- start failover ----------*****-----------')
        self.StopNetCatHost()
        time.sleep(15)
        output=subprocess.check_output(['ip', 'netns', 'exec', 'test-keepalived-ns', 'ip', 'address', 'show', 'vethpeer']).rstrip().decode('utf-8')
        print('----------------------------------------------------')
        print('Floated IP 192.168.111.52 should be at slave which is now master vethpeer')
        print('----------------------------------------------------')

        print(output)
        self.assertRegex(output, "192.168.111.52")

        print('---------------Floated ip 192.168.111.52 should not be at vethtest-----------')
        print('----------------------------------------------------------------------------')
        output=subprocess.check_output(['ip','address', 'show', 'vethtest']).rstrip().decode('utf-8')
        print(output)

        self.assertNotRegex(output, "192.168.111.52")

        self.StopSocatNameSpace()

if __name__ == '__main__':
    unittest.main(testRunner=unittest.TextTestRunner(stream=sys.stdout,
                                                     verbosity=3))