From cb247e9fd597b39a60601537efb1924886132557 Mon Sep 17 00:00:00 2001 From: Leonardo Parente <23251360+leoparente@users.noreply.github.com> Date: Tue, 21 May 2024 14:58:28 -0300 Subject: [PATCH] add docstrings --- diode-napalm-agent/diode_napalm/__init__.py | 2 +- diode-napalm-agent/diode_napalm/cli/cli.py | 62 +++++++++++++++----- diode-napalm-agent/diode_napalm/client.py | 48 ++++++++++++++- diode-napalm-agent/diode_napalm/parser.py | 21 ++++++- diode-napalm-agent/diode_napalm/translate.py | 53 +++++++++++++---- diode-napalm-agent/pyproject.toml | 10 +--- 6 files changed, 160 insertions(+), 36 deletions(-) diff --git a/diode-napalm-agent/diode_napalm/__init__.py b/diode-napalm-agent/diode_napalm/__init__.py index be027bd..9549369 100644 --- a/diode-napalm-agent/diode_napalm/__init__.py +++ b/diode-napalm-agent/diode_napalm/__init__.py @@ -1,3 +1,3 @@ #!/usr/bin/env python # Copyright 2024 NetBox Labs Inc -"""NetBox Labs - Diode NAPAML namespace.""" +"""NetBox Labs - Diode NAPALM namespace.""" diff --git a/diode-napalm-agent/diode_napalm/cli/cli.py b/diode-napalm-agent/diode_napalm/cli/cli.py index d3e3e71..ebd1b7a 100644 --- a/diode-napalm-agent/diode_napalm/cli/cli.py +++ b/diode-napalm-agent/diode_napalm/cli/cli.py @@ -1,10 +1,11 @@ #!/usr/bin/env python # Copyright 2024 NetBox Labs Inc -"""Diode NAPALM Agent CLI""" +"""Diode NAPALM Agent CLI.""" import argparse import asyncio +import logging import sys from importlib.metadata import version from pathlib import Path @@ -15,8 +16,24 @@ from diode_napalm.version import version_semver import netboxlabs.diode.sdk.version as SdkVersion +# Set up logging +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + def parse_config_file(cfile: Path): + """ + Parse the configuration file and return the diode configuration. + + Args: + cfile (Path): Path to the configuration file. + + Returns: + cfg.diode: Parsed configuration data. + + Raises: + SystemExit: If the configuration file cannot be opened or parsed. + """ try: with open(cfile, "r") as f: cfg = parse_config(f.read()) @@ -28,21 +45,34 @@ def parse_config_file(cfile: Path): async def start_policy(cfg, client): + """ + Start the policy for the given configuration and client. + + Args: + cfg: Configuration data for the policy. + client: Client instance for data ingestion. + """ for info in cfg.data: - print(f"Get driver '{info.driver}'") + logger.info(f"Get driver '{info.driver}'") np_driver = get_network_driver(info.driver) - print(f"Getting information from '{info.hostname}'") + logger.info(f"Getting information from '{info.hostname}'") with np_driver(info.hostname, info.username, info.password, info.timeout, info.optional_args) as device: - data = {} - data["driver"] = info.driver - data["site"] = cfg.config.netbox.get("site", None) - data["device"] = device.get_facts() - data["interface"] = device.get_interfaces() + data = { + "driver": info.driver, + "site": cfg.config.netbox.get("site", None), + "device": device.get_facts(), + "interface": device.get_interfaces(), + } client.ingest(data) async def start_agent(cfg): - # start diode client + """ + Start the diode client and execute policies. + + Args: + cfg: Configuration data containing policies. + """ client = Client() client.init_client(target=cfg.config.target, api_key=cfg.config.api_key, tls_verify=cfg.config.tls_verify) @@ -53,14 +83,18 @@ async def start_agent(cfg): raise Exception(f"Unable to start policy {policy_name}: {e}") -def agent_main(): - parser = argparse.ArgumentParser(description="Diode Agent for SuzieQ") +def main(): + """ + Main entry point for the Diode NAPALM Agent CLI. + Parses command-line arguments and starts the agent. + """ + parser = argparse.ArgumentParser(description="Diode Agent for NAPALM") parser.add_argument( "-V", "--version", action="version", - version=f"Diode Agent version: {version_semver()}, SuzieQ version: {version('napalm')}, Diode SDK version: {SdkVersion.version_semver()}", - help='Display Diode Agent, SuzieQ and Diode SDK versions' + version=f"Diode Agent version: {version_semver()}, NAPALM version: {version('napalm')}, Diode SDK version: {SdkVersion.version_semver()}", + help='Display Diode Agent, NAPALM and Diode SDK versions' ) parser.add_argument( "-c", @@ -80,4 +114,4 @@ def agent_main(): if __name__ == '__main__': - agent_main() + main() diff --git a/diode-napalm-agent/diode_napalm/client.py b/diode-napalm-agent/diode_napalm/client.py index 43bb980..d5727a3 100644 --- a/diode-napalm-agent/diode_napalm/client.py +++ b/diode-napalm-agent/diode_napalm/client.py @@ -1,7 +1,8 @@ #!/usr/bin/env python # Copyright 2024 NetBox Labs Inc -"""Diode SDK Client for NAPALM""" +"""Diode SDK Client for NAPALM.""" +import logging from typing import Optional from diode_napalm.version import version_semver from diode_napalm.translate import translate_data @@ -10,16 +11,38 @@ APP_NAME = "diode-napalm-agent" APP_VERSION = version_semver() +# Set up logging +logging.basicConfig(level=logging.INFO) +logger = logging.getLogger(__name__) + class Client: + """ + Singleton class for managing the Diode client for NAPALM. + + This class ensures only one instance of the Diode client is created and provides methods + to initialize the client and ingest data. + + Attributes: + diode_client (DiodeClient): Instance of the DiodeClient. + """ _instance = None def __new__(cls): + """ + Create a new instance of the Client if one does not already exist. + + Returns: + Client: The singleton instance of the Client. + """ if cls._instance is None: cls._instance = super(Client, cls).__new__(cls) return cls._instance def __init__(self): + """ + Initialize the Client instance with no Diode client. + """ self.diode_client = None def init_client( @@ -28,12 +51,31 @@ def init_client( api_key: Optional[str] = None, tls_verify: bool = None ): + """ + Initialize the Diode client with the specified target, API key, and TLS verification. + + Args: + target (str): The target endpoint for the Diode client. + api_key (Optional[str]): The API key for authentication (default is None). + tls_verify (bool): Whether to verify TLS certificates (default is None). + """ self.diode_client = DiodeClient( target=target, app_name=APP_NAME, app_version=APP_VERSION, api_key=api_key, tls_verify=tls_verify) def ingest(self, data: dict): + """ + Ingest data using the Diode client after translating it. + + Args: + data (dict): The data to be ingested. + + Raises: + ValueError: If the Diode client is not initialized. + """ if self.diode_client is None: raise ValueError("diode client defined") ret = self.diode_client.ingest(translate_data(data)) - print(ret) - + if not len(ret.errors): + logger.info("successful ingestion") + else: + logger.error(ret) diff --git a/diode-napalm-agent/diode_napalm/parser.py b/diode-napalm-agent/diode_napalm/parser.py index ae72083..8cd07d6 100644 --- a/diode-napalm-agent/diode_napalm/parser.py +++ b/diode-napalm-agent/diode_napalm/parser.py @@ -1,6 +1,6 @@ #!/usr/bin/env python # Copyright 2024 NetBox Labs Inc -"""Parse Diode Agent Config file""" +"""Parse Diode Agent Config file.""" import yaml from typing import Any, Dict, List, Optional @@ -8,10 +8,12 @@ class ParseException(Exception): + """Custom exception for parsing errors.""" pass class Napalm(BaseModel): + """Model for NAPALM configuration.""" driver: str hostname: str username: str @@ -22,30 +24,47 @@ class Napalm(BaseModel): class DiscoveryConfig(BaseModel): + """Model for discovery configuration.""" netbox: Dict[str, str] class Policy(BaseModel): + """Model for a policy configuration.""" config: DiscoveryConfig data: List[Napalm] class DiodeConfig(BaseModel): + """Model for Diode configuration.""" target: str api_key: str tls_verify: bool class Diode(BaseModel): + """Model for Diode containing configuration and policies.""" config: DiodeConfig policies: Dict[str, Policy] class Config(BaseModel): + """Top-level model for the entire configuration.""" diode: Diode def parse_config(config_data: str): + """ + Parse the YAML configuration data into a Config object. + + Args: + config_data (str): The YAML configuration data as a string. + + Returns: + Config: The parsed configuration object. + + Raises: + ParseException: If there is an error in parsing the YAML or validating the data. + """ try: config = Config(**yaml.safe_load(config_data)) return config diff --git a/diode-napalm-agent/diode_napalm/translate.py b/diode-napalm-agent/diode_napalm/translate.py index 3d77b92..d3befd3 100644 --- a/diode-napalm-agent/diode_napalm/translate.py +++ b/diode-napalm-agent/diode_napalm/translate.py @@ -1,6 +1,6 @@ #!/usr/bin/env python # Copyright 2024 NetBox Labs Inc -"""Translate from NAPALM output format to Diode SDK entites""" +"""Translate from NAPALM output format to Diode SDK entities.""" from typing import Iterable from netboxlabs.diode.sdk.diode.v1.ingester_pb2 import ( @@ -15,6 +15,15 @@ def translate_device(device_info: dict) -> Device: + """ + Translate device information from NAPALM format to Diode SDK Device entity. + + Args: + device_info (dict): Dictionary containing device information. + + Returns: + Device: Translated Device entity. + """ manufacturer = Manufacturer(name=device_info["vendor"]) device_type = DeviceType( model=device_info["model"], @@ -37,6 +46,17 @@ def translate_device(device_info: dict) -> Device: def translate_interface(device: Device, if_name: str, interface_info: dict) -> Interface: + """ + Translate interface information from NAPALM format to Diode SDK Interface entity. + + Args: + device (Device): The device to which the interface belongs. + if_name (str): The name of the interface. + interface_info (dict): Dictionary containing interface information. + + Returns: + Interface: Translated Interface entity. + """ interface = Interface( device=device, name=if_name, @@ -51,17 +71,30 @@ def translate_interface(device: Device, if_name: str, interface_info: dict) -> I def translate_data(data: dict) -> Iterable[Entity]: + """ + Translate data from NAPALM format to Diode SDK entities. + + Args: + data (dict): Dictionary containing data to be translated. + + Returns: + Iterable[Entity]: Iterable of translated entities. + """ entities = [] - if "device" in data: - data["device"]["driver"] = data["driver"] - data["device"]["site"] = data["site"] - device = translate_device(data["device"]) + device_info = data.get("device") + if device_info: + device_info["driver"] = data.get("driver") + device_info["site"] = data.get("site") + device = translate_device(device_info) entities.append(Entity(device=device)) - if "interface" in data: - for if_name in data["interface"]: - if if_name in data["device"]["interface_list"]: - entities.append(Entity(interface=translate_interface( - device, if_name, data["interface"][if_name]))) + + interfaces = data.get("interface", {}) + interface_list = device_info.get("interface_list", []) + for if_name, interface_info in interfaces.items(): + if if_name in interface_list: + interface = translate_interface( + device, if_name, interface_info) + entities.append(Entity(interface=interface)) return entities diff --git a/diode-napalm-agent/pyproject.toml b/diode-napalm-agent/pyproject.toml index f006219..cfe4d7a 100644 --- a/diode-napalm-agent/pyproject.toml +++ b/diode-napalm-agent/pyproject.toml @@ -1,10 +1,10 @@ [project] name = "diode-napalm-agent" version = "0.0.1" # Overwritten during the build process -description = "NetBox Labs, Diode SuzieQ Agent" +description = "NetBox Labs, Diode NAPALM Agent" readme = "README.md" requires-python = ">=3.7" -license = {file = "LICENSE.txt"} +license = {file = "../LICENSE.txt"} authors = [ {name = "NetBox Labs", email = "support@netboxlabs.com" } # Optional ] @@ -40,7 +40,7 @@ test = ["coverage", "pytest", "pytest-cov"] "Homepage" = "https://netboxlabs.com/" [project.scripts] # Optional -diode-napalm-agent = "diode_napalm.cli.cli:agent_main" +diode-napalm-agent = "diode_napalm.cli.cli:main" [tool.setuptools] packages = [ @@ -55,10 +55,6 @@ build-backend = "setuptools.build_meta" [tool.ruff] line-length = 140 -exclude = [ - "netboxlabs/diode/sdk/diode/*", - "netboxlabs/diode/sdk/validate/*", -] [tool.ruff.format] quote-style = "double"