From 81688cd8b2c124d6147211b3a3b615cbdf90a3e7 Mon Sep 17 00:00:00 2001 From: Leonardo Parente <23251360+leoparente@users.noreply.github.com> Date: Fri, 27 Dec 2024 15:45:48 -0300 Subject: [PATCH] feat: add more defaults support for network and device discovery (#44) --- .../device_discovery/policy/models.py | 21 ++++- .../device_discovery/policy/runner.py | 4 +- .../device_discovery/translate.py | 82 ++++++++++++++++--- device-discovery/tests/policy/test_runner.py | 4 +- device-discovery/tests/test_client.py | 3 +- device-discovery/tests/test_translate.py | 58 ++++++++++--- network-discovery/config/config.go | 13 ++- network-discovery/policy/manager_test.go | 4 +- network-discovery/policy/runner.go | 15 +++- network-discovery/policy/runner_test.go | 7 +- 10 files changed, 168 insertions(+), 43 deletions(-) diff --git a/device-discovery/device_discovery/policy/models.py b/device-discovery/device_discovery/policy/models.py index f23c98a..821a291 100644 --- a/device-discovery/device_discovery/policy/models.py +++ b/device-discovery/device_discovery/policy/models.py @@ -30,14 +30,29 @@ class Napalm(BaseModel): default=None, description="Optional arguments" ) +class ObjectParameters(BaseModel): + """Model for object parameters.""" + + comments: str | None = Field(default=None, description="Comments, optional") + description: str | None = Field(default=None, description="Description, optional") + tags: list[str] | None = Field(default=None, description="Tags, optional") + +class Defaults(BaseModel): + """Model for default configuration.""" + + site: str | None = Field(default=None, description="Site name, optional") + role: str | None = Field(default=None, description="Device Role name, optional") + tags: list[str] | None = Field(default=None, description="Tags, optional") + device: ObjectParameters | None = Field(default=None, description="Device parameters, optional") + interface: ObjectParameters | None = Field(default=None, description="Interface parameters, optional") + ipaddress: ObjectParameters | None = Field(default=None, description="IP Address parameters, optional") + prefix: ObjectParameters | None = Field(default=None, description="Prefix parameters, optional") class Config(BaseModel): """Model for discovery configuration.""" schedule: str | None = Field(default=None, description="cron interval, optional") - defaults: dict[str, str] | None = Field( - default=None, description="NetBox configuration" - ) + defaults: Defaults | None = Field(default=None, description="Default configuration, optional") @field_validator("schedule") @classmethod diff --git a/device-discovery/device_discovery/policy/runner.py b/device-discovery/device_discovery/policy/runner.py index 064bcc3..ee135e8 100644 --- a/device-discovery/device_discovery/policy/runner.py +++ b/device-discovery/device_discovery/policy/runner.py @@ -13,7 +13,7 @@ from device_discovery.client import Client from device_discovery.discovery import discover_device_driver, supported_drivers -from device_discovery.policy.models import Config, Napalm, Status +from device_discovery.policy.models import Config, Defaults, Napalm, Status # Set up logging logging.basicConfig(level=logging.INFO) @@ -125,10 +125,10 @@ def run(self, id: str, scope: Napalm, config: Config): ) as device: data = { "driver": scope.driver, - "site": config.defaults.get("site", None), "device": device.get_facts(), "interface": device.get_interfaces(), "interface_ip": device.get_interfaces_ip(), + "defaults": config.defaults, } Client().ingest(scope.hostname, data) except Exception as e: diff --git a/device-discovery/device_discovery/translate.py b/device-discovery/device_discovery/translate.py index 34f7779..d1f2880 100644 --- a/device-discovery/device_discovery/translate.py +++ b/device-discovery/device_discovery/translate.py @@ -15,6 +15,8 @@ Prefix, ) +from device_discovery.policy.models import Defaults + def int32_overflows(number: int) -> bool: """ @@ -34,19 +36,29 @@ def int32_overflows(number: int) -> bool: return not (INT32_MIN <= number <= INT32_MAX) -def translate_device(device_info: dict) -> Device: +def translate_device(device_info: dict, defaults: Defaults) -> Device: """ Translate device information from NAPALM format to Diode SDK Device entity. Args: ---- device_info (dict): Dictionary containing device information. + defaults (Defaults): Default configuration. Returns: ------- Device: Translated Device entity. """ + tags = list(defaults.tags) if defaults.tags else [] + description = None + comments = None + + if defaults.device: + tags.extend(defaults.device.tags) + description = defaults.device.description + comments = defaults.device.comments + device = Device( name=device_info.get("hostname"), device_type=DeviceType( @@ -55,15 +67,19 @@ def translate_device(device_info: dict) -> Device: platform=Platform( name=device_info.get("driver"), manufacturer=device_info.get("vendor") ), + role=defaults.role, serial=device_info.get("serial_number"), status="active", - site=device_info.get("site"), + site=defaults.site, + tags=tags, + description=description, + comments=comments, ) return device def translate_interface( - device: Device, if_name: str, interface_info: dict + device: Device, if_name: str, interface_info: dict, defaults: Defaults ) -> Interface: """ Translate interface information from NAPALM format to Diode SDK Interface entity. @@ -73,18 +89,29 @@ def translate_interface( device (Device): The device to which the interface belongs. if_name (str): The name of the interface. interface_info (dict): Dictionary containing interface information. + defaults (Defaults): Default configuration. Returns: ------- Interface: Translated Interface entity. """ + tags = list(defaults.tags) if defaults.tags else [] + description = None + + if defaults.interface: + tags.extend(defaults.interface.tags) + description = defaults.interface.description + + description = interface_info.get("description", description) + interface = Interface( device=device, name=if_name, enabled=interface_info.get("is_enabled"), mac_address=interface_info.get("mac_address"), - description=interface_info.get("description"), + description=description, + tags=tags, ) # Convert napalm interface speed from Mbps to Netbox Kbps @@ -100,7 +127,7 @@ def translate_interface( def translate_interface_ips( - interface: Interface, interfaces_ip: dict + interface: Interface, interfaces_ip: dict, defaults: Defaults ) -> Iterable[Entity]: """ Translate IP address and Prefixes information for an interface. @@ -110,12 +137,32 @@ def translate_interface_ips( interface (Interface): The interface entity. if_name (str): The name of the interface. interfaces_ip (dict): Dictionary containing interface IP information. + defaults (Defaults): Default configuration. Returns: ------- Iterable[Entity]: Iterable of translated IP address and Prefixes entities. """ + tags = defaults.tags if defaults.tags else [] + ip_tags = list(tags) + ip_comments = None + ip_description = None + + prefix_tags = list(tags) + prefix_comments = None + prefix_description = None + + if defaults.ipaddress: + ip_tags.extend(defaults.ipaddress.tags) + ip_comments = defaults.ipaddress.comments + ip_description = defaults.ipaddress.description + + if defaults.prefix: + prefix_tags.extend(defaults.prefix.tags) + prefix_comments = defaults.prefix.comments + prefix_description = defaults.prefix.description + ip_entities = [] for if_ip_name, ip_info in interfaces_ip.items(): @@ -127,14 +174,22 @@ def translate_interface_ips( ip_entities.append( Entity( prefix=Prefix( - prefix=str(network), site=interface.device.site + prefix=str(network), + site=interface.device.site, + tags=prefix_tags, + comments=prefix_comments, + description=prefix_description, ) ) ) ip_entities.append( Entity( ip_address=IPAddress( - address=ip_address, interface=interface + address=ip_address, + interface=interface, + tags=ip_tags, + comments=ip_comments, + description=ip_description, ) ) ) @@ -157,20 +212,25 @@ def translate_data(data: dict) -> Iterable[Entity]: """ entities = [] + defaults = data.get("defaults", Defaults()) + device_info = data.get("device", {}) interfaces = data.get("interface", {}) interfaces_ip = data.get("interface_ip", {}) if device_info: device_info["driver"] = data.get("driver") - device_info["site"] = data.get("site") - device = translate_device(device_info) + device = translate_device(device_info, defaults) entities.append(Entity(device=device)) 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) + interface = translate_interface( + device, if_name, interface_info, defaults + ) entities.append(Entity(interface=interface)) - entities.extend(translate_interface_ips(interface, interfaces_ip)) + entities.extend( + translate_interface_ips(interface, interfaces_ip, defaults) + ) return entities diff --git a/device-discovery/tests/policy/test_runner.py b/device-discovery/tests/policy/test_runner.py index 6ba38a5..a02b040 100644 --- a/device-discovery/tests/policy/test_runner.py +++ b/device-discovery/tests/policy/test_runner.py @@ -7,7 +7,7 @@ import pytest from apscheduler.triggers.date import DateTrigger -from device_discovery.policy.models import Config, Napalm, Status +from device_discovery.policy.models import Config, Defaults, Napalm, Status from device_discovery.policy.runner import PolicyRunner @@ -20,7 +20,7 @@ def policy_runner(): @pytest.fixture def sample_config(): """Fixture for a sample Config object.""" - return Config(schedule="0 * * * *", defaults={"site": "New York"}) + return Config(schedule="0 * * * *", defaults=Defaults(site="New York")) @pytest.fixture diff --git a/device-discovery/tests/test_client.py b/device-discovery/tests/test_client.py index 3fae9a5..7be429d 100644 --- a/device-discovery/tests/test_client.py +++ b/device-discovery/tests/test_client.py @@ -2,6 +2,7 @@ # Copyright 2024 NetBox Labs Inc """NetBox Labs - Client Unit Tests.""" +from types import SimpleNamespace from unittest.mock import patch import pytest @@ -35,7 +36,7 @@ def sample_data(): "GigabitEthernet0/0": {"ipv4": {"192.0.2.1": {"prefix_length": 24}}} }, "driver": "ios", - "site": "New York", + "defaults": SimpleNamespace(site="New York", role=None, tags = None, device = None), } diff --git a/device-discovery/tests/test_translate.py b/device-discovery/tests/test_translate.py index 7e8b544..e1074dd 100644 --- a/device-discovery/tests/test_translate.py +++ b/device-discovery/tests/test_translate.py @@ -3,7 +3,9 @@ """NetBox Labs - Translate Unit Tests.""" import pytest +from netboxlabs.diode.sdk.ingester import Tag +from device_discovery.policy.models import Defaults, ObjectParameters from device_discovery.translate import ( translate_data, translate_device, @@ -60,21 +62,42 @@ def sample_interfaces_ip(): return {"GigabitEthernet0/0": {"ipv4": {"192.0.2.1": {"prefix_length": 24}}}} -def test_translate_device(sample_device_info): +@pytest.fixture +def sample_defaults(): + """Sample defaults for testing.""" + return Defaults( + site="New York", + tags=["tag1", "tag2"], + device=ObjectParameters(comments="testing", tags=["devtag"]), + interface=ObjectParameters(description="testing", tags=["inttag"]), + ipaddress=ObjectParameters(description="ip test", tags=["iptag"]), + prefix=ObjectParameters(description="prefix test",tags=["prefixtag"]), + role="router", + ) + + +def test_translate_device(sample_device_info, sample_defaults): """Ensure device translation is correct.""" - device = translate_device(sample_device_info) + device = translate_device(sample_device_info, sample_defaults) assert device.name == "router1" assert device.device_type.model == "ISR4451" assert device.platform.name == "ios" assert device.serial == "123456789" assert device.site.name == "New York" + assert device.comments == "testing" + assert device.tags == [Tag(name="tag1"), Tag(name="tag2"), Tag(name="devtag")] -def test_translate_interface(sample_device_info, sample_interface_info): +def test_translate_interface( + sample_device_info, sample_interface_info, sample_defaults +): """Ensure interface translation is correct.""" - device = translate_device(sample_device_info) + device = translate_device(sample_device_info, sample_defaults) interface = translate_interface( - device, "GigabitEthernet0/0", sample_interface_info["GigabitEthernet0/0"] + device, + "GigabitEthernet0/0", + sample_interface_info["GigabitEthernet0/0"], + sample_defaults, ) assert interface.device.name == "router1" assert interface.name == "GigabitEthernet0/0" @@ -83,17 +106,19 @@ def test_translate_interface(sample_device_info, sample_interface_info): assert interface.mac_address == "00:1C:58:29:4A:71" assert interface.speed == 1000000 assert interface.description == "Uplink Interface" + assert interface.tags == [Tag(name="tag1"), Tag(name="tag2"), Tag(name="inttag")] def test_translate_interface_with_overflow_data( - sample_device_info, sample_interface_overflows_info + sample_device_info, sample_interface_overflows_info, sample_defaults ): """Ensure interface translation is correct.""" - device = translate_device(sample_device_info) + device = translate_device(sample_device_info, sample_defaults) interface = translate_interface( device, "GigabitEthernet0/0", sample_interface_overflows_info["GigabitEthernet0/0"], + sample_defaults, ) assert interface.device.name == "router1" assert interface.name == "GigabitEthernet0/0" @@ -102,20 +127,30 @@ def test_translate_interface_with_overflow_data( assert interface.mac_address == "00:1C:58:29:4A:71" assert interface.speed == 0 assert interface.description == "Uplink Interface" + assert interface.tags == [Tag(name="tag1"), Tag(name="tag2"), Tag(name="inttag")] def test_translate_interface_ips( - sample_device_info, sample_interface_info, sample_interfaces_ip + sample_device_info, sample_interface_info, sample_interfaces_ip, sample_defaults ): """Ensure interface IPs translation is correct.""" - device = translate_device(sample_device_info) + device = translate_device(sample_device_info, sample_defaults) interface = translate_interface( - device, "GigabitEthernet0/0", sample_interface_info["GigabitEthernet0/0"] + device, + "GigabitEthernet0/0", + sample_interface_info["GigabitEthernet0/0"], + sample_defaults, + ) + ip_entities = list( + translate_interface_ips(interface, sample_interfaces_ip, sample_defaults) ) - ip_entities = list(translate_interface_ips(interface, sample_interfaces_ip)) assert len(ip_entities) == 2 assert ip_entities[0].prefix.prefix == "192.0.2.0/24" assert ip_entities[1].ip_address.address == "192.0.2.1/24" + assert ip_entities[0].prefix.description == "prefix test" + assert ip_entities[1].ip_address.description == "ip test" + assert ip_entities[0].prefix.tags == [Tag(name="tag1"), Tag(name="tag2"), Tag(name="prefixtag")] + assert ip_entities[1].ip_address.tags == [Tag(name="tag1"), Tag(name="tag2"), Tag(name="iptag")] def test_translate_data( @@ -127,7 +162,6 @@ def test_translate_data( "interface": sample_interface_info, "interface_ip": sample_interfaces_ip, "driver": "ios", - "site": "New York", } entities = list(translate_data(data)) assert len(entities) == 4 diff --git a/network-discovery/config/config.go b/network-discovery/config/config.go index c1854ef..f77485a 100644 --- a/network-discovery/config/config.go +++ b/network-discovery/config/config.go @@ -14,11 +14,18 @@ type Scope struct { Targets []string `yaml:"targets"` } +// Defaults represents the supported default values for a policy +type Defaults struct { + Description string `yaml:"description,omitempty"` + Comments string `yaml:"comments,omitempty"` + Tags []string `yaml:"tags,omitempty"` +} + // PolicyConfig represents the configuration of a policy type PolicyConfig struct { - Schedule *string `yaml:"schedule"` - Defaults map[string]string `yaml:"defaults"` - Timeout int `yaml:"timeout"` + Schedule *string `yaml:"schedule,omitempty"` + Defaults Defaults `yaml:"defaults"` + Timeout int `yaml:"timeout"` } // Policy represents a network-discovery policy diff --git a/network-discovery/policy/manager_test.go b/network-discovery/policy/manager_test.go index 2da326f..f8b580e 100644 --- a/network-discovery/policy/manager_test.go +++ b/network-discovery/policy/manager_test.go @@ -42,7 +42,7 @@ func TestManagerParsePolicies(t *testing.T) { policy1: config: defaults: - site: New York NY + comments: test scope: targets: - 192.168.1.1/24 @@ -51,7 +51,7 @@ func TestManagerParsePolicies(t *testing.T) { policies, err := manager.ParsePolicies(yamlData) assert.NoError(t, err) assert.Contains(t, policies, "policy1") - assert.Equal(t, "New York NY", policies["policy1"].Config.Defaults["site"]) + assert.Equal(t, "test", policies["policy1"].Config.Defaults.Comments) }) t.Run("No Policies", func(t *testing.T) { diff --git a/network-discovery/policy/runner.go b/network-discovery/policy/runner.go index 0616a6b..e34bc26 100644 --- a/network-discovery/policy/runner.go +++ b/network-discovery/policy/runner.go @@ -97,11 +97,18 @@ func (r *Runner) run() { ip := &diode.IPAddress{ Address: diode.String(host.Addresses[0].Addr + "/32"), } - if r.config.Defaults["description"] != "" { - ip.Description = diode.String(r.config.Defaults["description"]) + if r.config.Defaults.Description != "" { + ip.Description = diode.String(r.config.Defaults.Description) } - if r.config.Defaults["comments"] != "" { - ip.Comments = diode.String(r.config.Defaults["comments"]) + if r.config.Defaults.Comments != "" { + ip.Comments = diode.String(r.config.Defaults.Comments) + } + if len(r.config.Defaults.Tags) > 0 { + var tags []*diode.Tag + for _, tag := range r.config.Defaults.Tags { + tags = append(tags, &diode.Tag{Name: diode.String(tag)}) + } + ip.Tags = tags } entities = append(entities, ip) } diff --git a/network-discovery/policy/runner_test.go b/network-discovery/policy/runner_test.go index 8127a1e..dfbb8de 100644 --- a/network-discovery/policy/runner_test.go +++ b/network-discovery/policy/runner_test.go @@ -79,9 +79,10 @@ func TestRunnerRun(t *testing.T) { policyConfig := config.Policy{ Config: config.PolicyConfig{ Schedule: nil, - Defaults: map[string]string{ - "description": "Test", - "comments": "This is a test", + Defaults: config.Defaults{ + Description: "Test", + Comments: "This is a test", + Tags: []string{"test", "ip"}, }, }, Scope: config.Scope{