Fixed examples and container improvements

docs/tutorial_01_discovering_containers.rst

@@ -138,15 +138,15 @@ instance.
 
    # To work only with the first 10 samples:
    batch.leakage_section = range(10)
-   print(trace_batch)
+   print(batch)
 
    # To work with only with one tenth of the sample:
    batch.leakage_section = range(0, 100, 10)
-   print(trace_batch)
+   print(batch)
 
    # To cancel `leakage_section`:
    batch.leakage_section = None  # cancelling leakage_section
-   print(trace_batch)
+   print(batch)
 
 This will output:
 
@@ -185,7 +185,7 @@ See :code:`lascar/tools/processing` for a list of existing processing.
    )
 
    # Principal component analysis on leakage with 3 components
-   batch.leakage_processing = PcaProcessing(trace_batch, 3)
+   batch.leakage_processing = PcaProcessing(batch, 3)
 
    # No leakage processing
    batch.leakage_processing = None

docs/tutorial_04_acquisition_setup_example.rst

@@ -45,9 +45,9 @@ oscilloscope.
            """
            :param number_of_traces: Number of traces in the container
            """
-           super().__init__(self, number_of_traces)
            self.dut = Dut()
            self.oscilloscope = Oscilloscope()
+           super().__init__(number_of_traces)
 
        def generate_trace(self, index: int):
            """

docs/tutorial_07_dpa_example.rst

@@ -43,7 +43,7 @@ conditioned by a single key byte (256 guesses).
        return sbox[value["plaintext"][3] ^ guess] & 1
 
    guess_range = range(256)
-   dpa_engine = DpaEngine("dpa", selection_function, guess_range)
+   dpa_engine = DpaEngine(selection_function, guess_range)
 
 We can now create a :class:`Session <lascar.session.Session>`, register the
 :code:`dpa_lsb_engine`, and run it:

docs/tutorial_08_session_outputs.rst

@@ -69,10 +69,10 @@ of the output of the 3rd sbox, the other on the MSB.
        return (sbox[value["plaintext"][3] ^ guess] >> 7) & 1
 
    dpa_lsb_engine = DpaEngine(
-       "dpa_lsb", selection_function_lsb, guess_range, solution=container.key[3]
+       selection_function_lsb, guess_range, solution=container.key[3], name="dpa_lsb"
    )
    dpa_msb_engine = DpaEngine(
-       "dpa_msb", selection_function_lsb, guess_range, solution=container.key[3]
+       selection_function_msb, guess_range, solution=container.key[3], name="dpa_msb"
    )
 
 Dictionnary output

examples/ascad/02-snr.py

@@ -36,7 +36,7 @@
 container = Hdf5Container(
     filename, leakages_dataset_name="traces", values_dataset_name="metadata"
 )
-container.number_of_traces = 5000  # only 5000 traces used over the 60000 available
+container = container.limited(5000)  # Only 5000 traces used over the 60000 available
 
 
 """
@@ -47,33 +47,33 @@
 - the partition_values (range(256) for all of them, since we look at byte values)
 """
 snr_1_engine = SnrEngine(
-    "SNR1: unmasked sbox output",
     lambda value: sbox[value["plaintext"][3] ^ value["key"][3]],  # sbox(p[3] ⊕ k[3])
     range(256),
+    name="SNR1: unmasked sbox output"
 )
 
 snr_2_engine = SnrEngine(
-    "SNR2: masked sbox output",
     lambda value: sbox[value["plaintext"][3] ^ value["key"][3]]
     ^ value["masks"][15],  # sbox(p[3] ⊕ k[3]) ⊕ rout
     range(256),
+    name="SNR2: masked sbox output"
 )
 
 snr_3_engine = SnrEngine(
-    "SNR3: common output mask out", lambda value: value["masks"][15], range(256)  # rout
+    lambda value: value["masks"][15], range(256), name="SNR3: common output mask out"  # rout
 )
 
 snr_4_engine = SnrEngine(
-    "SNR4: masked sbox output in linear parts",
     lambda value: sbox[value["plaintext"][3] ^ value["key"][3]]
     ^ value["masks"][1],  # sbox(p[3] ⊕ k[3]) ⊕ r[3]
     range(256),
+    name="SNR4: masked sbox output in linear parts"
 )
 
 snr_5_engine = SnrEngine(
-    "SNR5: sbox output mask in linear parts",
     lambda value: value["masks"][1],  # r[3]
     range(256),
+    name="SNR5: sbox output mask in linear parts"
 )
 
 engines = [snr_1_engine, snr_2_engine, snr_3_engine, snr_4_engine, snr_5_engine]

examples/ascad/03-keras-train.py

@@ -70,13 +70,12 @@ def partition_function(value):
 
 # An engine has to be created, dedicated to the profiling of a classifier (here a keras neural network) with leakages/labels.
 nn_profile_engine = ProfileEngine(
-    "nn_profile",
     nn,
     partition_function,
     partition_values,
     epochs=5,
     batch_size=200,
-    test_size=0.1,
+    test_size=0.1
 )
 
 Session(profiling_container, engine=nn_profile_engine).run()

examples/ascad/04-keras-test.py

@@ -71,15 +71,15 @@ def guess_function(value, guess):
 guess_range = range(256)  # The possbles value for the guess
 
 # An engine has to be created, dedicated to use a classifier (here a keras neural network nn) with Side-Channel Traces.
-nn_match_engine = MatchEngine(
-    "nn_match", nn, guess_function, guess_range, solution=solution
-)
+nn_match_engine = MatchEngine(nn, guess_function, guess_range, solution=solution, name="nn_match")
 
 
 # Now, 5 times in a row we randomly take 2000 traces from attack_container, and compute the mean rank of the correct key, every 10 traces.
 ranks = []
 
-for i,container in enumerate(split_container(attack_container, size_of_splits=2000)[:5]):
+for i, container in enumerate(
+    split_container(attack_container, size_of_splits=2000)[:5]
+):
 
     session = Session(
         container,

examples/ascad/05-cpa-high-order.py

@@ -40,18 +40,18 @@
     leakage_section=poi,
 )
 
-container.number_of_traces = 500
+container = container.limited(500)
 
 # Then we make a container that will apply a CenteredProduct to recombine all the points of interest.
 container.leakage_processing = CenteredProductProcessing(container)
 
 
 # Now a classical CPA, targetting the output of the 3rd Sbox:
 cpa_engine = CpaEngine(
-    "cpa-high-order",
     lambda value, guess: hamming(sbox[value["plaintext"][3] ^ guess]),
     range(256),
     solution=container[0].value["key"][3],
+    name="cpa-high-order",
     jit=False,
 )
 

examples/base/lra.py

@@ -18,14 +18,8 @@ def guess_function(
 guess_range = range(256)
 
 leakage_model = HammingPrecomputedModel()
-
-container = BasicAesSimulationContainer(
-    5000, 0
-)  
-
-attack = LraEngine(
-    "lra", partition, partition_size, guess_function, guess_range
-)
+container = BasicAesSimulationContainer(5000, 0)
+attack = LraEngine("lra", partition, partition_size, guess_function, guess_range)
 
 session = Session(container)
 session.add_engine(attack)

examples/base/ttest.py

@@ -41,7 +41,6 @@ def partition_function(value):
     return partition_function
 
 
-number_of_partitions = 2  # number of possible classes (~output of the partiton_function) for the partition_function
 ttest_engines = [
     TTestEngine(get_partition_function(i)) for i in range(16)
 ]

examples/rainbow/rainbow_stm32.py

@@ -111,7 +111,6 @@ def generate_trace(self, idx):
 
     container = RainbowSubBytesContainer(250)
     engine = CpaEngine(
-        "cpa byte 5",
         lambda value, secret: sbox[value["plaintext"][5] ^ secret],
         range(256),
         solution=42,

lascar/container/__init__.py

@@ -24,6 +24,7 @@
 from .container import AbstractContainer
 from .container import Container
 from .container import TraceBatchContainer
+from .container import Slice 
 from .filtered_container import FilteredContainer
 from .filtered_container import RandomizedContainer
 from .filtered_container import split_container

lascar/container/container.py

@@ -109,10 +109,10 @@ class Container:
 
     Where both leakage and value can be represented as a numpy array.
 
-    The role of the Container class is to be overloaded so that it can deliver traces, stored as a specified format.
-    Mostly, the __getitem__/__setitem__ have to be overloaded when user want to write its own Container format class.
-
-    :param number_of_traces:
+    The role of the Container class is to be overloaded so that it can deliver traces,
+    stored as a specified format. Mostly, the `__getitem__`, `__setitem__` and
+    `__len__` methods have to be overloaded when user want to write its own Container
+    format class.
     """
 
     def __init__(self, **kwargs):
@@ -149,9 +149,6 @@ def __init__(self, **kwargs):
         )
         self._value_abstract = AbstractArray(self.values.shape[1:], self.values.dtype)
 
-        self.number_of_traces = kwargs.get(
-            "number_of_traces", self.number_of_traces_max
-        )
         self.leakage_section = kwargs.get("leakage_section", None)
         self.value_section = kwargs.get("value_section", None)
         self.leakage_processing = kwargs.get("leakage_processing", None)
@@ -345,9 +342,6 @@ def plot_leakage(self, key):
 
         plot([self[i].leakage for i in key])
 
-    def __len__(self):
-        return self.number_of_traces
-
     def __iter__(self):
         """
         Container is iterable
@@ -359,7 +353,7 @@ def __iter__(self):
             yield self[i]
 
     def __str__(self):
-        res = "Container with %d traces. " % (self.number_of_traces)
+        res = f"Container with {len(self)} traces. "
         res += "leakages: %s, values: %s. " % (
             self._leakage_abstract,
             self._value_abstract,
@@ -396,6 +390,25 @@ def get_leakage_mean_var(self):
         session = Session(self).run()
         return session["mean"].finalize(), session["var"].finalize()
 
+    def sliced(self, start: int, stop: int) -> "Slice":
+        """
+        Wraps the current container in a :class:`Slice` to select a slice of the
+        traces.
+
+        :param start: Start index
+        :param stop: End index, excluded.
+        """
+        return Slice(self, start, stop)
+
+    def limited(self, count: int) -> "Slice":
+        """
+        Wraps the current container in a :class:`Slice` to limit the number of
+        traces.
+
+        :param count: Number of traces limit.
+        """
+        return Slice(self, 0, count)
+
 
 # def to_trace(func):
 #     def wrapper(*args, **kwargs):
@@ -423,16 +436,15 @@ class AbstractContainer(Container):
 
     def __init__(self, number_of_traces, **kwargs):
         self.logger = logging.getLogger(__name__)
-
+        self.number_of_traces = number_of_traces
         trace = self.generate_trace(0)
         self.leakages = AbstractArray(
             (number_of_traces,) + trace.leakage.shape, trace.leakage.dtype
         )
         self.values = AbstractArray(
             (number_of_traces,) + trace.value.shape, trace.value.dtype
         )
-
-        Container.__init__(self, **kwargs)
+        super().__init__(**kwargs)
 
     def generate_trace(self, idx):
         """
@@ -466,6 +478,9 @@ def generate_trace_batch(self, idx_begin, idx_end):
 
         return TraceBatchContainer(leakages, values)
 
+    def __len__(self):
+        return self.number_of_traces
+
     def __getitem__(self, key):
         """
         :rtype: Trace or TraceBatch depending on key
@@ -577,6 +592,9 @@ def __setitem__(self, key, value):
         self.leakages[key] = value.leakage if isinstance(key, int) else value.leakages
         self.values[key] = value.value if isinstance(key, int) else value.values
 
+    def __len__(self):
+        return len(self.leakages)
+
     def save(self, filename):
         """
         Save the current TraceBatchContainer to a file using np.save
@@ -722,3 +740,38 @@ def generate_trace(self, idx):
         leakage = self.get_leakage()
 
         return Trace(leakage, value)
+
+
+class Slice:
+    def __init__(self, base: Container, start: int, stop: int):
+        """
+        Container which is a slice of a parent container, in order to reduce the
+        number of traces.
+
+        :param base: Base container
+        :param start: Start index
+        :param stop: End index, excluded.
+        """
+        self.base = base
+        if stop < start:
+            raise ValueError("Slice stop must be >= start")
+        self.start = start
+        self.stop = stop
+
+    def __getitem__(self, key):
+        if isinstance(key, int):
+            if key < self.start or key >= self.stop:
+                raise IndexError()
+            return self.base[key]
+        elif isinstance(key, slice):
+            if key.start < self.start or key.stop > self.stop:
+                raise IndexError()
+            return self.base[key]
+        elif isinstance(key, list):
+            for i in key:
+                if i < self.start or i >= self.stop:
+                    raise IndexError()
+            return self.base[key]
+
+    def __len__(self):
+        return self.stop - self.start

lascar/container/filtered_container.py

@@ -73,7 +73,7 @@ class RandomizedContainer(FilteredContainer):
 
     def __init__(self, container, **kwargs):
         FilteredContainer.__init__(
-            self, container, np.random.permutation(container.number_of_traces), **kwargs
+            self, container, np.random.permutation(len(container)), **kwargs
         )
 
 
@@ -84,7 +84,7 @@ def split_container(container, random=True, **kwargs):
     :param kwargs: specify either the number of splits (number_of_splits), or the size of each split (size_of_splits)
     :return: a list of number_of_splits containers OR a list of containers of size_of_splits each.
     """
-    n = container.number_of_traces
+    n = len(container)
     if random:
         indexes = np.random.permutation(n)
     else:

lascar/container/hdf5_container.py

@@ -71,6 +71,9 @@ def __getitem__(self, key):
     def __setitem__(self, key, value):
         TraceBatchContainer.__setitem__(self, key, value)
 
+    def __len__(self):
+        return len(self.leakages)
+
     @staticmethod
     def void_container(
         filename,
@@ -143,7 +146,7 @@ def export(
         leakage, value = container[0]
         out = Hdf5Container.void_container(
             filename,
-            container.number_of_traces,
+            len(container),
             leakage.shape,
             leakage.dtype,
             value.shape,