diff --git a/docs/source/api-reference/drivers/index.md b/docs/source/api-reference/drivers/index.md index cc3aa3b92..63829bae7 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 tftp.md ustreamer.md yepkit.md diff --git a/docs/source/api-reference/drivers/snmp.md b/docs/source/api-reference/drivers/snmp.md new file mode 100644 index 000000000..4e429a234 --- /dev/null +++ b/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 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 43e0508dd..f0f399595 100644 --- a/packages/jumpstarter-driver-composite/jumpstarter_driver_composite/driver_test.py +++ b/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 diff --git a/packages/jumpstarter-driver-power/jumpstarter_driver_power/client.py b/packages/jumpstarter-driver-power/jumpstarter_driver_power/client.py index ff6a572ef..9af2b281d 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,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 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/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/__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..331859dd8 --- /dev/null +++ b/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 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..677617bdf --- /dev/null +++ b/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() 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() 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 e784a9ba2..55db768f6 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -20,6 +20,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 }