converge/cmd/agent/agent.go

package main

import (
	"bufio"
	"converge/pkg/agent"
	"converge/pkg/comms"
	"converge/pkg/iowrappers"
	"converge/pkg/terminal"
	"converge/pkg/websocketutil"
	"crypto/tls"
	"flag"
	"fmt"
	"github.com/gliderlabs/ssh"
	"github.com/gorilla/websocket"
	"github.com/pkg/sftp"
	"io"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/exec"
	"runtime"
	"strings"
	"time"

	_ "embed"
)

//go:embed hostkey.pem
var hostPrivateKey []byte

func SftpHandler(sess ssh.Session) {
	uid := int(time.Now().UnixMilli())
	agent.Login(uid, sess)
	defer agent.LogOut(uid)

	debugStream := io.Discard
	serverOptions := []sftp.ServerOption{
		sftp.WithDebug(debugStream),
	}
	server, err := sftp.NewServer(
		sess,
		serverOptions...,
	)
	if err != nil {
		log.Printf("sftp tcpserver init error: %s\n", err)
		return
	}
	if err := server.Serve(); err == io.EOF {
		server.Close()
		fmt.Println("sftp client exited session.")
	} else if err != nil {
		fmt.Println("sftp tcpserver completed with error:", err)
	}
}

func passwordAuth(ctx ssh.Context, password string) bool {
	// Replace with your own logic to validate username and password
	return ctx.User() == "abc" && password == "123"
}

func sshServer(hostKeyFile string, shellCommand string) *ssh.Server {
	ssh.Handle(func(s ssh.Session) {
		workingDirectory, _ := os.Getwd()
		env := append(os.Environ(), fmt.Sprintf("agentdir=%s", workingDirectory))
		process, err := terminal.PtySpawner.Start(s, env, shellCommand)
		if err != nil {
			panic(err)
		}
		uid := int(time.Now().UnixMilli())
		agent.Login(uid, s)
		iowrappers.SynchronizeStreams(process.Pipe(), s)
		agent.LogOut(uid)
		process.Wait()
		process.Wait()
	})

	log.Println("starting ssh server, waiting for debug sessions")
	server := ssh.Server{
		PasswordHandler: passwordAuth,
		SubsystemHandlers: map[string]ssh.SubsystemHandler{
			"sftp": SftpHandler,
		},
	}
	//err := generateHostKey(hostKeyFile, 2048)
	//if err != nil {
	//	log.Printf("Could not create host key file '%s': %v", hostKeyFile, err)
	//}
	//option := ssh.HostKeyFile(hostKeyFile)

	option := ssh.HostKeyPEM(hostPrivateKey)
	option(&server)

	return &server
}

func echoServer(conn io.ReadWriter) {
	log.Println("Echo service started")
	io.Copy(conn, conn)
}

func netCatServer(conn io.ReadWriter) {
	stdio := bufio.NewReadWriter(
		bufio.NewReaderSize(os.Stdin, 0),
		bufio.NewWriterSize(os.Stdout, 0))
	iowrappers.SynchronizeStreams(conn, stdio)
}

type AgentService interface {
	Run(listener net.Listener)
}

type ListenerServer func() *ssh.Server

func (server ListenerServer) Run(listener net.Listener) {
	server().Serve(listener)
}

type ConnectionServer func(conn io.ReadWriter)

func (server ConnectionServer) Run(listener net.Listener) {
	for {
		conn, err := listener.Accept()
		if err != nil {
			panic(err)
		}
		go server(conn)
	}
}

type ReaderFunc func(p []byte) (n int, err error)

func (f ReaderFunc) Read(p []byte) (n int, err error) {
	return f(p)
}

func main() {

	usage := "agent [options] <wsUrl> \n" +
		"\n" +
		"Run agent with <wsUrl> of the form ws[s]://<host>[:port]/agent/<ID>\n" +
		"Here <ID> is the unique id of the agent that allows rendez-vous with an end-user.\n" +
		"The end-user must specify the same id when connecting using ssh.\n"

	flag.Usage = func() {
		fmt.Fprintf(os.Stderr, usage+"\n")
		flag.PrintDefaults()
	}
	advanceWarningTime := flag.Duration("warning-time", 5*time.Minute, "advance warning time before sessio ends")
	agentExpriryTime := flag.Duration("expiry-time", 10*time.Minute, "expiry time of the session")
	tickerInterval := flag.Duration("check-interval", 60*time.Second, "interval at which expiry is checked")
	insecure := flag.Bool("insecure", false, "allow invalid certificates")

	flag.Parse()
	if flag.NArg() != 1 {
		flag.Usage()
		os.Exit(1)
	}
	wsURL := flag.Arg(0)

	dialer := websocket.Dialer{
		Proxy:            http.ProxyFromEnvironment,
		HandshakeTimeout: 45 * time.Second,
	}
	if *insecure {
		dialer.TLSClientConfig = &tls.Config{InsecureSkipVerify: true}
	}
	conn, _, err := dialer.Dial(wsURL, nil)
	if err != nil {
		log.Println("WebSocket connection error:", err)
		return
	}
	wsConn := websocketutil.NewWebSocketConn(conn)
	defer wsConn.Close()

	commChannel, err := comms.NewCommChannel(comms.Agent, wsConn)
	if err != nil {
		panic(err)
	}

	var service AgentService
	shells := []string{"bash", "sh", "ash", "ksh", "zsh", "fish", "tcsh", "csh"}
	if runtime.GOOS == "windows" {
		shells = []string{"powershell", "bash"}
	}

	shell := ""
	for _, candidate := range shells {
		shell, err = exec.LookPath(candidate)
		if err == nil {
			break
		}
	}
	if shell == "" {
		log.Printf("Cannot find a shell in %v", shells)
		os.Exit(1)
	}
	log.Printf("Using shell %s for remote sessions", shell)
	service = ListenerServer(func() *ssh.Server {
		return sshServer("hostkey.pem", shell)
	})
	//service = ConnectionServer(netCatServer)
	//service = ConnectionServer(echoServer)
	log.Println()
	log.Printf("Clients should use the following commands to connect to this agent:")
	log.Println()
	clientUrl := strings.ReplaceAll(wsURL, "/agent/", "/client/")
	sshCommand := fmt.Sprintf("ssh -oServerAliveInterval=10 -oProxyCommand=\"wsproxy %s\"  abc@localhost",
		clientUrl)
	sftpCommand := fmt.Sprintf("sftp -oServerAliveInterval=10 -oProxyCommand=\"wsproxy %s\" abc@localhost",
		clientUrl)
	log.Println("  # For SSH")
	log.Println("  " + sshCommand)
	log.Println()
	log.Println("  # for SFTP")
	log.Println("  " + sftpCommand)
	log.Println()
	urlObject, _ := url.Parse(wsURL)

	log.Printf("wsproxy can be downloaded from %s",
		strings.ReplaceAll(urlObject.Scheme, "ws", "http")+
			"://"+urlObject.Host+"/docs/wsproxy")
	log.Println()

	agent.ConfigureAgent(commChannel, *advanceWarningTime, *agentExpriryTime, *tickerInterval)
	service.Run(commChannel.Session)
}