overhaul to properly support terminals

Author: coalaura

elevate.go

@@ -2,8 +2,6 @@ package main
 
 import (
 	"fmt"
-	"os"
-	"os/exec"
 	"syscall"
 	"unsafe"
 
@@ -83,15 +81,3 @@ func elevate(app, params, cwd string) error {
 
 	return nil
 }
-
-func regular(app, params, cwd string) error {
-	cmd := exec.Command(app, params)
-
-	cmd.Dir = cwd
-
-	cmd.Stdout = os.Stdout
-	cmd.Stderr = os.Stderr
-	cmd.Stdin = os.Stdin
-
-	return cmd.Run()
-}

go.mod

@@ -1,16 +1,22 @@
 module sudo
 
-go 1.23.4
+go 1.24.2
 
-toolchain go1.23.5
+toolchain go1.24.4
 
 require (
-	github.com/coalaura/logger v1.2.3
-	golang.org/x/sys v0.29.0
-	golang.org/x/term v0.28.0
+	github.com/coalaura/logger v1.4.5
+	github.com/shirou/gopsutil v3.21.11+incompatible
+	golang.org/x/sys v0.34.0
 )
 
 require (
+	github.com/go-ole/go-ole v1.3.0 // indirect
 	github.com/gookit/color v1.5.4 // indirect
+	github.com/stretchr/testify v1.9.0 // indirect
+	github.com/tklauser/go-sysconf v0.3.15 // indirect
+	github.com/tklauser/numcpus v0.10.0 // indirect
 	github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e // indirect
+	github.com/yusufpapurcu/wmi v1.2.4 // indirect
+	golang.org/x/term v0.33.0 // indirect
 )

go.sum

@@ -1,20 +1,44 @@
 github.com/coalaura/logger v1.2.3 h1:uy6XjA+NjQKddjsY9FbwQplM4Rv9gJ43mVcE3YTyT9E=
 github.com/coalaura/logger v1.2.3/go.mod h1:4lGnkI2dT7d0GSO3l1ci6dhxAJy3SWqjOi03k+Ugwxc=
+github.com/coalaura/logger v1.4.5 h1:xXazOab4qXaltUbD4TrQdSs2TtLB+k6t0t6y/M8LR3Q=
+github.com/coalaura/logger v1.4.5/go.mod h1:3HCYCWmsWmYW175e2/fZL9BWjJutr2W+7adeh1BPHkg=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
+github.com/go-ole/go-ole v1.2.6 h1:/Fpf6oFPoeFik9ty7siob0G6Ke8QvQEuVcuChpwXzpY=
+github.com/go-ole/go-ole v1.2.6/go.mod h1:pprOEPIfldk/42T2oK7lQ4v4JSDwmV0As9GaiUsvbm0=
+github.com/go-ole/go-ole v1.3.0 h1:Dt6ye7+vXGIKZ7Xtk4s6/xVdGDQynvom7xCFEdWr6uE=
+github.com/go-ole/go-ole v1.3.0/go.mod h1:5LS6F96DhAwUc7C+1HLexzMXY1xGRSryjyPPKW6zv78=
 github.com/gookit/color v1.5.4 h1:FZmqs7XOyGgCAxmWyPslpiok1k05wmY3SJTytgvYFs0=
 github.com/gookit/color v1.5.4/go.mod h1:pZJOeOS8DM43rXbp4AZo1n9zCU2qjpcRko0b6/QJi9w=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
-github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
-github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
+github.com/shirou/gopsutil v3.21.11+incompatible h1:+1+c1VGhc88SSonWP6foOcLhvnKlUeu/erjjvaPEYiI=
+github.com/shirou/gopsutil v3.21.11+incompatible/go.mod h1:5b4v6he4MtMOwMlS0TUMTu2PcXUg8+E1lC7eC3UO/RA=
+github.com/stretchr/testify v1.9.0 h1:HtqpIVDClZ4nwg75+f6Lvsy/wHu+3BoSGCbBAcpTsTg=
+github.com/stretchr/testify v1.9.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY=
+github.com/tklauser/go-sysconf v0.3.12 h1:0QaGUFOdQaIVdPgfITYzaTegZvdCjmYO52cSFAEVmqU=
+github.com/tklauser/go-sysconf v0.3.12/go.mod h1:Ho14jnntGE1fpdOqQEEaiKRpvIavV0hSfmBq8nJbHYI=
+github.com/tklauser/go-sysconf v0.3.15 h1:VE89k0criAymJ/Os65CSn1IXaol+1wrsFHEB8Ol49K4=
+github.com/tklauser/go-sysconf v0.3.15/go.mod h1:Dmjwr6tYFIseJw7a3dRLJfsHAMXZ3nEnL/aZY+0IuI4=
+github.com/tklauser/numcpus v0.6.1 h1:ng9scYS7az0Bk4OZLvrNXNSAO2Pxr1XXRAPyjhIx+Fk=
+github.com/tklauser/numcpus v0.6.1/go.mod h1:1XfjsgE2zo8GVw7POkMbHENHzVg3GzmoZ9fESEdAacY=
+github.com/tklauser/numcpus v0.10.0 h1:18njr6LDBk1zuna922MgdjQuJFjrdppsZG60sHGfjso=
+github.com/tklauser/numcpus v0.10.0/go.mod h1:BiTKazU708GQTYF4mB+cmlpT2Is1gLk7XVuEeem8LsQ=
 github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e h1:JVG44RsyaB9T2KIHavMF/ppJZNG9ZpyihvCd0w101no=
 github.com/xo/terminfo v0.0.0-20220910002029-abceb7e1c41e/go.mod h1:RbqR21r5mrJuqunuUZ/Dhy/avygyECGrLceyNeo4LiM=
+github.com/yusufpapurcu/wmi v1.2.4 h1:zFUKzehAFReQwLys1b/iSMl+JQGSCSjtVqQn9bBrPo0=
+github.com/yusufpapurcu/wmi v1.2.4/go.mod h1:SBZ9tNy3G9/m5Oi98Zks0QjeHVDvuK0qfxQmPyzfmi0=
 golang.org/x/exp v0.0.0-20220909182711-5c715a9e8561 h1:MDc5xs78ZrZr3HMQugiXOAkSZtfTpbJLDr/lwfgO53E=
 golang.org/x/exp v0.0.0-20220909182711-5c715a9e8561/go.mod h1:cyybsKvd6eL0RnXn6p/Grxp8F5bW7iYuBgsNCOHpMYE=
+golang.org/x/sys v0.0.0-20190916202348-b4ddaad3f8a3/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
+golang.org/x/sys v0.1.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.8.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.11.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.29.0 h1:TPYlXGxvx1MGTn2GiZDhnjPA9wZzZeGKHHmKhHYvgaU=
 golang.org/x/sys v0.29.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
-golang.org/x/term v0.28.0 h1:/Ts8HFuMR2E6IP/jlo7QVLZHggjKQbhu/7H0LJFr3Gg=
-golang.org/x/term v0.28.0/go.mod h1:Sw/lC2IAUZ92udQNf3WodGtn4k/XoLyZoh8v/8uiwek=
+golang.org/x/sys v0.34.0 h1:H5Y5sJ2L2JRdyv7ROF1he/lPdvFsd0mJHFw2ThKHxLA=
+golang.org/x/sys v0.34.0/go.mod h1:BJP2sWEmIv4KK5OTEluFJCKSidICx8ciO85XgH3Ak8k=
+golang.org/x/term v0.33.0 h1:NuFncQrRcaRvVmgRkvM3j/F00gWIAlcmlB8ACEKmGIg=
+golang.org/x/term v0.33.0/go.mod h1:s18+ql9tYWp1IfpV9DmCtQDDSRBUjKaw9M1eAv5UeF0=
 gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
 gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=

main.go

@@ -1,9 +1,7 @@
 package main
 
 import (
-	"fmt"
 	"os"
-	"regexp"
 	"strings"
 
 	"github.com/coalaura/logger"
@@ -16,7 +14,7 @@ var log = logger.New().WithOptions(logger.Options{
 func main() {
 	isAdmin, err := isElevated()
 	if err != nil {
-		log.FatalF("Unable to resolve elevation status: %v", err)
+		log.Fatalf("Unable to resolve elevation status: %v\n", err)
 
 		os.Exit(1)
 	}
@@ -33,69 +31,38 @@ func main() {
 		log.Note("Already elevated.")
 	}
 
-	cwd, _ := os.Getwd()
-
-	var (
-		app  string
-		keep bool
-	)
-
-	if args[0] == "su" {
-		// Don't need to do anything here if we are already elevated
-		if isAdmin {
-			return
-		}
-
-		// --keep or -k to not close current window
-		for _, arg := range args[0:] {
-			if arg == "-k" || arg == "--keep" {
-				keep = true
-
-				break
-			}
-		}
-	} else {
-		app = strings.Join(args, " ")
-	}
-
-	cmd, arg := build(app, cwd, isAdmin)
+	terminal, shell, err := GetCurrentTerminalAndShell()
+	if err != nil {
+		log.Fatalf("Unable to resolve terminal/shell: %v\n", err)
 
-	if isAdmin {
-		regular(cmd, arg, cwd)
-	} else {
-		elevate(cmd, arg, cwd)
+		os.Exit(1)
 	}
 
-	if app == "" && !keep {
-		terminateParent()
-	}
-}
+	command := args[0]
+	args = args[1:]
 
-func build(app, cwd string, isAdmin bool) (string, string) {
-	var cmd string
+	isSu := command == "su"
+	isWho := command == "who"
 
-	if isAdmin {
-		cmd = app
-	} else {
-		if app == "" {
-			cmd = fmt.Sprintf("%s /K \"cd /d %s\"", shellWithWT(), cwd)
-		} else {
-			rgx := regexp.MustCompile(`\s*;\s*`)
-			app = rgx.ReplaceAllString(app, " && ")
+	// special "sudo su" handling
+	if isSu || isWho {
+		command, args = terminal.Build(shell, "", nil)
+		if len(command) == 0 {
+			log.Fatalf("Unable to resolve terminal/shell: %v\n", err)
 
-			var sh string
+			os.Exit(1)
+		}
 
-			if isShell(app) {
-				sh = shell()
-			} else {
-				sh = shellWithWT()
-			}
+		if isWho {
+			log.Infof("%s %s\n", command, strings.Join(args, " "))
 
-			cmd = fmt.Sprintf("%s /K \"cd /d %s && %s && exit\"", sh, cwd, app)
+			return
 		}
+	} else {
+		command, args = terminal.Build(shell, command, args)
 	}
 
-	parts := strings.Split(cmd, " ")
+	cwd, _ := os.Getwd()
 
-	return parts[0], strings.Join(parts[1:], " ")
+	elevate(command, strings.Join(args, " "), cwd)
 }

run.cmd

@@ -1,5 +1,10 @@
 @echo off
 
+echo Building...
 go build -o temp.exe
-temp.exe su --keep
-rm temp.exe
+
+echo Running...
+temp.exe who
+
+echo Done.
+del temp.exe

shell.go

@@ -2,100 +2,39 @@ package main
 
 import (
 	"fmt"
-	"os"
-	"os/exec"
 	"path/filepath"
 	"strings"
-	"syscall"
-
-	"golang.org/x/sys/windows"
 )
 
-const PROCESS_QUERY_LIMITED_INFORMATION = 0x1000
-
-func getParentProcessName() (string, error) {
-	ppid := uint32(os.Getppid())
-
-	hProcess, err := windows.OpenProcess(PROCESS_QUERY_LIMITED_INFORMATION, false, ppid)
-	if err != nil {
-		return "", err
-	}
-
-	defer windows.CloseHandle(hProcess)
-
-	buffer := make([]uint16, windows.MAX_PATH)
-	size := uint32(len(buffer))
-
-	err = windows.QueryFullProcessImageName(hProcess, 0, &buffer[0], &size)
-	if err != nil {
-		return "", err
-	}
-
-	processName := syscall.UTF16ToString(buffer)
-
-	return filepath.Base(processName), nil
+type Shell struct {
+	Path string
 }
 
-func shell() string {
-	parent, err := getParentProcessName()
-	if err != nil {
-		return "cmd.exe"
+func (s *Shell) Build(cmd string, args []string) []string {
+	if cmd == "" {
+		return []string{s.Path}
 	}
 
-	path, err := exec.LookPath(parent)
-	if err != nil {
-		return "cmd.exe"
-	}
-
-	return path
-}
-
-func shellWithWT() string {
-	sh := shell()
-
-	if wt, err := exec.LookPath("wt"); err == nil {
-		if exists, rect := processExists("WindowsTerminal.exe"); exists {
-			var (
-				size     string
-				position string
-			)
-
-			if w, h := terminalSize(); w != 0 {
-				size = fmt.Sprintf(" --size %d,%d ", w, h)
-			}
-
-			if rect != nil {
-				position = fmt.Sprintf(" --pos %d,%d", rect.Left+10, rect.Top)
-			}
-
-			sh = fmt.Sprintf("%s%s%s %s", wt, size, position, sh)
-		}
-	}
-
-	return sh
-}
+	var command []string
 
-func isShell(app string) bool {
-	test := filepath.Base(app)
+	exe := strings.ToLower(filepath.Base(s.Path))
 
-	if !strings.HasSuffix(app, ".exe") {
-		test += ".exe"
+	switch exe {
+	case "cmd.exe":
+		command = []string{s.Path, "/c"}
+	case "powershell.exe", "pwsh.exe":
+		command = []string{s.Path, "-NoProfile", "-Command"}
+	case "bash.exe", "zsh.exe", "fish.exe", "git-bash.exe", "nu.exe":
+		command = []string{s.Path, "-c"}
 	}
 
-	shells := []string{
-		"cmd.exe",
-		"powershell.exe",
-		"pwsh.exe",
-		"wsl.exe",
-		"bash.exe",
-		"zsh.exe",
+	if len(command) == 0 {
+		return append([]string{cmd}, args...)
 	}
 
-	for _, sh := range shells {
-		if strings.EqualFold(sh, test) {
-			return true
-		}
+	if len(args) > 0 {
+		cmd = fmt.Sprintf("%s %s", cmd, strings.Join(args, " "))
 	}
 
-	return false
+	return append(command, cmd)
 }