chiark / gitweb /
make-secnet-sites: Use argparse rather than ad-hoc parser
[secnet.git] / 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