Mark-VIII/tests/test_robustness.py

import os
import sys
import unittest
from unittest import mock

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


def setUpModule():
    __import__('importlib').reload(server)


def ctx(body=None):
    return type('C', (), {'body': body, 'args': (), 'query': {}})()


class DaemonRetryTest(unittest.TestCase):
    def setUp(self):
        self.old_sleep = server.DAEMON_SOCK_RETRY_SLEEP
        self.old_retries = server.DAEMON_SOCK_RETRIES
        server.DAEMON_SOCK_RETRY_SLEEP = 0
        server.DAEMON_SOCK_RETRIES = 2

    def tearDown(self):
        server.DAEMON_SOCK_RETRY_SLEEP = self.old_sleep
        server.DAEMON_SOCK_RETRIES = self.old_retries

    def test_sock_retries_then_succeeds(self):
        attempts = {'n': 0}

        class Sock:
            def __init__(self):
                self.chunks = [b'HTTP/1.1 200 OK\r\n\r\n{"ok":true}', b'']

            def settimeout(self, t):
                pass

            def connect(self, addr):
                attempts['n'] += 1
                if attempts['n'] < 2:
                    raise OSError('busy')

            def sendall(self, data):
                pass

            def recv(self, n):
                return self.chunks.pop(0) if self.chunks else b''

            def close(self):
                pass

        with mock.patch.object(server.socket, 'socket', lambda *a, **k: Sock()):
            status, data = server.daemon_sock_call('GET', '/api/pineap/get_config')
        self.assertEqual(attempts['n'], 2)
        self.assertEqual(status, 200)
        self.assertEqual(data, {'ok': True})

    def test_call_retries_on_oserror(self):
        attempts = {'n': 0}

        def boom(*a, **k):
            attempts['n'] += 1
            raise OSError('down')

        with mock.patch.object(server.socket, 'socket', boom):
            status, data = server.daemon_call('GET', '/api/api_ping')
        self.assertEqual(attempts['n'], 3)
        self.assertEqual(status, 0)
        self.assertIsNone(data)


class ConfigWriteSafetyTest(unittest.TestCase):
    def setUp(self):
        self.old_sock = server.daemon_sock_call

    def tearDown(self):
        server.daemon_sock_call = self.old_sock
    def test_set_config_refuses_when_get_fails(self):
        calls = []

        def fake(method, path, body=None, timeout=10):
            calls.append((method, path))
            if method == 'GET':
                return 0, None
            return 200, {'success': True}

        server.daemon_sock_call = fake
        status, payload = server.h_pineap_set_config(ctx({'loghandshake': True}))
        self.assertEqual(status, 502)
        self.assertIn('could not read', payload['error'])
        self.assertFalse(any(c[0] == 'PUT' for c in calls))

    def test_hostapd_set_refuses_when_get_fails(self):
        calls = []

        def fake(method, path, body=None, timeout=10):
            calls.append((method, path))
            if method == 'GET':
                return 0, None
            return 200, {'success': True}

        server.daemon_sock_call = fake
        status, payload = server.h_pineap_hostapd_set(ctx({'pineape_auth_pass': True}))
        self.assertEqual(status, 502)
        self.assertFalse(any(c[0] == 'PUT' for c in calls))


class Hak5RetryTest(unittest.TestCase):
    def setUp(self):
        self.old_sleep = server.HAK5_RETRY_SLEEP
        self.old_run = server.device_run
        server.HAK5_RETRY_SLEEP = 0

    def tearDown(self):
        server.HAK5_RETRY_SLEEP = self.old_sleep
        server.device_run = self.old_run

    def test_hak5_raises_after_retries(self):
        calls = []

        def fake(args, timeout=20, input_data=None):
            calls.append(args)
            return 1, '', 'busy'

        server.device_run = fake
        with self.assertRaises(RuntimeError):
            server.hak5('PINEAPPLE_SSID_POOL_LIST')
        self.assertEqual(len(calls), 3)

    def test_hak5_succeeds_on_retry(self):
        calls = []

        def fake(args, timeout=20, input_data=None):
            calls.append(args)
            if len(calls) < 2:
                return 1, '', 'busy'
            return 0, 'ok\n', ''

        server.device_run = fake
        out = server.hak5('PINEAPPLE_SSID_POOL_LIST')
        self.assertEqual(out, 'ok\n')
        self.assertEqual(len(calls), 2)

    def test_hak5_treats_error_text_as_failure(self):
        calls = []

        def fake(args, timeout=20, input_data=None):
            calls.append(args)
            return 0, '', 'ERROR: invalid time (expected number of seconds)'

        server.device_run = fake
        with self.assertRaises(RuntimeError):
            server.hak5('PINEAPPLE_EXAMINE_CHANNEL', '140')
        self.assertEqual(len(calls), 3)


class SsidPoolFailureTest(unittest.TestCase):
    def setUp(self):
        self.old_hak5 = server.hak5

    def tearDown(self):
        server.hak5 = self.old_hak5
    def test_ssids_post_add_returns_502_on_hak5_failure(self):
        server.hak5 = lambda *a, **k: (_ for _ in ()).throw(RuntimeError('busy'))
        status, payload = server.h_ssids_post(ctx({'action': 'add', 'ssid': 'NewNet'}))
        self.assertEqual(status, 502)
        self.assertIn('ssid pool update failed', payload['error'])

    def test_examine_returns_502_on_hak5_failure(self):
        server.hak5 = lambda *a, **k: (_ for _ in ()).throw(RuntimeError('busy'))
        status, payload = server.h_recon_examine(type('C', (), {
            'args': (), 'body': {'bssid': 'AA:BB:CC:DD:EE:FF'}})())
        self.assertEqual(status, 502)
        self.assertEqual(payload['error'], 'examine failed')


class AdvertiseBlockTest(unittest.TestCase):
    def setUp(self):
        self.old_proxy = server._daemon_proxy
        self.old_uci = server._uci_section

    def tearDown(self):
        server._daemon_proxy = self.old_proxy
        server._uci_section = self.old_uci
    def test_advertise_enable_always_refused(self):
        server._uci_section = lambda name: {'disable': '0'}
        calls = []
        server._daemon_proxy = lambda method, path, body=None, timeout=15: (
            calls.append((method, path, body)) or (200, {'success': True}))
        status, payload = server.h_pineap_advertise(ctx({'enable': True}))
        self.assertEqual(status, 400)
        self.assertIn('cannot be re-enabled', payload['error'])
        self.assertEqual(calls, [])

    def test_advertise_disable_still_proxies(self):
        calls = []
        server._daemon_proxy = lambda method, path, body=None, timeout=15: (
            calls.append(path) or (200, {'success': True}))
        status, payload = server.h_pineap_advertise(ctx({'enable': False}))
        self.assertEqual(status, 200)
        self.assertEqual(calls, ['ssidpool/disable'])


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