converge/cmd/agent/agent.go

package main

import (
	"bufio"
	"cidebug/pkg/agent"
	"cidebug/pkg/iowrappers"
	"cidebug/pkg/websocketutil"
	"fmt"
	"github.com/gorilla/websocket"
	"io"
	"log"
	"net"
	"os"
	"os/exec"
	"syscall"
	"time"
	"unsafe"

	"github.com/creack/pty"
	"github.com/gliderlabs/ssh"
	"github.com/hashicorp/yamux"
	"github.com/pkg/sftp"

	_ "embed"
)

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

func SftpHandler(sess ssh.Session) {
	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 setWinsize(f *os.File, w, h int) {
	syscall.Syscall(syscall.SYS_IOCTL, f.Fd(), uintptr(syscall.TIOCSWINSZ),
		uintptr(unsafe.Pointer(&struct{ h, w, x, y uint16 }{uint16(h), uint16(w), 0, 0})))
}

func sshServer(hostKeyFile string) *ssh.Server {
	ssh.Handle(func(s ssh.Session) {
		// TODO shell should  be made configurable
		cmd := exec.Command("bash")
		ptyReq, winCh, isPty := s.Pty()
		if isPty {
			workingDirectory, _ := os.Getwd()
			cmd.Env = append(os.Environ(),
				fmt.Sprintf("TERM=%s", ptyReq.Term),
				fmt.Sprintf("agentdir=%s", workingDirectory))
			f, err := pty.Start(cmd)
			if err != nil {
				panic(err)
			}
			go func() {
				for win := range winCh {
					setWinsize(f, win.Width, win.Height)
				}
			}()
			uid := int(time.Now().UnixMilli())
			agent.Login(uid, s)

			go func() {
				io.Copy(f, s) // stdin
			}()
			io.Copy(s, f) // stdout
			cmd.Wait()
			agent.LogOut(uid)
		} else {
			io.WriteString(s, "No PTY requested.\n")
			s.Exit(1)
		}
	})

	log.Println("starting ssh server")
	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() {
	wsURL := os.Args[1]

	advanceWarningTime := 10 * time.Minute
	agentExpriryTime := 30 * time.Minute
	tickerInterval := 60 * time.Second
	agent.ConfigureAgent(advanceWarningTime, agentExpriryTime, tickerInterval)

	conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil)
	if err != nil {
		log.Println("WebSocket connection error:", err)
		return
	}
	wsConn := websocketutil.NewWebSocketConn(conn)
	defer wsConn.Close()

	listener, err := yamux.Server(wsConn, nil)
	if err != nil {
		panic(err)
	}

	// Need to create listener implementation that aactually listens for websocket connections.
	var service AgentService
	service = ListenerServer(func() *ssh.Server {
		return sshServer("hostkey.pem")
	})
	//service = ConnectionServer(netCatServer)
	//service = ConnectionServer(echoServer)
	service.Run(listener)
}