fingerbot-cli/tests/test_lan_session.py

import asyncio
import hashlib
import hmac as _hmac
import struct

import pytest

from tuya_re import crypto, lan_protocol as lp
from tuya_re.const import LanCommand
from tuya_re.lan_session import LanSwitch

KEY = b"0123456789abcdef"
KEY_STR = KEY.decode()  # "0123456789abcdef" is ascii


async def _read_frame(reader):
    header = await reader.readexactly(16)
    length = struct.unpack(">4I", header)[3]
    return header + await reader.readexactly(length)


class FakePlug:
    def __init__(self, state):
        self.state = state
        self.server = None
        self.host = self.port = None

    async def start(self):
        self.server = await asyncio.start_server(self._handle, "127.0.0.1", 0)
        sock = self.server.sockets[0].getsockname()
        self.host, self.port = sock[0], sock[1]

    async def _handle(self, reader, writer):
        try:
            while True:
                frame = await _read_frame(reader)
                msg = lp.unpack_frame(frame)
                req = lp.decode_v33_payload(KEY, msg.payload)
                if msg.command == LanCommand.DP_QUERY:
                    body = lp.encode_v33_command(
                        KEY, LanCommand.DP_QUERY, {"dps": {"1": self.state}}
                    )
                    writer.write(
                        lp.pack_frame(msg.seqno, LanCommand.DP_QUERY, body, retcode=0)
                    )
                elif msg.command == LanCommand.CONTROL:
                    self.state = bool(req["dps"]["1"])
                    body = lp.encode_v33_command(KEY, LanCommand.CONTROL, {})
                    writer.write(
                        lp.pack_frame(msg.seqno, LanCommand.CONTROL, body, retcode=0)
                    )
                await writer.drain()
        except (asyncio.IncompleteReadError, ConnectionResetError):
            pass


@pytest.fixture
async def plug():
    p = FakePlug(state=False)
    await p.start()
    yield p
    p.server.close()


async def test_status_reads_dp1(plug):
    plug.state = True
    s = LanSwitch(
        device_id="d", local_key=KEY_STR, version="3.3", host=plug.host, port=plug.port
    )
    assert await s.status() is True


async def test_on_off_toggle(plug):
    s = LanSwitch("d", KEY_STR, "3.3", plug.host, plug.port)
    await s.turn_on()
    assert plug.state is True
    await s.turn_off()
    assert plug.state is False
    await s.toggle()  # was off -> on
    assert plug.state is True


async def test_wrong_key_raises(plug):
    s = LanSwitch("d", "ffffffffffffffff", "3.3", plug.host, plug.port)
    with pytest.raises(lp.LanFrameError):
        await s.status()


async def test_unsupported_version_raises(plug):
    s = LanSwitch("d", KEY_STR, "3.5", plug.host, plug.port)
    with pytest.raises(NotImplementedError):
        await s.status()


# ---------------------------------------------------------------------------
# v3.4 fake server and tests
# ---------------------------------------------------------------------------

REMOTE_NONCE = b"abcdef0123456789"


class FakePlugV34:
    def __init__(self, state):
        self.state = state
        self.server = None
        self.host = self.port = None

    async def start(self):
        self.server = await asyncio.start_server(self._handle, "127.0.0.1", 0)
        s = self.server.sockets[0].getsockname()
        self.host, self.port = s[0], s[1]

    async def _handle(self, reader, writer):
        try:
            # negotiation
            start = lp.unpack_frame(await _read_frame(reader), hmac_key=KEY)
            local_nonce = crypto.aes_ecb_decrypt(KEY, start.payload)[:16]
            resp_plain = (
                REMOTE_NONCE + _hmac.new(KEY, local_nonce, hashlib.sha256).digest()
            )
            writer.write(
                lp.pack_frame(
                    2,
                    LanCommand.SESS_KEY_NEG_RESP,
                    crypto.aes_ecb_encrypt(KEY, resp_plain),
                    hmac_key=KEY,
                )
            )
            await writer.drain()
            _finish = lp.unpack_frame(await _read_frame(reader), hmac_key=KEY)
            session = lp.derive_session_key_v34(KEY, local_nonce, REMOTE_NONCE)
            # commands
            while True:
                msg = lp.unpack_frame(await _read_frame(reader), hmac_key=session)
                req = lp.decode_payload(session, msg.payload, version="3.4")
                if msg.command == LanCommand.DP_QUERY:
                    body = lp.encode_command(
                        session,
                        LanCommand.DP_QUERY,
                        {"dps": {"1": self.state}},
                        version="3.4",
                    )
                    writer.write(
                        lp.pack_frame(
                            msg.seqno,
                            LanCommand.DP_QUERY,
                            body,
                            retcode=0,
                            hmac_key=session,
                        )
                    )
                elif msg.command == LanCommand.CONTROL_NEW:
                    self.state = bool(req["data"]["dps"]["1"])
                    body = lp.encode_command(
                        session, LanCommand.CONTROL_NEW, {}, version="3.4"
                    )
                    writer.write(
                        lp.pack_frame(
                            msg.seqno,
                            LanCommand.CONTROL_NEW,
                            body,
                            retcode=0,
                            hmac_key=session,
                        )
                    )
                await writer.drain()
        except (asyncio.IncompleteReadError, ConnectionResetError, lp.LanFrameError):
            pass


@pytest.fixture
async def plug34():
    p = FakePlugV34(state=False)
    await p.start()
    yield p
    p.server.close()


async def test_v34_status_and_toggle(plug34):
    s = LanSwitch("d", KEY_STR, "3.4", plug34.host, plug34.port)
    plug34.state = True
    assert await s.status() is True
    await s.turn_off()
    assert plug34.state is False
    await s.turn_on()
    assert plug34.state is True
    await s.toggle()
    assert plug34.state is False