snmp: initial SNMP driver for power control

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

@@ -10,6 +10,7 @@ Drivers packages from the [drivers](https://github.com/jumpstarter-dev/jumpstart
 can.md
 pyserial.md
 sdwire.md
+snmp.md
 tftp.md
 ustreamer.md
 yepkit.md

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

@@ -0,0 +1,67 @@
+# 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
+
+```{eval-rst}
+.. autoclass:: jumpstarter_driver_snmp.client.SNMPServerClient()
+   :members:
+   :show-inheritance:
+```
+
+## Examples
+
+Power cycling a device:
+```python
+snmp_client.cycle(wait=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-composite/jumpstarter_driver_composite/driver_test.py

@@ -1,3 +1,4 @@
+
 from jumpstarter_driver_power.driver import MockPower
 
 from .driver import Composite

packages/jumpstarter-driver-power/jumpstarter_driver_power/client.py

@@ -1,3 +1,4 @@
+import time
 from collections.abc import Generator
 
 import asyncclick as click
@@ -13,6 +14,15 @@ def on(self) -> None:
     def off(self) -> None:
         self.call("off")
 
+    def cycle(self, wait: int = 2):
+        """Power cycle the device"""
+        self.logger.info("Starting power cycle sequence")
+        self.off()
+        self.logger.info(f"Waiting {wait} seconds...")
+        time.sleep(wait)
+        self.on()
+        self.logger.info("Power cycle sequence complete")
+
     def read(self) -> Generator[PowerReading, None, None]:
         for v in self.streamingcall("read"):
             yield PowerReading.model_validate(v, strict=True)
@@ -33,4 +43,10 @@ def off():
             """Power off"""
             self.off()
 
+        @base.command()
+        @click.option('--wait', '-w', default=2, help='Wait time in seconds between off and on')
+        def cycle(wait):
+            """Power cycle"""
+            click.echo(f"Power cycling with {wait} seconds wait time...")
+            self.cycle(wait)
         return base

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/client.py

@@ -0,0 +1,28 @@
+from dataclasses import dataclass
+
+import asyncclick as click
+from jumpstarter_driver_power.client import PowerClient
+
+
+@dataclass(kw_only=True)
+class SNMPServerClient(PowerClient):
+    """Client interface for SNMP Power Control"""
+
+    def on(self):
+        """Turn power on"""
+        self.call("on")
+
+    def off(self):
+        """Turn power off"""
+        self.call("off")
+
+    def cli(self):
+        @click.group()
+        def snmp():
+            """SNMP power control commands"""
+            pass
+
+        for cmd in super().cli().commands.values():
+            snmp.add_command(cmd)
+
+        return snmp

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

@@ -0,0 +1,209 @@
+import asyncio
+import socket
+from dataclasses import dataclass, field
+from enum import Enum, IntEnum
+
+from pysnmp.carrier.asyncio.dgram import udp
+from pysnmp.entity import config, engine
+from pysnmp.entity.rfc3413 import cmdgen
+from pysnmp.proto import rfc1902
+
+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
+
+class SNMPError(Exception):
+    """Base exception for SNMP errors"""
+    pass
+
+@dataclass(kw_only=True)
+class SNMPServer(Driver):
+    """SNMP Power Control Driver"""
+    host: str = field()
+    user: str = field()
+    port: int = field(default=161)
+    plug: int = field()
+    oid: str = field(default="1.3.6.1.4.1.13742.6.4.1.2.1.2.1")
+    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__"):
+            super().__post_init__()
+
+        try:
+            self.ip_address = socket.gethostbyname(self.host)
+            self.logger.debug(f"Resolved {self.host} to {self.ip_address}")
+        except socket.gaierror as e:
+            raise SNMPError(f"Failed to resolve hostname {self.host}: {e}") from e
+
+        self.full_oid = tuple(int(x) for x in self.oid.split('.')) + (self.plug,)
+
+    def _setup_snmp(self):
+        snmp_engine = engine.SnmpEngine()
+
+        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 = "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,
+                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(
+            snmp_engine,
+            "my-creds",
+            self.user,
+            security_level
+        )
+
+        config.add_target_address(
+            snmp_engine,
+            "my-target",
+            udp.DOMAIN_NAME,
+            (self.ip_address, self.port),
+            "my-creds",
+            timeout=int(self.timeout * 100),
+        )
+
+        config.add_transport(
+            snmp_engine,
+            udp.DOMAIN_NAME,
+            udp.UdpAsyncioTransport().open_client_mode()
+        )
+
+        return snmp_engine
+
+    @classmethod
+    def client(cls) -> str:
+        return "jumpstarter_driver_snmp.client.SNMPServerClient"
+
+    def _snmp_set(self, state: PowerState):
+        result = {"success": False, "error": None}
+
+        def callback(snmpEngine, sendRequestHandle, errorIndication,
+                    errorStatus, errorIndex, varBinds, cbCtx):
+            self.logger.debug(f"Callback {errorIndication} {errorStatus} {errorIndex} {varBinds}")
+            if errorIndication:
+                self.logger.error(f"SNMP error: {errorIndication}")
+                result["error"] = f"SNMP error: {errorIndication}"
+            elif errorStatus:
+                self.logger.error(f"SNMP status: {errorStatus}")
+                result["error"] = (
+                    f"SNMP error: {errorStatus.prettyPrint()} at "
+                    f"{varBinds[int(errorIndex) - 1][0] if errorIndex else '?'}"
+                )
+            else:
+                result["success"] = True
+                for oid, val in varBinds:
+                    self.logger.debug(f"{oid.prettyPrint()} = {val.prettyPrint()}")
+            self.logger.debug(f"SNMP set result: {result}")
+
+        try:
+            self.logger.info(f"Sending power {state.name} command to {self.host}")
+            created_loop = False
+
+            try:
+                asyncio.get_running_loop()
+            except RuntimeError:
+              loop = asyncio.new_event_loop()
+              asyncio.set_event_loop(loop)
+              created_loop = True
+
+            snmp_engine = self._setup_snmp()
+
+            cmdgen.SetCommandGenerator().send_varbinds(
+                snmp_engine,
+                "my-target",
+                None,
+                "",
+                [(self.full_oid, rfc1902.Integer(state.value))],
+                callback,
+            )
+
+            snmp_engine.open_dispatcher(self.timeout)
+            snmp_engine.close_dispatcher()
+
+            if not result["success"]:
+                raise SNMPError(result["error"])
+
+            return f"Power {state.name} command sent successfully"
+
+        except Exception as e:
+            error_msg = f"SNMP set failed: {str(e)}"
+            self.logger.error(error_msg)
+            raise SNMPError(error_msg) from e
+        finally:
+            if created_loop:
+                loop.close()
+
+    @export
+    def on(self):
+        """Turn power on"""
+        return self._snmp_set(PowerState.ON)
+
+    @export
+    def off(self):
+        """Turn power off"""
+        return self._snmp_set(PowerState.OFF)
+
+    def close(self):
+        """No cleanup needed since engines are created per operation"""
+        if hasattr(super(), "close"):
+            super().close()