[ota] Offer encryption with the api key so enabling it works over OTA

tests/unit_tests/test_espota2_noise.py

@@ -10,6 +10,7 @@ when the installed aioesphomeapi predates the noise module.
 from __future__ import annotations
 
 import base64
+from collections.abc import Callable
 import hashlib
 import io
 from pathlib import Path
@@ -109,8 +110,17 @@ class FakeEncryptedDevice(threading.Thread):
             return
         server_flags = espota2.SERVER_FEATURE_SUPPORTS_NOISE if self.offer_noise else 0
         sock.sendall(bytes([espota2.RESPONSE_FEATURE_FLAGS, server_flags]))
-        if not (self.offer_noise and noise_negotiated):
+        if noise_negotiated and not self.offer_noise:
             return  # the client fails closed; nothing further arrives
+        if not noise_negotiated:
+            # A device that does not require encryption lets a plaintext
+            # client through
+            self._transfer(
+                lambda byte: sock.sendall(bytes([byte])),
+                lambda length: _recv_exact(sock, length),
+                lambda remaining: sock.recv(min(remaining, 4096)),
+            )
+            return
 
         from cryptography.exceptions import InvalidTag
         from noise.connection import NoiseConnection
@@ -149,6 +159,20 @@ class FakeEncryptedDevice(threading.Thread):
             assert len(plaintext) == length, "control units must be one per frame"
             return plaintext
 
+        def recv_data(remaining: int) -> bytes:
+            plaintext = proto.decrypt(_recv_frame(sock))
+            assert 0 < len(plaintext) <= espota2.NOISE_MAX_PLAINTEXT
+            return plaintext
+
+        self._transfer(send_byte, recv_unit, recv_data)
+
+    def _transfer(
+        self,
+        send_byte: Callable[[int], None],
+        recv_unit: Callable[[int], bytes],
+        recv_data: Callable[[int], bytes],
+    ) -> None:
+        """The post-handshake exchange, identical over both transports."""
         send_byte(espota2.RESPONSE_AUTH_OK)
         recv_unit(1)  # ota type
         size = int.from_bytes(recv_unit(4), "big")
@@ -159,9 +183,9 @@ class FakeEncryptedDevice(threading.Thread):
         received = b""
         acked = 0
         while len(received) < size:
-            plaintext = proto.decrypt(_recv_frame(sock))
-            assert 0 < len(plaintext) <= espota2.NOISE_MAX_PLAINTEXT
-            received += plaintext
+            chunk = recv_data(size - len(received))
+            assert chunk, "client closed mid-transfer"
+            received += chunk
             if self.version >= espota2.OTA_VERSION_2_0:
                 while acked + espota2.UPLOAD_BLOCK_SIZE <= len(received) or (
                     len(received) == size and acked < size
@@ -240,6 +264,29 @@ def test_client_fails_closed_when_device_lacks_encryption() -> None:
     device.join_and_check()
 
 
+def test_plaintext_client_accepted_by_offering_device() -> None:
+    """A device that offers but does not require encryption still takes a
+    plaintext upload from a client with no key configured."""
+    firmware = bytes(range(256)) * 40
+    device = FakeEncryptedDevice(offer_noise=True, require_noise=False)
+    with patch("time.sleep"):
+        _upload(device, firmware, None)
+    device.join_and_check()
+    assert device.received == firmware
+
+
+def test_keyed_client_encrypts_with_offering_device() -> None:
+    """The upload that turns on `ota: encryption:` is already encrypted when
+    the running firmware offers it."""
+    pytest.importorskip("aioesphomeapi.noise")
+    firmware = bytes(range(256)) * 40
+    device = FakeEncryptedDevice(offer_noise=True, require_noise=False)
+    with patch("time.sleep"):
+        _upload(device, firmware, PSK)
+    device.join_and_check()
+    assert device.received == firmware
+
+
 def test_plaintext_client_gets_encryption_required_error() -> None:
     """A client without a key gets the device's 0x94 error message."""
     device = FakeEncryptedDevice()