ipaddrset.py

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