NetworkVisualizer/ip.py

from ast import Str
import re

class net:
    type = 0
    start = None
    end = None
    cidr = 0

    def __init__(self):
        pass

    def parseNet(self, netstr: str):
        netstr = netstr.strip()
        # basic check
        if re.fullmatch("^[\.:0-9a-f]*\/[0-9]*$", netstr) is None:
            raise ValueError("Net {} is malformed".format(netstr))
        
        adr, cidr = netstr.split("/")

        start = ip(adr)
        start.numerical = start._mask(int(cidr))
        end = ip(adr)
        end.numerical = end._mask(int(cidr), True)

        self.type = start.type
        self.start = start
        self.end = end
        self.cidr = int(cidr)

    def getSize(self):
        if self.type == 0:
            return 32-self.cidr
        return 128-self.cidr

def _genMask(n: int):
    x = 0
    for i in range(n):
        x = (x << 1) + 1
    return ~x

class ip:
    type = 0
    numerical = 0

    def __init__(self):
        pass

    def __init__(self, octets, type):
        self.fromOctets(octets, type)

    def __init__(self, adr: Str):
        self.parseAdrString(adr)

    def toStr(self, mask=0):
        if self.type == 0:
            o = [str(self.getOctet(i, mask)) for i in range(4)]
            return ".".join(o)
        else:
            raise NotImplementedError()

    def getOctet(self, octet, mask=0):
        h = 4 if self.type == 0 else 8
        o = 8 if self.type == 0 else 16
        n = self._mask(mask)
        return (n >> (o * (h - octet - 1))) & ~_genMask(o)

    def fromOctets(self, octets, type):
        if type == 0:
            c = 4
            o = 8
        else:
            c = 8
            o = 16

        self.type = type
        self.numerical = 0
        for i in range(c):
            self.numerical = (self.numerical << o) + octets[i]

    def fromNumerical(self, num: int, type):
        self.num = num
        self.type = type

    def _mask(self, n: int, sethigh: bool=False):
        h = 32 if self.type == 0 else 128
        if sethigh:
            return self.numerical | ~_genMask(h-n)
        else:
            return self.numerical & _genMask(h-n)

    def _splitMergedOctet(self, octet: str, splits: int):
        n = int(octet)
        r = [0]*splits

        if n > 1 << (splits*8):
            raise ValueError("Invalid Octet {}".format(octet))

        for i in range(splits-1, -1, -1):
            r[i] = n % 256
            n = n >> 8
        
        return r

    def toInt(self, trunc_bits = 0):
        return self.numerical >> trunc_bits

    def _parseV4(self, adr: str):
        # basic validation
        if re.fullmatch("^[\.0-9]*$", adr) is None:
            raise ValueError("Address {} contains illegal symbols".format(adr))

        parts = adr.split(".")

        r = []
        if len(parts) < 4:
            for i in range(len(parts)-1):
                r += [int(parts[i])]
            r += self._splitMergedOctet(parts[-1], 4 - len(parts) + 1)
        elif len(parts) == 4:
            for i in range(len(parts)):
                r += [int(parts[i])]
        else:
            raise ValueError("Invalid IPv4 {}, to many octets".format(adr))

        # validate each octet
        for p in r:
            n = int(p)
            if (n < 0) or (n > 255):
                raise ValueError("Invalid octet {} in address {}".format(p, adr))

        self.fromOctets(r, 0)

    def _parseV6(self, adr: str):
        # ignore case
        adr = adr.lower()
        # basic validation
        if re.fullmatch("^[:0-9a-f]*$", adr) is None:
            raise ValueError("Address {} contains illegal symbols".format(adr))

        # handle null case, because it breaks the rest
        if adr == "::":
            self.numerical = 0
            self.type = 1
            return

        # handle leading and trailing ellipses
        # error out, when there is a ":" without an ellipsis
        if adr[0] == ":":
            adr = adr[1:]
            if adr[0] != ":":
                raise ValueError("Address :{} starts with : that is not part of an ellipsis".format(adr))

        if adr[-1] == ":":
            adr = adr[:-1]
            if adr[-1] != ":":
                raise ValueError("Address {}: ends with : that is not part of an ellipsis".format(adr))

        parts = adr.split(":")

        # validate each group
        empties = 0
        for p in parts:
            if p == "":
                empties += 1
                if empties > 1:
                    raise ValueError("Address can contain only one ellipsis! Illegal address {}".format(adr))
            else:
                n = int(p, 16)
                if (n < 0) or (n > 65535):
                    raise ValueError("Invalid octet {} in address {}".format(p, adr))

        if len(parts) < 8 and empties == 0:
            raise ValueError("Invalid address {}: not enough octets without ellipsis".format(p, adr))

        if len(parts) > 8:
            raise ValueError("Invalid IPv6 {}, to many octets".format(adr))

        r = [0]*8

        # assign octets until we reach :: or the end of the address
        for i in range(len(parts)):
            if parts[i] == "":
                break
            else:
                r[i] = int(parts[i], 16)
        
        # assign remaining octets backwards (filling :: with 0s)
        if i < 7:
            i = 7
            while (parts[i - 8] != ""):
                r[i] = int(parts[i - 8], 16)
                i -= 1

        self.fromOctets(r, 1)

    def parseAdrString(self, adr: str):
        if adr.count(":") == 0:
            self._parseV4(adr)
        else:
            self._parseV6(adr)