diff --git a/pyroute2/conntrack.py b/pyroute2/conntrack.py index 36e4d26cd789e8727d91aaf862540ba58dfca8cb..0f380be64fa5e9b9cb98875468e2930094021848 100644 --- a/pyroute2/conntrack.py +++ b/pyroute2/conntrack.py @@ -9,6 +9,9 @@ from pyroute2.netlink.nfnetlink.nfctsocket import NFCTSocket class NFCTATcpProtoInfo(object): + __slots__ = ('state', 'wscale_orig', 'wscale_reply', 'flags_orig', + 'flags_reply') + def __init__(self, state, wscale_orig=None, wscale_reply=None, flags_orig=None, flags_reply=None): self.state = state @@ -49,6 +52,9 @@ class NFCTATcpProtoInfo(object): class ConntrackEntry(object): + __slots__ = ('tuple_orig', 'tuple_reply', 'status', 'timeout', + 'protoinfo', 'mark', 'id', 'use') + def __init__(self, family, tuple_orig, tuple_reply, cta_status, cta_timeout, cta_protoinfo, cta_mark, cta_id, cta_use): self.tuple_orig = NFCTAttrTuple.from_netlink(family, tuple_orig) @@ -86,6 +92,8 @@ class Conntrack(NFCTSocket): """ High level conntrack functions """ + def __init__(self, nlm_generator=True, **kwargs): + super(Conntrack, self).__init__(nlm_generator=nlm_generator, **kwargs) def stat(self): """ Return current statistics per CPU @@ -107,16 +115,16 @@ class Conntrack(NFCTSocket): Same result than /proc/sys/net/netfilter/nf_conntrack_count file or conntrack -C command """ - ndmsg = super(Conntrack, self).count() - return ndmsg[0].get_attr('CTA_STATS_GLOBAL_ENTRIES') + for ndmsg in super(Conntrack, self).count(): + return ndmsg.get_attr('CTA_STATS_GLOBAL_ENTRIES') def conntrack_max_size(self): """ Return the max size of connection tracking table /proc/sys/net/netfilter/nf_conntrack_max """ - ndmsg = super(Conntrack, self).conntrack_max_size() - return ndmsg[0].get_attr('CTA_STATS_GLOBAL_MAX_ENTRIES') + for ndmsg in super(Conntrack, self).conntrack_max_size(): + return ndmsg.get_attr('CTA_STATS_GLOBAL_MAX_ENTRIES') def delete(self, entry): if isinstance(entry, ConntrackEntry): @@ -125,7 +133,12 @@ class Conntrack(NFCTSocket): tuple_orig = entry else: raise NotImplementedError() - self.entry('del', tuple_orig=tuple_orig) + for ndmsg in self.entry('del', tuple_orig=tuple_orig): + return ndmsg + + def entry(self, cmd, **kwargs): + for res in super(Conntrack, self).entry(cmd, **kwargs): + return res def dump_entries(self, mark=None, mark_mask=0xffffffff, tuple_orig=None, tuple_reply=None): diff --git a/pyroute2/netlink/nfnetlink/nfctsocket.py b/pyroute2/netlink/nfnetlink/nfctsocket.py index e8abd331540cd9acbd9182dd72e90e6463320f27..949183431d9f9850a83a2aa6b762422d469acfda 100644 --- a/pyroute2/netlink/nfnetlink/nfctsocket.py +++ b/pyroute2/netlink/nfnetlink/nfctsocket.py @@ -367,6 +367,11 @@ class NFCTAttr(object): class NFCTAttrTuple(NFCTAttr): + + __slots__ = ('saddr', 'daddr', 'proto', 'sport', 'dport', + 'icmp_id', 'icmp_type', 'family', '_attr_ip', + '_attr_icmp') + def __init__(self, family=socket.AF_INET, saddr=None, daddr=None, proto=None, sport=None, dport=None, icmp_id=None, icmp_type=None, icmp_code=None): @@ -587,8 +592,8 @@ class NFCTSocket(NetlinkSocket): IPCTNL_MSG_CT_GET_UNCONFIRMED: nfct_msg, }.items()) - def __init__(self, nfgen_family=socket.AF_INET): - super(NFCTSocket, self).__init__(family=NETLINK_NETFILTER) + def __init__(self, nfgen_family=socket.AF_INET, **kwargs): + super(NFCTSocket, self).__init__(family=NETLINK_NETFILTER, **kwargs) self.register_policy(self.policy) self._nfgen_family = nfgen_family diff --git a/tests/general/test_conntrack.py b/tests/general/test_conntrack.py index 4a271133aaae082ca3002a0ca9f75aaa9bf51728..32a800fa1adb7a786a5581a669ee3dd414c93e63 100644 --- a/tests/general/test_conntrack.py +++ b/tests/general/test_conntrack.py @@ -86,7 +86,7 @@ class TestConntrack(BasicSetup): # These values should be pretty the same, but the call is not atomic # so some sessions may end or begin that time. Thus the difference # may occur, but should not be significant. - assert abs(len(self.ct.dump()) - self.ct.count()) < 3 + assert abs(len(list(self.ct.dump())) - self.ct.count()) < 3 class TestNFCTSocket(BasicSetup):