From c55af948570dc1a4bd57512fb6c02b269ade915d Mon Sep 17 00:00:00 2001
From: Erik Brakkee <erik@brakkee.org>
Date: Sat, 20 Jul 2024 18:40:22 +0200
Subject: [PATCH] generating key automatically on the agent side. Should be
 done later at the rendez-vous server since there will be many agents running
 on different servers

---
 cmd/agent/agent.go | 51 ++++++++++++++++++++++++++++++++++++++++++----
 1 file changed, 47 insertions(+), 4 deletions(-)

diff --git a/cmd/agent/agent.go b/cmd/agent/agent.go
index 27a43b6..b98637b 100755
--- a/cmd/agent/agent.go
+++ b/cmd/agent/agent.go
@@ -3,6 +3,10 @@ package main
 import (
 	"bufio"
 	"cidebug/pkg/iowrappers"
+	"crypto/rand"
+	"crypto/rsa"
+	"crypto/x509"
+	"encoding/pem"
 	"fmt"
 	"github.com/gorilla/websocket"
 	"io"
@@ -50,7 +54,7 @@ func setWinsize(f *os.File, w, h int) {
 		uintptr(unsafe.Pointer(&struct{ h, w, x, y uint16 }{uint16(h), uint16(w), 0, 0})))
 }
 
-func sshServer() *ssh.Server {
+func sshServer(hostKeyFile string) *ssh.Server {
 	ssh.Handle(func(s ssh.Session) {
 		cmd := exec.Command("bash")
 		ptyReq, winCh, isPty := s.Pty()
@@ -78,12 +82,18 @@ func sshServer() *ssh.Server {
 
 	log.Println("starting ssh server")
 	server := ssh.Server{
-		//Addr:            ":2222",
 		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
 }
 
@@ -116,10 +126,41 @@ func (server ConnectionServer) Run(listener net.Listener) {
 		if err != nil {
 			panic(err)
 		}
-		go echoServer(conn)
+		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)
@@ -138,7 +179,9 @@ func main() {
 	log.Println("Connection established to rendez-vous server, waiting for debug sessions")
 
 	var service AgentService
-	service = ListenerServer(sshServer)
+	service = ListenerServer(func() *ssh.Server {
+		return sshServer("hostkey.pem")
+	})
 	//service = ConnectionServer(echoServer)
 	//service := ConnectionServer(netCatServer)
 	service.Run(listener)