Add COMMAND env variable support for shell command

Dockerfile

@@ -14,6 +14,8 @@ RUN apk add --no-cache ca-certificates
 
 COPY --from=builder /app/sshserver /usr/local/bin/sshserver
 
+ENV COMMAND /app/tui
+
 EXPOSE 22
 
 CMD ["/usr/local/bin/sshserver"]

docker-compose.yml

@@ -1,11 +0,0 @@
-services:
-  ssh-server:
-    image: ssh-server
-    build: .
-    ports:
-      - "22:22"
-    volumes:
-      - ../tui:/app/tui:ro
-    cap_add:
-      - SYS_CHROOT
-    restart: unless-stopped

main.go

@@ -81,7 +81,11 @@ func handleChannel(channel ssh.Channel, requests <-chan *ssh.Request) {
 			req.Reply(true, nil)
 		case "shell":
 			req.Reply(true, nil)
-			cmd := exec.Command("/app/tui")
+			command := os.Getenv("COMMAND")
+			if command == "" {
+				command = "/app/tui"
+			}
+			cmd := exec.Command(command)
 			cmd.Env = []string{"PATH=/bin"}
 			cmd.Dir = "/"
 			var err error