From fb5e863d8c3e1829736c0b66cedeb7aa6adc91c3 Mon Sep 17 00:00:00 2001
From: Connor Nelson <Connor@ConnorNelson.com>
Date: Sat, 7 Sep 2024 14:35:10 -0700
Subject: [PATCH] Workspace: Add sudo

---
 workspace/Dockerfile    |  2 +-
 workspace/core/init.nix |  8 ++------
 workspace/core/sudo.nix | 21 +++++++++++++++++++++
 workspace/core/sudo.py  | 35 +++++++++++++++++++++++++++++++++++
 workspace/flake.nix     |  2 ++
 5 files changed, 61 insertions(+), 7 deletions(-)
 create mode 100644 workspace/core/sudo.nix
 create mode 100644 workspace/core/sudo.py

diff --git a/workspace/Dockerfile b/workspace/Dockerfile
index 9e961ffe8..560233195 100644
--- a/workspace/Dockerfile
+++ b/workspace/Dockerfile
@@ -32,7 +32,7 @@ fi
 
 nix build --print-build-logs ".#$DOJO_WORKSPACE" --out-link /out/nix/var/nix/profiles/default
 nix copy --to /out --no-require-sigs ".#$DOJO_WORKSPACE"
-for suid_binary in "python-suid" "bash-suid" "sh-suid" "windows"; do
+for suid_binary in "python-suid" "bash-suid" "sh-suid" "sudo"; do
     suid_binary_path=$(realpath "/out/nix/var/nix/profiles/default/bin/${suid_binary}")
     FILE="/out/${suid_binary_path}"
     if [ -e $FILE ]; then
diff --git a/workspace/core/init.nix b/workspace/core/init.nix
index 91d943b33..74aeb8127 100644
--- a/workspace/core/init.nix
+++ b/workspace/core/init.nix
@@ -46,12 +46,8 @@ let
     exec > /run/dojo/var/root/init.log 2>&1
     chmod 600 /run/dojo/var/root/init.log
 
-    # TODO: Better support privileged mode
-    if [ "$DOJO_MODE" = "privileged" ] && [ -f /usr/bin/sudo ]; then
-      chmod 4755 /usr/bin/sudo
-      usermod -aG sudo hacker
-      echo 'hacker ALL=(ALL) NOPASSWD:ALL' >> /etc/sudoers
-      passwd -d root
+    if [ "$DOJO_MODE" = "privileged" ]; then
+      touch /run/dojo/var/root/privileged
     fi
 
     if [ -x "/challenge/.init" ]; then
diff --git a/workspace/core/sudo.nix b/workspace/core/sudo.nix
new file mode 100644
index 000000000..909f06ac8
--- /dev/null
+++ b/workspace/core/sudo.nix
@@ -0,0 +1,21 @@
+{ pkgs }:
+
+pkgs.stdenv.mkDerivation {
+  name = "sudo";
+  src = ./sudo.py;
+
+  unpackPhase = ''
+    runHook preUnpack
+    cp $src $PWD
+    runHook postUnpack
+  '';
+
+  installPhase = ''
+    runHook preInstall
+    mkdir -p $out/bin
+    echo "#!/usr/bin/env python-suid" > $out/bin/sudo
+    cat ${./sudo.py} >> $out/bin/sudo
+    chmod +x $out/bin/sudo
+    runHook postInstall
+  '';
+}
diff --git a/workspace/core/sudo.py b/workspace/core/sudo.py
new file mode 100644
index 000000000..1e7de3a86
--- /dev/null
+++ b/workspace/core/sudo.py
@@ -0,0 +1,35 @@
+import os
+import shutil
+import sys
+
+
+def error(message):
+    print(message, file=sys.stderr)
+    sys.exit(1)
+
+
+def main():
+    program = os.path.basename(sys.argv[0])
+
+    if not os.path.exists("/run/dojo/var/root/priviliged"):
+        error(f"{program}: workspace is not privileged")
+
+    os.setuid(os.geteuid())
+
+    if len(sys.argv) < 2:
+        error(f"Usage: {program} <command> [args...]")
+
+    command = sys.argv[1]
+    command_path = shutil.which(sys.argv[1])
+    if not command_path:
+        error(f"{program}: {command}: command not found")
+    argv = sys.argv[1:]
+
+    try:
+        os.execve(command_path, argv, os.environ)
+    except:
+        os.exit(1)
+
+
+if __name__ == "__main__":
+    main()
\ No newline at end of file
diff --git a/workspace/flake.nix b/workspace/flake.nix
index 0f69cdb36..503ee019f 100644
--- a/workspace/flake.nix
+++ b/workspace/flake.nix
@@ -18,6 +18,7 @@
 
           init = import ./core/init.nix { inherit pkgs; };
           suid-interpreter = import ./core/suid-interpreter.nix { inherit pkgs; };
+          sudo = import ./core/sudo.nix { inherit pkgs; };
           ssh-entrypoint = import ./core/ssh-entrypoint.nix { inherit pkgs; };
           service = import ./services/service.nix { inherit pkgs; };
           code-service = import ./services/code.nix { inherit pkgs; };
@@ -66,6 +67,7 @@
             (lib.hiPrio ldd)
 
             init
+            sudo
             ssh-entrypoint
             service
             code-service