6 files changed, 68 insertions, 27 deletions

diff --git a/pyecsca/sca/scope/base.py b/pyecsca/sca/scope/base.py
index 78a6431..5b2cde2 100644
--- a/pyecsca/sca/scope/base.py
+++ b/pyecsca/sca/scope/base.py
@@ -1,8 +1,16 @@
+from enum import Enum, auto
 from typing import Tuple, Sequence, Optional
 
-import numpy as np
 from public import public
 
+from ..trace import Trace
+
+
+@public
+class SampleType(Enum):
+    Raw = auto()
+    Volt = auto()
+
 
 @public
 class Scope(object):
@@ -30,13 +38,15 @@ class Scope(object):
         """
         raise NotImplementedError
 
-    def setup_channel(self, channel: str, coupling: str, range: float, enable: bool) -> None:
+    def setup_channel(self, channel: str, coupling: str, range: float, offset: float,
+                      enable: bool) -> None:
         """
         Setup a channel to use the coupling method and measure the given voltage range.
 
         :param channel: The channel to measure.
         :param coupling: The coupling method ("AC" or "DC).
         :param range: The voltage range to measure.
+        :param offset: The analog voltage offset added to the input. Not supported on many scopes.
         :param enable: Whether to enable or disable the channel.
         """
         raise NotImplementedError
@@ -80,11 +90,12 @@ class Scope(object):
         """
         raise NotImplementedError
 
-    def retrieve(self, channel: str) -> Optional[np.ndarray]:
+    def retrieve(self, channel: str, type: SampleType) -> Optional[Trace]:
         """
         Retrieve a captured trace of a channel.
 
         :param channel: The channel to retrieve the trace from.
+        :param type: The type of returned samples.
         :return: The captured trace (if any).
         """
         raise NotImplementedError
diff --git a/pyecsca/sca/scope/chipwhisperer.py b/pyecsca/sca/scope/chipwhisperer.py
index f8a8b7b..cb12c91 100644
--- a/pyecsca/sca/scope/chipwhisperer.py
+++ b/pyecsca/sca/scope/chipwhisperer.py
@@ -4,7 +4,8 @@ import numpy as np
 from chipwhisperer.capture.scopes.OpenADC import OpenADC
 from public import public
 
-from .base import Scope
+from .base import Scope, SampleType
+from ..trace import Trace
 
 
 @public
@@ -30,7 +31,7 @@ class ChipWhispererScope(Scope):  # pragma: no cover
         self.scope.samples = posttrig
         return self.scope.clock.freq_ctr, self.scope.samples
 
-    def setup_channel(self, channel: str, coupling: str, range: float, enable: bool) -> None:
+    def setup_channel(self, channel: str, coupling: str, range: float, offset: float, enable: bool) -> None:
         pass
 
     def setup_trigger(self, channel: str, threshold: float, direction: str, delay: int,
@@ -51,8 +52,11 @@ class ChipWhispererScope(Scope):  # pragma: no cover
     def capture(self, timeout: Optional[int] = None) -> bool:
         return not self.scope.capture()
 
-    def retrieve(self, channel: str) -> Optional[np.ndarray]:
-        return self.scope.get_last_trace()
+    def retrieve(self, channel: str, type: SampleType) -> Optional[Trace]:
+        data = self.scope.get_last_trace()
+        if data is None:
+            return None
+        return Trace(data, {"sampling_frequency": self.scope.clock.clkgen_freq, "channel": channel})
 
     def stop(self) -> None:
         pass
diff --git a/pyecsca/sca/scope/picoscope_alt.py b/pyecsca/sca/scope/picoscope_alt.py
index dfa69a9..f93c997 100644
--- a/pyecsca/sca/scope/picoscope_alt.py
+++ b/pyecsca/sca/scope/picoscope_alt.py
@@ -1,13 +1,13 @@
 from time import time_ns, sleep
 from typing import Optional, Tuple, Sequence, Union
 
-import numpy as np
 from picoscope.ps4000 import PS4000
 from picoscope.ps5000 import PS5000
 from picoscope.ps6000 import PS6000
 from public import public
 
-from .base import Scope
+from .base import Scope, SampleType
+from ..trace import Trace
 
 
 @public
@@ -17,6 +17,7 @@ class PicoScopeAlt(Scope):  # pragma: no cover
         super().__init__()
         self.ps = ps
         self.trig_ratio: float = 0.0
+        self.frequency: Optional[float] = None
 
     def open(self) -> None:
         self.ps.open()
@@ -31,10 +32,11 @@ class PicoScopeAlt(Scope):  # pragma: no cover
         if max_samples < samples:
             self.trig_ratio = (pretrig / samples)
             samples = max_samples
+        self.frequency = actual_frequency
         return actual_frequency, samples
 
-    def setup_channel(self, channel: str, coupling: str, range: float, enable: bool) -> None:
-        self.ps.setChannel(channel, coupling, range, 0.0, enable)
+    def setup_channel(self, channel: str, coupling: str, range: float, offset: float, enable: bool) -> None:
+        self.ps.setChannel(channel, coupling, range, offset, enable)
 
     def setup_trigger(self, channel: str, threshold: float, direction: str, delay: int,
                       timeout: int, enable: bool) -> None:
@@ -54,8 +56,14 @@ class PicoScopeAlt(Scope):  # pragma: no cover
                 return False
         return True
 
-    def retrieve(self, channel: str) -> Optional[np.ndarray]:
-        return self.ps.getDataV(channel)
+    def retrieve(self, channel: str, type: SampleType) -> Optional[Trace]:
+        if type == SampleType.Raw:
+            data = self.ps.getDataRaw(channel)
+        else:
+            data = self.ps.getDataV(channel)
+        if data is None:
+            return None
+        return Trace(data, {"sampling_frequency": self.frequency, "channel": channel})
 
     def stop(self) -> None:
         self.ps.stop()
diff --git a/pyecsca/sca/scope/picoscope_sdk.py b/pyecsca/sca/scope/picoscope_sdk.py
index e1d4a52..014680c 100644
--- a/pyecsca/sca/scope/picoscope_sdk.py
+++ b/pyecsca/sca/scope/picoscope_sdk.py
@@ -11,7 +11,8 @@ from picosdk.ps5000 import ps5000
 from picosdk.ps6000 import ps6000
 from public import public
 
-from .base import Scope
+from .base import Scope, SampleType
+from ..trace import Trace
 
 
 def adc2volt(adc: Union[np.ndarray, ctypes.c_int16],
@@ -75,14 +76,16 @@ class PicoScopeSdk(Scope):  # pragma: no cover
     def setup_frequency(self, frequency: int, pretrig: int, posttrig: int) -> Tuple[int, int]:
         return self.set_frequency(frequency, pretrig, posttrig)
 
-    def set_channel(self, channel: str, enabled: bool, coupling: str, range: float):
+    def set_channel(self, channel: str, enabled: bool, coupling: str, range: float, offset: float):
+        if offset != 0.0:
+            raise ValueError("Offset not supported.")
         assert_pico_ok(
                 self.__dispatch_call("SetChannel", self.handle, self.CHANNELS[channel], enabled,
                                      self.COUPLING[coupling], self.RANGES[range]))
         self.ranges[channel] = range
 
-    def setup_channel(self, channel: str, coupling: str, range: float, enable: bool):
-        self.set_channel(channel, enable, coupling, range)
+    def setup_channel(self, channel: str, coupling: str, range: float,  offset: float, enable: bool):
+        self.set_channel(channel, enable, coupling, range, offset)
 
     def _set_freq(self, frequency: int, pretrig: int, posttrig: int, period_bound: float,
                   timebase_bound: int,
@@ -167,7 +170,7 @@ class PicoScopeSdk(Scope):  # pragma: no cover
                 return False
         return True
 
-    def retrieve(self, channel: str) -> Optional[np.ndarray]:
+    def retrieve(self, channel: str, type: SampleType) -> Optional[Trace]:
         if self.samples is None:
             raise ValueError
         actual_samples = ctypes.c_int32(self.samples)
@@ -176,7 +179,11 @@ class PicoScopeSdk(Scope):  # pragma: no cover
                 self.__dispatch_call("GetValues", self.handle, 0, ctypes.byref(actual_samples), 1,
                                      0, 0, ctypes.byref(overflow)))
         arr = np.array(self.buffers[channel], dtype=np.int16)
-        return adc2volt(arr, self.ranges[channel], self.MAX_ADC_VALUE)
+        if type == SampleType.Raw:
+            data = arr
+        else:
+            data = adc2volt(arr, self.ranges[channel], self.MAX_ADC_VALUE)
+        return Trace(data, {"sampling_frequency": self.frequency, "channel": channel})
 
     def stop(self):
         assert_pico_ok(self.__dispatch_call("Stop"))
@@ -306,15 +313,16 @@ class PS6000Scope(PicoScopeSdk):  # pragma: no cover
 
     COUPLING = {
         "AC": ps6000.PS6000_COUPLING["PS6000_AC"],
-        "DC": ps6000.PS6000_COUPLING["PS6000_DC_1M"]
+        "DC": ps6000.PS6000_COUPLING["PS6000_DC_1M"],
+        "DC_50": ps6000.PS6000_COUPLING["PS6000_DC_50R"]
     }
 
     def open(self):
         assert_pico_ok(ps6000.ps6000OpenUnit(ctypes.byref(self.handle), None))
 
-    def set_channel(self, channel: str, enabled: bool, coupling: str, range: float):
+    def set_channel(self, channel: str, enabled: bool, coupling: str, range: float, offset: float):
         assert_pico_ok(ps6000.ps6000SetChannel(self.handle, self.CHANNELS[channel], enabled,
-                                               self.COUPLING[coupling], self.RANGES[range], 0,
+                                               self.COUPLING[coupling], self.RANGES[range], offset,
                                                ps6000.PS6000_BANDWIDTH_LIMITER["PS6000_BW_FULL"]))
 
     def set_buffer(self, channel: str, enable: bool):
diff --git a/pyecsca/sca/trace/trace.py b/pyecsca/sca/trace/trace.py
index 82d1ee4..ca20021 100644
--- a/pyecsca/sca/trace/trace.py
+++ b/pyecsca/sca/trace/trace.py
@@ -9,11 +9,13 @@ from public import public
 
 @public
 class Trace(object):
-    """A trace, which has an optional title, optional data bytes and mandatory samples."""
+    """A trace, which has some samples and metadata."""
     meta: Mapping[str, Any]
     samples: ndarray
 
     def __init__(self, samples: ndarray, meta: Mapping[str, Any] = None, trace_set: Any = None):
+        if meta is None:
+            meta = {}
         self.meta = meta
         self.samples = samples
         self.trace_set = trace_set
diff --git a/test/ec/test_point.py b/test/ec/test_point.py
index 76e1103..a8b9fd7 100644
--- a/test/ec/test_point.py
+++ b/test/ec/test_point.py
@@ -64,14 +64,14 @@ class PointTests(TestCase):
                       X=Mod(0x2, self.secp128r1.curve.prime),
                       Y=Mod(0x3, self.secp128r1.curve.prime),
                       Z=Mod(1, self.secp128r1.curve.prime))
-        assert pt.equals(other)
+        self.assertTrue(pt.equals(other))
         self.assertNotEqual(pt, other)
-        assert not pt.equals(2)
+        self.assertFalse(pt.equals(2))
         self.assertNotEqual(pt, 2)
 
         infty_one = InfinityPoint(self.coords)
         infty_other = InfinityPoint(self.coords)
-        assert infty_one.equals(infty_other)
+        self.assertTrue(infty_one.equals(infty_other))
         self.assertEqual(infty_one, infty_other)
 
         mont = MontgomeryModel()
@@ -79,5 +79,13 @@ class PointTests(TestCase):
                           X=Mod(0x64daccd2656420216545e5f65221eb,
                                 0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa),
                           Z=Mod(1, 0xaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa))
-        assert not pt.equals(different)
+        self.assertFalse(pt.equals(different))
         self.assertNotEqual(pt, different)
+
+    def test_bytes(self):
+        pt = Point(self.coords,
+                   X=Mod(0x4, self.secp128r1.curve.prime),
+                   Y=Mod(0x6, self.secp128r1.curve.prime),
+                   Z=Mod(2, self.secp128r1.curve.prime))
+        self.assertEqual(bytes(pt), b"\x04\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x04\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x06\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x02")
+        self.assertEqual(bytes(InfinityPoint(self.coords)), b"\x00")
\ No newline at end of file