From 58e2e18467c54ee6714d5d47bc36ee916cad1972 Mon Sep 17 00:00:00 2001 From: Benny Zlotnik Date: Wed, 5 Feb 2025 12:18:37 +0200 Subject: [PATCH 1/8] snmp: initial SNMP driver for power control Signed-off-by: Benny Zlotnik --- packages/jumpstarter-driver-snmp/README.md | 0 .../jumpstarter_driver_snmp/__init__.py | 0 .../jumpstarter_driver_snmp/client.py | 48 +++++ .../jumpstarter_driver_snmp/driver.py | 166 ++++++++++++++++++ .../jumpstarter-driver-snmp/pyproject.toml | 44 +++++ pyproject.toml | 1 + 6 files changed, 259 insertions(+) create mode 100644 packages/jumpstarter-driver-snmp/README.md create mode 100644 packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/__init__.py create mode 100644 packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/client.py create mode 100644 packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/driver.py create mode 100644 packages/jumpstarter-driver-snmp/pyproject.toml diff --git a/packages/jumpstarter-driver-snmp/README.md b/packages/jumpstarter-driver-snmp/README.md new file mode 100644 index 000000000..e69de29bb diff --git a/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/__init__.py b/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/client.py b/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/client.py new file mode 100644 index 000000000..8cc212aeb --- /dev/null +++ b/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/client.py @@ -0,0 +1,48 @@ +from dataclasses import dataclass + +import asyncclick as click + +from jumpstarter.client import DriverClient + + +@dataclass(kw_only=True) +class SNMPServerClient(DriverClient): + """Client interface for SNMP Power Control""" + + def power_on(self): + """Turn power on""" + return self.call("power_on") + + def power_off(self): + """Turn power off""" + return self.call("power_off") + + def power_cycle(self): + """Power cycle the device""" + return self.call("power_cycle") + + def cli(self): + @click.group() + def snmp(): + """SNMP power control commands""" + pass + + @snmp.command() + def on(): + """Turn power on""" + result = self.power_on() + click.echo(result) + + @snmp.command() + def off(): + """Turn power off""" + result = self.power_off() + click.echo(result) + + @snmp.command() + def cycle(): + """Power cycle the device""" + result = self.power_cycle() + click.echo(result) + + return snmp diff --git a/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/driver.py b/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/driver.py new file mode 100644 index 000000000..ca5e8f410 --- /dev/null +++ b/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/driver.py @@ -0,0 +1,166 @@ +import asyncio +import socket +import time +from dataclasses import dataclass, field +from enum import 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 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) + 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") + + def __post_init__(self): + if hasattr(super(), "__post_init__"): + super().__post_init__() + + try: + self.ip_address = socket.gethostbyname(self.host) + self.logger.info(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): + try: + # TODO: switch to anyio? + asyncio.get_running_loop() + except RuntimeError: + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + + + snmp_engine = engine.SnmpEngine() + + config.add_v3_user( + snmp_engine, + self.user, + config.USM_AUTH_NONE, + None + ) + config.add_target_parameters( + snmp_engine, + "my-creds", + self.user, + "noAuthNoPriv" + ) + + config.add_transport( + snmp_engine, + udp.DOMAIN_NAME, + udp.UdpAsyncioTransport().open_client_mode() + ) + + config.add_target_address( + snmp_engine, + "my-target", + udp.DOMAIN_NAME, + (self.ip_address, self.port), + "my-creds" + ) + + 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.info(f"Callback {errorIndication} {errorStatus} {errorIndex} {varBinds}") + if errorIndication: + self.logger.info(f"SNMP error: {errorIndication}") + result["error"] = f"SNMP error: {errorIndication}" + elif errorStatus: + self.logger.info(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.info(f"SNMP set result: {result}") + + try: + self.logger.info(f"Sending power {state.name} command to {self.host}") + + 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 + + @export + def power_on(self): + """Turn power on""" + return self._snmp_set(PowerState.ON) + + @export + def power_off(self): + """Turn power off""" + return self._snmp_set(PowerState.OFF) + + @export + def power_cycle(self): + """Power cycle the device""" + try: + self.logger.info("Starting power cycle sequence") + self.power_off() + self.logger.info(f"Waiting {self.quiescent_period} seconds...") + time.sleep(self.quiescent_period) + self.power_on() + return "Power cycle completed successfully" + except Exception as e: + error_msg = f"Power cycle failed: {str(e)}" + self.logger.error(error_msg) + raise SNMPError(error_msg) from e + + def close(self): + """No cleanup needed since engines are created per operation""" + if hasattr(super(), "close"): + super().close() diff --git a/packages/jumpstarter-driver-snmp/pyproject.toml b/packages/jumpstarter-driver-snmp/pyproject.toml new file mode 100644 index 000000000..5b46388d2 --- /dev/null +++ b/packages/jumpstarter-driver-snmp/pyproject.toml @@ -0,0 +1,44 @@ +[project] +name = "jumpstarter-driver-snmp" +version = "0.1.0" +description = "SNMP driver" +readme = "README.md" +requires-python = ">=3.12.3" +license = { text = "Apache-2.0" } +authors = [ + { name = "Benny Zlotnik", email = "bzlotnik@redhat.com" } +] + +dependencies = [ + "jumpstarter", + "pysnmp==7.1.16" +] + + +[dependency-groups] +dev = [ + "pytest>=8.3.2", + "pytest-cov>=6.0.0", + "pytest-anyio>=0.0.0", + "pytest-asyncio>=0.0.0", + "jumpstarter-testing" +] + + +[tool.pytest.ini_options] +log_cli = true +log_cli_level = "INFO" +testpaths = ["jumpstarter_driver_snmp"] +asyncio_mode = "auto" + +[tool.hatch.metadata.hooks.vcs.urls] +Homepage = "https://jumpstarter.dev" +source_archive = "https://github.com/jumpstarter-dev/repo/archive/{commit_hash}.zip" + +[tool.hatch.version] +source = "vcs" +raw-options = { 'root' = '../../'} + +[build-system] +requires = ["hatchling", "hatch-vcs"] +build-backend = "hatchling.build" diff --git a/pyproject.toml b/pyproject.toml index db5adcabf..27d3f5302 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -19,6 +19,7 @@ jumpstarter-driver-power = { workspace = true } jumpstarter-driver-pyserial = { workspace = true } jumpstarter-driver-sdwire = { workspace = true } jumpstarter-driver-tftp = { workspace = true } +jumpstarter-driver-snmp = { workspace = true } jumpstarter-driver-ustreamer = { workspace = true } jumpstarter-imagehash = { workspace = true } jumpstarter-kubernetes = { workspace = true } From 63ecbc5a5ce37102231871fa8b8d874b635288e0 Mon Sep 17 00:00:00 2001 From: Benny Zlotnik Date: Thu, 6 Feb 2025 13:32:34 +0200 Subject: [PATCH 2/8] inherit power client and standartize method names Signed-off-by: Benny Zlotnik --- .../jumpstarter_driver_power/client.py | 17 ++++ .../jumpstarter_driver_snmp/client.py | 32 +++----- .../jumpstarter_driver_snmp/driver.py | 77 +++++++++++-------- 3 files changed, 71 insertions(+), 55 deletions(-) diff --git a/packages/jumpstarter-driver-power/jumpstarter_driver_power/client.py b/packages/jumpstarter-driver-power/jumpstarter_driver_power/client.py index 924f59675..23797fab9 100644 --- a/packages/jumpstarter-driver-power/jumpstarter_driver_power/client.py +++ b/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,16 @@ def on(self) -> str: def off(self) -> str: return self.call("off") + def cycle(self, quiescent_period: int = 2) -> str: + """Power cycle the device""" + self.logger.info("Starting power cycle sequence") + self.off() + self.logger.info(f"Waiting {quiescent_period} seconds...") + time.sleep(quiescent_period) + self.on() + self.logger.info("Power cycle sequence complete") + return "Power cycled" + def read(self) -> Generator[PowerReading, None, None]: for v in self.streamingcall("read"): yield PowerReading.model_validate(v, strict=True) @@ -33,4 +44,10 @@ def off(): """Power off""" click.echo(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(self.cycle(wait)) + return base diff --git a/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/client.py b/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/client.py index 8cc212aeb..517957281 100644 --- a/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/client.py +++ b/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/client.py @@ -1,25 +1,20 @@ from dataclasses import dataclass import asyncclick as click - -from jumpstarter.client import DriverClient +from jumpstarter_driver_power.client import PowerClient @dataclass(kw_only=True) -class SNMPServerClient(DriverClient): +class SNMPServerClient(PowerClient): """Client interface for SNMP Power Control""" - def power_on(self): + def on(self) -> str: """Turn power on""" - return self.call("power_on") + return self.call("on") - def power_off(self): + def off(self) -> str: """Turn power off""" - return self.call("power_off") - - def power_cycle(self): - """Power cycle the device""" - return self.call("power_cycle") + return self.call("off") def cli(self): @click.group() @@ -27,22 +22,13 @@ def snmp(): """SNMP power control commands""" pass - @snmp.command() - def on(): - """Turn power on""" - result = self.power_on() - click.echo(result) - - @snmp.command() - def off(): - """Turn power off""" - result = self.power_off() - click.echo(result) + for cmd in super().cli().commands.values(): + snmp.add_command(cmd) @snmp.command() def cycle(): """Power cycle the device""" - result = self.power_cycle() + result = self.cycle() click.echo(result) return snmp diff --git a/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/driver.py b/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/driver.py index ca5e8f410..458b9c767 100644 --- a/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/driver.py +++ b/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/driver.py @@ -1,6 +1,5 @@ import asyncio import socket -import time from dataclasses import dataclass, field from enum import IntEnum @@ -30,6 +29,10 @@ class SNMPServer(Driver): 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) def __post_init__(self): if hasattr(super(), "__post_init__"): @@ -37,7 +40,7 @@ def __post_init__(self): try: self.ip_address = socket.gethostbyname(self.host) - self.logger.info(f"Resolved {self.host} to {self.ip_address}") + 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 @@ -45,26 +48,51 @@ def __post_init__(self): def _setup_snmp(self): try: - # TODO: switch to anyio? asyncio.get_running_loop() except RuntimeError: loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) - snmp_engine = engine.SnmpEngine() - config.add_v3_user( - snmp_engine, - self.user, - config.USM_AUTH_NONE, - None - ) + 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 + ) + else: + security_level = 'noAuthNoPriv' + config.add_v3_user( + snmp_engine, + self.user, + config.USM_AUTH_NONE, + None + ) + config.add_target_parameters( snmp_engine, "my-creds", self.user, - "noAuthNoPriv" + security_level ) config.add_transport( @@ -92,12 +120,12 @@ def _snmp_set(self, state: PowerState): def callback(snmpEngine, sendRequestHandle, errorIndication, errorStatus, errorIndex, varBinds, cbCtx): - self.logger.info(f"Callback {errorIndication} {errorStatus} {errorIndex} {varBinds}") + self.logger.debug(f"Callback {errorIndication} {errorStatus} {errorIndex} {varBinds}") if errorIndication: - self.logger.info(f"SNMP error: {errorIndication}") + self.logger.error(f"SNMP error: {errorIndication}") result["error"] = f"SNMP error: {errorIndication}" elif errorStatus: - self.logger.info(f"SNMP status: {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 '?'}" @@ -106,7 +134,7 @@ def callback(snmpEngine, sendRequestHandle, errorIndication, result["success"] = True for oid, val in varBinds: self.logger.debug(f"{oid.prettyPrint()} = {val.prettyPrint()}") - self.logger.info(f"SNMP set result: {result}") + self.logger.debug(f"SNMP set result: {result}") try: self.logger.info(f"Sending power {state.name} command to {self.host}") @@ -136,30 +164,15 @@ def callback(snmpEngine, sendRequestHandle, errorIndication, raise SNMPError(error_msg) from e @export - def power_on(self): + def on(self): """Turn power on""" return self._snmp_set(PowerState.ON) @export - def power_off(self): + def off(self): """Turn power off""" return self._snmp_set(PowerState.OFF) - @export - def power_cycle(self): - """Power cycle the device""" - try: - self.logger.info("Starting power cycle sequence") - self.power_off() - self.logger.info(f"Waiting {self.quiescent_period} seconds...") - time.sleep(self.quiescent_period) - self.power_on() - return "Power cycle completed successfully" - except Exception as e: - error_msg = f"Power cycle failed: {str(e)}" - self.logger.error(error_msg) - raise SNMPError(error_msg) from e - def close(self): """No cleanup needed since engines are created per operation""" if hasattr(super(), "close"): From e08fc881938005ce72454f0bea9f3a73911e433e Mon Sep 17 00:00:00 2001 From: Benny Zlotnik Date: Thu, 6 Feb 2025 17:18:36 +0200 Subject: [PATCH 3/8] add tests and doc Signed-off-by: Benny Zlotnik --- docs/source/api-reference/drivers/snmp.md | 82 ++++++++++ .../examples/exporter.yaml | 18 ++ .../jumpstarter_driver_snmp/driver.py | 116 +++++++------ .../jumpstarter_driver_snmp/driver_test.py | 154 ++++++++++++++++++ 4 files changed, 324 insertions(+), 46 deletions(-) create mode 100644 docs/source/api-reference/drivers/snmp.md create mode 100644 packages/jumpstarter-driver-snmp/examples/exporter.yaml create mode 100644 packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/driver_test.py diff --git a/docs/source/api-reference/drivers/snmp.md b/docs/source/api-reference/drivers/snmp.md new file mode 100644 index 000000000..c42f52fea --- /dev/null +++ b/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 diff --git a/packages/jumpstarter-driver-snmp/examples/exporter.yaml b/packages/jumpstarter-driver-snmp/examples/exporter.yaml new file mode 100644 index 000000000..31a59aec5 --- /dev/null +++ b/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: +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" diff --git a/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/driver.py b/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/driver.py index 458b9c767..3f5cd857a 100644 --- a/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/driver.py +++ b/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() diff --git a/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/driver_test.py b/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/driver_test.py new file mode 100644 index 000000000..880b3e2c6 --- /dev/null +++ b/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/driver_test.py @@ -0,0 +1,154 @@ +from unittest.mock import MagicMock, patch + +import pytest +from pysnmp.entity import config as snmp_config + +from jumpstarter_driver_snmp.driver import AuthProtocol, PrivProtocol, SNMPServer + + +class MockMibObject: + def getInstIdFromIndices(self, *args): + return (1, 3, 6) + +def setup_mock_snmp_engine(): + mock_engine = MagicMock() + mock_builder = MagicMock() + + mock_entry = MockMibObject() + mock_builder.import_symbols.return_value = [mock_entry] + mock_engine.get_mib_builder.return_value = mock_builder + + mock_engine.transport_dispatcher = MagicMock() + mock_engine.transport_dispatcher.start = MagicMock() + mock_engine.transport_dispatcher.stop = MagicMock() + + return mock_engine + +@pytest.mark.parametrize("auth_config", [ + { + "user": "usr-no-auth", + "auth_protocol": AuthProtocol.NONE, + "auth_key": None, + "priv_protocol": PrivProtocol.NONE, + "priv_key": None, + "expected_args_len": 2 # only user and engine args for noAuth + }, + { + "user": "usr-md5-none", + "auth_protocol": AuthProtocol.MD5, + "auth_key": "authkey1", + "priv_protocol": PrivProtocol.NONE, + "priv_key": None, + "expected_args_len": 4 # engine, user, auth_protocol, auth_key + }, + { + "user": "usr-sha-des", + "auth_protocol": AuthProtocol.SHA, + "auth_key": "authkey1", + "priv_protocol": PrivProtocol.DES, + "priv_key": "privkey1", + "expected_args_len": 6 # engine, user, auth_protocol, auth_key, priv_protocol, priv_key + }, +]) +def test_snmp_auth_configurations(auth_config): + """Test different SNMP authentication configurations""" + with patch('pysnmp.entity.config.add_v3_user') as mock_add_user, \ + patch('pysnmp.entity.engine.SnmpEngine', return_value=setup_mock_snmp_engine()), \ + patch('pysnmp.entity.config.add_target_parameters'), \ + patch('pysnmp.entity.config.add_target_address'), \ + patch('pysnmp.entity.config.add_transport'): + + server = SNMPServer( + host="localhost", + user=auth_config["user"], + plug=1, + auth_protocol=auth_config["auth_protocol"], + auth_key=auth_config["auth_key"], + priv_protocol=auth_config["priv_protocol"], + priv_key=auth_config["priv_key"] + ) + + server._setup_snmp() + + args, _ = mock_add_user.call_args + + assert len(args) == auth_config["expected_args_len"] + + assert args[1] == auth_config["user"] + + if auth_config["auth_protocol"] != AuthProtocol.NONE: + if auth_config["auth_protocol"] == AuthProtocol.MD5: + expected_auth = snmp_config.USM_AUTH_HMAC96_MD5 + else: + expected_auth = snmp_config.USM_AUTH_HMAC96_SHA + + assert args[2] == expected_auth + assert args[3] == auth_config["auth_key"] + + if auth_config["priv_protocol"] != PrivProtocol.NONE: + if auth_config["priv_protocol"] == PrivProtocol.DES: + expected_priv = snmp_config.USM_PRIV_CBC56_DES + else: + expected_priv = snmp_config.USM_PRIV_CFB128_AES + + assert args[4] == expected_priv + assert args[5] == auth_config["priv_key"] + +@patch('pysnmp.entity.config.add_v3_user') +@patch('pysnmp.entity.engine.SnmpEngine') +def test_power_on_command(mock_engine, mock_add_user): + """Test power on command execution""" + mock_engine.return_value = setup_mock_snmp_engine() + + with patch('pysnmp.entity.rfc3413.cmdgen.SetCommandGenerator.send_varbinds') as mock_send, \ + patch('asyncio.get_running_loop', side_effect=RuntimeError), \ + patch('asyncio.new_event_loop'), \ + patch('asyncio.set_event_loop'), \ + patch('pysnmp.entity.config.add_target_parameters'), \ + patch('pysnmp.entity.config.add_target_address'), \ + patch('pysnmp.entity.config.add_transport'): + + server = SNMPServer( + host="localhost", + user="testuser", + plug=1 + ) + + def side_effect(*args): + callback = args[-1] + callback(None, None, None, None, None, [], None) + mock_send.side_effect = side_effect + + result = server.on() + assert "Power ON command sent successfully" in result + mock_send.assert_called_once() + + +@patch('pysnmp.entity.config.add_v3_user') +@patch('pysnmp.entity.engine.SnmpEngine') +def test_power_off_command(mock_engine, mock_add_user): + """Test power off command execution""" + mock_engine.return_value = setup_mock_snmp_engine() + + with patch('pysnmp.entity.rfc3413.cmdgen.SetCommandGenerator.send_varbinds') as mock_send, \ + patch('asyncio.get_running_loop', side_effect=RuntimeError), \ + patch('asyncio.new_event_loop'), \ + patch('asyncio.set_event_loop'), \ + patch('pysnmp.entity.config.add_target_parameters'), \ + patch('pysnmp.entity.config.add_target_address'), \ + patch('pysnmp.entity.config.add_transport'): + + server = SNMPServer( + host="localhost", + user="testuser", + plug=1 + ) + + def side_effect(*args): + callback = args[-1] + callback(None, None, None, None, None, [], None) + mock_send.side_effect = side_effect + + result = server.off() + assert "Power OFF command sent successfully" in result + mock_send.assert_called_once() From 0412944b6f4cfde672d2f5bd5f4954abf47af3d7 Mon Sep 17 00:00:00 2001 From: Benny Zlotnik Date: Thu, 6 Feb 2025 18:03:32 +0200 Subject: [PATCH 4/8] snmp: remove cycle from snmp client it's already in the base power client Signed-off-by: Benny Zlotnik --- .../jumpstarter_driver_snmp/client.py | 6 ------ 1 file changed, 6 deletions(-) diff --git a/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/client.py b/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/client.py index 517957281..97740388f 100644 --- a/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/client.py +++ b/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/client.py @@ -25,10 +25,4 @@ def snmp(): for cmd in super().cli().commands.values(): snmp.add_command(cmd) - @snmp.command() - def cycle(): - """Power cycle the device""" - result = self.cycle() - click.echo(result) - return snmp From a6b9ca198dce7d529bbf3923bf1198d2350328c3 Mon Sep 17 00:00:00 2001 From: Benny Zlotnik Date: Thu, 6 Feb 2025 19:26:33 +0200 Subject: [PATCH 5/8] snmp: drop return values for power control methods Signed-off-by: Benny Zlotnik --- docs/source/api-reference/drivers/snmp.md | 6 +-- .../driver_test.py | 38 ++++++++++--------- .../jumpstarter_driver_power/client.py | 25 ++++++------ .../jumpstarter_driver_power/client_test.py | 9 ++--- .../jumpstarter_driver_power/driver.py | 28 +++++++++----- .../jumpstarter_driver_power/driver_test.py | 14 +++++-- .../jumpstarter_driver_snmp/client.py | 8 ++-- .../jumpstarter_testing/pytest_test.py | 2 +- .../jumpstarter/config/exporter_test.py | 2 +- 9 files changed, 74 insertions(+), 58 deletions(-) diff --git a/docs/source/api-reference/drivers/snmp.md b/docs/source/api-reference/drivers/snmp.md index c42f52fea..8332997e9 100644 --- a/docs/source/api-reference/drivers/snmp.md +++ b/docs/source/api-reference/drivers/snmp.md @@ -53,11 +53,11 @@ Turn power off for the configured outlet. Returns: - str: Confirmation message -#### cycle(quiescent_period: int = 2) +#### cycle(wait: 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 +- wait: Time to wait in seconds between power off and power on Returns: - str: Confirmation message @@ -66,7 +66,7 @@ Returns: Power cycling a device: ```python -snmp_client.cycle(quiescent_period=3) +snmp_client.cycle(wait=3) ``` Basic power control: diff --git a/packages/jumpstarter-driver-composite/jumpstarter_driver_composite/driver_test.py b/packages/jumpstarter-driver-composite/jumpstarter_driver_composite/driver_test.py index 750ad902f..20ec6acf4 100644 --- a/packages/jumpstarter-driver-composite/jumpstarter_driver_composite/driver_test.py +++ b/packages/jumpstarter-driver-composite/jumpstarter_driver_composite/driver_test.py @@ -1,21 +1,23 @@ -from jumpstarter_driver_power.driver import MockPower -from .driver import Composite -from jumpstarter.common.utils import serve +from jumpstarter_driver_composite.driver import Composite +from jumpstarter.driver import Driver -def test_drivers_composite(): - with serve( - Composite( - children={ - "power0": MockPower(), - "composite1": Composite( - children={ - "power1": MockPower(), - }, - ), - }, - ) - ) as client: - assert client.power0.on() == "ok" - assert client.composite1.power1.on() == "ok" + +def test_composite_basic(): + class SimpleDriver(Driver): + @classmethod + def client(cls) -> str: + return "test.client.SimpleClient" + + child1 = SimpleDriver() + child2 = SimpleDriver() + + composite = Composite(children={ + "child1": child1, + "child2": child2 + }) + + assert len(composite.children) == 2 + assert composite.children["child1"] == child1 + assert composite.children["child2"] == child2 diff --git a/packages/jumpstarter-driver-power/jumpstarter_driver_power/client.py b/packages/jumpstarter-driver-power/jumpstarter_driver_power/client.py index 23797fab9..96f64ffdc 100644 --- a/packages/jumpstarter-driver-power/jumpstarter_driver_power/client.py +++ b/packages/jumpstarter-driver-power/jumpstarter_driver_power/client.py @@ -8,21 +8,20 @@ class PowerClient(DriverClient): - def on(self) -> str: - return self.call("on") + def on(self): + self.call("on") - def off(self) -> str: - return self.call("off") + def off(self): + self.call("off") - def cycle(self, quiescent_period: int = 2) -> str: + def cycle(self, wait: int = 2): """Power cycle the device""" self.logger.info("Starting power cycle sequence") self.off() - self.logger.info(f"Waiting {quiescent_period} seconds...") - time.sleep(quiescent_period) + self.logger.info(f"Waiting {wait} seconds...") + time.sleep(wait) self.on() self.logger.info("Power cycle sequence complete") - return "Power cycled" def read(self) -> Generator[PowerReading, None, None]: for v in self.streamingcall("read"): @@ -37,17 +36,19 @@ def base(): @base.command() def on(): """Power on""" - click.echo(self.on()) + self.on() + click.echo("Powered on") @base.command() def off(): """Power off""" - click.echo(self.off()) + self.off() + click.echo("Powered 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(self.cycle(wait)) - + click.echo(f"Power cycling with {wait} seconds wait time...") + self.cycle(wait) return base diff --git a/packages/jumpstarter-driver-power/jumpstarter_driver_power/client_test.py b/packages/jumpstarter-driver-power/jumpstarter_driver_power/client_test.py index 0fd9ae8e6..d1baa4f96 100644 --- a/packages/jumpstarter-driver-power/jumpstarter_driver_power/client_test.py +++ b/packages/jumpstarter-driver-power/jumpstarter_driver_power/client_test.py @@ -5,19 +5,18 @@ def test_client_mock_power(): with serve(MockPower()) as client: - assert client.on() == "ok" - assert client.off() == "ok" + client.on() + client.off() assert list(client.read()) == [ PowerReading(voltage=0.0, current=0.0), PowerReading(voltage=5.0, current=2.0), ] - def test_client_sync_mock_power(): with serve(SyncMockPower()) as client: - assert client.on() == "ok" - assert client.off() == "ok" + client.on() + client.off() assert list(client.read()) == [ PowerReading(voltage=0.0, current=0.0), diff --git a/packages/jumpstarter-driver-power/jumpstarter_driver_power/driver.py b/packages/jumpstarter-driver-power/jumpstarter_driver_power/driver.py index b5326a578..1be8a470e 100644 --- a/packages/jumpstarter-driver-power/jumpstarter_driver_power/driver.py +++ b/packages/jumpstarter-driver-power/jumpstarter_driver_power/driver.py @@ -11,23 +11,27 @@ def client(cls) -> str: return "jumpstarter_driver_power.client.PowerClient" @abstractmethod - async def on(self) -> str: ... + async def on(self): ... @abstractmethod - async def off(self) -> str: ... + async def off(self): ... @abstractmethod async def read(self) -> AsyncGenerator[PowerReading, None]: ... class MockPower(PowerInterface, Driver): + def __init__(self, children=None): + self._power_state = None + super().__init__() + @export - async def on(self) -> str: - return "ok" + async def on(self): + self._power_state = "on" @export - async def off(self) -> str: - return "ok" + async def off(self): + self._power_state = "off" @export async def read(self) -> AsyncGenerator[PowerReading, None]: @@ -36,13 +40,17 @@ async def read(self) -> AsyncGenerator[PowerReading, None]: class SyncMockPower(PowerInterface, Driver): + def __init__(self, children=None): + self._power_state = None + super().__init__() + @export - def on(self) -> str: - return "ok" + def on(self): + self._power_state = "on" @export - def off(self) -> str: - return "ok" + def off(self): + self._power_state = "off" @export def read(self) -> Generator[PowerReading, None]: diff --git a/packages/jumpstarter-driver-power/jumpstarter_driver_power/driver_test.py b/packages/jumpstarter-driver-power/jumpstarter_driver_power/driver_test.py index f0824ff47..63a2b506d 100644 --- a/packages/jumpstarter-driver-power/jumpstarter_driver_power/driver_test.py +++ b/packages/jumpstarter-driver-power/jumpstarter_driver_power/driver_test.py @@ -9,8 +9,11 @@ async def test_driver_mock_power(): driver = MockPower() - assert await driver.on() == "ok" - assert await driver.off() == "ok" + await driver.on() + assert driver._power_state == "on" + + await driver.off() + assert driver._power_state == "off" assert [v async for v in driver.read()] == [ PowerReading(voltage=0.0, current=0.0), @@ -21,8 +24,11 @@ async def test_driver_mock_power(): def test_driver_sync_mock_power(): driver = SyncMockPower() - assert driver.on() == "ok" - assert driver.off() == "ok" + driver.on() + assert driver._power_state == "on" + + driver.off() + assert driver._power_state == "off" assert list(driver.read()) == [ PowerReading(voltage=0.0, current=0.0), diff --git a/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/client.py b/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/client.py index 97740388f..331859dd8 100644 --- a/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/client.py +++ b/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/client.py @@ -8,13 +8,13 @@ class SNMPServerClient(PowerClient): """Client interface for SNMP Power Control""" - def on(self) -> str: + def on(self): """Turn power on""" - return self.call("on") + self.call("on") - def off(self) -> str: + def off(self): """Turn power off""" - return self.call("off") + self.call("off") def cli(self): @click.group() diff --git a/packages/jumpstarter-testing/jumpstarter_testing/pytest_test.py b/packages/jumpstarter-testing/jumpstarter_testing/pytest_test.py index b026ec5d3..f1697cc2c 100644 --- a/packages/jumpstarter-testing/jumpstarter_testing/pytest_test.py +++ b/packages/jumpstarter-testing/jumpstarter_testing/pytest_test.py @@ -12,7 +12,7 @@ def test_env(pytester: Pytester, monkeypatch): class TestSample(JumpstarterTest): def test_simple(self, client): - assert client.on() == "ok" + client.on() """ ) diff --git a/packages/jumpstarter/jumpstarter/config/exporter_test.py b/packages/jumpstarter/jumpstarter/config/exporter_test.py index 33348cc0a..d648c039f 100644 --- a/packages/jumpstarter/jumpstarter/config/exporter_test.py +++ b/packages/jumpstarter/jumpstarter/config/exporter_test.py @@ -53,7 +53,7 @@ async def test_exporter_serve(mock_controller): with start_blocking_portal() as portal: async with client.lease_async(metadata_filter=MetadataFilter(), lease_name=None, portal=portal) as lease: async with lease.connect_async() as client: - assert await client.power.call_async("on") == "ok" + await client.power.call_async("on") assert hasattr(client.nested, "tcp") tg.cancel_scope.cancel() From 255bfc892b05f2818d3e305236e9102658350985 Mon Sep 17 00:00:00 2001 From: Benny Zlotnik Date: Fri, 7 Feb 2025 13:47:57 +0200 Subject: [PATCH 6/8] snmp: close event loop if created Signed-off-by: Benny Zlotnik --- .../jumpstarter_driver_snmp/driver.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/driver.py b/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/driver.py index 3f5cd857a..677617bdf 100644 --- a/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/driver.py +++ b/packages/jumpstarter-driver-snmp/jumpstarter_driver_snmp/driver.py @@ -157,11 +157,14 @@ def callback(snmpEngine, sendRequestHandle, errorIndication, 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) + loop = asyncio.new_event_loop() + asyncio.set_event_loop(loop) + created_loop = True snmp_engine = self._setup_snmp() @@ -186,6 +189,9 @@ def callback(snmpEngine, sendRequestHandle, errorIndication, 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): From 66241a9b89d2e3966b4d924c5e7f785585f9c9a9 Mon Sep 17 00:00:00 2001 From: Benny Zlotnik Date: Sun, 9 Feb 2025 09:40:31 +0200 Subject: [PATCH 7/8] snmp: add doc to index Signed-off-by: Benny Zlotnik --- docs/source/api-reference/drivers/index.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/source/api-reference/drivers/index.md b/docs/source/api-reference/drivers/index.md index 59d7467b5..509407769 100644 --- a/docs/source/api-reference/drivers/index.md +++ b/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 ustreamer.md yepkit.md ``` From 135610efae46b9c0147b946290f2fb4db26f3454 Mon Sep 17 00:00:00 2001 From: Benny Zlotnik Date: Wed, 12 Feb 2025 13:42:51 +0200 Subject: [PATCH 8/8] snmp: use doc reference properly Signed-off-by: Benny Zlotnik --- docs/source/api-reference/drivers/snmp.md | 25 ++++--------------- .../jumpstarter_driver_power/client_test.py | 1 + 2 files changed, 6 insertions(+), 20 deletions(-) diff --git a/docs/source/api-reference/drivers/snmp.md b/docs/source/api-reference/drivers/snmp.md index 8332997e9..4e429a234 100644 --- a/docs/source/api-reference/drivers/snmp.md +++ b/docs/source/api-reference/drivers/snmp.md @@ -41,26 +41,11 @@ export: ### 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(wait: int = 2) -Power cycle the device with a configurable wait period between off and on states. - -Parameters: -- wait: Time to wait in seconds between power off and power on - -Returns: -- str: Confirmation message +```{eval-rst} +.. autoclass:: jumpstarter_driver_snmp.client.SNMPServerClient() + :members: + :show-inheritance: +``` ## Examples diff --git a/packages/jumpstarter-driver-power/jumpstarter_driver_power/client_test.py b/packages/jumpstarter-driver-power/jumpstarter_driver_power/client_test.py index d1baa4f96..9f204dd8a 100644 --- a/packages/jumpstarter-driver-power/jumpstarter_driver_power/client_test.py +++ b/packages/jumpstarter-driver-power/jumpstarter_driver_power/client_test.py @@ -13,6 +13,7 @@ def test_client_mock_power(): PowerReading(voltage=5.0, current=2.0), ] + def test_client_sync_mock_power(): with serve(SyncMockPower()) as client: client.on()