ipaddrset.py

   1 """IP address set manipulation, built on top of ipaddress.py"""
   2
   3 # This file is Free Software.  It was originally written for secnet.
   4 #
   5 # Copyright 2014 Ian Jackson
   6 #
   7 # You may redistribute secnet as a whole and/or modify it under the
   8 # terms of the GNU General Public License as published by the Free
   9 # Software Foundation; either version 3, or (at your option) any
  10 # later version.
  11 #
  12 # You may redistribute this file and/or modify it under the terms of
  13 # the GNU General Public License as published by the Free Software
  14 # Foundation; either version 2, or (at your option) any later version.
  15 # Note however that this version of ipaddrset.py uses the Python
  16 # ipaddr library from Google, which is licenced only under the Apache
  17 # Licence, version 2.0, which is only compatible with the GNU GPL v3
  18 # (or perhaps later versions), and not with the GNU GPL v2.
  19 #
  20 # This software is distributed in the hope that it will be useful,
  21 # but WITHOUT ANY WARRANTY; without even the implied warranty of
  22 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
  23 # GNU General Public License for more details.
  24 #
  25 # You should have received a copy of the GNU General Public License
  26 # along with this software; if not, see
  27 # https://www.gnu.org/licenses/gpl.html.
  28
  29 import ipaddress
  30
  31 _vsns = [6,4]
  32
  33 class IPAddressSet:
  34         "A set of IP addresses"
  35
  36         # constructors
  37         def __init__(self,l=[]):
  38                 "New set contains each IP*Network in the sequence l"
  39                 self._v = {}
  40                 for v in _vsns:
  41                         self._v[v] = [ ]
  42                 self.append(l)
  43
  44         # housekeeping and representation
  45         def _compact(self):
  46                 for v in _vsns:
  47                         self._v[v] = list(
  48                                 ipaddress.collapse_addresses(self._v[v]))
  49         def __repr__(self):
  50                 return "IPAddressSet(%s)" % self.networks()
  51         def str(self,comma=",",none="-"):
  52                 "Human-readable string with controllable delimiters"
  53                 if self:
  54                         return comma.join(map(str, self.networks()))
  55                 else:
  56                         return none
  57         def __str__(self):
  58                 return self.str()
  59
  60         # mutators
  61         def append(self,l):
  62                 "Appends each IP*Network in the sequence l to self"
  63                 self._append(l)
  64                 self._compact()
  65
  66         def _append(self,l):
  67                 "Appends each IP*Network in the sequence l to self"
  68                 for a in l:
  69                         self._v[a.version].append(a)
  70
  71         # enquirers including standard comparisons
  72         def __bool__(self):
  73                 for v in _vsns:
  74                         if self._v[v]:
  75                                 return True
  76                 return False
  77         __nonzero__=__bool__ # for python2
  78
  79         def __eq__(self,other):
  80                 for v in _vsns:
  81                         if self._v[v] != other._v[v]:
  82                                 return False
  83                 return True
  84         def __ne__(self,other): return not self.__eq__(other)
  85         def __ge__(self,other):
  86                 """True iff self completely contains IPAddressSet other"""
  87                 for o in other:
  88                         if not self._contains_net(o):
  89                                 return False
  90                 return True
  91         def __le__(self,other): return other.__ge__(self)
  92         def __gt__(self,other): return self!=other and other.__ge__(self)
  93         def __lt__(self,other): return other.__gt__(self)
  94
  95         def __cmp__(self,other):
  96                 if self==other: return 0
  97                 if self>=other: return +1
  98                 if self<=other: return -1
  99                 return NotImplemented
 100
 101         def __iter__(self):
 102                 "Iterates over minimal list of distinct IPNetworks in this set"
 103                 for v in _vsns:
 104                         for i in self._v[v]:
 105                                 yield i
 106
 107         def networks(self):
 108                 "Returns miminal list of distinct IPNetworks in this set"
 109                 return [i for i in self]
 110
 111         # set operations
 112         def intersection(self,other):
 113                 "Returns the intersection; does not modify self"
 114                 r = IPAddressSet()
 115                 for v in _vsns:
 116                         for i in self._v[v]:
 117                                 for j in other._v[v]:
 118                                         if i.overlaps(j):
 119                                                 if i.prefixlen > j.prefixlen:
 120                                                         r._append([i])
 121                                                 else:
 122                                                         r._append([j])
 123                 return r
 124         def union(self,other):
 125                 "Returns the union; does not modify self"
 126                 r = IPAddressSet()
 127                 r._append(self.networks())
 128                 r._append(other.networks())
 129                 r._compact()
 130                 return r
 131
 132         def _contains_net(self,n):
 133                 """True iff self completely contains IPNetwork n"""
 134                 for i in self:
 135                         if i.overlaps(n) and n.prefixlen >= i.prefixlen:
 136                                 return True
 137                 return False
 138
 139         def contains(self,thing):
 140                 """Returns True iff self completely contains thing.
 141                    thing may be an IPNetwork or an IPAddressSet"""
 142                 try:
 143                         v = [thing.version]
 144                 except KeyError:
 145                         v = None
 146                 if v:
 147                         return self._contains_net(ipaddress.ip_network(thing))
 148                 else:
 149                         return self.__ge__(thing)
 150
 151 def complete_set():
 152         "Returns a set containing all addresses"
 153         s=IPAddressSet()
 154         for v in _vsns:
 155                 if v==6: a=ipaddress.IPv6Address(0)
 156                 elif v==4: a=ipaddress.IPv4Address(0)
 157                 else: raise "internal error"
 158                 n=ipaddress.ip_network("%s/0" % a)
 159                 s.append([n])
 160         return s