Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

fix(anta.tests): Cleaning up STUN tests module #934

Merged
merged 16 commits into from
Dec 5, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
49 changes: 48 additions & 1 deletion anta/decorators.py
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,8 @@
F = TypeVar("F", bound=Callable[..., Any])


def deprecated_test(new_tests: list[str] | None = None) -> Callable[[F], F]:
# TODO: Remove this decorator in ANTA v2.0.0 in favor of deprecated_test_class
def deprecated_test(new_tests: list[str] | None = None) -> Callable[[F], F]: # pragma: no cover
"""Return a decorator to log a message of WARNING severity when a test is deprecated.

Parameters
Expand Down Expand Up @@ -62,6 +63,52 @@ async def wrapper(*args: Any, **kwargs: Any) -> Any:
return decorator


def deprecated_test_class(new_tests: list[str] | None = None) -> Callable[[type[AntaTest]], type[AntaTest]]:
"""Return a decorator to log a message of WARNING severity when a test is deprecated.

Parameters
----------
new_tests
A list of new test classes that should replace the deprecated test.

Returns
-------
Callable[[type], type]
A decorator that can be used to wrap test functions.

"""

def decorator(cls: type[AntaTest]) -> type[AntaTest]:
"""Actual decorator that logs the message.

Parameters
----------
cls
The cls to be decorated.

Returns
-------
cls
The decorated cls.
"""
orig_init = cls.__init__

def new_init(*args: Any, **kwargs: Any) -> None:
"""Overload __init__ to generate a warning message for deprecation."""
if new_tests:
new_test_names = ", ".join(new_tests)
logger.warning("%s test is deprecated. Consider using the following new tests: %s.", cls.name, new_test_names)
else:
logger.warning("%s test is deprecated.", cls.name)
orig_init(*args, **kwargs)

# NOTE: we are ignoring mypy warning as we want to assign to a method here
cls.__init__ = new_init # type: ignore[method-assign]
return cls

return decorator


def skip_on_platforms(platforms: list[str]) -> Callable[[F], F]:
"""Return a decorator to skip a test based on the device's hardware model.

Expand Down
35 changes: 35 additions & 0 deletions anta/input_models/stun.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
# Copyright (c) 2023-2024 Arista Networks, Inc.
# Use of this source code is governed by the Apache License 2.0
# that can be found in the LICENSE file.
"""Module containing input models for services tests."""

from __future__ import annotations

from ipaddress import IPv4Address

from pydantic import BaseModel, ConfigDict

from anta.custom_types import Port


class StunClientTranslation(BaseModel):
"""STUN (Session Traversal Utilities for NAT) model represents the configuration of an IPv4-based client translations."""

model_config = ConfigDict(extra="forbid")
source_address: IPv4Address
"""The IPv4 address of the STUN client"""
source_port: Port = 4500
"""The port number used by the STUN client for communication. Defaults to 4500."""
public_address: IPv4Address | None = None
"""The public-facing IPv4 address of the STUN client, discovered via the STUN server."""
public_port: Port | None = None
"""The public-facing port number of the STUN client, discovered via the STUN server."""

def __str__(self) -> str:
"""Return a human-readable string representation of the StunClientTranslation for reporting.

Examples
--------
Client 10.0.0.1 Port: 4500
"""
return f"Client {self.source_address} Port: {self.source_port}"
114 changes: 60 additions & 54 deletions anta/tests/stun.py
Original file line number Diff line number Diff line change
Expand Up @@ -7,29 +7,36 @@
# mypy: disable-error-code=attr-defined
from __future__ import annotations

from ipaddress import IPv4Address
from typing import ClassVar

from pydantic import BaseModel

from anta.custom_types import Port
from anta.decorators import deprecated_test_class
from anta.input_models.stun import StunClientTranslation
from anta.models import AntaCommand, AntaTemplate, AntaTest
from anta.tools import get_failed_logs, get_value
from anta.tools import get_value


class VerifyStunClientTranslation(AntaTest):
gmuloc marked this conversation as resolved.
Show resolved Hide resolved
"""Verifies the translation for a source address on a STUN client.

class VerifyStunClient(AntaTest):
"""Verifies STUN client settings, including local IP/port and optionally public IP/port.
This test performs the following checks for each specified address family:

1. Validates that there is a translation for the source address on the STUN client.
2. If public IP and port details are provided, validates their correctness against the configuration.

Expected Results
----------------
* Success: The test will pass if the STUN client is correctly configured with the specified IPv4 source address/port and public address/port.
* Failure: The test will fail if the STUN client is not configured or if the IPv4 source address, public address, or port details are incorrect.
* Success: If all of the following conditions are met:
- The test will pass if the source address translation is present.
- If public IP and port details are provided, they must also match the translation information.
* Failure: If any of the following occur:
- There is no translation for the source address on the STUN client.
- The public IP or port details, if specified, are incorrect.

Examples
--------
```yaml
anta.tests.stun:
- VerifyStunClient:
- VerifyStunClientTranslation:
stun_clients:
- source_address: 172.18.3.2
public_address: 172.18.3.21
Expand All @@ -43,73 +50,72 @@ class VerifyStunClient(AntaTest):
"""

categories: ClassVar[list[str]] = ["stun"]
commands: ClassVar[list[AntaCommand | AntaTemplate]] = [AntaTemplate(template="show stun client translations {source_address} {source_port}")]
commands: ClassVar[list[AntaCommand | AntaTemplate]] = [AntaTemplate(template="show stun client translations {source_address} {source_port}", revision=1)]

class Input(AntaTest.Input):
"""Input model for the VerifyStunClient test."""
"""Input model for the VerifyStunClientTranslation test."""

stun_clients: list[ClientAddress]

class ClientAddress(BaseModel):
"""Source and public address/port details of STUN client."""

source_address: IPv4Address
"""IPv4 source address of STUN client."""
source_port: Port = 4500
"""Source port number for STUN client."""
public_address: IPv4Address | None = None
"""Optional IPv4 public address of STUN client."""
public_port: Port | None = None
"""Optional public port number for STUN client."""
stun_clients: list[StunClientTranslation]
"""List of STUN clients."""
StunClientTranslation: ClassVar[type[StunClientTranslation]] = StunClientTranslation

def render(self, template: AntaTemplate) -> list[AntaCommand]:
"""Render the template for each STUN translation."""
return [template.render(source_address=client.source_address, source_port=client.source_port) for client in self.inputs.stun_clients]

@AntaTest.anta_test
def test(self) -> None:
"""Main test function for VerifyStunClient."""
"""Main test function for VerifyStunClientTranslation."""
self.result.is_success()

# Iterate over each command output and corresponding client input
for command, client_input in zip(self.instance_commands, self.inputs.stun_clients):
bindings = command.json_output["bindings"]
source_address = str(command.params.source_address)
source_port = command.params.source_port
input_public_address = client_input.public_address
input_public_port = client_input.public_port

# If no bindings are found for the STUN client, mark the test as a failure and continue with the next client
if not bindings:
self.result.is_failure(f"STUN client transaction for source `{source_address}:{source_port}` is not found.")
self.result.is_failure(f"{client_input} - STUN client translation not found.")
continue

# Extract the public address and port from the client input
public_address = client_input.public_address
public_port = client_input.public_port

# Extract the transaction ID from the bindings
transaction_id = next(iter(bindings.keys()))

# Prepare the actual and expected STUN data for comparison
actual_stun_data = {
"source ip": get_value(bindings, f"{transaction_id}.sourceAddress.ip"),
"source port": get_value(bindings, f"{transaction_id}.sourceAddress.port"),
}
expected_stun_data = {"source ip": source_address, "source port": source_port}

# If public address is provided, add it to the actual and expected STUN data
if public_address is not None:
actual_stun_data["public ip"] = get_value(bindings, f"{transaction_id}.publicAddress.ip")
expected_stun_data["public ip"] = str(public_address)

# If public port is provided, add it to the actual and expected STUN data
if public_port is not None:
actual_stun_data["public port"] = get_value(bindings, f"{transaction_id}.publicAddress.port")
expected_stun_data["public port"] = public_port

# If the actual STUN data does not match the expected STUN data, mark the test as failure
if actual_stun_data != expected_stun_data:
failed_log = get_failed_logs(expected_stun_data, actual_stun_data)
self.result.is_failure(f"For STUN source `{source_address}:{source_port}`:{failed_log}")
# Verifying the public address if provided
if input_public_address and str(input_public_address) != (actual_public_address := get_value(bindings, f"{transaction_id}.publicAddress.ip")):
self.result.is_failure(f"{client_input} - Incorrect public-facing address - Expected: {input_public_address} Actual: {actual_public_address}")

# Verifying the public port if provided
if input_public_port and input_public_port != (actual_public_port := get_value(bindings, f"{transaction_id}.publicAddress.port")):
self.result.is_failure(f"{client_input} - Incorrect public-facing port - Expected: {input_public_port} Actual: {actual_public_port}")


@deprecated_test_class(new_tests=["VerifyStunClientTranslation"])
class VerifyStunClient(VerifyStunClientTranslation):
"""(Deprecated) Verifies the translation for a source address on a STUN client.

Alias for the VerifyStunClientTranslation test to maintain backward compatibility.
When initialized, it will emit a deprecation warning and call the VerifyStunClientTranslation test.

TODO: Remove this class in ANTA v2.0.0.

Examples
--------
```yaml
anta.tests.stun:
- VerifyStunClient:
stun_clients:
- source_address: 172.18.3.2
public_address: 172.18.3.21
source_port: 4500
public_port: 6006
```
"""

# required to redefine name an description to overwrite parent class.
name = "VerifyStunClient"
description = "(Deprecated) Verifies the translation for a source address on a STUN client."


class VerifyStunServer(AntaTest):
Expand Down
17 changes: 17 additions & 0 deletions docs/api/tests.stun.md
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,8 @@ anta_title: ANTA catalog for STUN tests
~ that can be found in the LICENSE file.
-->

# Tests

::: anta.tests.stun
options:
show_root_heading: false
Expand All @@ -18,3 +20,18 @@ anta_title: ANTA catalog for STUN tests
filters:
- "!test"
- "!render"

# Input models

::: anta.input_models.stun

options:
show_root_heading: false
show_root_toc_entry: false
show_bases: false
merge_init_into_class: false
anta_hide_test_module_description: true
show_labels: true
filters:
- "!^__init__"
- "!^__str__"
9 changes: 8 additions & 1 deletion examples/tests.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -745,7 +745,14 @@ anta.tests.stp:
threshold: 10
anta.tests.stun:
- VerifyStunClient:
# Verifies STUN client settings, including local IP/port and optionally public IP/port.
# (Deprecated) Verifies the translation for a source address on a STUN client.
stun_clients:
- source_address: 172.18.3.2
public_address: 172.18.3.21
source_port: 4500
public_port: 6006
- VerifyStunClientTranslation:
# Verifies the translation for a source address on a STUN client.
stun_clients:
- source_address: 172.18.3.2
public_address: 172.18.3.21
Expand Down
Loading
Loading