controller-wisetopdataserver/python/biopro/util/inet.py

from typing import List, Optional

from .sys import system_call
from .text import part_suffix


class IPAddress:
    IPv4 = 'inet'
    IPv6 = 'inet6'

    def __init__(self):
        self.version: str = ''
        self.address: str = ''
        self.mask: int = 0
        self.bridge: Optional[str] = None
        self.scope: List[str] = []

    def __repr__(self):
        return '  %s %s/%d brd %s scope %s' % (
            self.version,
            self.address,
            self.mask,
            'NULL' if self.bridge is None else self.bridge,
            ' '.join(self.scope)
        )


class IPInterface:
    def __init__(self):
        self.interface: str = ''
        self.method: List[str] = []
        self.mtu: int = 0
        self.qdisc: str = ''
        self.state: str = ''
        self.group: str = ''
        self.qlen: int = 0
        self.link_type: str = ''
        self.link_addr: str = ''
        self.link_bridge: str = ''
        self.address: List[IPAddress] = []

    def __repr__(self):
        l1 = '%s <%s> mtu %d qdisc %s state %s group %s qlen %d' % (
            self.interface,
            ','.join(self.method),
            self.mtu,
            self.qdisc,
            self.state,
            self.state,
            self.qlen,
        )

        l2 = '  link/%s %s brd %s' % (
            self.link_type,
            self.link_addr,
            self.link_bridge
        )

        return '\n'.join([l1, l2, *map(repr, self.address)])


class Router:
    def __init__(self):
        self.address: str = ''
        self.dev: str = ''
        self.proto: str = ''
        self.metric: int = 0
        self.src_addr: Optional[str] = None
        self.scope: Optional[str] = None

        #

    def __repr__(self):
        return 'default via %s dev %s proto %s scope %s src %s metric %d' % (
            self.address,
            self.dev,
            self.proto,
            self.scope,
            'NULL' if self.src_addr is None else self.src_addr,
            self.metric
        )


def ip_addr(interface: str = None) -> List[IPInterface]:
    if interface is None:
        output = system_call('ip', 'addr')
    else:
        output = system_call('ip', 'addr', 'show', interface)

    ret: List[IPInterface] = []
    interface: IPInterface = None
    counter = 1

    for line in output.split('\n'):
        if line.startswith('%d:' % counter):
            counter += 1

            #
            interface = IPInterface()
            ret.append(interface)

            part = line.split(' ')

            #
            interface.interface = part[1]
            if interface.interface.endswith(':'):
                interface.interface = interface.interface[:-1]

            #
            method = part[2]
            if method.startswith('<') and method.endswith('>'):
                method = method[1:-1]

            interface.method = method.split(',')

            #
            i = 3
            while i < len(part):
                if part[i] == 'mtu':
                    interface.mtu = int(part[i + 1])
                    i += 2

                elif part[i] == 'qdisc':
                    interface.qdisc = part[i + 1]
                    i += 2

                elif part[i] == 'state':
                    interface.state = part[i + 1]
                    i += 2

                elif part[i] == 'group':
                    interface.group = part[i + 1]
                    i += 2

                elif part[i] == 'qlen':
                    interface.qlen = int(part[i + 1])
                    i += 2

                else:
                    i += 1

        else:
            line = line.strip()

            if line.startswith('link/'):
                part = line.split(' ')

                interface.link_type = part[0][5:]
                interface.link_addr = part[1]

                if len(part) >= 4 and part[2] == 'brd':
                    interface.link_bridge = part[3]

            elif line.startswith('inet'):
                part = line.split(' ')

                address = IPAddress()
                interface.address.append(address)

                #
                address.version = part[0]

                #
                address.address, address.mask = part_suffix(part[1], '/')
                address.mask = int(address.mask)

                #
                i = 2
                while i < len(part):
                    if part[i] == 'brd':
                        address.bridge = part[i + 1]
                        i += 2

                    elif part[i] == 'scope':
                        address.scope = part[i + 1:]
                        i = len(part)

                    else:
                        i += 1

    return ret


def ip_route() -> Router:
    output = system_call('ip', 'route')

    lines: List[str] = output.split('\n')
    part = lines[0].split(' ')

    router = Router()

    router.address = part[2]

    i = 3
    while i < len(part):
        if part[i] == 'dev':
            router.dev = part[i + 1]
            i += 2

        elif part[i] == 'proto':
            router.proto = part[i + 1]
            i += 2

        elif part[i] == 'metric':
            router.metric = int(part[i + 1])
            i += 2

        elif part[i] == 'scope':
            router.scope = part[i + 1]
            i += 2

        elif part[i] == 'src':
            router.src_addr = part[i + 1]
            i += 2

        else:
            i += 1

    return router


def allow_host() -> List[str]:
    ret = ['127.0.0.1']

    for ip in ip_addr():
        if ip.interface != 'lo':
            # not localhost

            for addr in ip.address:
                if addr.version == IPAddress.IPv4:
                    ret.append(addr.address)

    return ret


def main():
    import sys

    argv = sys.argv

    if len(argv) == 1 or argv[1] == 'addr':
        for _ret in ip_addr():
            print(_ret)

    elif argv[1] == 'route':
        print(ip_route())

    elif argv[1] == 'allow_host':
        print(allow_host())

    else:
        for ip in ip_addr():
            if ip.interface == argv[1]:
                for addr in ip.address:
                    if addr.version == IPAddress.IPv4:
                        print(addr.address)


if __name__ == '__main__':
    main()