A CIDR IP address calculator

subnet.py

"""A subnet IP address range calculator."""
import argparse
import copy
import math


class Subnet:
    """A subnet."""

    @staticmethod
    def get_routing_prefix_octets(mask):
        """Get the octets of a CIDR mask as a list of integers."""
        octets = [0, 0, 0, 0]
        elements = mask[:mask.index("/")].split(".")

        for i in range(4):
            if len(elements) > i:
                octets[i] = int(elements[i])

            if octets[i] < 0 or octets[i] > 255:
                raise ValueError(f"Invalid CIDR mask: {mask}")

        return octets

    @staticmethod
    def get_routing_prefix_length(mask):
        """Get the length of a CIDR mask as an integer."""
        length = int(mask[mask.index("/") + 1:])

        if length < 0 or length > 32:
            raise ValueError(f"Invalid CIDR mask: {mask}")

        return length

    @staticmethod
    def get_theoretical_ip_count(octets, length):
        """Get the theoretical number of available IP's in the subnet."""
        capacity = 1

        for i in range(int(length / 8), 4):
            if i < 3:
                capacity = capacity * (256 - octets[i])
            else:
                capacity = capacity * 256 - octets[i]

        return capacity

    @staticmethod
    def get_theoretical_size(octets, length):
        """Get the theoretical size of the subnet."""
        return int(math.pow(2, 32 - length))

    @classmethod
    def get_length(cls, octets, length):
        """Get the length of the subnet IP generator."""
        ip_count = cls.get_theoretical_ip_count(octets, length)
        size = cls.get_theoretical_size(octets, length)

        return min(ip_count, size)

    def __init__(self, mask="0.0.0.0/0"):
        """Override."""
        self.mask = mask

    def __repr__(self):
        """Override."""
        return f"{self.__class__.__name__}('{self.mask}')"

    def __str__(self):
        """Override."""
        return self.mask

    def __iter__(self):
        """Override."""
        return self.generate()

    def __next__(self):
        """Override."""
        return next(self.generator)

    def __len__(self):
        """Count the IP addresses in the subnet."""
        return self.length

    @property
    def mask(self):
        """Get mask."""
        return self._mask

    @mask.setter
    def mask(self, mask):
        """Set mask."""
        self.routing_prefix_octets = self.get_routing_prefix_octets(mask)
        self.routing_prefix_length = self.get_routing_prefix_length(mask)
        self._mask = mask
        self.generator = self.generate()
        self.length = self.get_length(self.routing_prefix_octets,
                                      self.routing_prefix_length)

    def generate(self):
        """Generate IP addresses in the subnet."""
        octets = copy.deepcopy(self.routing_prefix_octets)
        static_octets = int(self.routing_prefix_length / 8)

        for _ in range(len(self)):
            ip_address = ".".join([str(octet) for octet in octets])

            for i in reversed(range(static_octets, 4)):
                octets[i] += 1

                if octets[i] > 255:
                    octets[i] = 0
                else:
                    break

            yield ip_address


def parse_args(description=__doc__):
    """Parse input arguments."""
    parser = argparse.ArgumentParser(description=description)

    parser.add_argument(
        '-c', "--count",
        help="Count the IP addresses in the subnet",
        action='store_true',
        default=False
    )
    parser.add_argument(
        "cidr",
        help="A CIDR mask",
        nargs="*"
    )

    return parser.parse_args()


def main():
    """Do the CLI."""
    arguments = parse_args()

    for mask in arguments.cidr:
        subnet = Subnet(mask)

        if arguments.count:
            print(len(subnet))
        else:
            for ip in subnet:
                print(ip)


if __name__ == "__main__":
    main()