From 8c1de5941c22cd92166e6b8aa7fab3342673e8a6 Mon Sep 17 00:00:00 2001
From: Andrea Righi <arighi@nvidia.com>
Date: Mon, 23 Dec 2024 08:06:57 +0100
Subject: [PATCH] virtme: enable ssh support

Add support for specifying `--server ssh`, to automatically start and
configure sshd in the virtme-ng guest, and `--client ssh` to connect to
a virtme-ng guest via SSH.

The server's port can be customized with the `--port NUM` option
(default is 2222).

Example usage (start a vng instance with sshd enabled):

 $ vng --server ssh --port 2222

In another shell session on the host, connect to the running instance as
a client:

 $ vng --client ssh --port 2222

Note that it's also possible to use ssh or scp directly, the --client
option is provided for convenience and to be consistent with the vsock
option.

Signed-off-by: Andrea Righi <arighi@nvidia.com>
---
 README.md                       |  9 +++++++++
 setup.py                        |  1 +
 virtme/commands/run.py          | 34 ++++++++++++++++++++++++++++++---
 virtme/guest/virtme-init        |  4 ++++
 virtme/guest/virtme-sshd-script | 10 ++++++++++
 virtme_ng/run.py                |  2 +-
 virtme_ng_init                  |  2 +-
 7 files changed, 57 insertions(+), 5 deletions(-)
 create mode 100755 virtme/guest/virtme-sshd-script

diff --git a/README.md b/README.md
index dc13f1a..0919cfc 100644
--- a/README.md
+++ b/README.md
@@ -388,6 +388,15 @@ Examples
    $ vng --client
 ```
 
+ - Enable ssh in the vng guest:
+```
+   # Start the vng instance with ssh server support:
+   $ vng --server ssh
+
+   # Connect to the vng guest from the host via ssh:
+   $ vng --client ssh
+```
+
  - Run virtme-ng inside a docker container:
 ```
    $ docker run -it --privileged ubuntu:23.10 /bin/bash
diff --git a/setup.py b/setup.py
index d9d7a12..ebeb590 100755
--- a/setup.py
+++ b/setup.py
@@ -98,6 +98,7 @@ def run(self):
 package_files = [
     "virtme-init",
     "virtme-udhcpc-script",
+    "virtme-sshd-script",
     "virtme-snapd-script",
     "virtme-sound-script",
 ]
diff --git a/virtme/commands/run.py b/virtme/commands/run.py
index 6f7a926..ee3a5ff 100644
--- a/virtme/commands/run.py
+++ b/virtme/commands/run.py
@@ -330,7 +330,7 @@ def make_parser() -> argparse.ArgumentParser:
     )
 
     g = parser.add_argument_group(title="Remote Console")
-    cli_srv_choices = ["console"]
+    cli_srv_choices = ["console", "ssh"]
     g.add_argument(
         "--server",
         action="store",
@@ -968,6 +968,28 @@ def cleanup_console_script():
     qemuargs.extend(["-device", "vhost-vsock-pci,guest-cid=%d" % args.port])
 
 
+def ssh_client(args):
+    remote_cmd = args.remote_cmd if args.remote_cmd else ""
+    cmd = "ssh -p %d localhost %s" % (args.port, remote_cmd)
+    if args.dry_run:
+        print(cmd)
+    else:
+        os.system(cmd)
+
+
+def ssh_server(args, arch, qemuargs, kernelargs):
+    # Implicitly enable dhcp to automatically get an IP on the network
+    # interface and prevent interface renaming.
+    kernelargs.extend(["virtme.dhcp", "net.ifnames=0", "biosdevname=0"])
+
+    # Tell virtme-ng-init / virtme-init to start sshd.
+    kernelargs.extend(["virtme.ssh"])
+
+    # Setup a port forward network interface for the guest.
+    qemuargs.extend(["-device", "%s,netdev=ssh" % (arch.virtio_dev_type("net"))])
+    qemuargs.extend(["-netdev", "user,id=ssh,hostfwd=tcp::%d-:22" % args.port])
+
+
 # Allowed characters in mount paths.  We can extend this over time if needed.
 _SAFE_PATH_PATTERN = "[a-zA-Z0-9_+ /.-]+"
 _RWDIR_RE = re.compile("^(%s)(?:=(%s))?$" % (_SAFE_PATH_PATTERN, _SAFE_PATH_PATTERN))
@@ -980,7 +1002,10 @@ def do_it() -> int:
         if args.server is not None:
             arg_fail('--client cannot be used with --server.')
 
-        console_client(args)
+        if args.client == 'vsock':
+            console_client(args)
+        elif args.client == 'ssh':
+            ssh_client(args)
         sys.exit(0)
 
     arch = architectures.get(args.arch)
@@ -1526,7 +1551,10 @@ def get_net_mac(index):
         )
 
     if args.server is not None:
-        console_server(args, qemu, arch, qemuargs, kernelargs)
+        if args.server == "vsock":
+            console_server(args, qemu, arch, qemuargs, kernelargs)
+        elif args.server == "ssh":
+            ssh_server(args, arch, qemuargs, kernelargs)
 
     if args.pwd:
         rel_pwd = os.path.relpath(os.getcwd(), args.root)
diff --git a/virtme/guest/virtme-init b/virtme/guest/virtme-init
index f97da74..6e36203 100755
--- a/virtme/guest/virtme-init
+++ b/virtme/guest/virtme-init
@@ -257,6 +257,10 @@ if cat /proc/cmdline |grep -q -E '(^| )virtme.dhcp($| )'; then
     wait
 fi
 
+if cat /proc/cmdline |grep -q -E '(^| )virtme.ssh($| )'; then
+    $(dirname $0)/virtme-sshd-script
+fi
+
 if cat /proc/cmdline |grep -q -E '(^| )virtme.snapd($| )'; then
     # If snapd is present in the system try to start it, to properly support snaps.
     snapd_bin="/usr/lib/snapd/snapd";
diff --git a/virtme/guest/virtme-sshd-script b/virtme/guest/virtme-sshd-script
new file mode 100755
index 0000000..537e57d
--- /dev/null
+++ b/virtme/guest/virtme-sshd-script
@@ -0,0 +1,10 @@
+#!/bin/bash
+#
+# Initialize ssh server for remote connections (option `--server ssh`)
+
+HOME=$(getent passwd "${virtme_user}" | cut -d: -f6)
+cat ${HOME}/.ssh/id_rsa.pub >> ${HOME}/.ssh/authorized_keys
+chown ${virtme_user} ${HOME}/.ssh/authorized_keys
+mkdir -p /run/sshd
+rm -f /var/run/nologin
+/usr/sbin/sshd -h ~/.ssh/id_rsa
diff --git a/virtme_ng/run.py b/virtme_ng/run.py
index 846f67f..9d824fd 100644
--- a/virtme_ng/run.py
+++ b/virtme_ng/run.py
@@ -487,7 +487,7 @@ def make_parser():
     )
 
     g_remote = parser.add_argument_group(title="Remote Console")
-    cli_srv_choices = ["console"]
+    cli_srv_choices = ["console", "ssh"]
 
     g_remote.add_argument(
         "--server",
diff --git a/virtme_ng_init b/virtme_ng_init
index fe8484d..9b1e02a 160000
--- a/virtme_ng_init
+++ b/virtme_ng_init
@@ -1 +1 @@
-Subproject commit fe8484d502456131e5975c051077da2ff67c5aa3
+Subproject commit 9b1e02a0deb094a36741b6172fb7ea8dc7dd8ddb