"""CIDR blocks share the first part of the bit sequence that comprises the
binary representation of the IP address

We determine the network portion of the address by applying a bitwise AND
operation to between the address and the netmask. A bitwise AND operation will
basically save the networking portion of the address and discard the host
portion.

In the CIDR notation 192.0.1.0/24, the prefix is 192.0.1.0, and the total number
of bits in the address is 24.  The least significant 8 bits are the IP addresses
that are part of this CIDR block.

10.10.1.16/32 is an address prefix with 32 bits, which is the highest number of
bits allowed in IPv4.

For example, we could express the idea that the IP address 192.168.0.15 is
associated with the netmask 255.255.255.0 by using the CIDR notation of
92.168.0.15/24. This means that the first 24 bits of the IP address given are
considered significant for the network routing.
"""
import re

class InvalidIPError(Exception):
    pass

class Ip(object):
    def __init__(self, decimal=None, binary=None):
        """Initialize an IP address given either its decimal (192.168.0.0) or
        binary (32 characters, each of which is 0 or 1) form"""
        if decimal is not None:
            self.raw = decimal
            self.decimal = self.raw
            self.binary = self._binary_from_decimal()
        elif binary is not None:
            self.raw = binary
            self.binary = binary
            self.decimal = self._decimal_from_binary()

        self._verify()

    def _verify(self):
        """Perform a quick validation that the decimal representation of the
        IP address is valid
        """
        octets = [int(x) for x in self.decimal.split(".")]

        if len(octets) != 4:
            raise InvalidIPError("{} is not a valid IP".format(self.decimal))

        for octet in octets:
            if octet not in range(256):
                raise InvalidIPError("{} is not a valid IP".format(
                    self.decimal))

    def _decimal_from_binary(self):
        """Takes a binary IP address and translates to decimal"""
        octets = []
        for i in range(0, len(self.binary), 8):
            frac = self.binary[i:i + 8]
            octets.append(str(int(frac, 2)))
        return ".".join(octets)

    def _binary_from_decimal(self):
        """Takes a decimal IP address and translates to binary"""
        binrep = ""
        for octet in self.raw.split("."):
            binrep += "{0:b}".format(int(octet)).zfill(8)
        return binrep


class Netmask(object):
    """Object represents the part after the / in a CIDR block"""
    def __init__(self, input_string):
        self.raw = input_string
        self.decimal = int(input_string)
        self.binary = self._binary()

    def _binary(self):
        """Return a string with `netmask` 1's padded up to 32 characters with
        zeroes
        """
        return "1"*self.decimal + "0"*(32-self.decimal)


class Cidr(object):
    """Class representing a CIDR block"""

    def __init__(self, input_string):
        # Match a regex against x.x.x.x/y
        result = re.match(
            "(?P<ip>^[0-9]+.[0-9]+.[0-9]+.[0-9]+)/(?P<netmask>[0-9]{1,2})",
            input_string)

        self.ip = Ip(decimal=result.group("ip"))
        self.netmask = Netmask(result.group("netmask"))
        self.network_prefix = self._get_network_prefix()

    def _get_network_prefix(self):
        """Make a binary representation of the network prefix.  This will match
        the IP address of the CIDR block up to the first netmask bytes.

        Returns a string, 32 characters in length
        """
        np = ""
        for letter in zip(self.ip.binary, self.netmask.binary):
            # bitwise AND the binary IP with the binary netmask
            np += "1" if letter[0] == "1" and letter[1] == "1" else "0"
        return np

    def enumerate_ips(self):
        """Yield an Ip for each IP address inside of this CIDR block"""

        if self.netmask.decimal == 32:
            yield Ip(binary=self.network_prefix)
            raise StopIteration

        max_val = 2 ** (32-self.netmask.decimal)

        for i in range(max_val):
            binary_rep = (self.network_prefix[:self.netmask.decimal] +
                          "{0:b}".format(i).zfill(32-self.netmask.decimal))
            yield Ip(binary=binary_rep)


def test_enumeration():

    # Test a single IP
    single_ip = list(Cidr("128.128.128.128/32").enumerate_ips())
    assert len(single_ip) == 1
    assert single_ip[0].decimal == "128.128.128.128"

    # Do a non-power of two netmask
    cidr = Cidr("192.168.50.48/30")
    expected = ["192.168.50.48", "192.168.50.49", "192.168.50.50",
                "192.168.50.51"]
    for res in zip(expected, list(cidr.enumerate_ips())):
        assert res[1].decimal == res[0]

    # Check the length of a larger list
    assert len(list(Cidr("127.0.0.0/16").enumerate_ips())) == 65536


if __name__ == "__main__":
    print "Testing IP enumerator..."
    test_enumeration()
    print "Tests passed"