converge/cmd/agent/agent.go

package main

import (
	"bufio"
	"cidebug/pkg/iowrappers"
	"cidebug/pkg/websocketutil"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"github.com/gorilla/websocket"
	"io"
	"log"
	"net"
	"os"
	"os/exec"
	"syscall"
	"unsafe"

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

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) {
		cmd := exec.Command("bash")
		ptyReq, winCh, isPty := s.Pty()
		if isPty {
			cmd.Env = append(os.Environ(), fmt.Sprintf("TERM=%s", ptyReq.Term))
			f, err := pty.Start(cmd)
			if err != nil {
				panic(err)
			}
			go func() {
				for win := range winCh {
					setWinsize(f, win.Width, win.Height)
				}
			}()
			go func() {
				io.Copy(f, s) // stdin
			}()
			io.Copy(s, f) // stdout
			cmd.Wait()
		} 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(&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 generateHostKey(filename string, bitSize int) error {
	if _, err := os.Stat(filename); !os.IsNotExist(err) {
		log.Printf("Reusing key file '%s'", filename)
		return nil
	}
	privateKey, err := rsa.GenerateKey(rand.Reader, bitSize)
	if err != nil {
		return err
	}

	privateKeyPEM := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(privateKey),
	}

	privateKeyFile, err := os.Create(filename)
	if err != nil {
		return err
	}
	defer privateKeyFile.Close()

	log.Printf("Generating key '%s'", filename)
	return pem.Encode(privateKeyFile, privateKeyPEM)
}

func main() {
	wsURL := os.Args[1]
	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)
}