Fix lint issues and add doc file to the ssh-mitm driver

Signed-off-by: Bella Khizgiyaev <bkhizgiy@redhat.com>

Author: bkhizgiy

packages/jumpstarter-driver-ssh-mitm/README.md

@@ -0,0 +1,92 @@
+# SSH MITM Driver
+
+`jumpstarter-driver-ssh-mitm` provides secure SSH proxy functionality where private keys
+are stored on the exporter and never transmitted to clients.
+
+## Installation
+
+```{code-block} console
+:substitutions:
+$ pip3 install --extra-index-url {{index_url}} jumpstarter-driver-ssh-mitm
+```
+
+## Configuration
+
+Example configuration with inline key:
+
+```yaml
+export:
+  ssh_mitm:
+    type: jumpstarter_driver_ssh_mitm.driver.SSHMITM
+    config:
+      default_username: "root"
+      ssh_identity: |
+        -----BEGIN OPENSSH PRIVATE KEY-----
+        ...
+        -----END OPENSSH PRIVATE KEY-----
+    children:
+      tcp:
+        type: jumpstarter_driver_network.driver.TcpNetwork
+        config:
+          host: "192.168.1.100"
+          port: 22
+```
+
+Example configuration with key file:
+
+```yaml
+export:
+  ssh_mitm:
+    type: jumpstarter_driver_ssh_mitm.driver.SSHMITM
+    config:
+      default_username: "root"
+      ssh_identity_file: "/path/to/private/key"
+    children:
+      tcp:
+        type: jumpstarter_driver_network.driver.TcpNetwork
+        config:
+          host: "192.168.1.100"
+          port: 22
+```
+
+### Config parameters
+
+| Parameter         | Description                                              | Type | Required | Default |
+| ----------------- | -------------------------------------------------------- | ---- | -------- | ------- |
+| default_username  | Default SSH username                                     | str  | no       | ""      |
+| ssh_identity      | SSH private key content (inline)                         | str  | no*      | None    |
+| ssh_identity_file | Path to SSH private key file                             | str  | no*      | None    |
+
+\* Either `ssh_identity` or `ssh_identity_file` must be provided.
+
+### Required children
+
+- `tcp`: A `TcpNetwork` driver providing target host and port
+
+## Usage
+
+```bash
+# Execute a command
+j ssh_mitm whoami
+
+# Interactive shell (native SSH via port forwarding)
+j ssh_mitm shell
+
+# Interactive shell (gRPC REPL, no local SSH required)
+j ssh_mitm shell --repl
+
+# Port forward for ssh/scp/rsync
+j ssh_mitm forward -p 2222
+# Then: ssh -p 2222 localhost
+```
+
+## API Reference
+
+```{eval-rst}
+.. autoclass:: jumpstarter_driver_ssh_mitm.driver.SSHMITM()
+```
+
+```{eval-rst}
+.. autoclass:: jumpstarter_driver_ssh_mitm.client.SSHMITMClient()
+    :members: execute, run
+```

packages/jumpstarter-driver-ssh-mitm/jumpstarter_driver_ssh_mitm/client.py

@@ -16,8 +16,8 @@
 import shlex
 import shutil
 import subprocess
-import time
 import textwrap
+import time
 import uuid
 from dataclasses import dataclass
 from difflib import get_close_matches
@@ -60,6 +60,7 @@ def resolve_command(self, ctx, args):
 @dataclass
 class SSHMITMCommandRunResult:
     """Result of executing a command via SSH MITM."""
+
     return_code: int
     stdout: str
     stderr: str
@@ -69,15 +70,15 @@ class SSHMITMCommandRunResult:
 class SSHMITMClient(DriverClient):
     """
     Client for SSH MITM proxy driver.
-    
+
     Provides secure SSH access where the private key never leaves the exporter.
     Commands are executed via gRPC - the driver runs SSH on behalf of the client.
     """
 
-    def cli(self):
+    def cli(self):  # noqa: C901
         """Create CLI command for 'j ssh_mitm'."""
         client = self
-        
+
         @click.group(
             "ssh",
             cls=DefaultCommandGroup,
@@ -119,7 +120,7 @@ def run(ctx, args):
 
             if result.return_code != 0:
                 ctx.exit(result.return_code)
-        
+
         @ssh_cmd.command("shell")
         @click.option(
             "--repl",
@@ -131,7 +132,7 @@ def run(ctx, args):
         def shell(ctx, repl, ssh_args):
             """
             Launch an SSH session through the MITM proxy.
-            
+
             By default, spawns the system 'ssh' binary via port forwarding.
             Use --repl for the lightweight gRPC REPL shell.
             """
@@ -141,7 +142,7 @@ def shell(ctx, repl, ssh_args):
                 exit_code = client._launch_native_ssh(ssh_args)
                 if exit_code != 0:
                     ctx.exit(exit_code)
-        
+
         @ssh_cmd.command("forward")
         @click.option(
             "--host",
@@ -162,15 +163,15 @@ def shell(ctx, repl, ssh_args):
         def forward(local_host, local_port):
             """
             Expose the MITM proxy as a local TCP port for native SSH/scp/rsync.
-            
+
             Example:
                 j ssh_mitm forward -p 2222
                 ssh -p 2222 localhost
             """
             client._start_forward(local_host, local_port)
-        
+
         return ssh_cmd
-        
+
     def _ensure_ssh_binary(self) -> str:
         ssh_path = shutil.which("ssh")
         if not ssh_path:
@@ -198,59 +199,59 @@ def _launch_native_ssh(self, ssh_args: tuple[str, ...]) -> int:
             self.logger.debug("Launching native SSH: %s", shlex.join(ssh_command))
             return subprocess.call(ssh_command)
 
-    def _run_shell(self):
+    def _run_shell(self):  # noqa: C901
         """Run interactive shell via gRPC commands."""
         username = self.call("get_default_username") or "user"
         hostname = "dut"
-        
+
         try:
             result = self.execute(["hostname", "-s"])
             if result.return_code == 0 and result.stdout.strip():
                 hostname = result.stdout.strip()
         except Exception as e:
             self.logger.debug("Failed to get hostname: %s", e)
-        
+
         click.echo(f"Connected to {hostname} via SSH MITM proxy")
         click.echo("Type 'exit' or Ctrl+D to exit")
         click.echo()
-        
+
         cwd = "~"
-        
+
         while True:
             try:
                 prompt = click.style(f"{username}@{hostname}", fg="green", bold=True)
                 prompt += click.style(":", fg="white")
                 prompt += click.style(cwd, fg="blue", bold=True)
                 prompt += click.style("$ ", fg="white")
-                
+
                 cmd = input(prompt)
-                
+
                 if not cmd.strip():
                     continue
-                    
+
                 if cmd.strip() == "exit":
                     click.echo("Connection closed.")
                     break
-                
+
                 if cmd.strip().startswith("cd "):
                     new_dir = cmd.strip()[3:].strip()
                     result = self.execute(
                         [
                             "bash",
                             "-c",
-                            f'cd {shlex.quote(cwd)} 2>/dev/null; cd {shlex.quote(new_dir)} && pwd',
+                            f"cd {shlex.quote(cwd)} 2>/dev/null; cd {shlex.quote(new_dir)} && pwd",
                         ]
                     )
                     if result.return_code == 0 and result.stdout.strip():
                         cwd = result.stdout.strip()
                     else:
                         click.echo(f"cd: {new_dir}: No such file or directory", err=True)
                     continue
-                
+
                 if cmd.strip() == "cd":
                     cwd = "~"
                     continue
-                
+
                 # Execute command in current directory using newline-delimited heredoc to avoid interpolation
                 token = f"JSSHMITM_{uuid.uuid4().hex}"
                 script = (
@@ -265,20 +266,20 @@ def _run_shell(self):
                     + "\n"
                 )
                 result = self.execute(["bash", "-lc", script])
-                
+
                 if result.stdout:
                     click.echo(result.stdout, nl=False)
                 if result.stderr:
                     click.echo(result.stderr, nl=False, err=True)
-                    
+
             except EOFError:
                 click.echo()
                 click.echo("Connection closed.")
                 break
             except KeyboardInterrupt:
                 click.echo("^C")
                 continue
-    
+
     def _start_forward(self, local_host: str, local_port: int):
         """Expose the SSH MITM server on a local TCP port."""
         click.echo("Starting local forward (Ctrl+C to stop)...")
@@ -300,18 +301,18 @@ def _start_forward(self, local_host: str, local_port: int):
     def execute(self, args) -> SSHMITMCommandRunResult:
         """
         Execute command on DUT via gRPC.
-        
+
         The command is run on the exporter using the stored SSH key,
         then results are returned.
         """
         return_code, stdout, stderr = self.call("execute_command", *args)
-        
+
         return SSHMITMCommandRunResult(
             return_code=return_code,
             stdout=stdout,
             stderr=stderr,
         )
-    
+
     def run(self, args) -> SSHMITMCommandRunResult:
         """Alias for execute()."""
         return self.execute(args)