Commit | Line | Data |
---|---|---|
6437945a IJ |
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 |