diff --git a/at_client/atclient.py b/at_client/atclient.py index 4ca5ab0..f343b31 100644 --- a/at_client/atclient.py +++ b/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() diff --git a/test/shared_key_recovery_test.py b/test/shared_key_recovery_test.py new file mode 100644 index 0000000..a0dd4eb --- /dev/null +++ b/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()