#!/usr/bin/python3
# Test NetworkManager on simulated WIFI devices

__author__ = 'Martin Pitt <martin.pitt@ubuntu.com>'
__copyright__ = '(C) 2013 Canonical Ltd.'
__license__ = 'GPL v2 or later'

import sys
import os
import os.path
import time
import subprocess
import socket
import unittest

from gi.repository import NetworkManager, NMClient, GLib

sys.path.append(os.path.dirname(__file__))
import network_test_base

SSID = 'fake net'

# If True, NetworkManager logs directly to stdout, to watch logs in real time
NM_LOG_STDOUT = os.getenv('NM_LOG_STDOUT', False)

# avoid accidentally destroying any real config
os.environ['GSETTINGS_BACKEND'] = 'memory'

# we currently get a lot of WARNINGs/CRITICALs from GI (leaked objects from
# previous test runs/main loops?) Redirect them to stdout, to avoid failing
# autopkgtests
os.dup2(sys.stdout.fileno(), sys.stderr.fileno())


class NetworkManagerTest(network_test_base.NetworkTestBase):
    '''Provide common functionality for NM tests'''

    def start_nm(self, wait_dev_client=True):
        '''Start NetworkManager and initialize client object

        If wait_dev_client is True (default), wait until NM recognizes
        self.dev_client; otherwise, just wait until NM has initialized (for
        coldplug mode).
        '''
        # create local configuration; this allows us to have full control, and
        # we also need to blacklist the AP device so that NM does not tear it
        # down; we also blacklist any existing real interface to avoid
        # interfering with it, and for getting predictable results
        blacklist = ''
        for iface in os.listdir('/sys/class/net'):
            if iface != self.dev_client:
                with open('/sys/class/net/%s/address' % iface) as f:
                    blacklist += 'mac:%s;' % f.read().strip()

        conf = os.path.join(self.workdir, 'NetworkManager.conf')
        with open(conf, 'w') as f:
            f.write('[main]\nplugins=keyfile\n\n[keyfile]\nunmanaged-devices=%s\n' % blacklist)

        if NM_LOG_STDOUT:
            f_log = None
        else:
            log = os.path.join(self.workdir, 'NetworkManager.log')
            f_log = os.open(log, os.O_CREAT | os.O_WRONLY | os.O_SYNC)

        # don't depend on or change system-wide state file, use a fresh one for
        # each test run
        state = os.path.join(self.workdir, 'NetworkManager.state')
        try:
            os.unlink(state)
        except OSError:
            pass
        p = subprocess.Popen(['NetworkManager', '--log-level=DEBUG', '--no-daemon', '--config=' + conf,
                              '--state-file=%s' % state],
                             stdout=f_log, stderr=subprocess.STDOUT)
        # automatically terminate process at end of test case
        self.addCleanup(p.wait)
        self.addCleanup(p.terminate)

        if NM_LOG_STDOUT:
            # let it initialize, then print a marker
            time.sleep(1)
            print('******* NM initialized *********\n\n')
            self.addCleanup(print, '\n\n******* Shutting down NM *********')
        else:
            self.addCleanup(os.close, f_log)

            # this should be fast, give it 2 s to initialize
            if wait_dev_client:
                self.poll_text(log, '<info> (%s): exported as /org/freedesktop/' % self.dev_client, timeout=20)
            else:
                self.poll_text(log, '<info> Networking is ', timeout=20)

        self.nmclient = NMClient.Client()
        self.assertTrue(self.nmclient.networking_get_enabled())

        # FIXME: This certainly ought to be true, but isn't
        #self.assertTrue(self.nmclient.get_manager_running())

        # determine device objects
        for d in self.nmclient.get_devices():
            if d.props.interface == self.dev_ap:
                self.assertEqual(d.get_driver(), 'mac80211_hwsim')
                self.assertEqual(d.get_hw_address(), self.mac_ap)
                self.nmdev_ap = d
            elif d.props.interface == self.dev_client:
                self.assertEqual(d.get_driver(), 'mac80211_hwsim')
                self.assertEqual(d.get_hw_address(), self.mac_client)
                self.nmdev = d

        self.assertTrue(hasattr(self, 'nmdev_ap'), 'Could not determine AP NM device')
        self.assertTrue(hasattr(self, 'nmdev'), 'Could not determine client NM device')

        self.process_glib_events()

    def tearDown(self):
        if not hasattr(self, 'nmdev') or not self.nmdev:
            return

        # if we don't do this, NMClient's internal state becomes confused when
        # removing connections and throws CRITICALs/WARNINGs like mad, so stop
        # it from trying to fiddle with connections
        self.nmclient.networking_set_enabled(False)

        # remove all created connections
        cons = self.nmdev.get_available_connections()
        self.counter = len(cons)
        if self.counter > 0:
            ml = GLib.MainLoop()

            def cb(con, error, data):
                self.counter -= 1
                if self.counter == 0:
                    ml.quit()

            for con in cons:
                con.delete(cb, None)
            cons = None
            ml.run()

        # verify that NM properly deconfigures the devices
        out = subprocess.check_output(['iw', 'dev', self.dev_client, 'link'],
                                      universal_newlines=True).strip()
        self.assertEqual(out, 'Not connected.')

        out = subprocess.check_output(['ip', 'a', 'show', 'dev', self.dev_client],
                                      universal_newlines=True).strip()
        self.assertRegex(out, 'state DOWN')
        self.assertFalse('inet 192' in out)
        self.assertFalse('inet6 2600' in out)

    @classmethod
    def process_glib_events(klass):
        '''Process pending GLib main loop events'''

        context = GLib.MainContext.default()
        while context.iteration(False):
            pass

    def assertEventually(self, condition, message=None, timeout=50):
        '''Assert that condition function eventually returns True.

        timeout is in deciseconds, defaulting to 50 (5 seconds). message is
        printed on failure.
        '''
        while timeout >= 0:
            self.process_glib_events()
            if condition():
                break
            timeout -= 1
            time.sleep(0.1)
        else:
            self.fail(message or 'timed out waiting for ' + str(condition))

    def wait_ap(self, timeout):
        '''Wait for AccessPoint NM object to appear, and return it'''

        self.assertEventually(lambda: len(self.nmdev.get_access_points()) > 0,
                              'timed out waiting for AP to be detected',
                              timeout=timeout)

        return self.nmdev.get_access_points()[0]


class Coldplug(NetworkManagerTest):
    '''In these tests NM starts after setting up the AP'''

    def test_no_ap(self):
        '''no available access point'''

        self.start_nm()
        self.assertEventually(self.nmclient.networking_get_enabled, timeout=20)

        # state independent properties
        self.assertEqual(self.nmdev.props.device_type, NetworkManager.DeviceType.WIFI)
        self.assertTrue(self.nmdev.props.managed)
        self.assertFalse(self.nmdev.props.firmware_missing)
        self.assertTrue(self.nmdev.props.udi.startswith('/sys/devices/'), self.nmdev.props.udi)

        # get_version() plausibility check
        out = subprocess.check_output(['nmcli', '--version'], universal_newlines=True)
        cli_version = out.split()[-1]
        self.assertTrue(cli_version[0].isdigit())
        self.assertEqual(self.nmclient.get_version(), cli_version)

        # state dependent properties (disconnected)
        self.assertEqual(self.nmdev.get_state(), NetworkManager.DeviceState.DISCONNECTED)
        self.assertEqual(self.nmdev.get_access_points(), [])
        self.assertEqual(self.nmdev.get_available_connections(), [])

    def test_open_b_ip4(self):
        '''Open network, 802.11b, IPv4'''

        self.do_test('hw_mode=b\nchannel=1\nssid=' + SSID, None, 11000)

    def test_open_b_ip6_raonly(self):
        '''Open network, 802.11b, IPv6 with only RA'''

        self.do_test('hw_mode=b\nchannel=1\nssid=' + SSID, 'ra-only', 11000)

    def test_open_b_ip6_dhcp(self):
        '''Open network, 802.11b, IPv6 with DHCP'''

        self.do_test('hw_mode=b\nchannel=1\nssid=' + SSID, '', 11000)

    def test_open_g_ip4(self):
        '''Open network, 802.11g, IPv4'''

        self.do_test('hw_mode=g\nchannel=1\nssid=' + SSID, None, 54000)

    @unittest.skip('FIXME: this makes bitrates fail, and causes timeouts')
    def test_open_a_ip4(self):
        '''Open network, 802.11a, IPv4'''

        self.do_test('hw_mode=a\nchannel=36\nssid=' + SSID, None, 54000)

    @unittest.skip('FIXME: this makes bitrates fail, and causes timeouts')
    def test_open_a_ip6(self):
        '''Open network, 802.11a, IPv6 with only RA'''

        self.do_test('hw_mode=a\nchannel=36\nssid=' + SSID, 'ra-only', 54000)

    def test_wpa1_ip4(self):
        '''WPA1, 802.11g, IPv4'''

        self.do_test('''hw_mode=g
channel=1
ssid=%s
wpa=1
wpa_key_mgmt=WPA-PSK
wpa_pairwise=TKIP
wpa_passphrase=12345678
''' % SSID, None, 54000, '12345678')

    def test_wpa2_ip4(self):
        '''WPA2, 802.11g, IPv4'''

        self.do_test('''hw_mode=g
channel=1
ssid=%s
wpa=2
wpa_key_mgmt=WPA-PSK
wpa_pairwise=CCMP
wpa_passphrase=12345678
''' % SSID, None, 54000, '12345678')

    def test_wpa2_ip6(self):
        '''WPA2, 802.11g, IPv6 with only RA'''

        self.do_test('''hw_mode=g
channel=1
ssid=%s
wpa=2
wpa_key_mgmt=WPA-PSK
wpa_pairwise=CCMP
wpa_passphrase=12345678
''' % SSID, 'ra-only', 54000, '12345678')

    #
    # Common test code
    #

    def do_test(self, hostapd_conf, ipv6_mode, expected_max_bitrate, secret=None):
        self.setup_ap(hostapd_conf, ipv6_mode)
        self.start_nm()

        # on coldplug we expect the AP to be picked out fast
        ap = self.wait_ap(timeout=30)
        self.assertTrue(ap.get_path().startswith('/org/freedesktop/NetworkManager'))
        self.assertEqual(ap.get_mode(), getattr(NetworkManager, '80211Mode').INFRA)
        self.assertEqual(ap.get_max_bitrate(), expected_max_bitrate)
        #self.assertEqual(ap.get_flags(), )

        # should not auto-connect
        self.assertEqual(self.nmclient.get_active_connections(), [])

        # connect to that AP
        (conn, active_conn) = self.connect_to_ap(ap, secret)

        # check NMActiveConnection object
        self.assertEqual([c.get_uuid() for c in self.nmclient.get_active_connections()],
                         [active_conn.get_uuid()])
        self.assertEqual([d.get_udi() for d in active_conn.get_devices()], [self.nmdev.get_udi()])

        self.check_connected_device_config(ipv6_mode)

        # check corresponding NMConnection object
        wireless_setting = conn.get_setting_wireless()
        self.assertEqual(wireless_setting.get_ssid(), SSID.encode())
        self.assertEqual(wireless_setting.get_hidden(), False)
        if secret:
            self.assertEqual(wireless_setting.get_security(), NetworkManager.SETTING_WIRELESS_SECURITY_SETTING_NAME)
        else:
            self.assertEqual(wireless_setting.get_security(), None)
        # for debugging
        #conn.dump()

        self.check_low_level_config(ipv6_mode)

    #
    # Helper methods
    #

    def connect_to_ap(self, ap, secret):
        '''Connect to an NMAccessPoint.

        secret should be None for open networks, and a string with the password
        for WEP/WPA.

        Return (NMConnection, NMActiveConnection) objects.
        '''
        # If we have a secret, supply it to the new connection right away;
        # adding it afterwards with update_secrets() does not work, and we
        # can't implement a SecretAgent as get_secrets() would need to build a
        # map of a map of gpointers to gpointers which is too much for PyGI
        if secret:
            partial_conn = NetworkManager.Connection.new()
            partial_conn.add_setting(NetworkManager.SettingWirelessSecurity.new())
            # FIXME: needs update for other auth types
            partial_conn.update_secrets(NetworkManager.SETTING_WIRELESS_SECURITY_SETTING_NAME,
                                        {'psk': secret})
        else:
            partial_conn = None

        ml = GLib.MainLoop()
        self.cb_conn = None

        def add_activate_cb(client, conn, conn_path, error, data):
            self.cb_conn = conn
            self.cb_error = error
            ml.quit()
        self.nmclient.add_and_activate_connection(partial_conn, self.nmdev, ap.get_path(), add_activate_cb, None)
        ml.run()
        self.assertEqual(self.cb_error, None)
        active_conn = self.cb_conn
        self.cb_conn = None

        conn = self.conn_from_active_conn(active_conn)
        self.assertTrue(conn.verify())

        # verify need_secrets()
        needed_secrets = conn.need_secrets()
        if secret is None:
            self.assertEqual(needed_secrets, (None, []))
        else:
            self.assertEqual(needed_secrets[0], NetworkManager.SETTING_WIRELESS_SECURITY_SETTING_NAME)
            self.assertEqual(type(needed_secrets[1]), list)
            self.assertGreaterEqual(len(needed_secrets[1]), 1)
            # FIXME: needs update for other auth types
            self.assertIn(needed_secrets[1][0], [NetworkManager.SETTING_WIRELESS_SECURITY_PSK])

        # we are usually ACTIVATING at this point; wait for completion
        # TODO: 5s is not enough, argh slow DHCP client
        self.assertEventually(lambda: active_conn.get_state() == NetworkManager.ActiveConnectionState.ACTIVATED,
                              'timed out waiting for %s to get activated' % active_conn.get_connection(),
                              timeout=80)
        self.assertEqual(self.nmdev.get_state(), NetworkManager.DeviceState.ACTIVATED)
        return (conn, active_conn)

    def conn_from_active_conn(self, active_conn):
        '''Get NMConnection object for an NMActiveConnection object'''

        # this sometimes takes a second try, when the corresponding
        # NMConnection object is not yet available
        tries = 3
        while tries > 0:
            self.process_glib_events()
            path = active_conn.get_connection()
            for c in self.nmdev.get_available_connections():
                if c.get_path() == path:
                    return c
            tries -= 1

        self.fail('Could not find NMConnection object for %s' % path)

    def check_connected_device_config(self, ipv6_mode):
        '''Check NMDevice configuration state after being connected'''

        if ipv6_mode is not None:
            # FIXME: why do we need to wait here, if state is already ACTIVATED?
            self.assertEventually(lambda: self.nmdev.get_ip6_config() is not None, timeout=1)
            self.assertEqual(self.nmdev.get_ip4_config(), None)
            conf = self.nmdev.get_ip6_config()
            self.assertNotEqual(conf, None)
            # we expect at least a link-local and a RA prefix or DHCP assigned
            # address
            self.assertGreaterEqual(len(conf.get_addresses()), 2)
            # note, we cannot call IP6Address.get_address(), as that returns a
            # raw gpointer; check address with low-level tools only
        else:
            # FIXME: why do we need to wait here, if state is already ACTIVATED?
            self.assertEventually(lambda: self.nmdev.get_ip4_config() is not None, timeout=1)
            self.assertEqual(self.nmdev.get_ip6_config(), None)
            conf = self.nmdev.get_ip4_config()
            self.assertNotEqual(conf, None)
            self.assertEqual(len(conf.get_addresses()), 1)
            self.assertEqual(socket.ntohl(conf.get_addresses()[0].get_address()) & 0xFFFFFF00,
                             0xC0A80500)  # 192.168.5.x

    def check_low_level_config(self, ipv6_mode):
        '''Check actual hardware state with ip/iw after being connected'''

        # verify link device configuration with iw
        out = subprocess.check_output(['iw', 'dev', self.dev_client, 'link'],
                                      universal_newlines=True)
        self.assertRegex(out, 'Connected to ' + self.mac_ap)
        self.assertRegex(out, 'SSID: ' + SSID)

        # verify IP device configuration with ip
        out = subprocess.check_output(['ip', 'a', 'show', 'dev', self.dev_client],
                                      universal_newlines=True)
        self.assertRegex(out, 'state UP')
        if ipv6_mode is not None:
            if ipv6_mode in ('', 'slaac'):
                # has global address from our DHCP server
                self.assertRegex(out, 'inet6 2600::[0-9a-z]+/')
            else:
                # has address with our prefix and MAC
                self.assertRegex(out, 'inet6 2600::ff:fe00:[0-9a-z]+/64 scope global (?:tentative )?dynamic')
                # has address with our prefix and random IP (Privacy Extension)
                self.assertRegex(out, 'inet6 2600:[0-9a-z:]+/64 scope global temporary (?:tentative )?dynamic')

            # has a link-local address
            self.assertRegex(out, 'inet6 fe80::ff:fe00:[0-9a-z:]+/64 scope link')
        else:
            self.assertRegex(out, 'inet 192.168.5.\d+/24')


class Hotplug(NetworkManagerTest):
    '''In these tests APs are set up while NM is already running'''

    def setUp(self):
        self.start_nm(wait_dev_client=False)

    def test_auto_detect_new_ap(self):
        '''new AP is being detected automatically within 30s'''

        self.setup_ap('hw_mode=b\nchannel=1\nssid=' + SSID, None)
        self.start_nm()
        ap = self.wait_ap(timeout=300)
        # get_ssid returns a byte array
        self.assertEqual(ap.get_ssid(), SSID.encode())
        self.assertEqual(self.nmdev.get_active_access_point(), None)


# write to stdout, not stderr
runner = unittest.TextTestRunner(stream=sys.stdout, verbosity=2)
unittest.main(testRunner=runner)
