fix: replace a shared key that cannot be an RSA ciphertext

at_client/atclient.py

@@ -147,10 +147,38 @@ def get_encryption_key_shared_by_me(self, key: SharedKey):
             else:
                 raise response.get_exception()
 
+        stored_key = response.get_raw_data_response()
         try:
-            return EncryptionUtil.rsa_decrypt_from_base64(response.get_raw_data_response(), self.keys[KeysUtil.encryption_private_key_name])
+            return EncryptionUtil.rsa_decrypt_from_base64(stored_key, self.keys[KeysUtil.encryption_private_key_name])
         except Exception as e:
+            # A stored value that is not the right length cannot be an RSA ciphertext
+            # for our key at all, so the record is unusable rather than merely
+            # undecryptable. Left to stand, every later send to this recipient fails
+            # identically for ever, and restarting re-reads the same record. Treat it
+            # like a missing key and mint a fresh one, as the not-found branch above does.
+            if self._stored_shared_key_is_unusable(stored_key):
+                print(f"{to_lookup} is unusable ({e}); replacing it with a new shared key")
+                return self.create_shared_encryption_key(key)
             raise AtDecryptionException(f"Failed to decrypt {to_lookup} - {e}")
+
+    def _stored_shared_key_is_unusable(self, stored_key):
+        """True when a stored shared key cannot be an RSA ciphertext for our own key.
+
+        RSA ciphertext is exactly as long as the key, so a different length means the
+        record is damaged — for example a truncated or interrupted write.
+
+        Deliberately narrow: a value of the correct length that still fails to decrypt
+        is far more likely to mean the wrong keys are loaded, and replacing the shared
+        key in that case would rotate a working key for both parties. That case keeps
+        raising.
+        """
+        try:
+            ciphertext = base64.b64decode(stored_key)
+            private_key = EncryptionUtil.private_key_from_base64(
+                self.keys[KeysUtil.encryption_private_key_name])
+            return len(ciphertext) != private_key.key_size // 8
+        except Exception:
+            return False  # cannot tell: keep the existing behaviour
 
     def get_encryption_key_shared_by_other(self, shared_key: SharedKey):
         shared_shared_key_name = shared_key.get_shared_shared_key_name()

test/shared_key_recovery_test.py

@@ -0,0 +1,86 @@
+import base64
+import unittest
+from unittest.mock import MagicMock
+
+from at_client import AtClient
+from at_client.common import AtSign
+from at_client.common.keys import SharedKey
+from at_client.connections.response import Response
+from at_client.exception.atexception import AtDecryptionException
+from at_client.util import EncryptionUtil, KeysUtil
+
+
+class SharedKeyRecoveryTest(unittest.TestCase):
+    """Network-free tests for recovering from a damaged stored shared key.
+
+    A sender keeps its own copy of the AES key it shares with a recipient, encrypted
+    to its own public key. If that record is damaged, every later send to the
+    recipient fails on it, and restarting re-reads the same record.
+    """
+
+    @classmethod
+    def setUpClass(cls):
+        # generate_rsa_key_pair returns DER bytes; a keystore holds them base64-encoded,
+        # which is what the RSA helpers expect.
+        private_der, public_der = EncryptionUtil.generate_rsa_key_pair()
+        cls.private_key = base64.b64encode(private_der).decode()
+        cls.public_key = base64.b64encode(public_der).decode()
+        cls.key_size_bytes = EncryptionUtil.private_key_from_base64(
+            cls.private_key).key_size // 8
+
+    def _client(self, stored_key):
+        """A client whose llookup of the shared key returns `stored_key`."""
+        client = AtClient.__new__(AtClient)  # bypass the network-connecting __init__
+        client.atsign = AtSign("@alice")
+        client.keys = {
+            KeysUtil.encryption_private_key_name: self.private_key,
+            KeysUtil.encryption_public_key_name: self.public_key,
+        }
+        client.secondary_connection = MagicMock()
+        client.secondary_connection.execute_command.return_value = \
+            Response().set_raw_data_response(stored_key)
+        client.create_shared_encryption_key = MagicMock(return_value='replacement key')
+        return client
+
+    def _shared_key(self):
+        return SharedKey('demo', AtSign('@alice'), AtSign('@bob'))
+
+    def test_readable_stored_key_is_returned(self):
+        aes_key = EncryptionUtil.generate_aes_key_base64()
+        stored = EncryptionUtil.rsa_encrypt_to_base64(aes_key, self.public_key)
+        client = self._client(stored)
+        self.assertEqual(client.get_encryption_key_shared_by_me(self._shared_key()), aes_key)
+        client.create_shared_encryption_key.assert_not_called()
+
+    def test_wrong_length_record_is_replaced(self):
+        # Not an RSA ciphertext for this key: it cannot be decrypted by anyone.
+        damaged = base64.b64encode(b'truncated').decode()
+        client = self._client(damaged)
+        result = client.get_encryption_key_shared_by_me(self._shared_key())
+        self.assertEqual(result, 'replacement key')
+        client.create_shared_encryption_key.assert_called_once()
+
+    def test_correct_length_but_undecryptable_still_raises(self):
+        """Most likely the wrong keys are loaded — replacing would rotate a good key."""
+        _, other_public_der = EncryptionUtil.generate_rsa_key_pair()
+        stored = EncryptionUtil.rsa_encrypt_to_base64(
+            EncryptionUtil.generate_aes_key_base64(),
+            base64.b64encode(other_public_der).decode())
+        self.assertEqual(len(base64.b64decode(stored)), self.key_size_bytes)
+        client = self._client(stored)
+        with self.assertRaises(AtDecryptionException):
+            client.get_encryption_key_shared_by_me(self._shared_key())
+        client.create_shared_encryption_key.assert_not_called()
+
+    def test_unusable_check_is_length_based(self):
+        client = self._client('unused')
+        good = EncryptionUtil.rsa_encrypt_to_base64('x', self.public_key)
+        self.assertFalse(client._stored_shared_key_is_unusable(good))
+        self.assertTrue(client._stored_shared_key_is_unusable(
+            base64.b64encode(b'short').decode()))
+        # Undecidable input keeps the existing behaviour rather than replacing a key.
+        self.assertFalse(client._stored_shared_key_is_unusable('not base64 !!'))
+
+
+if __name__ == '__main__':
+    unittest.main()