add tests and doc

Signed-off-by: Benny Zlotnik <[email protected]>

docs/source/api-reference/drivers/snmp.md

@@ -0,0 +1,82 @@
+# SNMP
+
+**driver**: `jumpstarter_driver_snmp.driver.SNMPServer`
+
+A driver for controlling power via SNMP-enabled PDUs (Power Distribution Units).
+
+## Driver configuration
+```yaml
+export:
+  power:
+    type: "jumpstarter_driver_snmp.driver.SNMPServer"
+    config:
+      host: "pdu.mgmt.com"
+      user: "labuser"
+      plug: 32
+      port: 161
+      oid: "1.3.6.1.4.1.13742.6.4.1.2.1.2.1"
+      auth_protocol: "NONE"
+      auth_key: null
+      priv_protocol: "NONE"
+      priv_key: null
+      timeout: 5.0
+```
+
+### Config parameters
+
+| Parameter | Description | Type | Required | Default |
+|-----------|-------------|------|----------|---------|
+| host | Hostname or IP address of the SNMP-enabled PDU | str | yes | |
+| user | SNMP v3 username | str | yes | |
+| plug | PDU outlet number to control | int | yes | |
+| port | SNMP port number | int | no | 161 |
+| oid | Base OID for power control | str | no | "1.3.6.1.4.1.13742.6.4.1.2.1.2.1" |
+| auth_protocol | Authentication protocol ("NONE", "MD5", "SHA") | str | no | "NONE" |
+| auth_key | Authentication key when auth_protocol is not "NONE" | str | no | null |
+| priv_protocol | Privacy protocol ("NONE", "DES", "AES") | str | no | "NONE" |
+| priv_key | Privacy key when priv_protocol is not "NONE" | str | no | null |
+| timeout | SNMP timeout in seconds | float | no | 5.0 |
+
+## SNMPServerClient API
+
+### Methods
+
+#### on()
+Turn power on for the configured outlet.
+
+Returns:
+- str: Confirmation message
+
+#### off()
+Turn power off for the configured outlet.
+
+Returns:
+- str: Confirmation message
+
+#### cycle(quiescent_period: int = 2)
+Power cycle the device with a configurable wait period between off and on states.
+
+Parameters:
+- quiescent_period: Time to wait in seconds between power off and power on
+
+Returns:
+- str: Confirmation message
+
+## Examples
+
+Power cycling a device:
+```python
+snmp_client.cycle(quiescent_period=3)
+```
+
+Basic power control:
+```python
+snmp_client.off()
+snmp_client.on()
+```
+
+Using the CLI:
+```bash
+j power on
+j power off
+j power cycle --wait 3

packages/jumpstarter-driver-snmp/examples/exporter.yaml

@@ -0,0 +1,18 @@
+apiVersion: jumpstarter.dev/v1alpha1
+kind: ExporterConfig
+endpoint: grpc.jumpstarter.192.168.0.203.nip.io:8082
+metadata:
+  namespace: default
+  name: demo
+tls:
+  ca: ''
+  insecure: true
+token: <token>
+export:
+  power:
+    type: "jumpstarter_driver_snmp.driver.SNMPServer"
+    config:
+      host: "pdu.mgmt.com"
+      user: "labuser"
+      plug: 32
+      oid: "1.3.6.1.4.1.13742.6.4.1.2.1.2.1"

packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/driver.py

@@ -1,7 +1,7 @@
 import asyncio
 import socket
 from dataclasses import dataclass, field
-from enum import IntEnum
+from enum import Enum, IntEnum
 
 from pysnmp.carrier.asyncio.dgram import udp
 from pysnmp.entity import config, engine
@@ -11,6 +11,16 @@
 from jumpstarter.driver import Driver, export
 
 
+class AuthProtocol(str, Enum):
+    NONE = "NONE"
+    MD5 = "MD5"
+    SHA = "SHA"
+
+class PrivProtocol(str, Enum):
+    NONE = "NONE"
+    DES = "DES"
+    AES = "AES"
+
 class PowerState(IntEnum):
     OFF = 0
     ON = 1
@@ -25,14 +35,13 @@ class SNMPServer(Driver):
     host: str = field()
     user: str = field()
     port: int = field(default=161)
-    quiescent_period: int = field(default=5)
-    timeout: int = 3
     plug: int = field()
     oid: str = field(default="1.3.6.1.4.1.13742.6.4.1.2.1.2.1")
-    auth_protocol: str = field(default=None)  # 'MD5' or 'SHA'
-    auth_key: str = field(default=None)
-    priv_protocol: str = field(default=None)  # 'DES' or 'AES'
-    priv_key: str = field(default=None)
+    auth_protocol: AuthProtocol = field(default=AuthProtocol.NONE)
+    auth_key: str | None = field(default=None)
+    priv_protocol: PrivProtocol = field(default=PrivProtocol.NONE)
+    priv_key: str | None = field(default=None)
+    timeout: float = field(default=5.0)
 
     def __post_init__(self):
         if hasattr(super(), "__post_init__"):
@@ -47,45 +56,54 @@ def __post_init__(self):
         self.full_oid = tuple(int(x) for x in self.oid.split('.')) + (self.plug,)
 
     def _setup_snmp(self):
-        try:
-            asyncio.get_running_loop()
-        except RuntimeError:
-            loop = asyncio.new_event_loop()
-            asyncio.set_event_loop(loop)
-
         snmp_engine = engine.SnmpEngine()
 
-        if self.auth_protocol and self.auth_key:
-            if self.priv_protocol and self.priv_key:
-                security_level = 'authPriv'
-                auth_protocol = getattr(config, f'usmHMAC{self.auth_protocol}AuthProtocol')
-                priv_protocol = getattr(config, f'usmPriv{self.priv_protocol}Protocol')
-
-                config.add_v3_user(
-                    snmp_engine,
-                    self.user,
-                    auth_protocol,
-                    self.auth_key,
-                    priv_protocol,
-                    self.priv_key
-                )
-            else:
-                security_level = 'authNoPriv'
-                auth_protocol = getattr(config, f'usmHMAC{self.auth_protocol}AuthProtocol')
-
-                config.add_v3_user(
-                    snmp_engine,
-                    self.user,
-                    auth_protocol,
-                    self.auth_key
-                )
+        AUTH_PROTOCOLS = {
+            AuthProtocol.NONE: config.USM_AUTH_NONE,
+            AuthProtocol.MD5: config.USM_AUTH_HMAC96_MD5,
+            AuthProtocol.SHA: config.USM_AUTH_HMAC96_SHA,
+        }
+
+        PRIV_PROTOCOLS = {
+            PrivProtocol.NONE: config.USM_PRIV_NONE,
+            PrivProtocol.DES: config. USM_PRIV_CBC56_DES,
+            PrivProtocol.AES: config.USM_PRIV_CFB128_AES,
+        }
+
+        auth_protocol = AUTH_PROTOCOLS[self.auth_protocol]
+        priv_protocol = PRIV_PROTOCOLS[self.priv_protocol]
+
+        if self.auth_protocol == AuthProtocol.NONE:
+            security_level = "noAuthNoPriv"
+        elif self.priv_protocol == PrivProtocol.NONE:
+            security_level = "authNoPriv"
         else:
-            security_level = 'noAuthNoPriv'
+            security_level = "authPriv"
+
+        if security_level == "noAuthNoPriv":
+            config.add_v3_user(
+                snmp_engine,
+                self.user
+            )
+        elif security_level == "authNoPriv":
+            if not self.auth_key:
+                raise SNMPError("Authentication key required when auth_protocol is specified")
             config.add_v3_user(
                 snmp_engine,
                 self.user,
-                config.USM_AUTH_NONE,
-                None
+                auth_protocol,
+                self.auth_key
+            )
+        else:
+            if not self.auth_key or not self.priv_key:
+                raise SNMPError("Both auth_key and priv_key required for authenticated privacy")
+            config.add_v3_user(
+                snmp_engine,
+                self.user,
+                auth_protocol,
+                self.auth_key,
+                priv_protocol,
+                self.priv_key
             )
 
         config.add_target_parameters(
@@ -95,18 +113,19 @@ def _setup_snmp(self):
             security_level
         )
 
-        config.add_transport(
+        config.add_target_address(
             snmp_engine,
+            "my-target",
             udp.DOMAIN_NAME,
-            udp.UdpAsyncioTransport().open_client_mode()
+            (self.ip_address, self.port),
+            "my-creds",
+            timeout=int(self.timeout * 100),
         )
 
-        config.add_target_address(
+        config.add_transport(
             snmp_engine,
-            "my-target",
             udp.DOMAIN_NAME,
-            (self.ip_address, self.port),
-            "my-creds"
+            udp.UdpAsyncioTransport().open_client_mode()
         )
 
         return snmp_engine
@@ -138,6 +157,11 @@ def callback(snmpEngine, sendRequestHandle, errorIndication,
 
         try:
             self.logger.info(f"Sending power {state.name} command to {self.host}")
+            try:
+                asyncio.get_running_loop()
+            except RuntimeError:
+                loop = asyncio.new_event_loop()
+                asyncio.set_event_loop(loop)
 
             snmp_engine = self._setup_snmp()