Private
Public Access
131 lines
4.5 KiB
Python
131 lines
4.5 KiB
Python
"""tests/test_protocol.py"""
|
|
|
|
import asyncio
|
|
|
|
import pytest
|
|
|
|
from tuya_re import crypto, packet
|
|
from tuya_re.config import Credentials
|
|
from tuya_re.const import Code, DPType, SECURITY_FLAG_LOGIN, SECURITY_FLAG_SESSION
|
|
from tuya_re.protocol import Session
|
|
|
|
CREDS = Credentials(
|
|
device_id="testdevice000000000001",
|
|
uuid="uuidTESTtest0000",
|
|
local_key="abcdef0123456789",
|
|
product_id="mknd4lci",
|
|
)
|
|
SRAND = b"\x01\x02\x03\x04\x05\x06"
|
|
LOGIN_KEY = crypto.login_key(CREDS.local_key)
|
|
SESSION_KEY = crypto.session_key(CREDS.local_key, SRAND)
|
|
|
|
|
|
class FakeDevice:
|
|
"""Collects chunks Session writes; lets tests answer as the device."""
|
|
|
|
def __init__(self):
|
|
self.reassembler = packet.Reassembler()
|
|
self.frames: list[tuple[int, int, int, bytes]] = []
|
|
self.session: Session | None = None
|
|
self._device_seq = 100
|
|
|
|
async def send(self, chunk: bytes) -> None:
|
|
blob = self.reassembler.feed(chunk)
|
|
if blob is None:
|
|
return
|
|
key = LOGIN_KEY if blob[0] == SECURITY_FLAG_LOGIN else SESSION_KEY
|
|
frame = crypto.aes_cbc_decrypt(key, blob[1:17], blob[17:])
|
|
parsed = packet.parse_frame(frame)
|
|
self.frames.append(parsed)
|
|
self.reply(parsed)
|
|
|
|
def reply(self, parsed) -> None:
|
|
seq, _, code, _ = parsed
|
|
if code == Code.SENDER_DEVICE_INFO:
|
|
payload = bytes([1, 0, 3, 0, 0, 1]) + SRAND + bytes(2) + bytes(32)
|
|
self.respond(
|
|
seq, Code.SENDER_DEVICE_INFO, payload, LOGIN_KEY, SECURITY_FLAG_LOGIN
|
|
)
|
|
elif code == Code.SENDER_PAIR:
|
|
self.respond(
|
|
seq, Code.SENDER_PAIR, b"\x00", SESSION_KEY, SECURITY_FLAG_SESSION
|
|
)
|
|
elif code == Code.SENDER_DPS:
|
|
self.respond(
|
|
seq, Code.SENDER_DPS, b"\x00", SESSION_KEY, SECURITY_FLAG_SESSION
|
|
)
|
|
|
|
def respond(self, response_to, code, payload, key, flag) -> None:
|
|
self._device_seq += 1
|
|
frame = packet.build_frame(self._device_seq, response_to, code, payload)
|
|
blob = packet.encrypt_frame(frame, key, flag)
|
|
for chunk in packet.split_packets(blob):
|
|
self.session.handle_notification(chunk)
|
|
|
|
def push(self, code, payload, key=None, flag=SECURITY_FLAG_SESSION):
|
|
"""Device-initiated frame (response_to=0)."""
|
|
self._device_seq += 1
|
|
frame = packet.build_frame(self._device_seq, 0, code, payload)
|
|
blob = packet.encrypt_frame(frame, key or SESSION_KEY, flag)
|
|
for chunk in packet.split_packets(blob):
|
|
self.session.handle_notification(chunk)
|
|
|
|
|
|
@pytest.fixture
|
|
def linked():
|
|
device = FakeDevice()
|
|
session = Session(CREDS, device.send, response_timeout=1.0)
|
|
device.session = session
|
|
return device, session
|
|
|
|
|
|
async def test_handshake_derives_session_key_and_pairs(linked):
|
|
device, session = linked
|
|
await session.handshake()
|
|
assert session.paired
|
|
codes = [f[2] for f in device.frames]
|
|
assert codes == [Code.SENDER_DEVICE_INFO, Code.SENDER_PAIR]
|
|
# pair payload: uuid + key6 + device_id zero-padded to 44
|
|
pair_payload = device.frames[1][3]
|
|
assert pair_payload[:16] == CREDS.uuid.encode()
|
|
assert pair_payload[16:22] == CREDS.local_key[:6].encode()
|
|
assert pair_payload[22:44] == CREDS.device_id.encode()
|
|
assert len(pair_payload) == 44
|
|
|
|
|
|
async def test_set_dp_bool_sends_dps_frame(linked):
|
|
device, session = linked
|
|
await session.handshake()
|
|
await session.set_dp(108, DPType.BOOL, True)
|
|
code, payload = device.frames[-1][2], device.frames[-1][3]
|
|
assert code == Code.SENDER_DPS
|
|
assert payload == bytes([108, DPType.BOOL, 1, 1])
|
|
|
|
|
|
async def test_received_dps_are_recorded_and_acked(linked):
|
|
device, session = linked
|
|
await session.handshake()
|
|
# device pushes dp 108 = False (bool)
|
|
device.push(Code.RECEIVE_DP, bytes([108, DPType.BOOL, 1, 0]))
|
|
await asyncio.sleep(0) # let ack task run
|
|
assert session.dps[108] is False
|
|
assert device.frames[-1][2] == Code.RECEIVE_DP # our ack
|
|
assert device.frames[-1][1] != 0 # it is a response
|
|
|
|
|
|
async def test_time1_request_is_answered(linked):
|
|
device, session = linked
|
|
await session.handshake()
|
|
device.push(Code.RECEIVE_TIME1_REQ, b"")
|
|
await asyncio.sleep(0)
|
|
frame = device.frames[-1]
|
|
assert frame[2] == Code.RECEIVE_TIME1_REQ
|
|
assert len(frame[3]) == 13 + 2 # ascii ms timestamp + tz
|
|
|
|
|
|
async def test_handshake_timeout_raises(linked):
|
|
device, session = linked
|
|
device.reply = lambda parsed: None # device goes silent
|
|
with pytest.raises(TimeoutError):
|
|
await session.handshake()
|