Make netfilterqueue a package and add type hints

oremanj · oremanj · commit 1cbea513e6fe · 2022-01-14T01:04:51.000-07:00
diff --git a/CHANGES.txt b/CHANGES.txt
@@ -1,11 +1,13 @@
 v1.0.0, unreleased
     Propagate exceptions raised by the user's packet callback
-    Warn about exceptions raised by the packet callback during queue unbinding
+    Avoid calls to the packet callback during queue unbinding
     Raise an error if a packet verdict is set after its parent queue is closed
     set_payload() now affects the result of later get_payload()
     Handle signals received when run() is blocked in recv()
     Accept packets in COPY_META mode, only failing on an attempt to access the payload
     Add a parameter NetfilterQueue(sockfd=N) that uses an already-opened Netlink socket
+    Add type hints
+    Remove the Packet.payload attribute; it was never safe (treated as a char* but not NUL-terminated) nor documented, but was exposed in the API (perhaps inadvertently).
 
 v0.9.0, 12 Jan 2021
     Improve usability when Packet objects are retained past the callback
diff --git a/MANIFEST.in b/MANIFEST.in
@@ -1,6 +1,3 @@
-include *.txt
-include *.rst
-include *.c
-include *.pyx
-include *.pxd
-recursive-include tests/ *.py
+include LICENSE.txt README.rst CHANGES.txt
+recursive-include netfilterqueue *.py *.pyx *.pxd *.c *.pyi py.typed
+recursive-include tests *.py
diff --git a/ci.sh b/ci.sh
@@ -11,34 +11,42 @@ python setup.py sdist --formats=zip
 
 # ... but not to install it
 pip uninstall -y cython
+python setup.py build_ext
 pip install dist/*.zip
 
 pip install -Ur test-requirements.txt
 
 if [ "$CHECK_LINT" = "1" ]; then
     error=0
-    if ! black --check setup.py tests; then
+    black_files="setup.py tests netfilterqueue"
+    if ! black --check $black_files; then
+        error=$?
+        black --diff $black_files
+    fi
+    mypy --strict -p netfilterqueue || error=$?
+    ( mkdir empty; cd empty; python -m mypy.stubtest netfilterqueue ) || error=$?
+
+    if [ $error -ne 0 ]; then
         cat <<EOF
 !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 
-Formatting problems were found (listed above). To fix them, run
+Problems were found by static analysis (listed above).
+To fix formatting and see remaining errors, run:
 
    pip install -r test-requirements.txt
-   black setup.py tests
+   black $black_files
+   mypy --strict -p netfilterqueue
+   ( mkdir empty; cd empty; python -m mypy.stubtest netfilterqueue )
 
 in your local checkout.
 
-EOF
-        error=1
-    fi
-    if [ "$error" = "1" ]; then
-        cat <<EOF
 !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 !!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!
 EOF
+        exit 1
     fi
-    exit $error
+    exit 0
 fi
 
 cd tests
diff --git a/netfilterqueue/__init__.pxd b/netfilterqueue/__init__.pxd
@@ -201,7 +201,7 @@ cdef class Packet:
 
     # Packet details:
     cdef Py_ssize_t payload_len
-    cdef readonly unsigned char *payload
+    cdef unsigned char *payload
     cdef timeval timestamp
     cdef u_int8_t hw_addr[8]
 
diff --git a/netfilterqueue/__init__.pyi b/netfilterqueue/__init__.pyi
@@ -0,0 +1,45 @@
+import socket
+from enum import IntEnum
+from typing import Callable, Dict, Optional, Tuple
+
+__version__: str
+VERSION: Tuple[int, ...]
+
+COPY_NONE: int
+COPY_META: int
+COPY_PACKET: int
+
+class Packet:
+    hook: int
+    hw_protocol: int
+    id: int
+    mark: int
+    def get_hw(self) -> Optional[bytes]: ...
+    def get_payload(self) -> bytes: ...
+    def get_payload_len(self) -> int: ...
+    def get_timestamp(self) -> float: ...
+    def get_mark(self) -> int: ...
+    def set_payload(self, payload: bytes) -> None: ...
+    def set_mark(self, mark: int) -> None: ...
+    def retain(self) -> None: ...
+    def accept(self) -> None: ...
+    def drop(self) -> None: ...
+    def repeat(self) -> None: ...
+
+class NetfilterQueue:
+    def __new__(self, *, af: int = ..., sockfd: int = ...) -> NetfilterQueue: ...
+    def bind(
+        self,
+        queue_num: int,
+        user_callback: Callable[[Packet], None],
+        max_len: int = ...,
+        mode: int = COPY_PACKET,
+        range: int = ...,
+        sock_len: int = ...,
+    ) -> None: ...
+    def unbind(self) -> None: ...
+    def get_fd(self) -> int: ...
+    def run(self, block: bool = ...) -> None: ...
+    def run_socket(self, s: socket.socket) -> None: ...
+
+PROTOCOLS: Dict[int, str]
diff --git a/netfilterqueue/__init__.pyx b/netfilterqueue/__init__.pyx
@@ -5,7 +5,6 @@ function.
 Copyright: (c) 2011, Kerkhoff Technologies Inc.
 License: MIT; see LICENSE.txt
 """
-VERSION = (0, 9, 0)
 
 # Constants for module users
 COPY_NONE = 0
@@ -24,6 +23,10 @@ DEF SockCopySize = MaxCopySize + SockOverhead
 # Socket queue should hold max number of packets of copysize bytes
 DEF SockRcvSize = DEFAULT_MAX_QUEUELEN * SockCopySize // 2
 
+__package__ = "netfilterqueue"
+
+from ._version import __version__, VERSION
+
 from cpython.exc cimport PyErr_CheckSignals
 
 # A negative return value from this callback will stop processing and
diff --git a/netfilterqueue/_version.py b/netfilterqueue/_version.py
@@ -0,0 +1,4 @@
+# This file is imported from __init__.py and exec'd from setup.py
+
+__version__ = "0.9.0+dev"
+VERSION = (0, 9, 0)
diff --git a/netfilterqueue/py.typed b/netfilterqueue/py.typed
diff --git a/setup.py b/setup.py
@@ -1,7 +1,7 @@
 import os, sys
 from setuptools import setup, Extension
 
-VERSION = "0.9.0"  # Remember to change CHANGES.txt and netfilterqueue.pyx when version changes.
+exec(open("netfilterqueue/_version.py", encoding="utf-8").read())
 
 setup_requires = []
 try:
@@ -10,7 +10,9 @@
 
     ext_modules = cythonize(
         Extension(
-            "netfilterqueue", ["netfilterqueue.pyx"], libraries=["netfilter_queue"]
+            "netfilterqueue.__init__",
+            ["netfilterqueue/__init__.pyx"],
+            libraries=["netfilter_queue"],
         ),
         compiler_directives={"language_level": "3str"},
     )
@@ -21,7 +23,7 @@
         # setup_requires below.
         setup_requires = ["cython"]
     elif not os.path.exists(
-        os.path.join(os.path.dirname(__file__), "netfilterqueue.c")
+        os.path.join(os.path.dirname(__file__), "netfilterqueue/__init__.c")
     ):
         sys.stderr.write(
             "You must have Cython installed (`pip install cython`) to build this "
@@ -31,21 +33,27 @@
         )
         sys.exit(1)
     ext_modules = [
-        Extension("netfilterqueue", ["netfilterqueue.c"], libraries=["netfilter_queue"])
+        Extension(
+            "netfilterqueue.__init__",
+            ["netfilterqueue/__init__.c"],
+            libraries=["netfilter_queue"],
+        )
     ]
 
 setup(
-    ext_modules=ext_modules,
-    setup_requires=setup_requires,
-    python_requires=">=3.6",
     name="NetfilterQueue",
-    version=VERSION,
+    version=__version__,
     license="MIT",
     author="Matthew Fox",
     author_email="matt@tansen.ca",
     url="https://github.com/oremanj/python-netfilterqueue",
     description="Python bindings for libnetfilter_queue",
-    long_description=open("README.rst").read(),
+    long_description=open("README.rst", encoding="utf-8").read(),
+    packages=["netfilterqueue"],
+    ext_modules=ext_modules,
+    include_package_data=True,
+    setup_requires=setup_requires,
+    python_requires=">=3.6",
     classifiers=[
         "Development Status :: 5 - Production/Stable",
         "License :: OSI Approved :: MIT License",
diff --git a/test-requirements.in b/test-requirements.in
@@ -5,3 +5,4 @@ pytest-trio
 async_generator
 black
 platformdirs <= 2.4.0  # needed by black; 2.4.1+ don't support py3.6
+mypy; implementation_name == "cpython"
diff --git a/test-requirements.txt b/test-requirements.txt
@@ -22,8 +22,12 @@ idna==3.3
     # via trio
 iniconfig==1.1.1
     # via pytest
+mypy==0.931 ; implementation_name == "cpython"
+    # via -r test-requirements.in
 mypy-extensions==0.4.3
-    # via black
+    # via
+    #   black
+    #   mypy
 outcome==1.1.0
     # via
     #   pytest-trio
@@ -57,10 +61,14 @@ sortedcontainers==2.4.0
 toml==0.10.2
     # via pytest
 tomli==1.2.3
-    # via black
+    # via
+    #   black
+    #   mypy
 trio==0.19.0
     # via
     #   -r test-requirements.in
     #   pytest-trio
 typing-extensions==4.0.1
-    # via black
+    # via
+    #   black
+    #   mypy
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -5,12 +5,12 @@
 import subprocess
 import sys
 import trio
-import unshare
+import unshare  # type: ignore
 import netfilterqueue
 from functools import partial
-from typing import AsyncIterator, Callable, Optional, Tuple
+from typing import Any, AsyncIterator, Callable, Dict, Optional, Tuple
 from async_generator import asynccontextmanager
-from pytest_trio.enable_trio_mode import *
+from pytest_trio.enable_trio_mode import *  # type: ignore
 
 
 # We'll create three network namespaces, representing a router (which
@@ -45,8 +45,8 @@ def enter_netns() -> None:
     subprocess.run("/sbin/ip link set lo up".split(), check=True)
 
 
-@pytest.hookimpl(tryfirst=True)
-def pytest_runtestloop():
+@pytest.hookimpl(tryfirst=True)  # type: ignore
+def pytest_runtestloop() -> None:
     if os.getuid() != 0:
         # Create a new user namespace for the whole test session
         outer = {"uid": os.getuid(), "gid": os.getgid()}
@@ -93,7 +93,9 @@ async def peer_main(idx: int, parent_fd: int) -> None:
     await peer.connect((peer_ip, peer_port))
 
     # Enter the message-forwarding loop
-    async def proxy_one_way(src, dest):
+    async def proxy_one_way(
+        src: trio.socket.SocketType, dest: trio.socket.SocketType
+    ) -> None:
         while src.fileno() >= 0:
             try:
                 msg = await src.recv(4096)
@@ -121,13 +123,13 @@ def _default_capture_cb(
 
 
 class Harness:
-    def __init__(self):
-        self._received = {}
-        self._conn = {}
-        self.dest_addr = {}
+    def __init__(self) -> None:
+        self._received: Dict[int, trio.MemoryReceiveChannel[bytes]] = {}
+        self._conn: Dict[int, trio.socket.SocketType] = {}
+        self.dest_addr: Dict[int, Tuple[str, int]] = {}
         self.failed = False
 
-    async def _run_peer(self, idx: int, *, task_status):
+    async def _run_peer(self, idx: int, *, task_status: Any) -> None:
         their_ip = PEER_IP[idx]
         my_ip = ROUTER_IP[idx]
         conn, child_conn = trio.socket.socketpair(socket.AF_UNIX, socket.SOCK_SEQPACKET)
@@ -169,10 +171,10 @@ async def _run_peer(self, idx: int, *, task_status):
                 # and its netns goes away. check=False to suppress that error.
                 await trio.run_process(f"ip link delete veth{idx}".split(), check=False)
 
-    async def _manage_peer(self, idx: int, *, task_status):
+    async def _manage_peer(self, idx: int, *, task_status: Any) -> None:
         async with trio.open_nursery() as nursery:
             await nursery.start(self._run_peer, idx)
-            packets_w, packets_r = trio.open_memory_channel(math.inf)
+            packets_w, packets_r = trio.open_memory_channel[bytes](math.inf)
             self._received[idx] = packets_r
             task_status.started()
             async with packets_w:
@@ -183,7 +185,7 @@ async def _manage_peer(self, idx: int, *, task_status):
                     await packets_w.send(msg)
 
     @asynccontextmanager
-    async def run(self):
+    async def run(self) -> AsyncIterator[None]:
         async with trio.open_nursery() as nursery:
             async with trio.open_nursery() as start_nursery:
                 start_nursery.start_soon(nursery.start, self._manage_peer, 1)
@@ -258,14 +260,14 @@ async def capture_packets_to(
         **options: int,
     ) -> AsyncIterator["trio.MemoryReceiveChannel[netfilterqueue.Packet]"]:
 
-        packets_w, packets_r = trio.open_memory_channel(math.inf)
+        packets_w, packets_r = trio.open_memory_channel[netfilterqueue.Packet](math.inf)
         queue_num, nfq = self.bind_queue(partial(cb, packets_w), **options)
         try:
             async with self.enqueue_packets_to(idx, queue_num):
                 async with packets_w, trio.open_nursery() as nursery:
 
                     @nursery.start_soon
-                    async def listen_for_packets():
+                    async def listen_for_packets() -> None:
                         while True:
                             await trio.lowlevel.wait_readable(nfq.get_fd())
                             nfq.run(block=False)
@@ -275,7 +277,7 @@ async def listen_for_packets():
         finally:
             nfq.unbind()
 
-    async def expect(self, idx: int, *packets: bytes):
+    async def expect(self, idx: int, *packets: bytes) -> None:
         for expected in packets:
             with trio.move_on_after(5) as scope:
                 received = await self._received[idx].receive()
@@ -291,13 +293,13 @@ async def expect(self, idx: int, *packets: bytes):
                     f"received {received!r}"
                 )
 
-    async def send(self, idx: int, *packets: bytes):
+    async def send(self, idx: int, *packets: bytes) -> None:
         for packet in packets:
             await self._conn[3 - idx].send(packet)
 
 
 @pytest.fixture
-async def harness() -> Harness:
+async def harness() -> AsyncIterator[Harness]:
     h = Harness()
     async with h.run():
         yield h
