inherit power client and standartize method names

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

Author: bennyz

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

@@ -1,4 +1,5 @@
 from collections.abc import Generator
+import time
 
 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

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

@@ -1,48 +1,34 @@
 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()
         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

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,41 +29,70 @@ 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__"):
             super().__post_init__()
 
         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
 
         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
-        )
+        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"):