140 lines
4.0 KiB
Python
140 lines
4.0 KiB
Python
"""
|
|
AngelaMos | 2026
|
|
test_protocol.py
|
|
"""
|
|
|
|
import pytest
|
|
|
|
from core.models import CommandType
|
|
from core.protocol import Message, MessageType, pack, unpack
|
|
|
|
|
|
TEST_KEY = "test-protocol-key"
|
|
|
|
|
|
class TestPackUnpack:
|
|
"""
|
|
Verify protocol message serialization roundtrips
|
|
"""
|
|
|
|
def test_register_roundtrip(self) -> None:
|
|
"""
|
|
REGISTER message survives pack -> unpack
|
|
"""
|
|
original = Message(
|
|
type=MessageType.REGISTER,
|
|
payload={
|
|
"hostname": "desktop-01",
|
|
"os": "Linux",
|
|
"username": "root",
|
|
"pid": 1234,
|
|
"internal_ip": "10.0.0.5",
|
|
"arch": "x86_64",
|
|
},
|
|
)
|
|
packed = pack(original, TEST_KEY)
|
|
restored = unpack(packed, TEST_KEY)
|
|
assert restored.type == MessageType.REGISTER
|
|
assert restored.payload["hostname"] == "desktop-01"
|
|
assert restored.payload["pid"] == 1234
|
|
|
|
def test_task_roundtrip(self) -> None:
|
|
"""
|
|
TASK message survives pack -> unpack
|
|
"""
|
|
original = Message(
|
|
type=MessageType.TASK,
|
|
payload={
|
|
"id": "task-001",
|
|
"command": CommandType.SHELL,
|
|
"args": "whoami",
|
|
},
|
|
)
|
|
restored = unpack(pack(original, TEST_KEY), TEST_KEY)
|
|
assert restored.type == MessageType.TASK
|
|
assert restored.payload["command"] == "shell"
|
|
|
|
def test_result_roundtrip(self) -> None:
|
|
"""
|
|
RESULT message survives pack -> unpack
|
|
"""
|
|
original = Message(
|
|
type=MessageType.RESULT,
|
|
payload={
|
|
"task_id": "task-001",
|
|
"output": "root\n",
|
|
"error": None,
|
|
},
|
|
)
|
|
restored = unpack(pack(original, TEST_KEY), TEST_KEY)
|
|
assert restored.type == MessageType.RESULT
|
|
assert restored.payload["output"] == "root\n"
|
|
|
|
def test_heartbeat_roundtrip(self) -> None:
|
|
"""
|
|
HEARTBEAT message survives pack -> unpack
|
|
"""
|
|
original = Message(
|
|
type=MessageType.HEARTBEAT,
|
|
payload={"id": "beacon-abc"},
|
|
)
|
|
restored = unpack(pack(original, TEST_KEY), TEST_KEY)
|
|
assert restored.type == MessageType.HEARTBEAT
|
|
|
|
def test_error_roundtrip(self) -> None:
|
|
"""
|
|
ERROR message survives pack -> unpack
|
|
"""
|
|
original = Message(
|
|
type=MessageType.ERROR,
|
|
payload={"detail": "command not found"},
|
|
)
|
|
restored = unpack(pack(original, TEST_KEY), TEST_KEY)
|
|
assert restored.type == MessageType.ERROR
|
|
assert restored.payload["detail"] == "command not found"
|
|
|
|
|
|
class TestUnpackValidation:
|
|
"""
|
|
Verify unpack rejects malformed data
|
|
"""
|
|
|
|
def test_invalid_base64(self) -> None:
|
|
"""
|
|
Non-base64 input raises ValueError
|
|
"""
|
|
with pytest.raises(ValueError, match="Invalid protocol message"):
|
|
unpack("not-valid-base64!!!", TEST_KEY)
|
|
|
|
def test_invalid_json_after_decode(self) -> None:
|
|
"""
|
|
Valid base64 but invalid JSON after XOR raises ValueError
|
|
"""
|
|
from core.encoding import encode
|
|
|
|
encoded = encode("this is not json", TEST_KEY)
|
|
with pytest.raises(ValueError, match="Invalid protocol message"):
|
|
unpack(encoded, TEST_KEY)
|
|
|
|
def test_missing_type_field(self) -> None:
|
|
"""
|
|
JSON without 'type' field raises ValueError
|
|
"""
|
|
from core.encoding import encode
|
|
|
|
encoded = encode('{"payload": {}}', TEST_KEY)
|
|
with pytest.raises(ValueError, match="Invalid protocol message"):
|
|
unpack(encoded, TEST_KEY)
|
|
|
|
def test_invalid_message_type(self) -> None:
|
|
"""
|
|
Unknown message type raises ValueError
|
|
"""
|
|
from core.encoding import encode
|
|
|
|
encoded = encode(
|
|
'{"type": "BOGUS", "payload": {}}', TEST_KEY
|
|
)
|
|
with pytest.raises(ValueError, match="Invalid protocol message"):
|
|
unpack(encoded, TEST_KEY)
|