Mark-VIII/tests/test_mk8_rfplan.py

import os, sys, unittest
sys.path.insert(0, os.path.join(os.path.dirname(__file__), '..', 'payload',
                                'user', 'remote_access', 'pager-webui'))
import importlib
import server
import mk8_rfplan


def setUpModule():
    importlib.reload(server)


class CtxStub(object):
    def __init__(self, body=None):
        self.body = body or {}


class RfPlanTest(unittest.TestCase):
    def setUp(self):
        self.runs = []
        self.seq = []
        self.paused = 0
        self.resumed = 0
        self.uci_show = {}
        self._old = (server.device_run, server._pause_hop,
                     server._resume_hop, server._read_hop)

        def fake_run(args, timeout=20, input_data=None):
            args = list(args)
            self.seq.append('cmd:' + ' '.join(str(a) for a in args))
            self.runs.append((args, timeout))
            if args[:2] == ['uci', 'show']:
                cfg = self.uci_show.get(args[2])
                if cfg is None:
                    return 1, '', 'entry not found'
                out = ''.join("%s.%s='%s'\n" % (args[2], k, v)
                              for k, v in sorted(cfg.items()))
                return 0, out, ''
            if args[:2] == ['uci', 'get']:
                if args[2] == 'pineapd.wlan1mon.hop':
                    return 0, '0\n', ''
                return 1, '', 'entry not found'
            if args[0] == 'iw':
                return 0, ('Connected to aa:bb:cc:dd:ee:ff (on wlan1up)\n'
                           '\tSSID: Net\n'), ''
            return 0, '', ''

        def fake_pause():
            self.paused += 1
            self.seq.append('pause')

        def fake_resume():
            self.resumed += 1
            self.seq.append('resume')

        server.device_run = fake_run
        server._pause_hop = fake_pause
        server._resume_hop = fake_resume
        server._read_hop = lambda: '0'

    def tearDown(self):
        (server.device_run, server._pause_hop,
         server._resume_hop, server._read_hop) = self._old

    @property
    def cmds(self):
        return [s[len('cmd:'):] for s in self.seq if s.startswith('cmd:')]

    def test_current_role_reads_uci(self):
        cases = [
            ({'mode': 'sta', 'disabled': '0'}, 'uplink'),
            ({'mode': 'sta', 'disabled': '1'}, 'idle'),
            ({'mode': 'ap'}, 'idle'),
            ({}, 'idle'),
        ]
        for cfg, want in cases:
            self.uci_show.clear()
            self.uci_show['wireless.wlan1up'] = cfg
            self.assertEqual(mk8_rfplan.current_role(), want, cfg)
        self.uci_show.clear()
        self.assertEqual(mk8_rfplan.current_role(), 'idle')

    def test_uplink_sets_sta_section_and_pauses_hop(self):
        result = mk8_rfplan.set_role('uplink', ssid='Net', psk='key')
        self.assertTrue(result.get('ok'), result)
        cmds = self.cmds
        for expected in (
                'uci set wireless.wlan1up=wifi-iface',
                'uci set wireless.wlan1up.device=radio1',
                'uci set wireless.wlan1up.mode=sta',
                'uci set wireless.wlan1up.network=cli',
                'uci set wireless.wlan1up.ssid=Net',
                'uci set wireless.wlan1up.encryption=psk2',
                'uci set wireless.wlan1up.key=key',
                'uci set wireless.wlan1up.disabled=0',
                'uci set network.cli=interface',
                'uci set network.cli.proto=dhcp'):
            self.assertIn(expected, cmds)
        commit = cmds.index('uci commit wireless')
        for c in ('uci set wireless.wlan1up.mode=sta',
                  'uci set wireless.wlan1up.disabled=0',
                  'uci set network.cli=interface',
                  'uci set network.cli.proto=dhcp'):
            self.assertLess(cmds.index(c), commit)
        reload_at = cmds.index('wifi reload')
        self.assertLess(commit, reload_at)
        self.assertEqual(self.paused, 1)
        self.assertLess(self.seq.index('pause'),
                        self.seq.index('cmd:wifi reload'))
        self.assertEqual(result['assoc'], 'aa:bb:cc:dd:ee:ff')

    def test_set_role_uplink_requires_ssid(self):
        result = mk8_rfplan.set_role('uplink')
        self.assertFalse(result['ok'])
        self.assertIn('ssid', result['error'])
        self.assertEqual(self.runs, [])
        self.assertEqual(self.paused, 0)

    def test_exclusivity_switch(self):
        self.uci_show['wireless.wlan1up'] = {'mode': 'ap'}
        mk8_rfplan.ensure_attack()
        self.assertEqual([c for c in self.cmds
                          if c.startswith(('uci set', 'uci commit'))], [])
        self.assertEqual(self.resumed, 0)
        self.assertEqual(self.paused, 0)
        self.uci_show['wireless.wlan1up'] = {'mode': 'sta', 'disabled': '0'}
        mk8_rfplan.ensure_attack()
        cmds = self.cmds
        self.assertIn('uci set wireless.wlan1up.disabled=1', cmds)
        self.assertIn('uci commit wireless', cmds)
        self.assertEqual(self.resumed, 1)

    def test_hop_paused_helper(self):
        self.assertTrue(mk8_rfplan.hop_paused())
        server._read_hop = lambda: None
        self.assertFalse(mk8_rfplan.hop_paused())
        server._read_hop = lambda: '1'
        self.assertFalse(mk8_rfplan.hop_paused())

    def test_h_rfplan_get_and_post(self):
        self.uci_show['wireless.wlan1up'] = {'mode': 'sta', 'disabled': '0'}
        status, data = server.h_rfplan_get(CtxStub())
        self.assertEqual(status, 200)
        self.assertEqual(data['role'], 'uplink')
        self.assertEqual(data['assoc'], 'aa:bb:cc:dd:ee:ff')
        self.assertTrue(data['hop_paused'])
        status, data = server.h_rfplan_post(CtxStub({'role': 'nope'}))
        self.assertEqual(status, 400)
        status, data = server.h_rfplan_post(CtxStub({'role': 'idle'}))
        self.assertEqual(status, 200)
        self.assertEqual(data['role'], 'idle')


if __name__ == '__main__':
    unittest.main()