dataset/seed/base/funcs/base.py

from risotto.utils import multi_function as _multi_function
from typing import List as _List


@_multi_function
def get_ip(zones: dict,
           server_name: str,
           ) -> str:
    if server_name is None:
        return
    if isinstance(server_name, list):
        return_list = True
    else:
        return_list = False
        server_name = [server_name]
    lst = []
    for s_name in server_name:
        host_name, domain_name = s_name.split('.', 1)
        for zone in zones.values():
            if domain_name == zone['domain_name']:
                break
        else:
            raise ValueError(f'cannot find IP in domain name "{domain_name}" (for "{s_name}")')
        ret = zone['hosts'][host_name]
        if not return_list:
            return ret
        if ret not in lst:
            lst.append(ret)
    return lst


@_multi_function
def get_zones_info(zones: dict,
                   type: str,
                   zone_names: _List[str]=None,
                   zone_name: str=None,
                   index: int=None,
                   ) -> str:
    if type == 'host_ip' and index != 0:
        return
    if zone_name:
        if zone_name not in zones:
            raise ValueError(f"cannot get zone informations in unknown zone '{zone_name}'")
        if type == 'cidr':
            return zones[zone_name]['host_ip'] + '/' + zones[zone_name]['network'].split('/')[-1]
        return zones[zone_name][type]
    return [data[type] for zone_name, data in zones.items() if not zone_names or zone_name in zone_names]