1# Lint as: python2, python3
2# Copyright (c) 2014 The Chromium OS Authors. All rights reserved.
3# Use of this source code is governed by a BSD-style license that can be
4# found in the LICENSE file.
5
6from __future__ import absolute_import
7from __future__ import division
8from __future__ import print_function
9
10import logging
11from six.moves import map
12from six.moves import zip
13
14from autotest_lib.client.common_lib import error
15
16def from_addr(addr, prefix_len=None):
17    """Build a Netblock object.
18
19    @param addr: string IP address with optional prefix length
20            (e.g. '192.168.1.1' or '192.168.1.1/24'). If |addr| has no
21            prefix length, then use the |prefix_len| parameter.
22    @param prefix_len: int number of bits forming the IP subnet prefix for
23            |addr|. This value will be preferred to the parsed value if
24            |addr| has a prefix length as well. If |addr|
25            has no prefix length and |prefix_len| is None, then an error
26            will be thrown.
27
28    """
29    if addr is None:
30        raise error.TestError('netblock.from_addr() expects non-None addr '
31                              'parameter.')
32
33    prefix_sep_count = addr.count('/')
34    if prefix_sep_count > 1:
35        raise error.TestError('Invalid IP address found: "%s".' % addr)
36
37    if prefix_sep_count == 1:
38        addr_str, prefix_len_str = addr.split('/')
39    else:
40        # No prefix separator.  Assume addr looks like '192.168.1.1'
41        addr_str = addr
42        # Rely on passed in |prefix_len|
43        prefix_len_str = None
44
45    if prefix_len is not None and prefix_len_str is not None:
46        logging.warning('Ignoring parsed prefix length of %s in favor of '
47                        'passed in value %d', prefix_len_str, prefix_len)
48    elif prefix_len is not None and prefix_len_str is None:
49        pass
50    elif prefix_len is None and prefix_len_str is not None:
51        prefix_len = int(prefix_len_str)
52    else:
53        raise error.TestError('Cannot construct netblock without knowing '
54                              'prefix length for addr: "%s".' % addr)
55
56    return Netblock(addr_str, prefix_len)
57
58
59class Netblock(object):
60    """Utility class for transforming netblock address to related strings."""
61
62    @staticmethod
63    def _octets_to_addr(octets):
64        """Transform a list of bytes into a string IP address.
65
66        @param octets list of ints (e.g. [192.168.0.1]).
67        @return string IP address (e.g. '192.168.0.1.').
68
69        """
70        return '.'.join(map(str, octets))
71
72
73    @staticmethod
74    def _int_to_octets(num):
75        """Tranform a 32 bit number into a list of 4 octets.
76
77        @param num: number to convert to octets.
78        @return list of int values <= 8 bits long.
79
80        """
81        return [(num >> s) & 0xff for s in (24, 16, 8, 0)]
82
83
84    @property
85    def netblock(self):
86        """@return the IPv4 address/prefix, e.g., '192.168.0.1/24'."""
87        return '/'.join([self._octets_to_addr(self._octets),
88                         str(self.prefix_len)])
89
90
91    @property
92    def netmask(self):
93        """@return the IPv4 netmask, e.g., '255.255.255.0'."""
94        return self._octets_to_addr(self._mask_octets)
95
96
97    @property
98    def prefix_len(self):
99        """@return the IPv4 prefix len, e.g., 24."""
100        return self._prefix_len
101
102
103    @property
104    def subnet(self):
105        """@return the IPv4 subnet, e.g., '192.168.0.0'."""
106        octets = [a & m for a, m in zip(self._octets, self._mask_octets)]
107        return self._octets_to_addr(octets)
108
109
110    @property
111    def broadcast(self):
112        """@return the IPv4 broadcast address, e.g., '192.168.0.255'."""
113        octets = [a | (m ^ 0xff)
114                  for a, m in zip(self._octets, self._mask_octets)]
115        return self._octets_to_addr(octets)
116
117
118    @property
119    def addr(self):
120        """@return the IPv4 address, e.g., '192.168.0.1'."""
121        return self._octets_to_addr(self._octets)
122
123
124    def __init__(self, addr_str, prefix_len):
125        """Construct a Netblock.
126
127        @param addr_str: string IP address (e.g. '192.168.1.1').
128        @param prefix_len: int length of subnet prefix (e.g. 24).
129
130        """
131        self._octets = list(map(int, addr_str.split('.')))
132        mask_bits = (-1 << (32 - prefix_len)) & 0xffffffff
133        self._mask_octets = self._int_to_octets(mask_bits)
134        self._prefix_len = prefix_len
135
136
137    def get_addr_in_block(self, offset):
138        """Get an address in a subnet.
139
140        For instance if this netblock represents 192.168.0.1/24,
141        then get_addr_in_block(5) would return 192.168.0.5.
142
143        @param offset int offset in block, (e.g. 5).
144        @return string address (e.g. '192.168.0.5').
145
146        """
147        offset = self._int_to_octets(offset)
148        octets = [(a & m) + o
149                  for a, m, o in zip(self._octets, self._mask_octets, offset)]
150        return self._octets_to_addr(octets)
151