client-exec CLI

cmd/byre/clientexec.go

@@ -0,0 +1,88 @@
+package main
+
+import (
+	"bytes"
+	"encoding/json"
+	"flag"
+	"fmt"
+	"net/http"
+	"os"
+	"strings"
+	"time"
+
+	"g.fc.run/ai/byre/internal/api"
+)
+
+const (
+	execUnavailable = 75
+	execTimedOut    = 124
+	execGrace       = time.Minute
+)
+
+func clientExecCmd(args []string) {
+	fs := flag.NewFlagSet("client-exec", flag.ExitOnError)
+	client := fs.String("client", "", "the client's name or id; the most recently active one when absent")
+	timeout := fs.Int("timeout", 30, "kill the command after this many seconds")
+	fs.Parse(args)
+	if fs.NArg() == 0 || *timeout <= 0 {
+		usage()
+	}
+	token := os.Getenv("BYRE_TOKEN")
+	if token == "" {
+		execFail(1, "BYRE_TOKEN is not set; byre sets it for everything it runs")
+	}
+
+	words := make([]string, fs.NArg())
+	for i, w := range fs.Args() {
+		words[i] = "'" + strings.ReplaceAll(w, "'", `'\''`) + "'"
+	}
+	req := api.ClientExec{Command: strings.Join(words, " "), TimeoutMs: *timeout * 1000}
+	if *client != "" {
+		req.Client = client
+	}
+	body, err := json.Marshal(req)
+	if err != nil {
+		execFail(1, err.Error())
+	}
+	hr, err := http.NewRequest(http.MethodPost, loopbackURL+"/api/clients/exec", bytes.NewReader(body))
+	if err != nil {
+		execFail(1, err.Error())
+	}
+	hr.Header.Set("Authorization", "Bearer "+token)
+	hr.Header.Set("Content-Type", "application/json")
+	hc := &http.Client{Timeout: time.Duration(*timeout)*time.Second + execGrace}
+	resp, err := hc.Do(hr)
+	if err != nil {
+		execFail(execUnavailable, err.Error())
+	}
+	defer resp.Body.Close()
+
+	if resp.StatusCode != http.StatusOK {
+		e := api.Error{}
+		err = json.NewDecoder(resp.Body).Decode(&e)
+		if err != nil {
+			e.Message = resp.Status
+		}
+		code := 1
+		if resp.StatusCode == http.StatusServiceUnavailable {
+			code = execUnavailable
+		}
+		execFail(code, e.Message)
+	}
+	res := api.ClientExecResult{}
+	err = json.NewDecoder(resp.Body).Decode(&res)
+	if err != nil {
+		execFail(1, fmt.Sprintf("decode reply: %s", err))
+	}
+	os.Stdout.Write(res.Stdout)
+	os.Stderr.Write(res.Stderr)
+	if res.ExitCode == -1 {
+		execFail(execTimedOut, fmt.Sprintf("killed on %s after %ds", res.Client, *timeout))
+	}
+	os.Exit(res.ExitCode)
+}
+
+func execFail(code int, msg string) {
+	fmt.Fprintln(os.Stderr, "byre client-exec: "+msg)
+	os.Exit(code)
+}

cmd/byre/main.go

@@ -26,13 +26,17 @@ func main() {
 		grepCmd(os.Args[2:])
 	case "list":
 		listCmd(os.Args[2:])
+	case "client-exec":
+		clientExecCmd(os.Args[2:])
+	case "screen":
+		screenCmd()
 	default:
 		usage()
 	}
 }
 
 func usage() {
-	fmt.Fprintln(os.Stderr, "usage: byre serve | byre install | byre check | byre run [--resume ID] | byre signing CHECKOUT | byre grep [--glob GLOB] [--ignore-case] [--files-only] PATTERN PATH | byre list PATH")
+	fmt.Fprintln(os.Stderr, "usage: byre serve | byre install | byre check | byre run [--resume ID] | byre signing CHECKOUT | byre grep [--glob GLOB] [--ignore-case] [--files-only] PATTERN PATH | byre list PATH | byre client-exec [--client NAME] [--timeout SECS] -- CMD... | byre screen")
 	os.Exit(2)
 }
 

cmd/byre/screen.go

@@ -0,0 +1,39 @@
+package main
+
+import (
+	"fmt"
+	"os"
+	"os/exec"
+
+	"howett.net/plist"
+)
+
+type ioregRoot struct {
+	IOConsoleUsers []struct {
+		UID    int  `plist:"kCGSSessionUserIDKey"`
+		Locked bool `plist:"CGSSessionScreenIsLocked"`
+	} `plist:"IOConsoleUsers"`
+}
+
+func screenCmd() {
+	out, err := exec.Command("ioreg", "-n", "Root", "-d1", "-a").Output()
+	if err != nil {
+		fatal("ioreg", "err", err)
+	}
+	root := ioregRoot{}
+	_, err = plist.Unmarshal(out, &root)
+	if err != nil {
+		fatal("decode ioreg", "err", err)
+	}
+	state := "no-session"
+	for _, u := range root.IOConsoleUsers {
+		if u.UID != os.Getuid() {
+			continue
+		}
+		state = "unlocked"
+		if u.Locked {
+			state = "locked"
+		}
+	}
+	fmt.Println(state)
+}

cmd/byre/serve.go

@@ -73,6 +73,9 @@ func serveCmd(args []string) {
 	if err := os.WriteFile(filepath.Join(*dataDir, "loopback-token"), []byte(token), 0o600); err != nil {
 		fatal("write loopback credential", "err", err)
 	}
+	if err := os.Setenv("BYRE_TOKEN", token); err != nil {
+		fatal("export loopback credential", "err", err)
+	}
 	home, err := os.UserHomeDir()
 	if err != nil {
 		fatal("home directory", "err", err)
@@ -101,7 +104,7 @@ func serveCmd(args []string) {
 	}
 	defer ts.Close()
 
-	srv, err := server.New(*dataDir, *ports, ts)
+	srv, err := server.New(*dataDir, *ports, ts, token)
 	if err != nil {
 		fatal("start server", "err", err)
 	}
@@ -115,7 +118,7 @@ func serveCmd(args []string) {
 	go updateDaily(ctx)
 
 	base := func(net.Listener) context.Context { return ctx }
-	loopback := &http.Server{Addr: *listenAddr, Handler: srv.LoopbackHandler(token), BaseContext: base}
+	loopback := &http.Server{Addr: *listenAddr, Handler: srv.LoopbackHandler(), BaseContext: base}
 	tailnet := &http.Server{Handler: srv.Handler(), BaseContext: base}
 
 	go func() {