Mark-VIII/tests/test_misc.py

import os
import sys
import tempfile
import unittest

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)


class PayloadsProxyTest(unittest.TestCase):
    def test_index_proxies(self):
        server.daemon_call = lambda m, p, body=None, token=None, timeout=15: (200, {'payloads': []})
        server.current_token = lambda: 'tok'
        class Ctx:
            args = ()
        status, payload = server.h_payloads_index(Ctx())
        self.assertEqual(status, 200)
        self.assertEqual(payload, {'payloads': []})

    def test_install_uses_key(self):
        calls = []
        def fake(m, p, body=None, token=None, timeout=15):
            calls.append((m, p))
            return 200, {'ok': True}
        server.daemon_call = fake
        server.current_token = lambda: 'tok'
        server.h_payloads_install(type('C', (), {'args': (), 'body': {'key': 'nautilus'}})())
        self.assertTrue(any(m == 'POST' and '/api/payloads/portal/nautilus/install' in p for m, p in calls))

    def test_remove_uses_key(self):
        calls = []
        def fake(m, p, body=None, token=None, timeout=15):
            calls.append((m, p))
            return 200, {'ok': True}
        server.daemon_call = fake
        server.current_token = lambda: 'tok'
        server.h_payloads_remove(type('C', (), {'args': (), 'body': {'key': 'nautilus'}})())
        self.assertTrue(any(m == 'POST' and '/api/payloads/portal/nautilus/remove' in p for m, p in calls))

    def test_install_surfaces_daemon_error_detail(self):
        server.daemon_call = lambda m, p, body=None, token=None, timeout=15: (
            500, {'error': 'network error: Get "https://downloads.hak5.org/.../download": dial tcp: lookup downloads.hak5.org on [::1]:53: server misbehaving'})
        server.current_token = lambda: 'tok'
        status, payload = server.h_payloads_install(type('C', (), {
            'args': (), 'body': {'key': 'recon~client~recon_reporter'}})())
        self.assertEqual(status, 500)
        self.assertIn('downloads.hak5.org', payload['detail'])

    def test_install_unwraps_raw_json_daemon_error(self):
        server.daemon_call = lambda m, p, body=None, token=None, timeout=15: (
            500, '{"error":"network error: Get \\"https://downloads.hak5.org/...\\": dial tcp: lookup downloads.hak5.org on [::1]:53: server misbehaving"}\n')
        server.current_token = lambda: 'tok'
        status, payload = server.h_payloads_install(type('C', (), {
            'args': (), 'body': {'key': 'recon~client~recon_reporter'}})())
        self.assertEqual(status, 500)
        self.assertIn('network error', payload['detail'])
        self.assertIn('downloads.hak5.org', payload['detail'])
        self.assertNotIn('{', payload['detail'])

    def test_installed_inventory_flattens_firmware_records(self):
        old = server._payload_daemon
        server._payload_daemon = lambda method, path, body=None: (200, [{
            'installed': {'key': 'user~general~demo', 'title': 'Demo',
                          'path': '/root/payloads/user/general/demo',
                          'category': 'general', 'launchpoint': 'payload.sh'},
            'missingmanifest': False, 'update': {'version': '2.0'}}])
        try:
            status, payload = server.h_payloads_installed(type('C', (), {'args': ()})())
        finally:
            server._payload_daemon = old
        self.assertEqual(status, 200)
        self.assertEqual(payload['payloads'][0]['title'], 'Demo')
        self.assertEqual(payload['payloads'][0]['update']['version'], '2.0')

    def test_install_rejects_unsafe_key(self):
        status, payload = server.h_payloads_install(type('C', (), {
            'args': (), 'body': {'key': '../payload'}})())
        self.assertEqual(status, 400)
        self.assertIn('key', payload['error'])

    def test_run_launches_only_resolved_installed_payload(self):
        old_record = server._payload_record
        old_safe = server._payload_safe_launch
        old_popen = server.subprocess.Popen
        old_run_dir = server.PAYLOAD_RUN_DIR

        class FakeProcess:
            pid = 4321
            def poll(self):
                return None

        with tempfile.TemporaryDirectory() as tmp:
            launch = os.path.join(tmp, 'payload.sh')
            with open(launch, 'w') as handle:
                handle.write('#!/bin/sh\n')
            server._payload_record = lambda key: (200, {
                'key': key, 'title': 'Demo', 'path': tmp,
                'launchpoint': 'payload.sh', 'interpreter': 'sh', 'disabled': False})
            server._payload_safe_launch = lambda item: (tmp, launch)
            calls = []
            server.subprocess.Popen = lambda command, **kwargs: calls.append((command, kwargs)) or FakeProcess()
            server.PAYLOAD_RUN_DIR = os.path.join(tmp, 'runs')
            server._payload_runs.clear()
            try:
                status, payload = server.h_payloads_run(type('C', (), {
                    'body': {'key': 'user~general~demo'}})())
            finally:
                server._payload_record = old_record
                server._payload_safe_launch = old_safe
                server.subprocess.Popen = old_popen
                server.PAYLOAD_RUN_DIR = old_run_dir
                server._payload_runs.clear()
        self.assertEqual(status, 200)
        self.assertEqual(calls[0][0], ['sh', launch])
        self.assertTrue(calls[0][1]['start_new_session'])
        self.assertTrue(payload['run']['running'])

    def test_stop_rejects_untracked_process(self):
        server._payload_runs.clear()
        status, payload = server.h_payloads_stop(type('C', (), {
            'body': {'id': 'not-tracked'}})())
        self.assertEqual(status, 404)

    def test_remove_rejects_active_webui(self):
        status, payload = server.h_payloads_remove(type('C', (), {
            'args': (), 'body': {'key': server.SELF_PAYLOAD_KEY}})())
        self.assertEqual(status, 409)
        self.assertIn('cannot remove itself', payload['error'])


class LoggingTest(unittest.TestCase):
    def test_system_tails(self):
        server.device_run = lambda args, timeout=20: (0, '\n'.join('line%d' % i for i in range(20)), '')
        status, payload = server.h_logging_system(type('C', (), {'args': (), 'query': {'lines': '5'}})())
        self.assertEqual(status, 200)
        self.assertEqual(len(payload['lines']), 5)

    def test_invalid_line_count_uses_default(self):
        server.device_run = lambda args, timeout=20: (
            0, '\n'.join('line%d' % i for i in range(250)), '')
        status, payload = server.h_logging_system(type('C', (), {
            'args': (), 'query': {'lines': 'invalid'}})())
        self.assertEqual(status, 200)
        self.assertEqual(len(payload['lines']), 200)

    def test_line_count_is_bounded(self):
        self.assertEqual(server._line_count(
            type('C', (), {'query': {'lines': '999999'}})(), 200), 2000)
        self.assertEqual(server._line_count(
            type('C', (), {'query': {'lines': '-10'}})(), 200), 0)

    def test_system_uses_bounded_logread(self):
        calls = []

        def fake(args, timeout=20, input_data=None):
            calls.append(list(args))
            return (0, '\n'.join('line%d' % i for i in range(20)), '')

        server.device_run = fake
        status, payload = server.h_logging_system(type('C', (), {
            'args': (), 'query': {'lines': '5'}})())
        self.assertEqual(status, 200)
        self.assertEqual(len(payload['lines']), 5)
        self.assertTrue(calls)
        self.assertEqual(calls[0][:2], ['logread', '-l'])

    def test_json_or_and_pool_list_tolerate_empty(self):
        self.assertIsNone(server._json_or(None))
        self.assertEqual(server._parse_pool_list(None), [])
        self.assertEqual(server._unique_keep_order(['a', '', 'a', 'b']), ['a', 'b'])


class SettingsTest(unittest.TestCase):
    def test_hostname_get(self):
        class H:
            command = 'GET'
        server.device_run = lambda args, timeout=20: (0, 'pager\n', '')
        status, payload = server.h_settings_hostname(type('C', (), {'args': (), 'h': H()})())
        self.assertEqual(status, 200)
        self.assertEqual(payload['hostname'], 'pager')

    def test_internet_status_uses_bounded_connectivity_check(self):
        calls = []
        old = server.device_run
        server.device_run = lambda args, timeout=20: (
            calls.append((args, timeout)) or (0, '', ''))
        try:
            status, payload = server.h_settings_internet(type('C', (), {})())
        finally:
            server.device_run = old
        self.assertEqual(status, 200)
        self.assertTrue(payload['online'])
        self.assertEqual(calls[0][0][-1], '1.1.1.1')
        self.assertEqual(calls[0][1], 5)

    def test_reboot_is_scheduled_after_response(self):
        scheduled = []
        class FakeTimer:
            def __init__(self, delay, fn):
                scheduled.append((delay, fn))
                self.daemon = False
            def start(self):
                scheduled.append('started')
        old = server.threading.Timer
        server.threading.Timer = FakeTimer
        try:
            status, payload = server.h_settings_reboot(type('C', (), {})())
        finally:
            server.threading.Timer = old
        self.assertEqual(status, 200)
        self.assertTrue(payload['ok'])
        self.assertEqual(scheduled[0][0], 1.0)
        self.assertEqual(scheduled[1], 'started')

    def test_password_uses_passwd_stdin_and_invalidates_session(self):
        calls = []
        def fake(args, timeout=20, input_data=None):
            calls.append((args, timeout, input_data))
            return 0, '', ''
        old_run = server.device_run
        old_session = server.SESSION_FILE
        fd, session = tempfile.mkstemp()
        os.close(fd)
        server.device_run = fake
        server.SESSION_FILE = session
        try:
            status, payload = server.h_settings_password(
                type('C', (), {'args': (), 'body': {'password': 'newpw'}})())
        finally:
            server.device_run = old_run
            server.SESSION_FILE = old_session
        self.assertEqual((status, payload), (200, {'ok': True}))
        self.assertEqual(calls, [(['/bin/passwd', 'root'], 15, b'newpw\nnewpw\n')])
        self.assertFalse(os.path.exists(session))

    def test_password_reports_passwd_failure(self):
        old_run = server.device_run
        server.device_run = lambda args, timeout=20, input_data=None: (1, '', 'bad password')
        try:
            status, payload = server.h_settings_password(
                type('C', (), {'args': (), 'body': {'password': 'newpw'}})())
        finally:
            server.device_run = old_run
        self.assertEqual(status, 500)
        self.assertEqual(payload['error'], 'bad password')

    def test_ntp_get_parses(self):
        class H:
            command = 'GET'
        server.device_run = lambda args, timeout=20: (0, "system.ntp.enabled='1'\nsystem.ntp.server='0.pool.ntp.org'\n", '')
        status, payload = server.h_settings_ntp(type('C', (), {'args': (), 'h': H()})())
        self.assertEqual(status, 200)
        self.assertEqual(payload['enabled'], True)
        self.assertIn('0.pool.ntp.org', payload['servers'])

    def test_ntp_get_parses_multiple_values_on_one_uci_line(self):
        class H:
            command = 'GET'
        server.device_run = lambda args, timeout=20: (
            0,
            "system.ntp.enabled='1'\n"
            "system.ntp.server='0.openwrt.pool.ntp.org' '1.openwrt.pool.ntp.org'\n",
            '')
        status, payload = server.h_settings_ntp(type('C', (), {'args': (), 'h': H()})())
        self.assertEqual(status, 200)
        self.assertEqual(payload['servers'], [
            '0.openwrt.pool.ntp.org', '1.openwrt.pool.ntp.org'])

    def test_service_running_detection(self):
        server.device_run = lambda args, timeout=20: (0, '', '')
        status, payload = server.h_settings_service(type('C', (), {'args': ()})())
        self.assertEqual(status, 200)
        self.assertIn('running', payload)
        self.assertIn('background', payload)

    def test_resources_parses_mark_style_filesystem_rows(self):
        sample = ("Filesystem Type 1024-blocks Used Available Capacity Mounted on\n"
                  "/dev/root squashfs 30720 30720 0 100% /rom\n"
                  "/dev/mmcblk0p1 ext4 3724208 21876 3492668 1% /mmc\n")
        rows = server._parse_df_rows(sample)
        self.assertEqual(len(rows), 2)
        self.assertEqual(rows[1]['format'], 'ext4')
        self.assertEqual(rows[1]['mount'], '/mmc')
        self.assertEqual(rows[1]['size'], 3724208 * 1024)

    def test_usb_parser_returns_structured_devices(self):
        server.device_run = lambda args, timeout=20: (
            0, 'Bus 001 Device 004: ID 0e8d:7961 MediaTek Inc. Wireless_Device\n', '')
        status, payload = server.h_settings_usb(type('C', (), {'args': ()})())
        self.assertEqual(status, 200)
        self.assertEqual(payload['devices'][0]['id'], '0e8d:7961')
        self.assertIn('MediaTek', payload['devices'][0]['name'])

    def test_password_rejects_mismatched_confirmation(self):
        status, payload = server.h_settings_password(type('C', (), {'body': {
            'new_password': 'newpass', 'repeat_password': 'different'}})())
        self.assertEqual(status, 400)
        self.assertIn('match', payload['error'])

    def test_management_wifi_get_hides_existing_password(self):
        old = server._uci_wifi_iface
        server._uci_wifi_iface = lambda name: {
            'ssid': 'Pager Mgmt', 'key': 'secretpass', 'disabled': '0', 'hidden': '1'}
        try:
            status, payload = server.h_settings_management_wifi(
                type('C', (), {'body': {}, 'args': ()})())
        finally:
            server._uci_wifi_iface = old
        self.assertEqual(status, 200)
        self.assertTrue(payload['has_password'])
        self.assertNotIn('password', payload)


class WifiClientModeTest(unittest.TestCase):
    @staticmethod
    def fake_device_run(calls, responses):
        def run(args, timeout=20, input_data=None):
            calls.append((list(args), timeout))
            key = (args[0], args[1])
            return responses.get(key, (0, '', ''))
        return run

    def test_client_state_parses_disabled(self):
        calls = []
        run = self.fake_device_run(calls, {
            ('uci', 'show'): (0,
                "wireless.wlan0cli=wifi-iface\n"
                "wireless.wlan0cli.ssid='OldNet'\n"
                "wireless.wlan0cli.disabled='1'\n"
                "wireless.wlan0cli.routed='0'\n", ''),
            ('iw', 'dev'): (0, '', ''),
            ('ip', '-4'): (0, '', ''),
        })
        old = server.device_run
        server.device_run = run
        try:
            status, payload = server.h_settings_wifi_client(type('C', (), {})())
        finally:
            server.device_run = old
        self.assertEqual(status, 200)
        self.assertFalse(payload['enabled'])
        self.assertFalse(payload['connected'])
        self.assertEqual(payload['ssid'], 'OldNet')
        self.assertFalse(payload['routed'])
        self.assertEqual(payload['ip'], '')

    def test_client_state_parses_connected(self):
        calls = []
        run = self.fake_device_run(calls, {
            ('uci', 'show'): (0,
                "wireless.wlan0cli=wifi-iface\n"
                "wireless.wlan0cli.ssid='All RPH Guest WIFI'\n"
                "wireless.wlan0cli.disabled='0'\n"
                "wireless.wlan0cli.routed='0'\n", ''),
            ('iw', 'dev'): (0,
                "Connected to 02:18:4a:a7:6a:d9 (on wlan0cli)\n"
                "\tSSID: All RPH Guest WIFI\n"
                "\tfreq: 2462\n"
                "\tsignal: -57 dBm\n", ''),
            ('ip', '-4'): (0,
                "6: wlan0cli: <BROADCAST,MULTICAST,UP,LOWER_UP> mtu 1500\n"
                "    inet 10.10.10.5/24 brd 10.10.10.255 scope global wlan0cli\n", ''),
        })
        old = server.device_run
        server.device_run = run
        try:
            status, payload = server.h_settings_wifi_client(type('C', (), {})())
        finally:
            server.device_run = old
        self.assertEqual(status, 200)
        self.assertTrue(payload['enabled'])
        self.assertTrue(payload['connected'])
        self.assertEqual(payload['connected_ssid'], 'All RPH Guest WIFI')
        self.assertEqual(payload['ip'], '10.10.10.5')
        self.assertEqual(payload['signal'], -57)
        self.assertEqual(payload['freq'], 2462)

    def test_scan_parses_networks_and_deduplicates(self):
        sample = (
            "BSS 02:18:4a:a7:6a:d2(on wlan0)\n"
            "\tfreq: 2462.0\n"
            "\tsignal: -63.00 dBm\n"
            "\tSSID: Riverwalk Plaza Staff\n"
            "\tRSN:\t * Version: 1\n"
            "\t\t * Group cipher: TKIP\n"
            "\t\t * Pairwise ciphers: CCMP TKIP\n"
            "\t\t * Authentication suites: PSK\n"
            "BSS 02:18:4a:a7:6a:d9(on wlan0)\n"
            "\tfreq: 2462.0\n"
            "\tsignal: -61.00 dBm\n"
            "\tSSID: All RPH Guest WIFI\n"
            "\tRSN:\t * Version: 1\n"
            "\t\t * Group cipher: CCMP\n"
            "\t\t * Pairwise ciphers: CCMP\n"
            "\t\t * Authentication suites: SAE\n"
            "BSS ea:cb:bc:8e:c5:0e(on wlan0)\n"
            "\tfreq: 2462.0\n"
            "\tsignal: -50.00 dBm\n"
            "\tSSID: OpenGuest\n"
            "BSS c6:cb:bc:8e:c5:0e(on wlan0)\n"
            "\tfreq: 2462.0\n"
            "\tsignal: -55.00 dBm\n"
            "\tSSID: \\x00\\x00\\x00\\x00\n"
            "BSS 42:18:4a:a7:6a:d2(on wlan0)\n"
            "\tfreq: 2462.0\n"
            "\tsignal: -65.00 dBm\n"
            "\tSSID: Riverwalk Plaza Staff\n"
            "\tWPA:\t * Version: 1\n"
            "\t\t * Group cipher: TKIP\n"
            "\t\t * Authentication suites: PSK\n")
        networks = server._parse_wifi_scan(sample)
        by_ssid = {n['ssid']: n for n in networks}
        self.assertIn('Riverwalk Plaza Staff', by_ssid)
        self.assertIn('All RPH Guest WIFI', by_ssid)
        self.assertIn('OpenGuest', by_ssid)
        self.assertEqual(by_ssid['Riverwalk Plaza Staff']['encryption'], 'WPA2')
        self.assertEqual(by_ssid['All RPH Guest WIFI']['encryption'], 'WPA3')
        self.assertEqual(by_ssid['OpenGuest']['encryption'], 'Open')
        self.assertEqual(by_ssid['OpenGuest']['channel'], 11)
        # hidden SSID entries are omitted
        self.assertNotIn('', by_ssid)
        # strongest BSS per SSID wins and results are signal-sorted (strongest first)
        self.assertEqual(networks[0]['ssid'], 'OpenGuest')
        self.assertEqual(networks[0]['signal'], -50)
        self.assertGreater(networks[0]['signal'], networks[1]['signal'])

    def test_scan_reports_device_failure(self):
        old = server.device_run
        server.device_run = lambda args, timeout=20: (1, '', 'scan not supported')
        try:
            status, payload = server.h_settings_wifi_client_scan(type('C', (), {})())
        finally:
            server.device_run = old
        self.assertEqual(status, 502)
        self.assertIn('scan', payload['error'])

    def test_connect_requires_ssid(self):
        class H:
            command = 'POST'
        old = server.device_run
        server.device_run = lambda args, timeout=20: (0, '', '')
        try:
            status, payload = server.h_settings_wifi_client_connect(
                type('C', (), {'h': H(), 'body': {'encryption': 'open'}})())
        finally:
            server.device_run = old
        self.assertEqual(status, 400)
        self.assertIn('SSID', payload['error'])

    def test_connect_rejects_short_password(self):
        class H:
            command = 'POST'
        calls = []
        run = self.fake_device_run(calls, {})
        old = server.device_run
        server.device_run = run
        try:
            status, payload = server.h_settings_wifi_client_connect(
                type('C', (), {'h': H(), 'body': {'ssid': 'X', 'encryption': 'wpa2', 'password': 'short'}})())
        finally:
            server.device_run = old
        self.assertEqual(status, 400)
        self.assertIn('password', payload['error'])
        self.assertEqual(calls, [])

    def test_connect_writes_uci_and_reloads(self):
        class H:
            command = 'POST'
        calls = []
        run = self.fake_device_run(calls, {})
        daemon_calls = []
        old_run = server.device_run
        old_sock = server.daemon_sock_call
        server.device_run = run
        server.daemon_sock_call = lambda m, p, body=None, timeout=10: (
            daemon_calls.append((m, p, body)) or (200, {'success': True}))
        try:
            status, payload = server.h_settings_wifi_client_connect(
                type('C', (), {'h': H(), 'body': {
                    'ssid': 'All RPH Guest WIFI', 'encryption': 'wpa2wpa3',
                    'password': 'Missions1', 'routed': True}})())
        finally:
            server.device_run = old_run
            server.daemon_sock_call = old_sock
        self.assertEqual(status, 200)
        sets = [args for args, _t in calls if args[:2] == ['uci', 'set']]
        expected = {
            'wireless.wlan0cli.ssid=All RPH Guest WIFI',
            'wireless.wlan0cli.encryption=sae-mixed',
            'wireless.wlan0cli.disabled=0',
            'wireless.wlan0cli.routed=1',
            'wireless.wlan0cli.key=Missions1',
        }
        got = {args[2] for args in sets}
        self.assertTrue(expected <= got)
        self.assertIn((['uci', 'commit', 'wireless'], 20), calls)
        self.assertIn((['wifi', 'reload'], 45), calls)
        self.assertEqual(daemon_calls,
            [('PUT', '/api/settings/wifi/set_client_route', {'routed': True})])

    def test_connect_open_clears_key(self):
        class H:
            command = 'POST'
        calls = []
        run = self.fake_device_run(calls, {})
        old_run = server.device_run
        old_sock = server.daemon_sock_call
        server.device_run = run
        server.daemon_sock_call = lambda m, p, body=None, timeout=10: (200, {'success': True})
        try:
            status, payload = server.h_settings_wifi_client_connect(
                type('C', (), {'h': H(), 'body': {'ssid': 'OpenGuest', 'encryption': 'open'}})())
        finally:
            server.device_run = old_run
            server.daemon_sock_call = old_sock
        self.assertEqual(status, 200)
        self.assertIn((['uci', 'delete', 'wireless.wlan0cli.key'], 20), calls)
        self.assertNotIn((['uci', 'set', 'wireless.wlan0cli.key='], 20), calls)

    def test_route_toggle_writes_uci_and_syncs_daemon(self):
        class H:
            command = 'POST'
        calls = []
        run = self.fake_device_run(calls, {})
        daemon_calls = []
        old_run = server.device_run
        old_sock = server.daemon_sock_call
        server.device_run = run
        server.daemon_sock_call = lambda m, p, body=None, timeout=10: (
            daemon_calls.append((m, p, body)) or (200, {'success': True}))
        try:
            status, payload = server.h_settings_wifi_client_route(
                type('C', (), {'h': H(), 'body': {'routed': True}})())
        finally:
            server.device_run = old_run
            server.daemon_sock_call = old_sock
        self.assertEqual(status, 200)
        self.assertIn((['uci', 'set', 'wireless.wlan0cli.routed=1'], 20), calls)
        self.assertIn((['uci', 'commit', 'wireless'], 20), calls)
        self.assertEqual(daemon_calls,
            [('PUT', '/api/settings/wifi/set_client_route', {'routed': True})])

    def test_disconnect_disables_and_reloads(self):
        class H:
            command = 'POST'
        calls = []
        run = self.fake_device_run(calls, {})
        old = server.device_run
        server.device_run = run
        try:
            status, payload = server.h_settings_wifi_client_disconnect(
                type('C', (), {'h': H(), 'body': {}})())
        finally:
            server.device_run = old
        self.assertEqual(status, 200)
        self.assertIn((['uci', 'set', 'wireless.wlan0cli.disabled=1'], 20), calls)
        self.assertIn((['uci', 'commit', 'wireless'], 20), calls)
        self.assertIn((['wifi', 'reload'], 45), calls)


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