Add parsing of ChipWhisperer trace sets.

3 files changed, 82 insertions, 13 deletions

diff --git a/pyecsca/__init__.py b/pyecsca/__init__.py
index 89149e4..a3f890f 100644
--- a/pyecsca/__init__.py
+++ b/pyecsca/__init__.py
@@ -8,3 +8,4 @@ from .tvla import *
 from .trace import *
 from .trace_set.base import *
 from .trace_set.inspector import *
+from .trace_set.chipwhisperer import *
\ No newline at end of file
diff --git a/pyecsca/align.py b/pyecsca/align.py
index 1c4b891..ff586ab 100644
--- a/pyecsca/align.py
+++ b/pyecsca/align.py
@@ -63,6 +63,7 @@ def align_correlation(reference: Trace, *traces: Trace,
         correlation = np.correlate(trace_part, reference_part, "same")
         max_correlation_offset = correlation.argmax(axis=0)
         max_correlation = correlation[max_correlation_offset]
+        del trace_part
         if max_correlation < min_correlation:
             return False, 0
         left_space = min(max_offset, reference_offset)
@@ -103,25 +104,26 @@ def align_peaks(reference: Trace, *traces: Trace,
 
 
 @public
-def align_sad(reference: Trace, *traces: Trace,
-              reference_offset: int, reference_length: int, max_offset: int) -> List[Trace]:
+def align_offset(reference: Trace, *traces: Trace,
+                 reference_offset: int, reference_length: int, max_offset: int,
+                 dist_func: Callable[[np.ndarray, np.ndarray], float], max_dist: float = float("inf")) -> List[Trace]:
     """
-    Align `traces` to the reference `trace` so that the Sum Of Absolute Differences between the
-    reference trace window from `reference_offset` of `reference_length` and the trace being aligned
-    within `max_offset` of the reference window is maximized.
+    Align `traces` to the reference `trace` so that the value of the `dist_func` is minimized
+    between the reference trace window from `reference_offset` of `reference_length` and the trace
+    being aligned within `max_offset` of the reference window.
 
     :param reference:
     :param traces:
     :param reference_offset:
     :param reference_length:
     :param max_offset:
+    :param dist_func:
     :return:
     """
     reference_part = reference.samples[reference_offset: reference_offset + reference_length]
-
     def align_func(trace):
         length = len(trace.samples)
-        best_sad = 0
+        best_distance = 0
         best_offset = 0
         for offset in range(-max_offset, max_offset):
             start = reference_offset + offset
@@ -129,17 +131,41 @@ def align_sad(reference: Trace, *traces: Trace,
             if start < 0 or stop >= length:
                 continue
             trace_part = trace.samples[start:stop]
-            # todo: add other distance functions here
-            sad = np.sum(np.abs(reference_part - trace_part))
-            if sad > best_sad:
-                best_sad = sad
+            distance = dist_func(reference_part, trace_part)
+            if distance < best_distance:
+                best_distance = distance
                 best_offset = offset
-        return True, best_offset
-
+        if best_distance < max_dist:
+            return True, best_offset
+        else:
+            return False, 0
     return align_reference(reference, *traces, align_func=align_func)
 
 
 @public
+def align_sad(reference: Trace, *traces: Trace,
+              reference_offset: int, reference_length: int, max_offset: int) -> List[Trace]:
+    """
+    Align `traces` to the reference `trace` so that the Sum Of Absolute Differences between the
+    reference trace window from `reference_offset` of `reference_length` and the trace being aligned
+    within `max_offset` of the reference window is maximized.
+
+    :param reference:
+    :param traces:
+    :param reference_offset:
+    :param reference_length:
+    :param max_offset:
+    :return:
+    """
+    def sad(reference_part, trace_part):
+        return float(np.sum(np.abs(reference_part - trace_part)))
+
+    return align_offset(reference, *traces,
+                        reference_offset=reference_offset, reference_length=reference_length,
+                        max_offset=max_offset, dist_func=sad)
+
+
+@public
 def align_dtw_scale(reference: Trace, *traces: Trace, radius: int = 1,
                     fast: bool = True) -> List[Trace]:
     """
diff --git a/pyecsca/trace_set/chipwhisperer.py b/pyecsca/trace_set/chipwhisperer.py
new file mode 100644
index 0000000..52300f6
--- /dev/null
+++ b/pyecsca/trace_set/chipwhisperer.py
@@ -0,0 +1,42 @@
+import numpy as np
+from .base import TraceSet
+from os.path import exists, isfile, join
+from configparser import ConfigParser
+from public import public
+
+from ..trace import Trace
+
+
+@public
+class ChipWhispererTraceSet(TraceSet):
+
+    def __init__(self, path: str = None, name: str = None):
+        if path is None and name is None:
+            super().__init__()
+        else:
+            data = self._read_data(path, name)
+            trace_data = data["traces"]
+            traces = [Trace(None, None, trace_samples, trace_set=self) for trace_samples in trace_data]
+            del data["traces"]
+            config = self._read_config(path, name)
+            super().__init__(*traces, **data, **config)
+
+    def _read_data(self, path, name):
+        types = {"keylist": None, "knownkey": None, "textin": None, "textout": None, "traces": None}
+        for type in types.keys():
+            type_path = join(path, name + "_" + type + ".npy")
+            if exists(type_path) and isfile(type_path):
+                types[type] = np.load(type_path)
+        return types
+
+    def _read_config(self, path, name):
+        config_path = join(path, "config_" + name + "_.cfg")
+        if exists(config_path) and isfile(config_path):
+            config = ConfigParser()
+            config.read(config_path)
+            return config["Trace Config"]
+        else:
+            return {}
+
+    def __repr__(self):
+        return "ChipWhispererTraceSet()"