To squash update to new API

create_resume_kti returns ResumeKTI
Use Dataset.get_dataset not to raise when reloading

scaaml/capture/aes/aes_capture.py

@@ -85,9 +85,7 @@ def capture_aes_dataset(firmware_sha256: str,
         raise ValueError('Holdout should not be captured with the same chip as'
                          'test or train')
 
-    # This over-writes the config file. The same happens with every call to
-    # Dataset.new_shard.
-    dataset = Dataset(
+    dataset = Dataset.get_dataset(
         root_path=root_path,
         shortname=shortname,
         architecture=architecture,

scaaml/capture/aes/crypto_alg.py

@@ -59,10 +59,6 @@ def __init__(self,
         self._progress_filename = full_progress_filename
         self._full_kt_filename = full_kt_filename
         self._full_progress_filename = full_progress_filename
-        # Key-plaintext iterator needs to be loaded from the saved file.
-        # The load happens during the first call of SCryptoAlgorithm.kti or
-        # directly by calling SCryptoAlgorithm.load_kti.
-        self._kti = None
 
         # Generate ans save key-text pairs.
         ktp = self._get_new_ktp()
@@ -80,7 +76,7 @@ def __init__(self,
         keys = np.array(keys, dtype=np.uint8)
         texts = np.array(texts, dtype=np.uint8)
         # Does not overwrite existing keys-texts file.
-        resume_kti.create_resume_kti(
+        self._kti = resume_kti.create_resume_kti(
             keys=keys,
             texts=texts,
             shard_length=self._examples_per_shard,
@@ -107,12 +103,6 @@ def attack_points(self, **kwargs: bytearray) -> Dict[str, bytearray]:
                 attack_point_name, **kwargs)
         return aps
 
-    def load_kti(self):
-        """Load kti. Reloads always."""
-        self._kti = resume_kti.ResumeKTI(
-            kt_filename=self._full_kt_filename,
-            progress_filename=self._full_progress_filename)
-
     def _get_new_ktp(self):
         ktp = ktp_scaaml()
         ktp.dataset = self._dataset
@@ -149,9 +139,7 @@ def __next__(self):
 
     @property
     def kti(self):
-        """Key-plaintext iterator. May need to load from file."""
-        if self._kti is None:
-            self.load_kti()
+        """Key-plaintext iterator."""
         return self._kti
 
     @property

tests/capture/aes/test_crypto_alg.py

@@ -130,37 +130,3 @@ def test_attack_points_info(mock_resumekti, mock_create_resume_kti):
         }
     }
     assert crypto_alg.attack_points_info() == api
-
-
-@patch.object(resume_kti, 'create_resume_kti')
-@patch.object(resume_kti, 'ResumeKTI')
-def test_load_kti(mock_resumekti, mock_create_resume_kti):
-    description = 'train'
-    implementation = 'MBEDTLS'
-    algorithm = 'simpleserial-aes'
-    keys = 3072
-    plaintexts = 256
-    repetitions = 1
-    examples_per_shard = 64
-    firmware_sha256 = 'TODO'
-    full_kt_filename = 'ktfilename.txt'
-    full_progress_filename = 'progressfilename.txt'
-
-    crypto_alg = SCryptoAlgorithm(
-        crypto_implementation=AESSBOX,
-        purpose=description,
-        implementation=implementation,
-        algorithm=algorithm,
-        keys=keys,
-        plaintexts=plaintexts,
-        repetitions=repetitions,
-        examples_per_shard=examples_per_shard,
-        firmware_sha256=firmware_sha256,
-        full_kt_filename=full_kt_filename,
-        full_progress_filename=full_progress_filename)
-    assert mock_resumekti.call_count == 0
-    assert crypto_alg.kti is not None
-    assert crypto_alg.kti is not None
-
-    mock_resumekti.assert_called_once_with(
-        kt_filename=full_kt_filename, progress_filename=full_progress_filename)