From ec55c26471482d2fb953cdbac74be40a58a4d07c Mon Sep 17 00:00:00 2001
From: Tulip <tulilirockz@outlook.com>
Date: Tue, 13 Aug 2024 00:31:01 -0300
Subject: [PATCH] feat: --nvidia flag for enter + export and status commands

---
 files/nopass-machinectl.rules | 17 ++++++++++
 src/lib/enter.nu              | 60 +++++++++++++++++++++++++++++++++++
 src/lib/export.nu             | 36 +++++++++++++++------
 src/lib/mod.nu                |  1 +
 src/lib/status.nu             | 22 +++++++++++++
 5 files changed, 126 insertions(+), 10 deletions(-)
 create mode 100644 files/nopass-machinectl.rules
 create mode 100644 src/lib/status.nu

diff --git a/files/nopass-machinectl.rules b/files/nopass-machinectl.rules
new file mode 100644
index 0000000..b527d8a
--- /dev/null
+++ b/files/nopass-machinectl.rules
@@ -0,0 +1,17 @@
+polkit.addRule(function(action, subject) {
+    if (
+        (
+            action.id == "org.freedesktop.machine1.shell" ||
+            action.id == "org.freedesktop.machine1.manage-machines" ||
+            action.id == "org.freedesktop.machine1.manage-images" ||
+            action.id == "org.freedesktop.machine1.login" ||
+            (
+                action.id == "org.freedesktop.systemd1.manage-units" &&
+                RegExp('systemd-nspawn@[A-Za-z0-9_-]+.service').test(action.lookup("unit")) === true
+            )
+        )
+        && subject.isInGroup("wheel")
+    ) {
+        return polkit.Result.YES;
+    }
+});
diff --git a/src/lib/enter.nu b/src/lib/enter.nu
index fde7529..698164d 100644
--- a/src/lib/enter.nu
+++ b/src/lib/enter.nu
@@ -27,6 +27,7 @@ export def --env "main enter" [
   --runner: string = "nspawn" # Runner of the machine: nspawn or vmspawn
   --no-kill # Do not kill machine preemptively before running launcher command
   --yes # Do not print any confirmation prompts 
+  --nvidia # Add binds for nvidia driver support
   machine: string # Name of the target to be logged into (either a directory, machine name, or image name)
   ...args: string
 ] {   
@@ -86,6 +87,7 @@ export def --env "main enter" [
     $"--private-users=(if $no_user_bind { "no" } else { "pick" })"
     $"--private-users-ownership=auto"
     "--keep-unit"
+    "--link-journal=try-guest"
     "--resolv-conf=replace-stub"
     (if $boot { "--boot" } else { $"--chdir=($env.PWD? | default "/home")" } )
     (if not $no_user_bind { $"--bind-user=($user)" } else { "--private-users=no" })
@@ -132,6 +134,64 @@ export def --env "main enter" [
     ] | each {|e| $"--bind=($e)"}))
   }  
 
+  if $nvidia {
+    # Taken from the arch wiki, not tested yet :(
+
+    $final_args = ($final_args | append ([
+      "/dev/nvidia0"
+      "/dev/nvidiactl"
+      "/dev/nvidia-modeset"
+      "/usr/bin/nvidia-bug-report.sh"
+      "/usr/bin/nvidia-cuda-mps-control"
+      "/usr/bin/nvidia-cuda-mps-server"
+      "/usr/bin/nvidia-debugdump"
+      "/usr/bin/nvidia-modprobe"
+      "/usr/bin/nvidia-ngx-updater"
+      "/usr/bin/nvidia-persistenced"
+      "/usr/bin/nvidia-powerd"
+      "/usr/bin/nvidia-sleep.sh"
+      "/usr/bin/nvidia-smi"
+      "/usr/bin/nvidia-xconfig"
+      "/usr/lib/gbm/nvidia-drm_gbm.so"
+      "/usr/lib/libEGL_nvidia.so"
+      "/usr/lib/libGLESv1_CM_nvidia.so"
+      "/usr/lib/libGLESv2_nvidia.so"
+      "/usr/lib/libGLX_nvidia.so"
+      "/usr/lib/libcuda.so"
+      "/usr/lib/libnvcuvid.so"
+      "/usr/lib/libnvidia-allocator.so"
+      "/usr/lib/libnvidia-cfg.so"
+      "/usr/lib/libnvidia-egl-gbm.so"
+      "/usr/lib/libnvidia-eglcore.so"
+      "/usr/lib/libnvidia-encode.so"
+      "/usr/lib/libnvidia-fbc.so"
+      "/usr/lib/libnvidia-glcore.so"
+      "/usr/lib/libnvidia-glsi.so"
+      "/usr/lib/libnvidia-glvkspirv.so"
+      "/usr/lib/libnvidia-ml.so"
+      "/usr/lib/libnvidia-ngx.so"
+      "/usr/lib/libnvidia-opticalflow.so"
+      "/usr/lib/libnvidia-ptxjitcompiler.so"
+      "/usr/lib/libnvidia-rtcore.so"
+      "/usr/lib/libnvidia-tls.so"
+      "/usr/lib/libnvidia-vulkan-producer.so"
+      "/usr/lib/libnvoptix.so"
+      "/usr/lib/modprobe.d/nvidia-utils.conf"
+      "/usr/lib/nvidia/wine/_nvngx.dll"
+      "/usr/lib/nvidia/wine/nvngx.dll"
+      "/usr/lib/nvidia/xorg/libglxserver_nvidia.so"
+      "/usr/lib/vdpau/libvdpau_nvidia.so"
+      "/usr/lib/xorg/modules/drivers/nvidia_drv.so"
+      "/usr/share/X11/xorg.conf.d/10-nvidia-drm-outputclass.conf"
+      "/usr/share/dbus-1/system.d/nvidia-dbus.conf"
+      "/usr/share/egl/egl_external_platform.d/15_nvidia_gbm.json"
+      "/usr/share/glvnd/egl_vendor.d/10_nvidia.json"
+      "/usr/share/licenses/nvidia-utils/LICENSE"
+      "/usr/share/vulkan/icd.d/nvidia_icd.json"
+      "/usr/share/vulkan/implicit_layer.d/nvidia_layers.json"
+    ] | each {|e| $"--bind=($e)"}))
+  }
+
   if ($extra_bind != "" and $extra_bind != null) {
     $final_args = ($final_args | append ($extra_bind | split row "," | each {|e| $"--bind(if $ro_binds {"-ro"} else {""})=($e)"}))
   }
diff --git a/src/lib/export.nu b/src/lib/export.nu
index 2e309b9..c77f900 100644
--- a/src/lib/export.nu
+++ b/src/lib/export.nu
@@ -1,16 +1,32 @@
+use logger.nu *
+use meta.nu NAME
+
 # Export programs from a machine by writing scripts and desktop files for exported apps
-export def "main export" [] {}
+export def "main export" [] {
+  $"Usage: ($NAME) export <command>"
+}
 
-export def "main export bin" [
+# Generate a script that will run a binary in a machine through nuspawn enter
+export def "main export binary" [
   --extra-flags: string # Flags for the exported binary
-  machine: string
+  --export-path: path = "~/.local/bin" # Path where the binary will be exported to
+  --nuspawn-args: string # Arguments that will be passed to nuspawn
+  --nuspawn-prefix: string # Any prefix that will run before the nuspawn caller runs (e.g.: bwrap nuspawn) 
+  --prefix: string # Any prefix that will run before the command gets ran (e.g.: gamescope steam) 
+  --nuspawn-binary: string = "nuspawn" # Path to the nuspawn script (must be executable) (defaults to $PATH)
+  machine: string # Machine that the program will be executed on
+  binary: string # Binary that will be exported to the host
 ] {
-  let script = "
-  #!/bin/sh
-  nuspawn enter ($machine) --
-  
-  
-  "
+  mkdir $export_path
+  let exported_binary_path = $"($export_path)/($binary | path basename)"
+  $"#!/bin/sh
+  ($nuspawn_prefix) ($nuspawn_binary) enter ($nuspawn_args) ($machine) /bin/sh -c \"($prefix) ($binary) ($extra_flags)\"
+  " | save -f $exported_binary_path
+  run-external "chmod" "+x" $exported_binary_path
+  logger success $"Succesfully exported binary to ($exported_binary_path)"
+}
 
+# Export desktop files and data from a program in a machine 
+export def "main export app" [] {
+  logger error "To be implemented"
 }
-export def "main export app" [] {}
diff --git a/src/lib/mod.nu b/src/lib/mod.nu
index 35f775d..ef33822 100644
--- a/src/lib/mod.nu
+++ b/src/lib/mod.nu
@@ -12,3 +12,4 @@ export use rename.nu *
 export use setup.nu *
 export use start.nu *
 export use oci.nu *
+export use status.nu *
diff --git a/src/lib/status.nu b/src/lib/status.nu
new file mode 100644
index 0000000..dbae3f3
--- /dev/null
+++ b/src/lib/status.nu
@@ -0,0 +1,22 @@
+use logger.nu *
+
+# Get status for all the machines or just a specific one
+export def "main status" [
+  machine?: string
+] {
+  if $machine != null {
+    let data = (machinectl status $machine --output json | complete | get stdout | from json)
+    if ($data | length) == 0 {
+      logger warning $"Machine ($machine) is not running"
+      return
+    }
+    return
+  }
+
+  let data = (machinectl list --output json | from json)
+  if ($data | length) == 0 {
+    logger warning "No machine is currently running"
+    return
+  }
+  $data
+}