diff --git a/.github/workflows/device-discovery-release.yaml b/.github/workflows/device-discovery-release.yaml index e970c4a..bed366e 100644 --- a/.github/workflows/device-discovery-release.yaml +++ b/.github/workflows/device-discovery-release.yaml @@ -20,6 +20,7 @@ env: permissions: contents: write + issues: write pull-requests: write jobs: diff --git a/.github/workflows/network-discovery-release.yaml b/.github/workflows/network-discovery-release.yaml index ec481cb..7cb8d38 100644 --- a/.github/workflows/network-discovery-release.yaml +++ b/.github/workflows/network-discovery-release.yaml @@ -19,6 +19,7 @@ env: permissions: contents: write + issues: write pull-requests: write jobs: diff --git a/README.md b/README.md index f71fdae..0b33774 100644 --- a/README.md +++ b/README.md @@ -2,5 +2,5 @@ Orb discovery backends collection -- [device-discovery](./device-discovery/README.md) - TBD -- [network-discovery](./network-discovery/README.md) - TBD \ No newline at end of file +- [device-discovery](./device-discovery/README.md) - Device Discovery Backend that uses [NAPALM](https://github.com/napalm-automation/napalm) Drivers. +- [network-discovery](./network-discovery/README.md) - Network Discovery Backend which is a wrapper over [NMAP](https://nmap.org/) scanner. \ No newline at end of file diff --git a/device-discovery/README.md b/device-discovery/README.md index f624d9c..e6540bf 100644 --- a/device-discovery/README.md +++ b/device-discovery/README.md @@ -1,40 +1,47 @@ # device-discovery Orb device discovery backend -### Config RFC -```yaml -discovery: - config: - target: grpc://localhost:8080/diode - api_key: ${DIODE_API_KEY} - host: 0.0.0.0 - port: 8072 +### Usage +```bash +usage: device-discovery [-h] [-V] [-s HOST] [-p PORT] -t DIODE_TARGET -k DIODE_API_KEY + +Orb Discovery Backend + +options: + -h, --help show this help message and exit + -V, --version Display Discovery, NAPALM and Diode SDK versions + -s HOST, --host HOST Server host + -p PORT, --port PORT Server port + -t DIODE_TARGET, --diode-target DIODE_TARGET + Diode target + -k DIODE_API_KEY, --diode-api-key DIODE_API_KEY + Diode API key. Environment variables can be used by wrapping them in ${} (e.g. + ${MY_API_KEY}) ``` ### Policy RFC ```yaml -discovery: - policies: - discovery_1: - config: - schedule: "* * * * *" #Cron expression - defaults: - site: New York NY - scope: - - hostname: 192.168.0.32 - username: ${USER} - password: admin - - driver: eos - hostname: 127.0.0.1 - username: admin - password: ${ARISTA_PASSWORD} - optional_args: - enable_password: ${ARISTA_PASSWORD} - discover_once: # will run only once - scope: - - hostname: 192.168.0.34 - username: ${USER} - password: ${PASSWORD} +policies: + discovery_1: + config: + schedule: "* * * * *" #Cron expression + defaults: + site: New York NY + scope: + - hostname: 192.168.0.32 + username: ${USER} + password: admin + - driver: eos + hostname: 127.0.0.1 + username: admin + password: ${ARISTA_PASSWORD} + optional_args: + enable_password: ${ARISTA_PASSWORD} + discover_once: # will run only once + scope: + - hostname: 192.168.0.34 + username: ${USER} + password: ${PASSWORD} ``` ## Run device-discovery device-discovery can be run by installing it with pip @@ -42,14 +49,16 @@ device-discovery can be run by installing it with pip git clone https://github.com/netboxlabs/orb-discovery.git cd orb-discovery/ pip install --no-cache-dir ./device-discovery/ -device-discovery -c config.yaml +device-discovery -t 'grpc://192.168.0.10:8080/diode' -k '${DIODE_API_KEY}' ``` ## Docker Image device-discovery can be build and run using docker: ```sh -docker build --no-cache -t device-discovery:develop -f device-discovery/docker/Dockerfile . -docker run -v /local/orb:/usr/local/orb/ -p 8072:8072 device-discovery:develop device-discovery -c /usr/local/orb/config.yaml +cd device-discovery +docker build --no-cache -t device-discovery:develop -f docker/Dockerfile . +docker run -e DIODE_API_KEY={YOUR_API_KEY} -p 8072:8072 device-discovery:develop \ + device-discovery -t 'grpc://192.168.0.10:8080/diode' -k '${DIODE_API_KEY}' ``` ### Routes (v1) diff --git a/device-discovery/device_discovery/discovery.py b/device-discovery/device_discovery/discovery.py index c086792..e14dcb6 100644 --- a/device-discovery/device_discovery/discovery.py +++ b/device-discovery/device_discovery/discovery.py @@ -16,9 +16,8 @@ def napalm_driver_list() -> list[str]: """ List the available NAPALM drivers. - This function scans the installed Python packages to identify NAPALM drivers, - appending their names (with the 'napalm-' prefix removed and hyphens replaced - with underscores) to a list of known drivers. + This function scans the installed Python modules to identify NAPALM drivers, + appending their names (with the 'napalm_' prefix removed) to a list of known drivers. Returns ------- @@ -28,11 +27,10 @@ def napalm_driver_list() -> list[str]: """ napalm_packages = ["ios", "eos", "junos", "nxos"] - prefix = "napalm-" - for dist in importlib_metadata.distributions(): - if dist.metadata["Name"].startswith(prefix): - package = dist.metadata["Name"][len(prefix) :].replace("-", "_") - napalm_packages.append(package) + prefix = "napalm_" + for dist in importlib_metadata.packages_distributions(): + if dist.startswith(prefix): + napalm_packages.append(dist[len(prefix) :]) return napalm_packages diff --git a/device-discovery/device_discovery/main.py b/device-discovery/device_discovery/main.py index 9e07c09..828d11c 100644 --- a/device-discovery/device_discovery/main.py +++ b/device-discovery/device_discovery/main.py @@ -3,6 +3,7 @@ """Orb Discovery entry point.""" import argparse +import os import sys from importlib.metadata import version @@ -10,7 +11,6 @@ import uvicorn from device_discovery.client import Client -from device_discovery.parser import parse_config_file from device_discovery.server import app from device_discovery.version import version_semver @@ -31,23 +31,50 @@ def main(): help="Display Discovery, NAPALM and Diode SDK versions", ) parser.add_argument( - "-c", - "--config", - metavar="config.yaml", - help="Yaml configuration file", + "-s", + "--host", + default="0.0.0.0", + help="Server host", + type=str, + required=False, + ) + parser.add_argument( + "-p", + "--port", + default=8072, + help="Server port", + type=str, + required=False, + ) + parser.add_argument( + "-t", + "--diode-target", + help="Diode target", + type=str, + required=True, + ) + + parser.add_argument( + "-k", + "--diode-api-key", + help="Diode API key. Environment variables can be used by wrapping them in ${} (e.g. ${MY_API_KEY})", type=str, required=True, ) - args = parser.parse_args() try: - cfg = parse_config_file(args.config) + args = parser.parse_args() + api_key = args.diode_api_key + if api_key.startswith("${") and api_key.endswith("}"): + env_var = api_key[2:-1] + api_key = os.getenv(env_var, api_key) + client = Client() - client.init_client(target=cfg.config.target, api_key=cfg.config.api_key) + client.init_client(target=args.diode_target, api_key=args.diode_api_key) uvicorn.run( app, - host=cfg.config.host, - port=cfg.config.port, + host=args.host, + port=args.port, ) except (KeyboardInterrupt, RuntimeError): pass diff --git a/device-discovery/device_discovery/parser.py b/device-discovery/device_discovery/parser.py deleted file mode 100644 index 3b7941d..0000000 --- a/device-discovery/device_discovery/parser.py +++ /dev/null @@ -1,120 +0,0 @@ -#!/usr/bin/env python -# Copyright 2024 NetBox Labs Inc -"""Parse Orb Discovery Config file.""" - -import os -from pathlib import Path - -import yaml -from pydantic import BaseModel, ValidationError - - -class ParseException(Exception): - """Custom exception for parsing errors.""" - - pass - -class DiscoveryConfig(BaseModel): - """Model for Diode configuration.""" - - host: str = "0.0.0.0" - port: int = 8072 - target: str - api_key: str - - -class Discovery(BaseModel): - """Model for Discovery containing configuration and policies.""" - - config: DiscoveryConfig - - -class Base(BaseModel): - """Top-level model for the entire configuration.""" - - discovery: Discovery - - -def resolve_env_vars(config): - """ - Recursively resolve environment variables in the configuration. - - Args: - ---- - config (dict): The configuration dictionary. - - Returns: - ------- - dict: The configuration dictionary with environment variables resolved. - - """ - if isinstance(config, dict): - return {k: resolve_env_vars(v) for k, v in config.items()} - if isinstance(config, list): - return [resolve_env_vars(i) for i in config] - if isinstance(config, str) and config.startswith("${") and config.endswith("}"): - env_var = config[2:-1] - return os.getenv(env_var, config) - return config - - -def parse_config(config_data: str) -> Base: - """ - 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: - # Parse the YAML configuration data - config_dict = yaml.safe_load(config_data) - # Resolve environment variables - resolved_config = resolve_env_vars(config_dict) - # Parse the data into the Config model - config = Base(**resolved_config) - return config - except yaml.YAMLError as e: - raise ParseException(f"YAML ERROR: {e}") - except ValidationError as e: - raise ParseException("Validation ERROR:", e) - - -def parse_config_file(file_path: Path) -> Discovery: - """ - Parse the Discovery configuration file and return the Discovery configuration object. - - This function reads the content of the specified YAML configuration file, - parses it into a `Config` object, and returns the `Discovery` part of the configuration. - - Args: - ---- - file_path (Path): The path to the YAML configuration file. - - Returns: - ------- - Discovery: The `Discovery` configuration object extracted from the parsed configuration. - - Raises: - ------ - ParseException: If there is an error parsing the YAML content or validating the data. - Exception: If there is an error opening the file or any other unexpected error. - - """ - try: - with open(file_path) as f: - cfg = parse_config(f.read()) - except ParseException: - raise - except Exception as e: - raise Exception(f"Unable to open config file {file_path}: {e.args[1]}") - return cfg.discovery diff --git a/device-discovery/device_discovery/policy/manager.py b/device-discovery/device_discovery/policy/manager.py index 0052fde..394624a 100644 --- a/device-discovery/device_discovery/policy/manager.py +++ b/device-discovery/device_discovery/policy/manager.py @@ -3,10 +3,10 @@ """Orb Discovery Policy Manager.""" import logging +import os import yaml -from device_discovery.parser import resolve_env_vars from device_discovery.policy.models import Policy, PolicyRequest from device_discovery.policy.runner import PolicyRunner @@ -15,6 +15,28 @@ logger = logging.getLogger(__name__) +def resolve_env_vars(config): + """ + Recursively resolve environment variables in the configuration. + + Args: + ---- + config (dict): The configuration dictionary. + + Returns: + ------- + dict: The configuration dictionary with environment variables resolved. + + """ + if isinstance(config, dict): + return {k: resolve_env_vars(v) for k, v in config.items()} + if isinstance(config, list): + return [resolve_env_vars(i) for i in config] + if isinstance(config, str) and config.startswith("${") and config.endswith("}"): + env_var = config[2:-1] + return os.getenv(env_var, config) + return config + class PolicyManager: """Policy Manager class.""" diff --git a/device-discovery/device_discovery/policy/models.py b/device-discovery/device_discovery/policy/models.py index 5536fc0..f23c98a 100644 --- a/device-discovery/device_discovery/policy/models.py +++ b/device-discovery/device_discovery/policy/models.py @@ -68,13 +68,7 @@ class Policy(BaseModel): scope: list[Napalm] -class Discovery(BaseModel): - """Model for Discovery containing configuration and policies.""" - - policies: dict[str, Policy] - - class PolicyRequest(BaseModel): """Model for a policy request.""" - discovery: Discovery + policies: dict[str, Policy] diff --git a/device-discovery/device_discovery/server.py b/device-discovery/device_discovery/server.py index 7024c33..75a3f0b 100644 --- a/device-discovery/device_discovery/server.py +++ b/device-discovery/device_discovery/server.py @@ -119,7 +119,7 @@ async def write_policy(request: PolicyRequest = Depends(parse_yaml_body)): """ started_policies = [] - policies = request.discovery.policies + policies = request.policies for name, policy in policies.items(): try: manager.start_policy(name, policy) diff --git a/device-discovery/docker/Dockerfile b/device-discovery/docker/Dockerfile index 70ecbc8..aad9437 100644 --- a/device-discovery/docker/Dockerfile +++ b/device-discovery/docker/Dockerfile @@ -2,11 +2,11 @@ FROM python:3.12-slim-bullseye AS builder WORKDIR /usr/src/app -COPY device-discovery/ ./device-discovery +COPY . . RUN python -m venv /opt/venv && \ . /opt/venv/bin/activate && \ - pip install --no-cache-dir ./device-discovery && \ + pip install --no-cache-dir . && \ deactivate FROM python:3.12-slim-bullseye diff --git a/device-discovery/tests/policy/test_manager.py b/device-discovery/tests/policy/test_manager.py index e81fe04..723b880 100644 --- a/device-discovery/tests/policy/test_manager.py +++ b/device-discovery/tests/policy/test_manager.py @@ -58,50 +58,47 @@ def test_start_existing_policy_raises_error(policy_manager, sample_policy): def test_parse_policy(policy_manager): """Test parsing YAML configuration into a PolicyRequest object.""" config_data = b""" - discovery: - policies: - policy1: - config: - schedule: "0 * * * *" - defaults: - site: "New York" - scope: - - driver: "ios" - hostname: "router1" - username: "admin" - password: "password" + policies: + policy1: + config: + schedule: "0 * * * *" + defaults: + site: "New York" + scope: + - driver: "ios" + hostname: "router1" + username: "admin" + password: "password" """ - with patch("device_discovery.parser.resolve_env_vars", side_effect=lambda x: x): - policy_request = policy_manager.parse_policy(config_data) + policy_request = policy_manager.parse_policy(config_data) # Verify structure of the parsed PolicyRequest - assert isinstance(policy_request, PolicyRequest) - assert "policy1" in policy_request.discovery.policies + assert isinstance(policy_request, PolicyRequest) + assert "policy1" in policy_request.policies def test_parse_policy_invalid_cron(policy_manager): """Test parsing YAML configuration with an invalid cron string.""" # Invalid cron string in schedule config_data = b""" - discovery: - policies: - policy1: - config: - schedule: "invalid cron string" - defaults: - site: "New York" - scope: - - driver: "ios" - hostname: "router1" - username: "admin" - password: "password" + policies: + policy1: + config: + schedule: "invalid cron string" + defaults: + site: "New York" + scope: + - driver: "ios" + hostname: "router1" + username: "admin" + password: "password" """ - with patch("device_discovery.parser.resolve_env_vars", side_effect=lambda x: x): - with pytest.raises(ValidationError) as exc_info: - policy_manager.parse_policy(config_data) - # Validate that the error is related to the invalid cron string - assert exc_info.match("Invalid cron schedule format.") + with pytest.raises(ValidationError) as exc_info: + policy_manager.parse_policy(config_data) + + # Validate that the error is related to the invalid cron string + assert exc_info.match("Invalid cron schedule format.") def test_policy_exists(policy_manager): diff --git a/device-discovery/tests/test_discovery.py b/device-discovery/tests/test_discovery.py index 331108e..33f1e06 100644 --- a/device-discovery/tests/test_discovery.py +++ b/device-discovery/tests/test_discovery.py @@ -22,11 +22,12 @@ def mock_get_network_driver(): with patch("device_discovery.discovery.get_network_driver") as mock: yield mock - @pytest.fixture -def mock_importlib_metadata_distributions(): - """Mock the importlib_metadata.distributions function.""" - with patch("device_discovery.discovery.importlib_metadata.distributions") as mock: +def mock_importlib_metadata_packages_distributions(): + """Mock the importlib_metadata.packages_distributions function.""" + with patch( + "device_discovery.discovery.importlib_metadata.packages_distributions" + ) as mock: yield mock @@ -148,20 +149,21 @@ def side_effect(driver_name): assert driver == "nxos", "Expected the 'ios' driver to be found" -def test_napalm_driver_list(mock_importlib_metadata_distributions): +def test_napalm_driver_list(mock_importlib_metadata_packages_distributions): """ Test the napalm_driver_list function to ensure it correctly lists available NAPALM drivers. Args: ---- - mock_importlib_metadata_distributions: Mocked importlib_metadata.distributions function. + mock_importlib_metadata_packages_distributions: Mocked + importlib_metadata.packages_distributions function. """ mock_distributions = [ - MagicMock(metadata={"Name": "napalm-srl"}), - MagicMock(metadata={"Name": "napalm-fake-driver"}), + "napalm_srl", + "napalm_fake_driver", ] - mock_importlib_metadata_distributions.return_value = mock_distributions + mock_importlib_metadata_packages_distributions.return_value = mock_distributions expected_drivers = ["ios", "eos", "junos", "nxos", "srl", "fake_driver"] drivers = napalm_driver_list() assert drivers == expected_drivers, f"Expected {expected_drivers}, got {drivers}" diff --git a/device-discovery/tests/test_main.py b/device-discovery/tests/test_main.py index d21f2e6..683f6ac 100644 --- a/device-discovery/tests/test_main.py +++ b/device-discovery/tests/test_main.py @@ -21,17 +21,6 @@ def mock_parse_args(): yield mock -@pytest.fixture -def mock_parse_config_file(): - """ - Fixture to mock the parse_config_file function. - - Mocks the parse_config_file method to simulate loading a configuration file. - """ - with patch("device_discovery.main.parse_config_file") as mock: - yield mock - - @pytest.fixture def mock_client(): """ @@ -54,18 +43,11 @@ def mock_uvicorn_run(): yield mock -def test_main_keyboard_interrupt(mock_parse_args, mock_parse_config_file): - """ - Test handling of KeyboardInterrupt in main. - - Args: - ---- - mock_parse_args: Mocked parse_args function. - mock_parse_config_file: Mocked parse_config_file function. - - """ - mock_parse_args.return_value = MagicMock(config="config.yaml") - mock_parse_config_file.side_effect = KeyboardInterrupt +def test_main_keyboard_interrupt(mock_parse_args): + """Test handling of KeyboardInterrupt in main.""" + mock_parse_args.return_value = MagicMock( + diode_target="grpc", diode_api_key="abc", host="0.0.0.0", port=1234 + ) with patch.object(sys, "exit", side_effect=Exception("Test Exit")): try: @@ -74,12 +56,11 @@ def test_main_keyboard_interrupt(mock_parse_args, mock_parse_config_file): assert str(e) == "Test Exit" -def test_main_with_config( - mock_parse_args, mock_parse_config_file, mock_client, mock_uvicorn_run -): +def test_main_with_config(mock_parse_args, mock_client, mock_uvicorn_run): """Test running the CLI with a configuration file and no environment file.""" - mock_parse_args.return_value = MagicMock(config="config.yaml") - mock_parse_config_file.return_value = MagicMock() + mock_parse_args.return_value = MagicMock( + diode_target="grpc", diode_api_key="abc", host="0.0.0.0", port=1234 + ) with patch.object(sys, "exit", side_effect=Exception("Test Exit")): try: @@ -87,17 +68,16 @@ def test_main_with_config( except Exception as e: assert str(e) == "Test Exit" - mock_parse_config_file.assert_called_once_with("config.yaml") + mock_parse_args.assert_called_once() mock_client.assert_called_once() mock_uvicorn_run.assert_called_once() -def test_main_start_server_failure( - mock_parse_args, mock_parse_config_file, mock_client, mock_uvicorn_run -): +def test_main_start_server_failure(mock_parse_args, mock_client, mock_uvicorn_run): """Test CLI failure when starting the agent.""" - mock_parse_args.return_value = MagicMock(config="config.yaml") - mock_parse_config_file.return_value = MagicMock() + mock_parse_args.return_value = MagicMock( + diode_target="grpc", diode_api_key="abc", host="0.0.0.0", port=1234 + ) mock_uvicorn_run.side_effect = Exception("Test Start Server Failure") with patch.object(sys, "exit", side_effect=Exception("Test Exit")) as mock_exit: @@ -106,7 +86,7 @@ def test_main_start_server_failure( except Exception as e: assert str(e) == "Test Exit" - mock_parse_config_file.assert_called_once_with("config.yaml") + mock_parse_args.assert_called_once() mock_client.assert_called_once() mock_uvicorn_run.assert_called_once() mock_exit.assert_called_once_with( @@ -135,20 +115,13 @@ def test_main_no_config_file(mock_parse_args): mock_exit.assert_called_once() -def test_main_missing_policy(mock_parse_args, mock_parse_config_file): - """ - Test handling of missing policy in start_agent. - - Args: - ---- - mock_parse_args: Mocked parse_args function. - mock_parse_config_file: Mocked parse_config_file function. - - """ - mock_parse_args.return_value = MagicMock(config="config.yaml", env=None, workers=2) +def test_main_missing_policy(mock_parse_args): + """Test handling of missing policy in start_agent.""" + mock_parse_args.return_value = MagicMock( + diode_target="grpc", diode_api_key="abc", host="0.0.0.0", port=1234 + ) mock_cfg = MagicMock() mock_cfg.policies = {"policy1": None} # Simulating a missing policy - mock_parse_config_file.return_value = mock_cfg with patch.object(sys, "exit", side_effect=Exception("Test Exit")): try: diff --git a/device-discovery/tests/test_parser.py b/device-discovery/tests/test_parser.py deleted file mode 100644 index b66b1d5..0000000 --- a/device-discovery/tests/test_parser.py +++ /dev/null @@ -1,97 +0,0 @@ -#!/usr/bin/env python -# Copyright 2024 NetBox Labs Inc -"""NetBox Labs - Parser Unit Tests.""" - -import os -from pathlib import Path -from unittest.mock import mock_open, patch - -import pytest - -from device_discovery.parser import ( - Base, - ParseException, - parse_config, - parse_config_file, - resolve_env_vars, -) - - -@pytest.fixture -def valid_yaml(): - """Valid Yaml Generator.""" - return """ - discovery: - config: - target: "target_value" - api_key: "api_key_value" - """ - - -@pytest.fixture -def invalid_yaml(): - """Invalid Yaml Generator.""" - return """ - discovery: - config: - api_key: "api_key_value" - host: "host_value" - """ - - -def test_parse_valid_config(valid_yaml): - """Ensure we can parse a valid configuration.""" - config = parse_config(valid_yaml) - assert isinstance(config, Base) - assert config.discovery.config.target == "target_value" - assert config.discovery.config.host == "0.0.0.0" - - -def test_parse_invalid_config(invalid_yaml): - """Ensure an invalid configuration raises a ParseException.""" - with pytest.raises(ParseException): - parse_config(invalid_yaml) - - -@patch("builtins.open", new_callable=mock_open, read_data="valid_yaml") -def test_parse_config_file(mock_file, valid_yaml): - """Ensure we can parse a configuration file.""" - with patch( - "device_discovery.parser.parse_config", return_value=parse_config(valid_yaml) - ): - config = parse_config_file(Path("fake_path.yaml")) - assert config.config.target == "target_value" - mock_file.assert_called_once_with(Path("fake_path.yaml")) - - -@patch("builtins.open", new_callable=mock_open, read_data="invalid_yaml") -def test_parse_config_file_parse_exception(mock_file): - """Ensure a ParseException in parse_config is propagated.""" - with patch( - "device_discovery.parser.parse_config", - side_effect=ParseException("Test Parse Exception"), - ): - with pytest.raises(ParseException): - parse_config_file(Path("fake_path.yaml")) - mock_file.assert_called_once_with(Path("fake_path.yaml")) - - -@patch.dict(os.environ, {"API_KEY": "env_api_key"}) -def test_resolve_env_vars(): - """Ensure environment variables are resolved correctly.""" - config_with_env_var = {"api_key": "${API_KEY}"} - resolved_config = resolve_env_vars(config_with_env_var) - assert resolved_config["api_key"] == "env_api_key" - - -def test_resolve_env_vars_no_env(): - """Ensure missing environment variables are handled correctly.""" - config_with_no_env_var = {"api_key": "${MISSING_KEY}"} - resolved_config = resolve_env_vars(config_with_no_env_var) - assert resolved_config["api_key"] == "${MISSING_KEY}" - - -def test_parse_config_file_exception(): - """Ensure file parsing errors are handled correctly.""" - with pytest.raises(Exception): - parse_config_file(Path("non_existent_file.yaml")) diff --git a/device-discovery/tests/test_server.py b/device-discovery/tests/test_server.py index 91934fa..8b614c1 100644 --- a/device-discovery/tests/test_server.py +++ b/device-discovery/tests/test_server.py @@ -7,7 +7,6 @@ import pytest import yaml from fastapi.testclient import TestClient -from pydantic import ValidationError from device_discovery.policy.models import PolicyRequest from device_discovery.server import app, manager @@ -23,18 +22,17 @@ def valid_policy_yaml(): Returns a YAML string that represents a valid PolicyRequest object. """ return """ - discovery: - policies: - policy1: - config: - schedule: "0 * * * *" - defaults: - site: "New York" - scope: - - driver: "ios" - hostname: "router1" - username: "admin" - password: "password" + policies: + policy1: + config: + schedule: "0 * * * *" + defaults: + site: "New York" + scope: + - driver: "ios" + hostname: "router1" + username: "admin" + password: "password" """ @@ -46,28 +44,27 @@ def multiple_policies_yaml(): Returns a YAML string that represents a valid PolicyRequest object. """ return """ - discovery: - policies: - policy1: - config: - schedule: "0 * * * *" - defaults: - site: "New York" - scope: - - driver: "ios" - hostname: "router1" - username: "admin" - password: "password" - policy2: - config: - schedule: "0 * * * *" - defaults: - site: "New York" - scope: - - driver: "ios" - hostname: "router1" - username: "admin" - password: "password" + policies: + policy1: + config: + schedule: "0 * * * *" + defaults: + site: "New York" + scope: + - driver: "ios" + hostname: "router1" + username: "admin" + password: "password" + policy2: + config: + schedule: "0 * * * *" + defaults: + site: "New York" + scope: + - driver: "ios" + hostname: "router1" + username: "admin" + password: "password" """ @@ -79,17 +76,16 @@ def invalid_policy_yaml(): Returns a YAML string that represents a valid PolicyRequest object. """ return """ - discovery: - policies: - policy1: - config: - schedule: "0 * * * *" - defaults: - site: "New York" - scope: - hostname: "router1" - username: "admin" - password: "password" + policies: + policy1: + config: + schedule: "0 * * * *" + defaults: + site: "New York" + scope: + hostname: "router1" + username: "admin" + password: "password" """ @@ -221,7 +217,7 @@ def test_write_policy_invalid_yaml(): response = client.post( "/api/v1/policies", headers={"Content-Type": "application/x-yaml"}, - json={"discovery": {"policies": {"policy1": {}}}}, + json={"policies": {"policy1": {}}}, ) assert response.status_code == 400 assert response.json() == {"detail": "Invalid YAML format"} @@ -238,7 +234,7 @@ def test_write_policy_validation_error(invalid_policy_yaml): assert response.json() == { "detail": [ { - "field": "discovery.policies.policy1.scope", + "field": "policies.policy1.scope", "type": "list_type", "error": "Input should be a valid list", } @@ -255,7 +251,7 @@ def test_write_policy_unexpected_parser_error(): response = client.post( "/api/v1/policies", headers={"Content-Type": "application/x-yaml"}, - json={"discovery": {"policies": {"policy1": {}}}}, + json={"policies": {"policy1": {}}}, ) assert response.status_code == 400 assert response.json() == {"detail": "unexpected error"} @@ -271,7 +267,7 @@ def test_write_policy_invalid_content_type(): response = client.post( "/api/v1/policies", headers={"content-type": "application/json"}, - json={"discovery": {"policies": {"policy1": {}}}}, + json={"policies": {"policy1": {}}}, ) assert response.status_code == 400 assert ( @@ -317,12 +313,12 @@ def test_write_policy_no_policy_error(): """ with patch( "device_discovery.server.parse_yaml_body", - return_value=PolicyRequest(discovery={"policies": {}}), + return_value=PolicyRequest(policies={}), ): response = client.post( "/api/v1/policies", headers={"Content-Type": "application/x-yaml"}, - json={"discovery": {"policies": {}}}, + json={"policies": {}}, ) assert response.status_code == 400 assert response.json()["detail"] == "no policies found in request" diff --git a/network-discovery/Makefile b/network-discovery/Makefile index b5989f6..b16a3fa 100644 --- a/network-discovery/Makefile +++ b/network-discovery/Makefile @@ -8,7 +8,7 @@ COMMIT_SHA := $(shell git rev-parse --short HEAD) .PHONY: build build: - CGO_ENABLED=$(CGO_ENABLED) GOOS=linux GOARCH=$(GOARCH) GOARM=$(GOARM) go build -mod=mod -o ${BUILD_DIR}/network-discovery cmd/main.go + CGO_ENABLED=$(CGO_ENABLED) GOOS=$(GOOS) GOARCH=$(GOARCH) GOARM=$(GOARM) go build -mod=mod -o ${BUILD_DIR}/network-discovery cmd/main.go .PHONY: lint lint: diff --git a/network-discovery/README.md b/network-discovery/README.md index 9b08a2a..228ada6 100644 --- a/network-discovery/README.md +++ b/network-discovery/README.md @@ -1,35 +1,39 @@ # network-discovery Orb network discovery backend -### Config RFC -```yaml -network: - config: - target: grpc://localhost:8080/diode - api_key: ${DIODE_API_KEY} - host: 0.0.0.0 - port: 8072 - log_level: info - log_format: json +### Usage +```sh +Usage of network-discovery: + -diode-api-key string + diode api key (REQUIRED). Environment variables can be used by wrapping them in ${} (e.g. ${MY_API_KEY}) + -diode-target string + diode target (REQUIRED) + -help + show this help + -host string + server host (default "0.0.0.0") + -log-format string + log format (default "TEXT") + -log-level string + log level (default "INFO") + -port int + server port (default 8073) ``` ### Policy RFC ```yaml -network: - policies: - network_1: - config: - schedule: "* * * * *" #Cron expression - defaults: - site: New York NY - scope: - targets: [192.168.1.0/24] - timeout: 5 #default 2 minutes - discover_once: # will run only once - scope: - targets: - - 92.168.0.34/24 - - google.com +policies: + network_1: + config: + schedule: "* * * * *" #Cron expression + timeout: 5 #default 2 minutes + scope: + targets: [192.168.1.0/24] + discover_once: # will run only once + scope: + targets: + - 192.168.0.34/24 + - google.com ``` ## Run device-discovery device-discovery can be run by installing it with pip @@ -37,7 +41,7 @@ device-discovery can be run by installing it with pip git clone https://github.com/netboxlabs/orb-discovery.git cd network-discovery/ make bin -build/network-discovery -c config.yaml +build/network-discovery --diode-target grpc://192.168.31.114:8080/diode --diode-api-key '${DIODE_API_KEY}' ``` ## Docker Image @@ -45,7 +49,10 @@ device-discovery can be build and run using docker: ```sh cd network-discovery/ docker build --no-cache -t network-discovery:develop -f docker/Dockerfile . -docker run -v /local/orb:/usr/local/orb/ --net=host device-discovery:develop network-discovery -c /usr/local/orb/config.yaml +docker run --net=host -e DIODE_API_KEY={YOUR_API_KEY} \ + network-discovery:develop network-discovery \ + --diode-target grpc://192.168.31.114:8080/diode \ + --diode-api-key '${DIODE_API_KEY}' ``` ### Routes (v1) diff --git a/network-discovery/cmd/main.go b/network-discovery/cmd/main.go index 30598d4..ef54682 100644 --- a/network-discovery/cmd/main.go +++ b/network-discovery/cmd/main.go @@ -6,10 +6,10 @@ import ( "fmt" "os" "os/signal" + "strings" "syscall" "github.com/netboxlabs/diode-sdk-go/diode" - "gopkg.in/yaml.v3" "github.com/netboxlabs/orb-discovery/network-discovery/config" "github.com/netboxlabs/orb-discovery/network-discovery/policy" @@ -20,47 +20,49 @@ import ( // AppName is the application name const AppName = "network-discovery" +func resolveEnv(value string) string { + // Check if the value starts with ${ and ends with } + if strings.HasPrefix(value, "${") && strings.HasSuffix(value, "}") { + // Extract the environment variable name + envVar := value[2 : len(value)-1] + // Get the value of the environment variable + envValue := os.Getenv(envVar) + if envValue != "" { + return envValue + } + fmt.Printf("error: environment variable %s is not set\n", envVar) + os.Exit(1) + } + // Return the original value if no substitution occurs + return value +} + func main() { - configPath := flag.String("config", "", "path to the configuration file (required)") + host := flag.String("host", "0.0.0.0", "server host") + port := flag.Int("port", 8073, "server port") + diodeTarget := flag.String("diode-target", "", "diode target (REQUIRED)") + diodeAPIKey := flag.String("diode-api-key", "", "diode api key (REQUIRED)."+ + " Environment variables can be used by wrapping them in ${} (e.g. ${MY_API_KEY})") + logLevel := flag.String("log-level", "INFO", "log level") + logFormat := flag.String("log-format", "TEXT", "log format") + help := flag.Bool("help", false, "show this help") flag.Parse() - if *configPath == "" { + if *help || *diodeTarget == "" || *diodeAPIKey == "" { fmt.Fprintf(os.Stderr, "Usage of network-discovery:\n") flag.PrintDefaults() - os.Exit(1) - - } - if _, err := os.Stat(*configPath); os.IsNotExist(err) { - fmt.Printf("configuration file '%s' does not exist", *configPath) - os.Exit(1) - } - fileData, err := os.ReadFile(*configPath) - if err != nil { - fmt.Printf("error reading configuration file: %v", err) - os.Exit(1) - } - - cfg := config.Config{ - Network: config.Network{ - Config: config.StartupConfig{ - Host: "0.0.0.0", - Port: 8073, - LogLevel: "INFO", - LogFormat: "TEXT", - }}, - } - - if err = yaml.Unmarshal(fileData, &cfg); err != nil { - fmt.Printf("error parsing configuration file: %v\n", err) + if *help { + os.Exit(0) + } os.Exit(1) } client, err := diode.NewClient( - cfg.Network.Config.Target, + *diodeTarget, AppName, version.GetBuildVersion(), - diode.WithAPIKey(cfg.Network.Config.APIKey), + diode.WithAPIKey(resolveEnv(*diodeAPIKey)), ) if err != nil { fmt.Printf("error creating diode client: %v\n", err) @@ -68,11 +70,10 @@ func main() { } ctx := context.Background() - logger := config.NewLogger(cfg.Network.Config.LogLevel, cfg.Network.Config.LogFormat) + logger := config.NewLogger(*logLevel, *logFormat) policyManager := policy.NewManager(ctx, logger, client) - server := server.Server{} - server.Configure(logger, policyManager, version.GetBuildVersion(), cfg.Network.Config) + server := server.NewServer(*host, *port, logger, policyManager, version.GetBuildVersion()) // handle signals done := make(chan bool, 1) diff --git a/network-discovery/config/config.go b/network-discovery/config/config.go index f69a664..c1854ef 100644 --- a/network-discovery/config/config.go +++ b/network-discovery/config/config.go @@ -12,13 +12,13 @@ type Status struct { // Scope represents the scope of a policy type Scope struct { Targets []string `yaml:"targets"` - Timeout int `yaml:"timeout"` } // PolicyConfig represents the configuration of a policy type PolicyConfig struct { Schedule *string `yaml:"schedule"` Defaults map[string]string `yaml:"defaults"` + Timeout int `yaml:"timeout"` } // Policy represents a network-discovery policy @@ -27,23 +27,7 @@ type Policy struct { Scope Scope `yaml:"scope"` } -// StartupConfig represents the configuration of the network-discovery service -type StartupConfig struct { - Target string `yaml:"target"` - APIKey string `yaml:"api_key"` - Host string `yaml:"host"` - Port int32 `yaml:"port"` - LogLevel string `yaml:"log_level"` - LogFormat string `yaml:"log_format"` -} - -// Network represents the network-discovery configuration -type Network struct { - Config StartupConfig `yaml:"config"` +// Policies represents a collection of network-discovery policies +type Policies struct { Policies map[string]Policy `mapstructure:"policies"` } - -// Config represents the configuration of the network-discovery service -type Config struct { - Network Network `mapstructure:"network"` -} diff --git a/network-discovery/go.mod b/network-discovery/go.mod index 53971ef..4347b80 100644 --- a/network-discovery/go.mod +++ b/network-discovery/go.mod @@ -40,12 +40,12 @@ require ( github.com/twitchyliquid64/golang-asm v0.15.1 // indirect github.com/ugorji/go/codec v1.2.12 // indirect golang.org/x/arch v0.8.0 // indirect - golang.org/x/crypto v0.26.0 // indirect + golang.org/x/crypto v0.31.0 // indirect golang.org/x/exp v0.0.0-20240613232115-7f521ea00fb8 // indirect golang.org/x/net v0.28.0 // indirect - golang.org/x/sync v0.8.0 // indirect - golang.org/x/sys v0.23.0 // indirect - golang.org/x/text v0.17.0 // indirect + golang.org/x/sync v0.10.0 // indirect + golang.org/x/sys v0.28.0 // indirect + golang.org/x/text v0.21.0 // indirect google.golang.org/genproto/googleapis/rpc v0.0.0-20240227224415-6ceb2ff114de // indirect google.golang.org/grpc v1.63.2 // indirect google.golang.org/protobuf v1.34.1 // indirect diff --git a/network-discovery/go.sum b/network-discovery/go.sum index cd03546..5ba1fa7 100644 --- a/network-discovery/go.sum +++ b/network-discovery/go.sum @@ -92,20 +92,20 @@ go.uber.org/goleak v1.3.0/go.mod h1:CoHD4mav9JJNrW/WLlf7HGZPjdw8EucARQHekz1X6bE= golang.org/x/arch v0.0.0-20210923205945-b76863e36670/go.mod h1:5om86z9Hs0C8fWVUuoMHwpExlXzs5Tkyp9hOrfG7pp8= golang.org/x/arch v0.8.0 h1:3wRIsP3pM4yUptoR96otTUOXI367OS0+c9eeRi9doIc= golang.org/x/arch v0.8.0/go.mod h1:FEVrYAQjsQXMVJ1nsMoVVXPZg6p2JE2mx8psSWTDQys= -golang.org/x/crypto v0.26.0 h1:RrRspgV4mU+YwB4FYnuBoKsUapNIL5cohGAmSH3azsw= -golang.org/x/crypto v0.26.0/go.mod h1:GY7jblb9wI+FOo5y8/S2oY4zWP07AkOJ4+jxCqdqn54= +golang.org/x/crypto v0.31.0 h1:ihbySMvVjLAeSH1IbfcRTkD/iNscyz8rGzjF/E5hV6U= +golang.org/x/crypto v0.31.0/go.mod h1:kDsLvtWBEx7MV9tJOj9bnXsPbxwJQ6csT/x4KIN4Ssk= golang.org/x/exp v0.0.0-20240613232115-7f521ea00fb8 h1:yixxcjnhBmY0nkL253HFVIm0JsFHwrHdT3Yh6szTnfY= golang.org/x/exp v0.0.0-20240613232115-7f521ea00fb8/go.mod h1:jj3sYF3dwk5D+ghuXyeI3r5MFf+NT2An6/9dOA95KSI= golang.org/x/net v0.28.0 h1:a9JDOJc5GMUJ0+UDqmLT86WiEy7iWyIhz8gz8E4e5hE= golang.org/x/net v0.28.0/go.mod h1:yqtgsTWOOnlGLG9GFRrK3++bGOUEkNBoHZc8MEDWPNg= -golang.org/x/sync v0.8.0 h1:3NFvSEYkUoMifnESzZl15y791HH1qU2xm6eCJU5ZPXQ= -golang.org/x/sync v0.8.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= +golang.org/x/sync v0.10.0 h1:3NQrjDixjgGwUOCaF8w2+VYHv0Ve/vGYSbdkTa98gmQ= +golang.org/x/sync v0.10.0/go.mod h1:Czt+wKu1gCyEFDUtn0jG5QVvpJ6rzVqr5aXyt9drQfk= golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg= -golang.org/x/sys v0.23.0 h1:YfKFowiIMvtgl1UERQoTPPToxltDeZfbj4H7dVUCwmM= -golang.org/x/sys v0.23.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= -golang.org/x/text v0.17.0 h1:XtiM5bkSOt+ewxlOE/aE/AKEHibwj/6gvWMl9Rsh0Qc= -golang.org/x/text v0.17.0/go.mod h1:BuEKDfySbSR4drPmRPG/7iBdf8hvFMuRexcpahXilzY= +golang.org/x/sys v0.28.0 h1:Fksou7UEQUWlKvIdsqzJmUmCX3cZuD2+P3XyyzwMhlA= +golang.org/x/sys v0.28.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA= +golang.org/x/text v0.21.0 h1:zyQAAkrwaneQ066sspRyJaG9VNi/YJ1NfzcGB3hZ/qo= +golang.org/x/text v0.21.0/go.mod h1:4IBbMaMmOPCJ8SecivzSH54+73PCFmPWxNTLm+vZkEQ= google.golang.org/genproto/googleapis/rpc v0.0.0-20240227224415-6ceb2ff114de h1:cZGRis4/ot9uVm639a+rHCUaG0JJHEsdyzSQTMX+suY= google.golang.org/genproto/googleapis/rpc v0.0.0-20240227224415-6ceb2ff114de/go.mod h1:H4O17MA/PE9BsGx3w+a+W2VOLLD1Qf7oJneAoU6WktY= google.golang.org/grpc v1.63.2 h1:MUeiw1B2maTVZthpU5xvASfTh3LDbxHd6IJ6QQVU+xM= diff --git a/network-discovery/policy/manager.go b/network-discovery/policy/manager.go index 3bc52b9..60568cc 100644 --- a/network-discovery/policy/manager.go +++ b/network-discovery/policy/manager.go @@ -32,16 +32,16 @@ func NewManager(ctx context.Context, logger *slog.Logger, client diode.Client) * // ParsePolicies parses the policies from the request func (m *Manager) ParsePolicies(data []byte) (map[string]config.Policy, error) { - var payload config.Config + var payload config.Policies if err := yaml.Unmarshal(data, &payload); err != nil { return nil, err } - if payload.Network.Policies == nil { + if len(payload.Policies) == 0 { return nil, errors.New("no policies found in the request") } - return payload.Network.Policies, nil + return payload.Policies, nil } // HasPolicy checks if the policy exists diff --git a/network-discovery/policy/manager_test.go b/network-discovery/policy/manager_test.go index 2ddeea7..2da326f 100644 --- a/network-discovery/policy/manager_test.go +++ b/network-discovery/policy/manager_test.go @@ -38,15 +38,14 @@ func TestManagerParsePolicies(t *testing.T) { t.Run("Valid Policies", func(t *testing.T) { yamlData := []byte(` - network: - policies: - policy1: - config: - defaults: - site: New York NY - scope: - targets: - - 192.168.1.1/24 + policies: + policy1: + config: + defaults: + site: New York NY + scope: + targets: + - 192.168.1.1/24 `) policies, err := manager.ParsePolicies(yamlData) @@ -67,19 +66,18 @@ func TestManagerPolicyLifecycle(t *testing.T) { logger := slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelDebug, AddSource: false})) manager := policy.NewManager(context.Background(), logger, nil) yamlData := []byte(` - network: - policies: - policy1: - scope: - targets: - - 192.168.1.1/24 - policy2: - scope: - targets: - - 192.168.2.1/24 - policy3: - scope: - targets: [] + policies: + policy1: + scope: + targets: + - 192.168.1.1/24 + policy2: + scope: + targets: + - 192.168.2.1/24 + policy3: + scope: + targets: [] `) policies, err := manager.ParsePolicies(yamlData) diff --git a/network-discovery/policy/runner.go b/network-discovery/policy/runner.go index adff50b..0616a6b 100644 --- a/network-discovery/policy/runner.go +++ b/network-discovery/policy/runner.go @@ -17,17 +17,21 @@ import ( type contextKey string // Define the policy key -const policyKey contextKey = "policy" +const ( + policyKey contextKey = "policy" + defaultTimeout = 2 * time.Minute +) // Runner represents the policy runner type Runner struct { - scanner *nmap.Scanner scheduler gocron.Scheduler ctx context.Context - cancel context.CancelFunc task gocron.Task client diode.Client logger *slog.Logger + timeout time.Duration + scope config.Scope + config config.PolicyConfig } // NewRunner returns a new policy runner @@ -53,35 +57,38 @@ func NewRunner(ctx context.Context, logger *slog.Logger, name string, policy con if err != nil { return nil, err } - timeout := time.Duration(policy.Scope.Timeout) * time.Minute - if timeout == 0 { - timeout = 2 * time.Minute + runner.timeout = time.Duration(policy.Config.Timeout) * time.Minute + if runner.timeout == 0 { + runner.timeout = defaultTimeout } runner.ctx = context.WithValue(ctx, policyKey, name) - runner.ctx, runner.cancel = context.WithTimeout(runner.ctx, timeout) - n, err := nmap.NewScanner( - runner.ctx, - nmap.WithTargets(policy.Scope.Targets...), + runner.scope = policy.Scope + runner.config = policy.Config + return runner, nil +} + +// run runs the policy +func (r *Runner) run() { + ctx, cancel := context.WithTimeout(r.ctx, r.timeout) + defer cancel() + scanner, err := nmap.NewScanner( + ctx, + nmap.WithTargets(r.scope.Targets...), nmap.WithPingScan(), nmap.WithNonInteractive(), ) if err != nil { - return nil, err + r.logger.Error("error creating scanner", slog.Any("error", err), slog.Any("policy", r.ctx.Value(policyKey))) + return } - runner.scanner = n - - return runner, nil -} - -// run runs the policy -func (r *Runner) run() error { - defer r.cancel() - result, warnings, err := r.scanner.Run() + r.logger.Info("running scanner", slog.Any("targets", r.scope.Targets), slog.Any("policy", r.ctx.Value(policyKey))) + result, warnings, err := scanner.Run() if len(*warnings) > 0 { r.logger.Warn("run finished with warnings", slog.String("warnings", fmt.Sprintf("%v", *warnings))) } if err != nil { - return err + r.logger.Error("error running scanner", slog.Any("error", err), slog.Any("policy", r.ctx.Value(policyKey))) + return } entities := make([]diode.Entity, 0, len(result.Hosts)) @@ -90,6 +97,12 @@ func (r *Runner) run() error { 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["comments"] != "" { + ip.Comments = diode.String(r.config.Defaults["comments"]) + } entities = append(entities, ip) } @@ -101,8 +114,6 @@ func (r *Runner) run() error { } else { r.logger.Info("entities ingested successfully", slog.Any("policy", r.ctx.Value(policyKey))) } - - return nil } // Start starts the policy runner diff --git a/network-discovery/policy/runner_test.go b/network-discovery/policy/runner_test.go index 7e70736..500d427 100644 --- a/network-discovery/policy/runner_test.go +++ b/network-discovery/policy/runner_test.go @@ -79,6 +79,13 @@ func TestRunnerRun(t *testing.T) { logger := slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{Level: slog.LevelDebug, AddSource: false})) mockClient := new(MockClient) policyConfig := config.Policy{ + Config: config.PolicyConfig{ + Schedule: nil, + Defaults: map[string]string{ + "description": "Test", + "comments": "This is a test", + }, + }, Scope: config.Scope{ Targets: []string{"localhost"}, }, @@ -103,7 +110,7 @@ func TestRunnerRun(t *testing.T) { select { case <-ingestCalled: // Ingest was called, proceed - case <-time.After(5 * time.Second): + case <-time.After(10 * time.Second): t.Fatal("Timeout: Ingest was not called") } diff --git a/network-discovery/server/server.go b/network-discovery/server/server.go index f73a600..1fd8d97 100644 --- a/network-discovery/server/server.go +++ b/network-discovery/server/server.go @@ -31,30 +31,37 @@ type Server struct { manager *policy.Manager stat config.Status logger *slog.Logger - config config.StartupConfig + host string + port int } func init() { gin.SetMode(gin.ReleaseMode) } -// Configure configures the network-discovery server -func (s *Server) Configure(logger *slog.Logger, manager *policy.Manager, version string, config config.StartupConfig) { - s.stat.Version = version - s.stat.StartTime = time.Now() - s.manager = manager - s.logger = logger - s.config = config - - s.router = gin.New() +// NewServer returns a new network-discovery server +func NewServer(host string, port int, logger *slog.Logger, manager *policy.Manager, version string) *Server { + server := &Server{ + router: gin.New(), + manager: manager, + stat: config.Status{ + Version: version, + StartTime: time.Now(), + }, + logger: logger, + host: host, + port: port, + } - v1 := s.router.Group("/api/v1") + v1 := server.router.Group("/api/v1") { - v1.GET("/status", s.getStatus) - v1.GET("/capabilities", s.getCapabilities) - v1.POST("/policies", s.createPolicy) - v1.DELETE("/policies/:policy", s.deletePolicy) + v1.GET("/status", server.getStatus) + v1.GET("/capabilities", server.getCapabilities) + v1.POST("/policies", server.createPolicy) + v1.DELETE("/policies/:policy", server.deletePolicy) } + + return server } // Router returns the router @@ -65,7 +72,7 @@ func (s *Server) Router() *gin.Engine { // Start starts the network-discovery server func (s *Server) Start() { go func() { - serv := fmt.Sprintf("%s:%d", s.config.Host, s.config.Port) + serv := fmt.Sprintf("%s:%d", s.host, s.port) s.logger.Info("starting network-discovery server at: " + serv) if err := s.router.Run(serv); err != nil { s.logger.Error("shutting down the server", "error", err) diff --git a/network-discovery/server/server_test.go b/network-discovery/server/server_test.go index 3c0caea..54b216c 100644 --- a/network-discovery/server/server_test.go +++ b/network-discovery/server/server_test.go @@ -15,7 +15,6 @@ import ( "github.com/stretchr/testify/assert" "github.com/stretchr/testify/mock" - "github.com/netboxlabs/orb-discovery/network-discovery/config" "github.com/netboxlabs/orb-discovery/network-discovery/policy" "github.com/netboxlabs/orb-discovery/network-discovery/server" ) @@ -40,15 +39,7 @@ func TestServerConfigureAndStart(t *testing.T) { client := new(MockClient) policyManager := policy.NewManager(ctx, logger, client) - srv := &server.Server{} - - config := config.StartupConfig{ - Host: "localhost", - Port: 8080, - } - version := "1.0.0" - - srv.Configure(logger, policyManager, version, config) + srv := server.NewServer("localhost", 8080, logger, policyManager, "1.0.0") srv.Start() // Check /status endpoint @@ -72,8 +63,7 @@ func TestServerGetCapabilities(t *testing.T) { client := new(MockClient) policyManager := policy.NewManager(ctx, logger, client) - srv := &server.Server{} - srv.Configure(logger, policyManager, "1.0.0", config.StartupConfig{}) + srv := server.NewServer("localhost", 8073, logger, policyManager, "1.0.0") // Check /capabilities endpoint w := httptest.NewRecorder() @@ -93,19 +83,17 @@ func TestServerCreateDeletePolicy(t *testing.T) { client := new(MockClient) policyManager := policy.NewManager(ctx, logger, client) - srv := &server.Server{} - srv.Configure(logger, policyManager, "1.0.0", config.StartupConfig{}) + srv := server.NewServer("localhost", 8073, logger, policyManager, "1.0.0") body := []byte(` - network: - policies: - test-policy: - config: - defaults: - site: New York NY - scope: - targets: - - 192.168.31.1/24 + policies: + test-policy: + config: + defaults: + site: New York NY + scope: + targets: + - 192.168.31.1/24 `) w := httptest.NewRecorder() @@ -120,16 +108,15 @@ func TestServerCreateDeletePolicy(t *testing.T) { // Try to create the same policy again body = []byte(` - network: - policies: - test-pol: - scope: - targets: - - 192.168.31.1/24 - test-policy: - scope: - targets: - - 192.168.31.1/24 + policies: + test-pol: + scope: + targets: + - 192.168.31.1/24 + test-policy: + scope: + targets: + - 192.168.31.1/24 `) w = httptest.NewRecorder() request, _ = http.NewRequest(http.MethodPost, "/api/v1/policies", bytes.NewReader(body)) @@ -183,8 +170,7 @@ func TestServerCreateInvalidPolicy(t *testing.T) { desc: "no policies found", contentType: "application/x-yaml", body: []byte(` - network: - config: {} + policies: {} `), returnCode: http.StatusBadRequest, returnMessage: `no policies found in the request`, @@ -193,18 +179,17 @@ func TestServerCreateInvalidPolicy(t *testing.T) { desc: "no targets found", contentType: "application/x-yaml", body: []byte(` - network: - policies: - test-policy: - scope: - targets: - - 192.168.31.1/24 - test-policy-invalid: - config: - defaults: - site: New York NY - scope: - ports: [80, 443] + policies: + test-policy: + scope: + targets: + - 192.168.31.1/24 + test-policy-invalid: + config: + defaults: + site: New York NY + scope: + ports: [80, 443] `), returnCode: http.StatusBadRequest, returnMessage: `test-policy-invalid : no targets found in the policy`, @@ -217,8 +202,7 @@ func TestServerCreateInvalidPolicy(t *testing.T) { client := new(MockClient) policyManager := policy.NewManager(ctx, logger, client) - srv := &server.Server{} - srv.Configure(logger, policyManager, "1.0.0", config.StartupConfig{}) + srv := server.NewServer("localhost", 8073, logger, policyManager, "1.0.0") // Create invalid policy request w := httptest.NewRecorder()