Mark-VIII/tests/test_mcp.py

import os
import shutil
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)


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


class McpDispatchTest(unittest.TestCase):
    def setUp(self):
        self.tmp = tempfile.mkdtemp(prefix='pager-mcp-')
        self.old_session = server.SESSION_FILE
        self.old_device_run = server.device_run
        self.old_daemon_sock_call = server.daemon_sock_call
        server.SESSION_FILE = os.path.join(self.tmp, 'session.json')
        with open(server.SESSION_FILE, 'w') as f:
            import json
            json.dump({'token': 't0k3n', 'serverid': 'x'}, f)
        server.device_run = lambda args, timeout=20, input_data=None: (0, '', '')
        server.daemon_sock_call = lambda method, path, body=None, timeout=10: (200, {
            'pineap_disabled': False, 'pineape_disabled': True,
            'autossidpool': True})

    def tearDown(self):
        server.SESSION_FILE = self.old_session
        server.device_run = self.old_device_run
        server.daemon_sock_call = self.old_daemon_sock_call
        shutil.rmtree(self.tmp)

    def msg(self, method, params=None, mid=1, jsonrpc='2.0'):
        m = {'jsonrpc': jsonrpc, 'id': mid, 'method': method}
        if params is not None:
            m['params'] = params
        return m

    def test_initialize_negotiates_protocol(self):
        status, body = server._mcp_dispatch(self.msg('initialize', {'protocolVersion': '2025-06-18'}))
        self.assertEqual(status, 200)
        self.assertEqual(body['result']['protocolVersion'], '2025-06-18')
        self.assertIn('tools', body['result']['capabilities'])
        self.assertEqual(body['result']['serverInfo']['name'], 'mark-viii')

    def test_tools_list_has_attack_and_recon_tools(self):
        status, body = server._mcp_dispatch(self.msg('tools/list'))
        names = [t['name'] for t in body['result']['tools']]
        self.assertIn('attack.deploy', names)
        self.assertIn('device.state', names)
        self.assertIn('recon.isearch', names)
        self.assertIn('loot.enterprise_creds', names)

    def test_tools_call_unknown_tool_errors(self):
        status, body = server._mcp_dispatch(self.msg('tools/call', {'name': 'nope', 'arguments': {}}))
        self.assertEqual(body['error']['code'], -32602)

    def test_ping(self):
        status, body = server._mcp_dispatch(self.msg('ping'))
        self.assertEqual(body['result'], {})

    def test_notifications_initialized_returns_202(self):
        status, body = server._mcp_dispatch({'jsonrpc': '2.0', 'method': 'notifications/initialized'})
        self.assertEqual(status, 202)
        self.assertIsNone(body)

    def test_resources_list_includes_skills(self):
        status, body = server._mcp_dispatch(self.msg('resources/list'))
        uris = [r['uri'] for r in body['result']['resources']]
        self.assertIn('skills://pineapple-control', uris)
        self.assertIn('device://state', uris)

    def test_prompts_list_includes_playbooks(self):
        status, body = server._mcp_dispatch(self.msg('prompts/list'))
        names = [p['name'] for p in body['result']['prompts']]
        self.assertIn('evil-wpa-attack', names)
        self.assertIn('evil-enterprise-attack', names)

    def test_bad_jsonrpc_rejected(self):
        status, body = server._mcp_dispatch({'jsonrpc': '1.0', 'id': 1, 'method': 'ping'})
        self.assertEqual(body['error']['code'], -32600)

    def test_endpoint_auth_accepts_bearer_token(self):
        status, body = server.h_mcp(ctx(self.msg('ping'), {'Authorization': 'Bearer t0k3n'}))
        self.assertEqual(status, 200)
        self.assertEqual(body['result'], {})

    def test_endpoint_auth_rejects_bad_token(self):
        status, body = server.h_mcp(ctx(self.msg('ping'), {'Authorization': 'Bearer wrong'}))
        self.assertEqual(status, 401)

    def test_deploy_tool_wraps_attack_handler(self):
        server._uci_wifi_iface = lambda name: {}
        server._uci_section = lambda name: {}
        server._verify_iface = lambda name, timeout=20: True
        server._allow_all_ssids = lambda: True
        server._deploy_enterprise = lambda args: {'ok': True, 'verified': True,
                                                  'iface': 'wlan1ent', 'band': '5'}
        status, body = server._mcp_dispatch(self.msg('tools/call', {
            'name': 'attack.deploy',
            'arguments': {'kind': 'enterprise', 'ssid': 'Corp', 'enctype': 'wpa2', 'channel': 36}}))
        self.assertEqual(status, 200)
        text = body['result']['content'][0]['text']
        self.assertIn('"verified": true', text)

    def test_capabilities_endpoint(self):
        status, body = server.h_harness_capabilities(ctx())
        self.assertEqual(status, 200)
        self.assertEqual(body['endpoint'], '/mcp')
        self.assertIn('attack.deploy', [t['name'] for t in body['tools']])
        self.assertIn('skills://wifi-deauth', [r['uri'] for r in body['resources']])


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