From ed922a235f8b73b84fb988b3d205966a4056f328 Mon Sep 17 00:00:00 2001
From: Erik Brakkee <erik@brakkee.org>
Date: Sun, 21 Jul 2024 11:22:05 +0200
Subject: [PATCH] agent now uses a fixed host key using the go embed package.
 Printing welcome message when user logs in.

---
 cmd/agent/agent.go  | 48 +++++++++++++--------------------------------
 pkg/sshutils/ssh.go | 35 +++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 34 deletions(-)
 create mode 100644 pkg/sshutils/ssh.go

diff --git a/cmd/agent/agent.go b/cmd/agent/agent.go
index 3df9cf5..ddc70c3 100755
--- a/cmd/agent/agent.go
+++ b/cmd/agent/agent.go
@@ -4,10 +4,6 @@ import (
 	"bufio"
 	"cidebug/pkg/iowrappers"
 	"cidebug/pkg/websocketutil"
-	"crypto/rand"
-	"crypto/rsa"
-	"crypto/x509"
-	"encoding/pem"
 	"fmt"
 	"github.com/gorilla/websocket"
 	"io"
@@ -22,8 +18,13 @@ import (
 	"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{
@@ -57,6 +58,8 @@ func setWinsize(f *os.File, w, h int) {
 
 func sshServer(hostKeyFile string) *ssh.Server {
 	ssh.Handle(func(s ssh.Session) {
+		hostname, _ := os.Hostname()
+		io.WriteString(s, fmt.Sprintf("Your are now on %s\n\n", hostname))
 		cmd := exec.Command("bash")
 		ptyReq, winCh, isPty := s.Pty()
 		if isPty {
@@ -88,11 +91,13 @@ func sshServer(hostKeyFile string) *ssh.Server {
 			"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)
+	//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
@@ -138,31 +143,6 @@ 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)
diff --git a/pkg/sshutils/ssh.go b/pkg/sshutils/ssh.go
new file mode 100644
index 0000000..3111c82
--- /dev/null
+++ b/pkg/sshutils/ssh.go
@@ -0,0 +1,35 @@
+package sshutils
+
+import (
+	"crypto/rand"
+	"crypto/rsa"
+	"crypto/x509"
+	"encoding/pem"
+	"log"
+	"os"
+)
+
+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)
+}