Refactor deduplication of walk oids
@@ -16,7 +16,17 @@ from asyncio import get_event_loop
from asyncio.events import AbstractEventLoop
from collections import OrderedDict
from ipaddress import ip_address
-from typing import Any, AsyncGenerator, Callable, Dict, List, Optional, Tuple
+from typing import (
+ Any,
+ AsyncGenerator,
+ Callable,
+ Dict,
+ Generator,
+ List,
+ Optional,
+ Set,
+ Tuple,
+)
from typing import Type as TType
from typing import TypeVar, cast
@@ -54,8 +64,6 @@ PyType = Any # TODO
TWalkResponse = AsyncGenerator[VarBind, None]
T = TypeVar("T", bound=TType[PyType]) # pylint: disable=invalid-name
-_set = set
-
LOG = logging.getLogger(__name__)
OID = ObjectIdentifier
@@ -72,6 +80,38 @@ class TFetcher(Protocol):
...
+def deduped_varbinds(
+ requested_oids: List[ObjectIdentifier],
+ grouped_oids: Dict[ObjectIdentifier, List[VarBind]],
+ yielded: Set[ObjectIdentifier],
+) -> Generator[VarBind, None, None]:
+ """
+ Generate grouped OIDs by ensuring they are contained in the original
+ request and have no duplicates.
+
+ :param requested_oids: A list of OIDs which were originally requested. If
+ any value from the grouped varbinds are not children of any of these
+ OIDs, the issue is logged and the value is skipped.
+ :param grouped_oids: The OIDs that need to be verified & deduped.
+ :param yielded: A set containing all OIDs that have already been
+ generated by this function. This set will be updated by this function
+ whenever a value is returned to detect duplicates.
+ """
+ for var in sorted(grouped_oids.values()):
+ for varbind in var:
+ containment = [varbind.oid in _ for _ in requested_oids]
+ if not any(containment) or varbind.oid in yielded:
+ LOG.debug(
+ "Unexpected device response: Returned VarBind %s "
+ "was either not contained in the requested tree or "
+ "appeared more than once. Skipping!",
+ varbind,
+ )
+ continue
+ yielded.add(varbind.oid)
+ yield varbind
+
+
class RawClient:
"""
A client to execute SNMP commands on a remote device.
@@ -284,19 +324,8 @@ class RawClient:
len(unfinished_oids),
len(oids),
)
- yielded = _set([])
- for var in sorted(grouped_oids.values()):
- for varbind in var:
- containment = [varbind.oid in _ for _ in requested_oids]
- if not any(containment) or varbind.oid in yielded:
- LOG.debug(
- "Unexpected device response: Returned VarBind %s "
- "was either not contained in the requested tree or "
- "appeared more than once. Skipping!",
- varbind,
- )
- continue
- yielded.add(varbind.oid)
+ yielded: Set[ObjectIdentifier] = set()
+ for varbind in deduped_varbinds(requested_oids, grouped_oids, yielded):
yield varbind
# As long as we have unfinished OIDs, we need to continue the walk for
@@ -332,12 +361,9 @@ class RawClient:
len(unfinished_oids),
len(oids),
)
- for var in sorted(grouped_oids.values()):
- for varbind in var:
- containment = [varbind.oid in _ for _ in requested_oids]
- if not any(containment) or varbind.oid in yielded:
- continue
- yielded.add(varbind.oid)
+ for varbind in deduped_varbinds(
+ requested_oids, grouped_oids, yielded
+ ):
yield varbind
async def multigetnext(
