ac107458/python-capport
3
0
Fork 0
Code Activity

flake8 linting

Browse Source
This commit is contained in:
Stefan Bühler
2023-01-12 13:16:51 +01:00
parent f9c9b98868
commit 9325950f51
21 changed files with 220 additions and 126 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/capport/utils/nft_socket.py
+16 -13
View File
@@ -2,13 +2,12 @@ from __future__ import annotations
import contextlib
import typing
import threading
from pyroute2.netlink.nfnetlink import nftsocket as _nftsocket
import pyroute2.netlink
import pyroute2.netlink.nlsocket
from pyroute2.netlink.nfnetlink import nfgen_msg
from pyroute2.netlink.nfnetlink import NFNL_SUBSYS_NFTABLES
import pyroute2.netlink # type: ignore
import pyroute2.netlink.nlsocket # type: ignore
from pyroute2.netlink.nfnetlink import NFNL_SUBSYS_NFTABLES # type: ignore
from pyroute2.netlink.nfnetlink import nfgen_msg # type: ignore
from pyroute2.netlink.nfnetlink import nftsocket as _nftsocket # type: ignore
NFPROTO_INET: int = 1 # nfgen_family "ipv4+ipv6"; strace decodes this as "AF_UNSPEC"
@@ -45,10 +44,12 @@ def _monkey_patch_pyroute2():
overwrite_methods(subcls)
overwrite_methods(pyroute2.netlink.nlmsg_base)
_monkey_patch_pyroute2()
def _build(msg_class: typing.Type[_NlMsgBase], /, attrs: typing.Dict={}, header: typing.Dict={}, **fields) -> _NlMsgBase:
def _build(msg_class: typing.Type[_NlMsgBase], /, attrs: dict = {}, header: dict = {}, **fields) -> _NlMsgBase:
msg = msg_class()
for key, value in header.items():
msg['header'][key] = value
@@ -66,7 +67,9 @@ def _build(msg_class: typing.Type[_NlMsgBase], /, attrs: typing.Dict={}, header:
# support passing nested attributes as dicts of subattributes (or lists of those)
if prime['nla_array']:
value = [
_build(nla_class, attrs=elem) if not isinstance(elem, pyroute2.netlink.nlmsg_base) and isinstance(elem, dict) else elem
_build(nla_class, attrs=elem)
if not isinstance(elem, pyroute2.netlink.nlmsg_base) and isinstance(elem, dict)
else elem
for elem in value
]
elif not isinstance(value, pyroute2.netlink.nlmsg_base) and isinstance(value, dict):
@@ -76,7 +79,7 @@ def _build(msg_class: typing.Type[_NlMsgBase], /, attrs: typing.Dict={}, header:
class NFTSocket(pyroute2.netlink.nlsocket.NetlinkSocket):
policy: typing.Dict[int, typing.Type[_nftsocket.nft_gen_msg]] = _nftsocket.NFTSocket.policy
policy: dict[int, typing.Type[_nftsocket.nft_gen_msg]] = _nftsocket.NFTSocket.policy
def __init__(self) -> None:
super().__init__(family=pyroute2.netlink.NETLINK_NETFILTER)
@@ -98,15 +101,15 @@ class NFTSocket(pyroute2.netlink.nlsocket.NetlinkSocket):
# abort does nothing if commit went through
tx.abort()
def nft_put(self, msg_type: int, msg_flags: int=0, /, *, attrs: typing.Dict={}, **fields) -> None:
def nft_put(self, msg_type: int, msg_flags: int = 0, /, *, attrs: dict = {}, **fields) -> None:
with self.begin() as tx:
tx.put(msg_type, msg_flags, attrs=attrs, **fields)
def nft_dump(self, msg_type: int, msg_flags: int=0, /, *, attrs: typing.Dict={}, **fields) -> None:
def nft_dump(self, msg_type: int, msg_flags: int = 0, /, *, attrs: dict = {}, **fields) -> None:
msg_flags |= pyroute2.netlink.NLM_F_DUMP
return self.nft_get(msg_type, msg_flags, attrs=attrs, **fields)
def nft_get(self, msg_type: int, msg_flags: int=0, /, *, attrs: typing.Dict={}, **fields) -> None:
def nft_get(self, msg_type: int, msg_flags: int = 0, /, *, attrs: dict = {}, **fields) -> None:
msg_class: typing.Type[_nftsocket.nft_gen_msg] = self.policy[msg_type]
msg_type = (NFNL_SUBSYS_NFTABLES << 8) | msg_type
msg_flags |= pyroute2.netlink.NLM_F_REQUEST
@@ -207,7 +210,7 @@ class NFTTransaction:
self._data += self._final_msg.data
self._final_msg = msg
def put(self, msg_type: int, msg_flags: int=0, /, *, attrs: typing.Dict={}, **fields) -> None:
def put(self, msg_type: int, msg_flags: int = 0, /, *, attrs: dict = {}, **fields) -> None:
msg_class: typing.Type[_nftsocket.nft_gen_msg] = self._socket.policy[msg_type]
msg_flags |= pyroute2.netlink.NLM_F_REQUEST # always set REQUEST
msg_flags &= ~pyroute2.netlink.NLM_F_ACK # make sure ACK is not set!