def wrap(*args, **kwargs):
callargs = getcallargs(fun, *args, **kwargs)
if callargs["sock"] is None:
# This variable is used only to debug leak in tests
COUNT['count'] += 1
with IPSet() as sock:
callargs["sock"] = sock
# We must pop kwargs here, else the function will receive
# a dict of dict
if "kwargs" in callargs:
callargs.update(callargs.pop("kwargs"))
return fun(**callargs) # pylint:disable=star-args
return fun(*args, **kwargs)
def open_netlink(self):
"""Open manually a netlink socket. You can use "with WiSet()" instead"""
if self.sock is None:
self.sock = IPSet()
def get_ipset_socket(**kwargs):
"""Get a socket that one can pass to several WiSet objects"""
return IPSet(**kwargs)
class WiSet(object):
"""Main high level ipset manipulation class.
Every high level ipset operation should be possible with this class,
you probably don't need other helpers of this module, except tools
to load data from kernel (:func:`load_all_ipsets` and :func:`load_ipset`)
For example, you can create and an entry in a ipset just with:
>>> with WiSet(name="mysuperipset") as myset:
>>> myset.create() # add the ipset in the kernel
>>> myset.add("198.51.100.1") # add one IP to the set
Netlink sockets are opened by __enter__ and __exit__ function, so you don't
have to manage it manually if you use the "with" keyword.
If you want to manage it manually (for example for long operation in
a daemon), you can do the following:
>>> myset = WiSet(name="mysuperipset")
>>> myset.open_netlink()
>>> # do stuff
>>> myset.close_netlink()
You can also don't initiate at all any netlink socket, this code will work:
>>> myset = WiSet(name="mysuperipset")
>>> myset.create()
>>> myset.destroy()
But do it very carefully. In that case, a netlink socket will be opened
in background for any operation. No socket will be leaked, but that
can consume resources.
You can also instantiate WiSet objects with :func:`load_all_ipsets` and
:func:`load_ipset`:
>>> all_sets_dict = load_all_ipsets()
>>> one_set = load_ipset(name="myset")
Have a look on content variable if you need list of entries in the Set.
"""
# pylint: disable=too-many-arguments
def __init__(
self,
name=None,
attr_type='hash:ip',
family=AF_INET,
sock=None,
timeout=None,
counters=False,
comment=False,
hashsize=None,
revision=None,
skbinfo=False,
):
self.name = name
self.hashsize = hashsize
self._attr_type = None
self.entry_type = None
self.attr_type = attr_type
self.family = family
self._content = None
self.sock = sock
self.timeout = timeout
self.counters = counters
self.comment = comment
self.revision = revision
self.index = None
self.skbinfo = skbinfo
def open_netlink(self):
"""Open manually a netlink socket. You can use "with WiSet()" instead"""
if self.sock is None:
self.sock = IPSet()
def close_netlink(self):
"""Clone any opened netlink socket"""
if self.sock is not None:
self.sock.close()
self.sock = None
@property
def attr_type(self):
return self._attr_type
@attr_type.setter
def attr_type(self, value):
self._attr_type = value
self.entry_type = value.split(":", 1)[1]
def __enter__(self):
self.open_netlink()
return self
def __exit__(self, exc_type, exc_value, traceback):
self.close_netlink()
@classmethod
def from_netlink(cls, ndmsg, content=False):
"""Create a ipset objects based on a parsed netlink message
:param ndmsg: the netlink message to parse
:param content: should we fill (and parse) entries info (can be slow
on very large set)
:type content: bool
"""
self = cls()
self.attr_type = ndmsg.get_attr("IPSET_ATTR_TYPENAME")
self.name = ndmsg.get_attr("IPSET_ATTR_SETNAME")
self.hashsize = ndmsg.get_attr("IPSET_ATTR_HASHSIZE")
self.family = ndmsg.get_attr("IPSET_ATTR_FAMILY")
self.revision = ndmsg.get_attr("IPSET_ATTR_REVISION")
self.index = ndmsg.get_attr("IPSET_ATTR_INDEX")
data = ndmsg.get_attr("IPSET_ATTR_DATA")
self.timeout = data.get_attr("IPSET_ATTR_TIMEOUT")
flags = data.get_attr("IPSET_ATTR_CADT_FLAGS")
if flags is not None:
self.counters = bool(flags & IPSET_FLAG_WITH_COUNTERS)
self.comment = bool(flags & IPSET_FLAG_WITH_COMMENT)
self.skbinfo = bool(flags & IPSET_FLAG_WITH_SKBINFO)
if content:
self.update_dict_content(ndmsg)
return self
def update_dict_content(self, ndmsg):
"""Update a dictionary statistics with values sent in netlink message
:param ndmsg: the netlink message
:type ndmsg: netlink message
"""
family = "IPSET_ATTR_IPADDR_IPV4"
ip_attr = "IPSET_ATTR_IP_FROM"
if self._content is None:
self._content = {}
timeout = None
entries = ndmsg.get_attr("IPSET_ATTR_ADT").get_attrs("IPSET_ATTR_DATA")
for entry in entries:
key = ""
for parse_type in self.entry_type.split(","):
if parse_type == "ip":
ip = entry.get_attr(ip_attr).get_attr(family)
key += ip
elif parse_type == "net":
ip = entry.get_attr(ip_attr).get_attr(family)
key += ip
cidr = entry.get_attr("IPSET_ATTR_CIDR")
if cidr is not None:
key += "/{0}".format(cidr)
elif parse_type == "iface":
key += entry.get_attr("IPSET_ATTR_IFACE")
elif parse_type == "set":
key += entry.get_attr("IPSET_ATTR_NAME")
elif parse_type == "mark":
key += str(hex(entry.get_attr("IPSET_ATTR_MARK")))
elif parse_type == "port":
proto = entry.get_attr('IPSET_ATTR_PROTO')
if proto is not None:
proto = IP_PROTOCOLS.get(proto, str(proto)).lower()
key += '{proto}:'.format(proto=proto)
key += str(entry.get_attr("IPSET_ATTR_PORT_FROM"))
elif parse_type == "mac":
key += entry.get_attr("IPSET_ATTR_ETHER")
key += ","
key = key.strip(",")
if self.timeout is not None:
timeout = entry.get_attr("IPSET_ATTR_TIMEOUT")
skbmark = entry.get_attr("IPSET_ATTR_SKBMARK")
if skbmark is not None:
# Convert integer to hex for mark/mask
# Only display mask if != 0xffffffff
if skbmark[1] != (2**32 - 1):
skbmark = "/".join([str(hex(mark)) for mark in skbmark])
else:
skbmark = str(hex(skbmark[0]))
entry_flag_parsed = {"physdev": False}
flags = entry.get_attr("IPSET_ATTR_CADT_FLAGS")
if flags is not None:
entry_flag_parsed["physdev"] = bool(flags & IPSET_FLAG_PHYSDEV)
entry_flag_parsed["wildcard"] = bool(
flags & IPSET_FLAG_IFACE_WILDCARD)
value = IPStats(packets=entry.get_attr("IPSET_ATTR_PACKETS"),
bytes=entry.get_attr("IPSET_ATTR_BYTES"),
comment=entry.get_attr("IPSET_ATTR_COMMENT"),
skbmark=skbmark,
timeout=timeout,
**entry_flag_parsed)
self._content[key] = value
def create(self, **kwargs):
"""Insert this Set in the kernel
Many options are set with python object attributes (like comments,
counters, etc). For non-supported type, kwargs are provided. See
:doc:`ipset` documentation for more information.
"""
create_ipset(self.name,
stype=self.attr_type,
family=self.family,
sock=self.sock,
timeout=self.timeout,
comment=self.comment,
counters=self.counters,
hashsize=self.hashsize,
skbinfo=self.skbinfo,
**kwargs)
def destroy(self):
"""Destroy this ipset in the kernel list.
It does not delete this python object (any content or other stored
values are keep in memory). This function will fail if the ipset is
still referenced (by example in iptables rules), you have been warned.
"""
destroy_ipset(self.name, sock=self.sock)
def add(self, entry, **kwargs):
"""Add an entry in this ipset.
If counters are enabled on the set, reset by default the value when
we add the element. Without this reset, kernel sometimes store old
values and can add very strange behavior on counters.
"""
if isinstance(entry, dict):
kwargs.update(entry)
entry = kwargs.pop("entry")
if self.counters:
kwargs["packets"] = kwargs.pop("packets", 0)
kwargs["bytes"] = kwargs.pop("bytes", 0)
skbmark = kwargs.get("skbmark")
if isinstance(skbmark, basestring):
skbmark = skbmark.split('/')
mark = int(skbmark[0], 16)
try:
mask = int(skbmark[1], 16)
except IndexError:
mask = int("0xffffffff", 16)
kwargs["skbmark"] = (mark, mask)
add_ipset_entry(self.name,
entry,
etype=self.entry_type,
sock=self.sock,
**kwargs)
def delete(self, entry, **kwargs):
"""Delete/remove an entry in this ipset"""
delete_ipset_entry(self.name,
entry,
etype=self.entry_type,
sock=self.sock,
**kwargs)
def test(self, entry, **kwargs):
"""Test if an entry is in this ipset"""
return test_ipset_entry(self.name,
entry,
etype=self.entry_type,
sock=self.sock,
**kwargs)
def test_list(self, entries, **kwargs):
"""Test if a list of a set of entries is in this ipset
Return a set of entries found in the IPSet
"""
return test_ipset_entries(self.name,
entries,
etype=self.entry_type,
sock=self.sock,
**kwargs)
def update_content(self):
"""Update the content dictionary with values from kernel"""
self._content = {}
update_wiset_content(self, sock=self.sock)
def flush(self):
"""Flush entries of the ipset"""
flush_ipset(self.name, sock=self.sock)
@property
def content(self):
"""Dictionary of entries in the set.
Keys are IP addresses (as string), values are IPStats tuples.
"""
if self._content is None:
self.update_content()
return self._content
def insert_list(self, entries):
"""Just a small helper to reduce the number of loops in main code."""
for entry in entries:
self.add(entry)
def replace_entries(self, new_list):
"""Replace the content of an ipset with a new list of entries.
This operation is like a flush() and adding all entries one by one. But
this call is atomic: it creates a temporary ipset and swap the content.
:param new_list: list of entries to add
:type new_list: list or :py:class:`set` of basestring or of
keyword arguments dict
"""
temp_name = str(uuid.uuid4())[0:8]
# Get a copy of ourself
temp = load_ipset(self.name, sock=self.sock)
temp.name = temp_name
temp.sock = self.sock
temp.create()
temp.insert_list(new_list)
swap_ipsets(self.name, temp_name, sock=self.sock)
temp.destroy()