start of test driver impl.
[python-rwhoisd.git] / rwhoisd / Cidr.py
1 # This file is part of python-rwhoisd
2 #
3 # Copyright (C) 2003, 2008 David E. Blacka
4 #
5 # This program is free software; you can redistribute it and/or modify
6 # it under the terms of the GNU General Public License as published by
7 # the Free Software Foundation; either version 2 of the License, or
8 # (at your option) any later version.
9 #
10 # This program is distributed in the hope that it will be useful, but
11 # WITHOUT ANY WARRANTY; without even the implied warranty of
12 # MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
13 # General Public License for more details.
14 #
15 # You should have received a copy of the GNU General Public License
16 # along with this program; if not, write to the Free Software
17 # Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307
18 # USA
19
20 import socket, types, copy, bisect, struct
21 import v6addr
22
23 def new(address, netlen = -1):
24     """Construct either a CidrV4 or CidrV6 object."""
25
26     # ints are probably v4 addresses.
27     if isinstance(address, int):
28         return CidrV4(address, netlen)
29     # longs could be v4 addresses, but we will only assume so if the
30     # value is small.
31     if isinstance(address, long):
32         if address <= pow(2, 32):
33             return CidrV4(address, netlen)
34         return CidrV6(address, netlen)
35     # otherwise, a colon in the address is a dead giveaway.
36     if ":" in address:
37         return CidrV6(address, netlen)
38     return CidrV4(address, netlen)
39
40 def valid_cidr(address):
41     """Returns the converted Cidr object if 'address' is valid CIDR
42     notation, False if not.  For the purposes of this module, valid
43     CIDR notation consists of a IPv4 or IPv6 address with an optional
44     trailing "/netlen"."""
45
46     if isinstance(address, Cidr): return address
47     try:
48         c = new(address)
49         return c
50     except (ValueError, socket.error):
51         return False
52
53
54 class Cidr:
55     """A class representing a generic CIDRized network value."""
56
57     def _initialize(self, address, netlen):
58         """This a common constructor that is used by the subclasses."""
59
60         if isinstance(address, int) or \
61                 isinstance(address, long) and netlen >= 0:
62             self.numaddr, self.netlen = address, netlen
63             self.addr = self._convert_ipaddr(address)
64             self.calc()
65             return
66
67         if not self.is_valid_cidr(address):
68             raise ValueError, \
69                 repr(address) + " is not a valid CIDR representation"
70
71         if netlen < 0:
72             if type(address) == types.StringType:
73                 if "/" in address:
74                     self.addr, self.netlen = address.split("/", 1)
75                 else:
76                     self.addr, self.netlen = address, self._max_netlen()
77             elif type(address) == types.TupleType:
78                 self.addr, self.netlen = address
79             else:
80                 raise TypeError, "address must be a string or a tuple"
81         else:
82             self.addr, self.netlen = address, netlen
83
84
85         # convert string network lengths to integer
86         if type(self.netlen) == types.StringType:
87             self.netlen = int(self.netlen)
88
89         self.calc()
90
91     def __str__(self):
92         return self.addr + "/" + str(self.netlen)
93
94     def __repr__(self):
95         return "<" + str(self) + ">"
96
97     def __cmp__(self, other):
98         """One CIDR network block is less than another if the start
99         address is numerically less or if the block is larger.  That
100         is, supernets will sort before subnets.  This ordering allows
101         for an efficient search for subnets of a given network."""
102
103         res = self._base_mask(self.numaddr) - other._base_mask(other.numaddr)
104         if res == 0: res = self.netlen - other.netlen
105         if res < 0: return -1
106         if res > 0: return 1
107         return 0
108
109     def calc(self):
110         """This method should be called after any change to the main
111         internal state: netlen or numaddr."""
112
113         # make sure the network length is valid
114         if not self.is_valid_netlen(self.netlen):
115             raise TypeError, "network length must be between 0 and %d" % \
116                 (self._max_netlen())
117
118         # convert the string ipv4 address to a 32bit number
119         self.numaddr = self._convert_ipstr(self.addr)
120         # calculate our netmask
121         self.mask = self._mask(self.netlen)
122         # force the cidr address into correct masked notation
123         self.numaddr &= self.mask
124
125         # convert the number back to a string to normalize the string
126         self.addr = self._convert_ipaddr(self.numaddr)
127
128     def is_supernet(self, other):
129         """returns True if the other Cidr object is a supernet (an
130         enclosing network block) of this one.  A Cidr object is a
131         supernet of itself."""
132         return other.numaddr & self.mask == self.numaddr
133
134     def is_subnet(self, other):
135         """returns True if the other Cidr object is a subnet (an
136         enclosednetwork block) of this one.  A Cidr object is a
137         subnet of itself."""
138         return self.numaddr & other.mask == other.numaddr
139
140     def netmask(self):
141         """return the netmask of this Cidr network"""
142         return self._convert_ipaddr(self.mask)
143
144     def length(self):
145         """return the length (in number of addresses) of this network block"""
146         return 1 << (self._max_netlen() - self.netlen);
147
148     def end(self):
149         """return the last IP address in this network block"""
150         return self._convert_ipaddr(self.numaddr + self.length() - 1)
151
152     def to_netblock(self):
153         return (self.addr, self.end())
154
155     def clone(self):
156         # we can get away with a shallow copy (so far)
157         return copy.copy(self)
158
159     def is_ipv6(self):
160         if isinstance(self, CidrV6): return True
161         return False
162
163     def is_valid_cidr(self, address):
164         if "/" in address:
165             addr, netlen = address.split("/", 1)
166             netlen = int(netlen)
167         else:
168             addr, netlen = address, 0
169         return self._is_valid_address(addr) and self.is_valid_netlen(netlen)
170
171     def is_valid_netlen(self, netlen):
172         if netlen < 0: return False
173         if netlen > self._max_netlen(): return False
174         return True
175
176
177 class CidrV4(Cidr):
178     """A class representing a CIDRized IPv4 network value.
179
180     Specifically, it is representing a contiguous IPv4 network block
181     that can be expressed as a ip-address/network-length pair."""
182
183     base_mask = 0xFFFFFFFF
184     msb_mask  = 0x80000000
185
186     def __init__(self, address, netlen = -1):
187         """This takes either a formatted string in CIDR notation:
188         (e.g., "127.0.0.1/32"), A tuple consisting of an formatting
189         string IPv4 address and a numeric network length, or the same
190         as two arguments."""
191
192         self._initialize(address, netlen)
193
194     def _is_valid_address(self, address):
195         """Returns True if the address is a legal IPv4 address."""
196         try:
197             self._convert_ipstr(address)
198             return True
199         except socket.error:
200             return False
201
202     def _base_mask(self, numaddr):
203         return numaddr & CidrV4.base_mask
204
205     def _max_netlen(self):
206         return 32
207
208     def _convert_ipstr(self, addr):
209         packed_numaddr = socket.inet_aton(addr)
210         return struct.unpack("!I", packed_numaddr)[0]
211
212     def _convert_ipaddr(self, numaddr):
213         packed_numaddr = struct.pack("!I", numaddr)
214         return socket.inet_ntoa(packed_numaddr)
215
216     def _mask(self, len):
217         return self._base_mask(CidrV4.base_mask << (32 - len))
218
219 class CidrV6(Cidr):
220     """A class representing a CIDRized IPv6 network value.
221
222     Specifically, it is representing a contiguous IPv6 network block
223     that can be expressed as a ipv6-address/network-length pair."""
224
225     base_mask  = 0xFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFFL # 128-bits of all ones.
226     msb_mask   = 0x80000000000000000000000000000000L
227     lower_mask = 0x0000000000000000FFFFFFFFFFFFFFFFL
228     upper_mask = 0xFFFFFFFFFFFFFFFF0000000000000000L
229
230     def __init__(self, address, netlen = -1):
231
232         self._initialize(address, netlen)
233
234     def _is_valid_address(self, address):
235         try:
236             self._convert_ipstr(address)
237             return True
238         except socket.error, e:
239             print "Failed to convert address string '%s': " + str(e) % (address)
240             return False
241
242     def _base_mask(self, numaddr):
243         return numaddr & CidrV6.base_mask
244
245     def _max_netlen(self):
246         return 128
247
248     def _convert_ipstr(self, addr):
249         packed_numaddr = socket.inet_pton(socket.AF_INET6, addr)
250         upper, lower = struct.unpack("!QQ", packed_numaddr);
251         return (upper << 64) | lower
252
253     def _convert_ipaddr(self, numaddr):
254         upper = (numaddr & CidrV6.upper_mask) >> 64;
255         lower = numaddr & CidrV6.lower_mask;
256         packed_numaddr = struct.pack("!QQ", upper, lower)
257         return socket.inet_ntop(socket.AF_INET6, packed_numaddr)
258
259     def _mask(self, len):
260         return self._base_mask(CidrV6.base_mask << (128 - len))
261
262
263 def netblock_to_cidr(start, end):
264     """Convert an arbitrary network block expressed as a start and end
265     address (inclusive) into a series of valid CIDR blocks."""
266
267     def largest_prefix(length, max_netlen, msb_mask):
268         # calculates the largest network length (smallest mask length)
269         # that can fit within the block length.
270         i = 1; v = length
271         while i <= max_netlen:
272             if v & msb_mask: break
273             i += 1; v <<= 1
274         return i
275     def netlen_to_mask(n, max_netlen, base_mask):
276         # convert the network length into its netmask
277         return ~((1 << (max_netlen - n)) - 1) & base_mask
278     def netlen_to_length(n, max_netlen, base_mask):
279         return 1 << (max_netlen - n) & base_mask
280
281     # convert the start and ending addresses of the netblock to Cidr
282     # object, mostly so we can get the numeric versions of their
283     # addresses.
284     cs = valid_cidr(start)
285     ce = valid_cidr(end)
286
287     # if either the start or ending addresses aren't valid addresses,
288     # quit now.
289     if not cs or not ce:
290         return None
291     # if the start and ending addresses aren't in the same family, quit now
292     if cs.is_ipv6() != ce.is_ipv6():
293         return None
294
295     max_netlen = cs._max_netlen()
296     msb_mask = cs.msb_mask
297     base_mask = cs.base_mask
298
299     # calculate the number of IP address in the netblock
300     block_len = ce.numaddr - cs.numaddr
301     # calcuate the largest CIDR block size that fits
302     netlen = largest_prefix(block_len + 1, max_netlen, msb_mask)
303
304     res = []; s = cs.numaddr
305     while block_len > 0:
306         mask = netlen_to_mask(netlen, max_netlen, base_mask)
307         # check to see if our current network length is valid
308         if (s & mask) != s:
309             # if not, shrink the network block size
310             netlen += 1
311             continue
312         # otherwise, we have a valid CIDR block, so add it to the list
313         res.append(new(s, netlen))
314         # and setup for the next round:
315         cur_len = netlen_to_length(netlen, max_netlen, base_mask)
316         s         += cur_len
317         block_len -= cur_len
318         netlen = largest_prefix(block_len + 1, max_netlen, msb_mask)
319     return res
320
321 # test driver
322 if __name__ == "__main__":
323     import sys
324     a = new("127.00.000.1/24")
325     b = new("127.0.0.1", 32)
326     c = new("24.232.119.192", 26)
327     d = new("24.232.119.0", 24)
328     e = new("24.224.0.0", 11)
329     f = new("216.168.111.0/27");
330     g = new("127.0.0.2/31");
331     h = new("127.0.0.16/32")
332     i = new("3ffe:4:201e:beef::0/64");
333     j = new("2001:3c01::/32")
334
335     print f.addr
336     print j.addr
337
338     try:
339         bad = new("24.261.119.0", 32)
340     except ValueError, x:
341         print "error:", x
342
343     print "cidr:", a, "num addresses:", a.length(), "ending address", \
344         a.end(), "netmask", a.netmask()
345
346     print "cidr:", j, "num addresses:", j.length(), "ending address", \
347         j.end(), "netmask", j.netmask()
348
349     clist = [a, b, c, d, e, f, g, h, i , j]
350     print "unsorted list of cidr objects:\n  ", clist
351
352
353     clist.sort()
354     print "sorted list of cidr object:\n  ", clist
355
356     k = new("2001:3c01::1:0", 120)
357     print "supernet: ", str(j), " supernet of ", str(k), "? ", \
358         str(j.is_supernet(k))
359     print "supernet: ", str(k), " supernet of ", str(j), "? ", \
360         str(k.is_supernet(j))
361     print "subnet: ", str(j), " subnet of ", str(k), "? ", \
362         str(j.is_subnet(k))
363     print "subnet: ", str(k), " subnet of ", str(j), "? ", \
364         str(k.is_subnet(j))
365
366     netblocks = [ ("192.168.10.0", "192.168.10.255"),
367                   ("192.168.10.0", "192.168.10.63"),
368                   ("172.16.0.0", "172.16.127.255"),
369                   ("24.33.41.22", "24.33.41.37"),
370                   ("196.11.1.0", "196.11.30.255"),
371                   ("192.247.1.0", "192.247.10.255"),
372                   ("10.131.43.3", "10.131.44.7"),
373                   ("3ffe:4:5::", "3ffe:4:5::ffff"),
374                   ("3ffe:4:5::", "3ffe:4:6::1")]
375
376     for start, end in netblocks:
377         print "netblock %s - %s:" % (start, end)
378         blocks = netblock_to_cidr(start, end)
379         print blocks