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"):