From e65625e32f6c5541d5b716b59a180ca0fb5c2d3c Mon Sep 17 00:00:00 2001
From: Tobias Waldekranz <tobias@waldekranz.com>
Date: Thu, 16 May 2024 08:54:27 +0000
Subject: [PATCH] test: Add support for attaching to devices via SSH

Support the familiar subprocess.run API, and runsh from netns.py.

This should only be used for debugging, or in the rare instances when
we want to test functionality that is not available via NETCONF.
---
 test/.env              |  2 +-
 test/docker/Dockerfile |  2 ++
 test/infamy/env.py     | 21 +++++++++++--------
 test/infamy/ssh.py     | 46 ++++++++++++++++++++++++++++++++++++++++++
 4 files changed, 62 insertions(+), 9 deletions(-)
 create mode 100644 test/infamy/ssh.py

diff --git a/test/.env b/test/.env
index 2c0fbef5..59be7453 100644
--- a/test/.env
+++ b/test/.env
@@ -2,7 +2,7 @@
 # shellcheck disable=SC2034,SC2154
 
 # Current container image
-INFIX_TEST=ghcr.io/kernelkit/infix-test:1.3
+INFIX_TEST=ghcr.io/kernelkit/infix-test:1.4
 
 ixdir=$(readlink -f "$testdir/..")
 logdir=$(readlink -f "$testdir/.log")
diff --git a/test/docker/Dockerfile b/test/docker/Dockerfile
index c9cc3073..c986de93 100644
--- a/test/docker/Dockerfile
+++ b/test/docker/Dockerfile
@@ -12,12 +12,14 @@ RUN apk add --no-cache \
     libc-dev \
     libyang-dev \
     linux-headers \
+    openssh-client \
     python3-dev \
     qemu-img \
     qemu-system-x86_64 \
     ruby-mustache \
     socat \
     squashfs-tools \
+    sshpass \
     tcpdump \
     tshark \
     make
diff --git a/test/infamy/env.py b/test/infamy/env.py
index be54ef60..44fb4553 100644
--- a/test/infamy/env.py
+++ b/test/infamy/env.py
@@ -5,7 +5,7 @@ import pydot
 import shlex
 import sys
 
-from . import neigh, netconf, tap, topology
+from . import neigh, netconf, ssh, tap, topology
 
 class NullEnv:
     def attr(self, _, default=None):
@@ -58,7 +58,7 @@ class Env(object):
 
         return val
 
-    def attach(self, node, port, factory_default=True):
+    def attach(self, node, port, protocol="netconf", factory_default=True):
         if self.ltop:
             mapping = self.ltop.mapping[node]
             node, port = self.ltop.xlate(node, port)
@@ -77,9 +77,14 @@ class Env(object):
         if not mgmtip:
             raise Exception(f"Failed, cannot find mgmt IP for {node}")
 
-        return netconf.Device(
-            location=netconf.Location(cport, mgmtip, password),
-            mapping=mapping,
-            yangdir=self.args.yangdir,
-            factory_default=factory_default
-        )
+        if protocol == "netconf":
+            return netconf.Device(
+                location=netconf.Location(cport, mgmtip, password),
+                mapping=mapping,
+                yangdir=self.args.yangdir,
+                factory_default=factory_default
+            )
+        elif protocol == "ssh":
+            return ssh.Device(ssh.Location(mgmtip, password))
+
+        raise Exception(f"Unsupported management procotol \"{protocol}\"")
diff --git a/test/infamy/ssh.py b/test/infamy/ssh.py
new file mode 100644
index 00000000..d44735aa
--- /dev/null
+++ b/test/infamy/ssh.py
@@ -0,0 +1,46 @@
+import subprocess
+
+from dataclasses import dataclass
+
+@dataclass
+class Location:
+    host: str
+    password: str
+    username: str = "admin"
+    port: int = 22
+
+class Device(object):
+    def __init__(self, location: Location):
+        self.location = location
+
+    def _mangle_subprocess_args(self, args, kwargs):
+        if not args:
+            return
+
+        args = list(args)
+        if type(args[0]) == str:
+            if kwargs.get("shell"):
+                args[0] = ["/bin/sh", "-c", args[0]]
+                kwargs["shell"] = False
+            else:
+                args[0] = [args[0]]
+
+        args[0] = [ "ssh",
+                    "-oStrictHostKeyChecking no",
+                    "-oUserKnownHostsFile /dev/null",
+	            "-oLogLevel QUIET",
+                    f"-l{self.location.username}",
+                    self.location.host ] + args[0]
+
+        if self.location.password:
+            args[0] = [ "sshpass" , f"-p{self.location.password}" ] + args[0]
+
+        return args, kwargs
+
+    def run(self, *args, **kwargs):
+        args, kwargs = self._mangle_subprocess_args(args, kwargs)
+        return subprocess.run(*args, **kwargs)
+
+    def runsh(self, script, *args, **kwargs):
+        return self.run("/bin/sh", text=True, input=script,
+                        stdout=subprocess.PIPE, stderr=subprocess.STDOUT, *args, **kwargs)