Implement interactive shell

cmd/wsh/main.go

@@ -8,6 +8,7 @@ import (
 	"strings"
 
 	"github.com/jessevdk/go-flags"
+	"golang.org/x/term"
 
 	"github.com/xconnio/berncrypt/go"
 	"github.com/xconnio/wamp-webrtc-go"
@@ -18,6 +19,8 @@ import (
 
 const (
 	defaultRealm             = "wampshell"
+	procedureInteractive     = "wampshell.shell.interactive"
+	procedureExec            = "wampshell.shell.exec"
 	procedureWebRTCOffer     = "wampshell.webrtc.offer"
 	topicOffererOnCandidate  = "wampshell.webrtc.offerer.on_candidate"
 	topicAnswererOnCandidate = "wampshell.webrtc.answerer.on_candidate"
@@ -65,11 +68,103 @@ func exchangeKeys(session *xconn.Session) (*keyPair, error) {
 	}, nil
 }
 
+func startInteractiveShell(session *xconn.Session, keys *keyPair) error {
+	fd := int(os.Stdin.Fd())
+	oldState, err := term.MakeRaw(fd)
+	if err != nil {
+		return fmt.Errorf("failed to set raw mode: %w", err)
+	}
+	defer func() { _ = term.Restore(fd, oldState) }()
+
+	firstProgress := true
+
+	call := session.Call(procedureInteractive).
+		ProgressSender(func(ctx context.Context) *xconn.Progress {
+			if firstProgress {
+				firstProgress = false
+				return xconn.NewProgress()
+			}
+
+			buf := make([]byte, 1024)
+			n, err := os.Stdin.Read(buf)
+			if err != nil {
+				return xconn.NewFinalProgress()
+			}
+
+			ciphertext, nonce, err := berncrypt.EncryptChaCha20Poly1305(buf[:n], keys.send)
+			if err != nil {
+				fmt.Printf("encryption error: %s", err)
+				os.Exit(1)
+			}
+			payload := append(nonce, ciphertext...)
+
+			return xconn.NewProgress(payload)
+		}).
+		ProgressReceiver(func(result *xconn.InvocationResult) {
+			if len(result.Args) > 0 {
+				encData := result.Args[0].([]byte)
+
+				if len(encData) < 12 {
+					fmt.Fprintln(os.Stderr, "invalid payload from server")
+					os.Exit(1)
+				}
+
+				plain, err := berncrypt.DecryptChaCha20Poly1305(encData[12:], encData[:12], keys.receive)
+				if err != nil {
+					_ = fmt.Errorf("decryption error: %w", err)
+				}
+
+				os.Stdout.Write(plain)
+			} else {
+				err = term.Restore(fd, oldState)
+				if err != nil {
+					return
+				}
+				os.Exit(0)
+			}
+		}).Do()
+
+	if call.Err != nil {
+		log.Fatalf("Shell error: %s", call.Err)
+	}
+	return nil
+}
+
+func runCommand(session *xconn.Session, keys *keyPair, args []string) error {
+	b := []byte(strings.Join(args, " "))
+
+	ciphertext, nonce, err := berncrypt.EncryptChaCha20Poly1305(b, keys.send)
+	if err != nil {
+		return fmt.Errorf("encryption error: %w", err)
+	}
+
+	payload := append(nonce, ciphertext...)
+
+	cmdResponse := session.Call(procedureExec).Args(payload).Do()
+	if cmdResponse.Err != nil {
+		return fmt.Errorf("command execution error: %w", cmdResponse.Err)
+	}
+
+	output, err := cmdResponse.Args.Bytes(0)
+	if err != nil {
+		fmt.Printf("Output parsing error: %v", err)
+		os.Exit(1)
+	}
+
+	plain, err := berncrypt.DecryptChaCha20Poly1305(output[12:], output[:12], keys.receive)
+	if err != nil {
+		return fmt.Errorf("decryption error: %w", err)
+	}
+	fmt.Print(string(plain))
+	return nil
+}
+
 type Options struct {
-	PeerToPeer bool `long:"p2p" description:"Use WebRTC for peer-to-peer connection"`
-	Args       struct {
+	Interactive bool `short:"i" long:"interactive" description:"Force interactive shell"`
+	PeerToPeer  bool `long:"p2p" description:"Use WebRTC for peer-to-peer connection"`
+	Args        struct {
 		Target string   `positional-arg-name:"host" required:"true"`
-		Cmd    []string `positional-arg-name:"command" required:"true"`
+		Cmd    []string `positional-arg-name:"command"`
 	} `positional-args:"yes"`
 }
 
@@ -105,11 +200,6 @@ func main() {
 		port = "8022"
 	}
 
-	anyArgs := make([]any, len(args))
-	for i, a := range args {
-		anyArgs[i] = a
-	}
-
 	privateKey, err := wampshell.ReadPrivateKeyFromFile()
 	if err != nil {
 		fmt.Printf("Error reading private key: %v\n", err)
@@ -118,8 +208,7 @@ func main() {
 
 	authenticator, err := auth.NewCryptoSignAuthenticator("", privateKey, nil)
 	if err != nil {
-		fmt.Printf("Error creating crypto sign authenticator: %v\n", err)
-		os.Exit(1)
+		log.Fatal("Error creating crypto sign authenticator:", err)
 	}
 
 	client := xconn.Client{
@@ -156,33 +245,15 @@ func main() {
 		panic(err)
 	}
 
-	b := []byte(strings.Join(args, " "))
-
-	ciphertext, nonce, err := berncrypt.EncryptChaCha20Poly1305(b, keys.send)
-	if err != nil {
-		panic(err)
-	}
-
-	payload := make([]byte, len(nonce)+len(ciphertext))
-	copy(payload, nonce)
-	copy(payload[len(nonce):], ciphertext)
-
-	cmdResponse := session.Call("wampshell.shell.exec").Args(payload).Do()
-	if cmdResponse.Err != nil {
-		fmt.Printf("Command execution error: %v\n", cmdResponse.Err)
-		os.Exit(1)
-	}
-
-	output, err := cmdResponse.Args.Bytes(0)
-	if err != nil {
-		fmt.Printf("Output parsing error: %v\n", err)
-		os.Exit(1)
+	if opts.Interactive || len(args) == 0 {
+		err := startInteractiveShell(session, keys)
+		if err != nil {
+			log.Fatal(err)
+		}
 	}
 
-	a, err := berncrypt.DecryptChaCha20Poly1305(output[12:], output[:12], keys.receive)
+	err = runCommand(session, keys, args)
 	if err != nil {
-		panic(err)
+		log.Fatal(err)
 	}
-	fmt.Print(string(a))
-
 }