annet 0.0__py3-none-any.whl

annet/lib.py

@@ -0,0 +1,128 @@
+import asyncio
+import os
+import shutil
+import sys
+from functools import lru_cache
+from pathlib import Path
+from typing import Awaitable, Optional
+
+import yaml
+from annet.annlib.lib import (  # pylint: disable=unused-import
+    ContextOrderedDict,
+    HuaweiNumBlock,
+    LMSMatcher,
+    add_annotation,
+    catch_ctrl_c,
+    cisco_collapse_vlandb,
+    cisco_expand_vlandb,
+    find_exc_in_stack,
+    find_modules,
+    first,
+    flatten,
+    huawei_collapse_vlandb,
+    huawei_expand_vlandb,
+    huawei_iface_ranges,
+    is_relative,
+    jinja_render,
+    jun_activate,
+    jun_is_inactive,
+    juniper_fmt_prefix_lists_acl,
+    juniper_port_split,
+    make_ip4_mask,
+    mako_render,
+    merge_dicts,
+    percentile,
+    uniq,
+)
+from contextlog import get_logger
+
+
+_TEMPLATE_CONTEXT_PATH: Optional[str] = None
+_DEFAULT_CONTEXT_PATH: Optional[str] = None
+
+
+def get_template_context_path() -> str:
+    if _TEMPLATE_CONTEXT_PATH is None:
+        set_template_context_path(str(Path(sys.modules["annet"].__file__).parent / "configs/context.yml"))
+    return _TEMPLATE_CONTEXT_PATH
+
+
+def set_template_context_path(path: str) -> None:
+    global _TEMPLATE_CONTEXT_PATH  # pylint: disable=global-statement
+    _TEMPLATE_CONTEXT_PATH = path
+
+
+def get_default_context_path() -> str:
+    if _DEFAULT_CONTEXT_PATH is None:
+        set_default_context_path("~/.annet/context.yml")
+    return _DEFAULT_CONTEXT_PATH
+
+
+def set_default_context_path(path: str) -> None:
+    global _DEFAULT_CONTEXT_PATH  # pylint: disable=global-statement
+    _DEFAULT_CONTEXT_PATH = path
+
+
+@lru_cache(maxsize=1)
+def _get_template_context():
+    with open(get_template_context_path()) as f:
+        return yaml.safe_load(f)
+
+
+def get_context_path(touch: Optional[bool] = False) -> str:
+    path = Path(os.getenv("ANN_CONTEXT_CONFIG_PATH", get_default_context_path())).expanduser().absolute()
+    if not path.exists():
+        src = get_template_context_path()
+        if not touch:
+            return str(src)
+        try:
+            # populate path with default configuration
+            path.parent.mkdir(parents=True, exist_ok=True)
+            shutil.copy(src, path)
+        except shutil.SameFileError:
+            pass
+    return str(path)
+
+
+def get_context() -> dict:
+    with open(get_context_path()) as f:
+        raw = yaml.safe_load(f)
+    _fill_in_default_generator_modules(raw)
+    context_name = os.getenv("ANN_SELECTED_CONTEXT", raw["selected_context"])
+    res = {k: raw[k][v] for k, v in raw["context"][context_name].items()}
+    if "ANN_GENERATORS_CONTEXT" in os.environ:  # an undocumented hack to maintain backwards compatibility; TODO: remove
+        res["generators"] = raw["generators"][os.getenv("ANN_GENERATORS_CONTEXT")]
+    return res
+
+
+@lru_cache(maxsize=1)
+def _warn_no_generators_in_context():
+    get_logger().warning(
+        "Older version of the context configuration found. Getting generators references from the template context"
+    )
+
+
+def _fill_in_default_generator_modules(raw: dict) -> bool:
+    """Backwards compatibility hack to add existing generators refs to context"""
+    if "generators" not in raw:
+        _warn_no_generators_in_context()
+        raw["generators"] = _get_template_context()["generators"]
+        for dst_context in raw["context"].values():
+            dst_context["generators"] = "default"
+        return True
+    return False
+
+
+def repair_context_file() -> None:
+    path = get_context_path()
+    with open(path) as f:
+        data = yaml.safe_load(f)
+    if _fill_in_default_generator_modules(data):
+        with open(path, "w") as f:
+            yaml.dump(data, f, sort_keys=False)
+
+
+def do_async(coro: Awaitable):
+    loop = asyncio.get_event_loop()
+    res = loop.run_until_complete(coro)
+    return res

annet/output.py

@@ -0,0 +1,167 @@
+import abc
+import os
+import posixpath
+import sys
+from typing import List, Optional, Tuple, Type, Dict
+
+import colorama
+from annet.annlib.output import (  # pylint: disable=unused-import
+    LABEL_NEW_PREFIX,
+    OutputWriter,
+    TextArgs,
+    capture_output,
+    dir_or_file_output,
+    format_file_diff,
+    print_as_json,
+    print_as_yaml,
+    print_err_label,
+    print_label,
+)
+from contextlog import get_logger
+
+from annet.cli_args import FileOutOptions, QueryOptions
+from annet.connectors import Connector
+from annet.storage import Device, storage_connector
+
+
+class _DriverConnector(Connector["OutputDriver"]):
+    name = "OutputDriver"
+    ep_name = "output"
+
+    def _get_default(self) -> Type["OutputDriver"]:
+        return OutputDriverBasic
+
+
+output_driver_connector = _DriverConnector()
+
+
+class OutputDriver(abc.ABC):
+    @abc.abstractmethod
+    def write_output(self, arg_out: FileOutOptions, items, query_result_count=1) -> None:
+        pass
+
+    @abc.abstractmethod
+    def format_fails(self, fail, fqdns: Optional[Dict[int, str]] = None) -> Tuple[str, str]:
+        pass
+
+    @abc.abstractmethod
+    def cfg_file_names(self, device: Device) -> List[str]:
+        pass
+
+    @abc.abstractmethod
+    def entire_config_dest_path(self, device: Device, config_path: str) -> str:
+        pass
+
+
+class OutputDriverBasic(OutputDriver):
+    def write_output(self, arg_out: FileOutOptions, items, query_result_count=1):
+        """
+        пишет результаты генерации в файл или директорию :dest
+        :dest - это директория в случаях:
+        - заканчивается на "/"
+        - существует директория с таким именем
+        - более одного устройства в результате запроса
+        - есть entire-генераторы (определяется по типу первого результата, если устройство одно)
+        """
+        logger = get_logger()
+
+        items_iter = iter(items)
+        try:
+            first_result = next(items_iter)
+        except StopIteration:
+            # нет результатов, ничего не пишем и не создаём
+            return
+
+        def _reassemble_items():
+            yield first_result
+            yield from items_iter
+
+        dest = arg_out.dest
+        dir_mode = dir_or_file_output(dest, query_result_count, suggest_dir=(os.sep in first_result[0]))
+        _reassembled_items = list(_reassemble_items())
+
+        for output_no, (label, output, is_fail) in enumerate(_reassembled_items):
+            writer = output if isinstance(output, OutputWriter) else OutputWriter(output)
+            label = os.path.normpath(label)
+            label_color = colorama.Back.RED if is_fail else colorama.Back.GREEN
+            if dest is None:
+                if hasattr(arg_out, "format") and arg_out.format == "json":
+                    if output_no > 0:
+                        sys.stdout.write(",")
+                    elif output_no == 0:
+                        sys.stdout.write("{")
+                    sys.stdout.write('"%s": ' % label)
+                    writer.write(sys.stdout)
+                    if output_no == len(_reassembled_items) - 1:
+                        sys.stdout.write("}")
+                else:
+                    if not arg_out.no_label:
+                        print_label(label, back_color=label_color)
+                    writer.write(sys.stdout)
+            elif dir_mode:
+                if label.startswith(LABEL_NEW_PREFIX):
+                    label = label[len(LABEL_NEW_PREFIX):]
+                if label.startswith(os.sep):
+                    # just in case.
+                    label = label.lstrip(os.sep)
+                if os.sep not in label:
+                    # vendor config
+                    label = os.path.basename(label)
+                else:
+                    # entire generated file
+                    parts = label.split(os.sep)
+                    hostname = parts[0]
+                    label = os.sep.join(parts[1:])
+                    if not arg_out.expand_path:
+                        label = os.path.basename(label)
+                    if query_result_count > 1:
+                        label = os.path.join(hostname, label)
+                file_dest = os.path.join(dest, "errors") if is_fail else dest
+                out_file = os.path.normpath(os.path.join(file_dest, label))
+                logger.info("writing '%s'", out_file)
+                dirname = os.path.dirname(out_file)
+                if not os.path.exists(dirname):
+                    os.makedirs(dirname)
+                with open(out_file, "w") as file:
+                    writer.write(file)
+            else:
+                logger.info("writing '%s'", dest)
+                with open(dest, "w") as file:
+                    writer.write(file)
+
+    def format_fails(self, fail, fqdns: Optional[Dict[int, str]] = None):
+        ret = []
+        fqdns = fqdns or {}
+        for (assignment, exc) in fail.items():
+            label = assignment
+            if assignment in fqdns:
+                label = fqdns[assignment]
+            elif isinstance(assignment, tuple):
+                label = assignment[0]
+            else:
+                ValueError("Failed to parse failed assignment %r" % assignment)
+            ret.append((label, getattr(exc, "formatted_output", f"{repr(exc)} (formatted_output is absent)"), True))
+        return ret
+
+    def cfg_file_names(self, device: Device) -> List[str]:
+        return [f"{device.hostname}.cfg"]
+
+    def entire_config_dest_path(self, device, config_path: str) -> str:
+        """Формирует путь к конфигу в директории destname.
+
+        Например, для устройства с hostname `my-device`:
+        ```
+        >>> device.entire_config_dest_path("/etc/frr/frr.conf")
+        'my-device.cfg/etc/frr/frr.conf'
+        >>>
+        ```
+        """
+        # NOTE: с полученным `config_path` работаем через `posixpath`, а не через `os.path`, потому что
+        #       entire-путь POSIX-специфичный; но в конце формируем путь через `os.path` для текущей платформы
+        if not posixpath.abspath(config_path):
+            raise RuntimeError(f"Want absolute config path, but relative received: {config_path}")
+        cfg_files = self.cfg_file_names(device)
+        # NOTE: получаем путь без "/" в начале, например, "etc/frr/frr.conf"
+        relative_config_path = posixpath.relpath(config_path, "/")
+        dest_config_path_parts = [cfg_files[0]] + relative_config_path.split(posixpath.sep)
+        return os.path.join(*dest_config_path_parts)