Replace Cython with pure Python

.github/workflows/tests.yml

@@ -18,7 +18,7 @@ jobs:
       - name: Install package
         run: poetry install
       - name: Insert test data
-        run: poetry run tests/data/insert.py
+        run: poetry run python tests/data/insert.py
       - name: Run tests
         run: poetry run pytest --cov=diamond_miner --cov-report=xml
       - uses: codecov/codecov-action@v3
@@ -41,30 +41,13 @@ jobs:
 
   pypi:
     needs: [test]
-    runs-on: ${{ matrix.os }}
-    strategy:
-      matrix:
-        os: [ubuntu-latest, macos-latest]
+    runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v3
       - uses: actions/setup-python@v4
         with:
           python-version: "3.x"
-      - name: Set up QEMU
-        if: runner.os == 'Linux'
-        uses: docker/setup-qemu-action@v2
-        with:
-          platforms: all
-      - name: Build wheels
-        uses: pypa/[email protected]
-        with:
-          output-dir: dist
-      - name: Upload artifacts
-        uses: actions/upload-artifact@v3
-        with:
-          path: dist/*.whl
       - uses: dioptra-io/publish-python-action@v1
         with:
           password: ${{ secrets.PYPI_TOKEN }}
           upload: ${{ startsWith(github.ref, 'refs/tags/v') }}
-          wheel: false

diamond_miner/format.py

@@ -0,0 +1,24 @@
+from ipaddress import IPv6Address
+
+
+def format_probe(
+    dst_addr_v6: int, src_port: int, dst_port: int, ttl: int, protocol: str
+) -> str:
+    """
+    Create a Caracal probe string.
+    Examples:
+        >>> from diamond_miner.format import format_probe
+        >>> format_probe(281470816487432, 24000, 33434, 1, "icmp")
+        '::ffff:808:808,24000,33434,1,icmp'
+    """
+    return f"{format_ipv6(dst_addr_v6)},{src_port},{dst_port},{ttl},{protocol}"
+
+
+def format_ipv6(addr: int) -> str:
+    """
+    Convert an IPv6 UInt128 to a string.
+        >>> from diamond_miner.format import format_ipv6
+        >>> format_ipv6(281470816487432)
+        '::ffff:808:808'
+    """
+    return str(IPv6Address(addr))

diamond_miner/grid.pyx → diamond_miner/grid.py

@@ -1,45 +1,45 @@
+from collections.abc import Iterator, Sequence
 from random import randint
+from typing import Any
 
 from pygfc import Permutation
 
 
 # The Cython version is approx. 2x faster on a M1 CPU.
-cdef class ParameterGrid:
-    cdef tuple parameters
-    cdef tuple size_
-    cdef len
-
-    def __init__(self, *parameters):
+class ParameterGrid:
+    def __init__(self, *parameters: Sequence):
         self.parameters = parameters
         self.size_ = tuple(len(p) for p in self.parameters)
         self.len = 1
         for dim in self.size_:
             self.len *= dim
 
-    def __getitem__(self, index):
+    def __getitem__(self, index: Sequence[int] | int) -> Sequence[Any]:
         if isinstance(index, int):
             if index >= len(self):
                 raise IndexError("index out of range")
             index = self.linear_to_subscript(index)
         return [p[i] for p, i in zip(self.parameters, index)]
 
-    def __iter__(self):
+    def __iter__(self) -> Iterator[Sequence[Any]]:
         return (self[index] for index in range(len(self)))
 
-    def __len__(self):
+    def __len__(self) -> int:
         return self.len
 
     @property
-    def size(self):
+    def size(self) -> Sequence[int]:
         return self.size_
 
-    def shuffled(self, int rounds = 6, seed = None):
-        seed = seed or randint(0, 2 ** 64 - 1)
+    def shuffled(
+        self, rounds: int = 6, seed: int | None = None
+    ) -> Iterator[Sequence[Any]]:
+        seed = seed or randint(0, 2**64 - 1)
         perm = Permutation(len(self), rounds, seed)
         return (self[index] for index in perm)
 
-    cdef list linear_to_subscript(self, index):
-        cdef list coordinates = []
+    def linear_to_subscript(self, index: int) -> Sequence[int]:
+        coordinates = []
         for dim in self.size_:
             index, coordinate = divmod(index, dim)
             coordinates.append(coordinate)

diamond_miner/mappers.pyx → diamond_miner/mappers.py

@@ -6,19 +6,12 @@
 """
 import random
 
-from libc.stdint cimport uint8_t, uint16_t
-
 from pygfc import Permutation
 
 from diamond_miner.defaults import DEFAULT_PREFIX_SIZE_V4
 
 
-cdef extern from *:
-    ctypedef long long int128_t "__int128_t"
-    ctypedef unsigned long long uint128_t "__uint128_t"
-
-
-cdef class SequentialFlowMapper:
+class SequentialFlowMapper:
     """
     Maps flow 0 to address 0, flow 1 to address 1, and so on until we have done
     the whole prefix. It then increases the port number in the same manner.
@@ -27,25 +20,23 @@
         >>> from diamond_miner.mappers import SequentialFlowMapper
         >>> mapper = SequentialFlowMapper()
         >>> mapper.offset(1)
-        >>> (1, 0)
+        (1, 0)
     """
 
-    cdef uint128_t prefix_size
-
-    def __init__(self, uint128_t prefix_size = DEFAULT_PREFIX_SIZE_V4):
+    def __init__(self, prefix_size: int = DEFAULT_PREFIX_SIZE_V4):
         assert prefix_size > 0, "prefix_size must be positive."
         self.prefix_size = prefix_size
 
-    cpdef uint128_t flow_id(self, uint128_t addr_offset, uint16_t port_offset, uint128_t prefix = 0):
+    def flow_id(self, addr_offset: int, port_offset: int, prefix: int = 0) -> int:
         return addr_offset + port_offset
 
-    cpdef (uint128_t, uint16_t) offset(self, uint128_t flow_id, uint128_t prefix = 0):
+    def offset(self, flow_id: int, prefix: int = 0) -> tuple[int, int]:
         if flow_id < self.prefix_size:
             return flow_id, 0
         return self.prefix_size - 1, flow_id - self.prefix_size + 1
 
 
-cdef class IntervalFlowMapper:
+class IntervalFlowMapper:
     """
     Similar to the `SequentialFlowMapper` but with an increment >= 1.
     This allows to target addresses .1, .33, .65, ... in priority,
@@ -55,14 +46,10 @@ def __init__(self, uint128_t prefix_size = DEFAULT_PREFIX_SIZE_V4):
         >>> from diamond_miner.mappers import IntervalFlowMapper
         >>> mapper = IntervalFlowMapper()
         >>> mapper.offset(1)
-        >>> (33, 0)
+        (33, 0)
     """
 
-    cdef uint128_t period
-    cdef uint128_t prefix_size
-    cdef uint128_t step
-
-    def __init__(self, uint128_t prefix_size = DEFAULT_PREFIX_SIZE_V4, uint128_t step = 32):
+    def __init__(self, prefix_size: int = DEFAULT_PREFIX_SIZE_V4, step: int = 32):
         assert prefix_size > 0, "prefix_size must be positive."
         assert prefix_size % 2 == 0, "prefix_size must be pair."
         assert step > 0, "step must be positive."
@@ -71,23 +58,23 @@ def __init__(self, uint128_t prefix_size = DEFAULT_PREFIX_SIZE_V4, uint128_t ste
         self.prefix_size = prefix_size
         self.step = step
 
-    cpdef uint128_t flow_id(self, uint128_t addr_offset, uint16_t port_offset, uint128_t prefix = 0):
+    def flow_id(self, addr_offset: int, port_offset: int, prefix: int = 0) -> int:
         if addr_offset == 0:
             return self.prefix_size - 1
         if port_offset != 0:
             return self.prefix_size + port_offset - 1
         q, r = divmod(addr_offset - 1, self.step)
         return r * self.period + q
 
-    cpdef (uint128_t, uint16_t) offset(self, uint128_t flow_id, uint128_t prefix = 0):
+    def offset(self, flow_id: int, prefix: int = 0) -> tuple[int, int]:
         if flow_id < self.prefix_size - 1:
             return ((flow_id * self.step) % (self.prefix_size - 1)) + 1, 0
         if flow_id == self.prefix_size - 1:
             return 0, 0
         return self.prefix_size - 1, flow_id - self.prefix_size + 1
 
 
-cdef class ReverseByteFlowMapper:
+class ReverseByteFlowMapper:
     """
     Maps flow `n` to address `reverse(n)` until we have done the whole prefix.
     It then increases the port number sequentially.
@@ -96,33 +83,33 @@ def __init__(self, uint128_t prefix_size = DEFAULT_PREFIX_SIZE_V4, uint128_t ste
         >>> from diamond_miner.mappers import ReverseByteFlowMapper
         >>> mapper = ReverseByteFlowMapper()
         >>> mapper.offset(1)
-        >>> (129, 0)
+        (129, 0)
     """
 
-    cpdef uint128_t flow_id(self, uint128_t addr_offset, uint16_t port_offset, uint128_t prefix = 0):
+    def flow_id(self, addr_offset: int, port_offset: int, prefix: int = 0) -> int:
         assert addr_offset < 256
         if addr_offset == 0:
             return 255
         if port_offset != 0:
             return 255 + port_offset
         return self.reverse_byte(addr_offset - 1)
 
-    cpdef (uint128_t, uint16_t) offset(self, uint128_t flow_id, uint128_t prefix = 0):
+    def offset(self, flow_id: int, prefix: int = 0) -> tuple[int, int]:
         if flow_id < 255:
             return self.reverse_byte(flow_id) + 1, 0
         if flow_id == 255:
             return 0, 0
         return 255, flow_id - 255
 
-    cdef uint8_t reverse_byte(self, uint8_t i):
+    def reverse_byte(self, i: int) -> int:
         # https://stackoverflow.com/a/2602885
         i = (i & 0xF0) >> 4 | (i & 0x0F) << 4
         i = (i & 0xCC) >> 2 | (i & 0x33) << 2
         i = (i & 0xAA) >> 1 | (i & 0x55) << 1
         return i
 
 
-cdef class RandomFlowMapper:
+class RandomFlowMapper:
     """
     Similar to the `SequentialFlowMapper` but with a random mapping between flow IDs and addresses.
     The mapping is randomized by prefix.
@@ -131,32 +118,29 @@ def __init__(self, uint128_t prefix_size = DEFAULT_PREFIX_SIZE_V4, uint128_t ste
         >>> from diamond_miner.mappers import RandomFlowMapper
         >>> mapper = RandomFlowMapper(seed=2022)
         >>> mapper.offset(1, prefix=1)
-        >>> (34, 0)
+        (34, 0)
         >>> mapper.offset(1, prefix=2)
-        >>> (145, 0)
+        (145, 0)
     """
 
-    cdef list permutations
-    cdef uint128_t prefix_size
-
-    def __init__(self, int seed, uint128_t prefix_size = DEFAULT_PREFIX_SIZE_V4):
+    def __init__(self, seed: int, prefix_size: int = DEFAULT_PREFIX_SIZE_V4):
         # We can generate a random permutation up to 2^64-1 only.
         assert prefix_size > 0, "prefix_size must be positive."
         self.permutations = []
-        self.prefix_size = min(prefix_size, (2 ** 64) - 1)
+        self.prefix_size = min(prefix_size, (2**64) - 1)
         random.seed(seed)
         for i in range(1024):
-            perm = Permutation(self.prefix_size, 3, random.randint(0, 2 ** 64))
+            perm = Permutation(self.prefix_size, 3, random.randint(0, 2**64))
             self.permutations.append(perm)
 
-    cpdef uint128_t flow_id(self, uint128_t addr_offset, uint16_t port_offset, uint128_t prefix):
+    def flow_id(self, addr_offset: int, port_offset: int, prefix: int) -> int:
         assert addr_offset < self.prefix_size
         if port_offset != 0:
             return self.prefix_size + port_offset - 1
         perm = self.permutations[prefix % len(self.permutations)]
-        return perm.inv(addr_offset)
+        return perm.inv(addr_offset)  # type: ignore
 
-    cpdef (uint128_t, uint16_t) offset(self, uint128_t flow_id, uint128_t prefix):
+    def offset(self, flow_id: int, prefix: int) -> tuple[int, int]:
         if flow_id < self.prefix_size:
             perm = self.permutations[prefix % len(self.permutations)]
             return perm[flow_id], 0