snmpagent: Implement an asyncio-based SNMP trap infrastructure supporting link up/down, PSU, and fan traps

Signed-off-by: Griffin Gao <griffin@micasnetworks.com>

Author: Griffin-micas

src/ax_interface/agent.py

@@ -2,13 +2,14 @@
 
 from .mib import MIBTable, MIBMeta
 from .socket_io import SocketManager
+from .trap import TrapInfra
 
 # how long to wait before forcibly killing background task(s) during the shutdown procedure.
 BACKGROUND_WAIT_TIMEOUT = 10  # seconds
 
 
 class Agent:
-    def __init__(self, mib_cls, enable_dynamic_frequency, update_frequency, loop):
+    def __init__(self, mib_cls, enable_dynamic_frequency, update_frequency, loop, trap_handlers=None):
         if not type(mib_cls) is MIBMeta:
             raise ValueError("Expected a class with type: {}".format(MIBMeta))
 
@@ -19,8 +20,14 @@ def __init__(self, mib_cls, enable_dynamic_frequency, update_frequency, loop):
         self.oid_updaters_enabled = asyncio.Event()
         self.stopped = asyncio.Event()
 
+        # Initialize our Trap infra
+        self.trap_infra = TrapInfra(loop, trap_handlers)
+
         # Initialize our MIB
-        self.mib_table = MIBTable(mib_cls, enable_dynamic_frequency, update_frequency)
+        if trap_handlers:
+            self.mib_table = MIBTable(mib_cls, enable_dynamic_frequency, update_frequency, self.trap_infra)
+        else:
+            self.mib_table = MIBTable(mib_cls, enable_dynamic_frequency, update_frequency)
 
         # containers
         self.socket_mgr = SocketManager(self.mib_table, self.run_enabled, self.loop)
@@ -54,6 +61,8 @@ async def run_in_event_loop(self):
     async def shutdown(self):
         # allow the agent to quit
         self.run_enabled.clear()
+        # shutdown trap infra
+        await self.trap_infra.shutdown()
         # close the socket
         self.socket_mgr.close()
         # wait for the agent to signal back

src/ax_interface/mib.py

@@ -313,14 +313,16 @@ class MIBTable(dict):
 
     def __init__(self, mib_cls,
                  enable_dynamic_frequency=DEFAULT_ENABLE_DYNAMIC_FREQUENCY,
-                 update_frequency=DEFAULT_UPDATE_FREQUENCY):
+                 update_frequency=DEFAULT_UPDATE_FREQUENCY,
+                 trap_infra_obj=None):
         if type(mib_cls) is not MIBMeta:
             raise ValueError("Supplied object is not a MIB class instance.")
         super().__init__(getattr(mib_cls, MIBMeta.KEYSTORE))
         self.enable_dynamic_frequency = enable_dynamic_frequency
         self.update_frequency = update_frequency
         self.updater_instances = getattr(mib_cls, MIBMeta.UPDATERS)
         self.prefixes = getattr(mib_cls, MIBMeta.PREFIXES)
+        self.trap_infra_obj = trap_infra_obj
 
     @staticmethod
     def _done_background_task_callback(fut):
@@ -340,6 +342,8 @@ def start_background_tasks(self, event):
             fut = asyncio.ensure_future(updater.start())
             fut.add_done_callback(MIBTable._done_background_task_callback)
             tasks.append(fut)
+        if self.trap_infra_obj:
+            asyncio.create_task(self.trap_infra_obj.db_listener())
         return asyncio.gather(*tasks)
 
     def _find_parent_prefix(self, item):

src/ax_interface/pdu_implementations.py

@@ -362,12 +362,26 @@ def __init__(self, *args, **kwargs):
         # header only.
 
 
-# class NotifyPDU(ContextOptionalPDU):
-#     """
-#         https://tools.ietf.org/html/rfc2741#section-6.2.10
-#         """
-#     header_type_ = PduTypes.NOTIFY
-# TODO: Impl
+class NotifyPDU(ContextOptionalPDU):
+    """
+        https://tools.ietf.org/html/rfc2741#section-6.2.10
+    """
+    header_type_ = PduTypes.NOTIFY
+
+    def __init__(self, header=None, context=None, payload=None, varBinds=[]):
+        super().__init__(header=header, context=context, payload=payload)
+
+        self.varBinds = varBinds
+        # Reducing the header length as per RFC 2741
+        # https://tools.ietf.org/html/rfc2741#section-6.1
+        payload_len = len(self.encode()) - constants.AGENTX_HEADER_LENGTH
+        self.header = self.header._replace(payload_length=payload_len)
+
+    def encode(self):
+        ret = super().encode()
+        for value in self.varBinds:
+            ret += value.to_bytes(self.header.endianness)
+        return ret
 
 
 

src/ax_interface/protocol.py

@@ -4,6 +4,7 @@
 from .encodings import ObjectIdentifier
 from .pdu import PDUHeader, PDUStream
 from .pdu_implementations import RegisterPDU, ResponsePDU, OpenPDU
+from .trap import TrapInfra
 
 
 class AgentX(asyncio.Protocol):
@@ -49,6 +50,7 @@ def opening_handshake(self):
 
     def register_subtrees(self, pdu):
         self.session_id = pdu.header.session_id
+        TrapInfra.protocol_obj = self
         logger.info("AgentX session starting with ID: {}".format(self.session_id))
 
         for idx, subtree in enumerate(self.mib_table.prefixes):