Playing with ssh

client/ssh/server.go

@@ -0,0 +1,123 @@
+package main
+
+import (
+	"crypto/rand"
+	"crypto/rsa"
+	"fmt"
+	"github.com/gliderlabs/ssh"
+	gossh "golang.org/x/crypto/ssh"
+	"io"
+	"net"
+	"sync"
+)
+
+type Server struct {
+	listener    net.Listener
+	allowedKeys []ssh.PublicKey
+	mu          sync.Mutex
+}
+
+func (srv *Server) UpdateKeys(newKeys []string) error {
+	srv.mu.Lock()
+	defer srv.mu.Unlock()
+
+	srv.allowedKeys = make([]ssh.PublicKey, len(newKeys))
+	for _, strKey := range newKeys {
+		parsedKey, _, _, _, err := ssh.ParseAuthorizedKey([]byte(strKey))
+		if err != nil {
+			return err
+		}
+		srv.allowedKeys = append(srv.allowedKeys, parsedKey)
+	}
+
+	return nil
+}
+
+func NewSSHServer() (*Server, error) {
+	ln, err := net.Listen("tcp", ":2222")
+	if err != nil {
+		return nil, err
+	}
+	return &Server{listener: ln, mu: sync.Mutex{}}, nil
+}
+
+// Stop stops SSH server. Blocking
+func (srv *Server) Stop() error {
+	err := srv.listener.Close()
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+// Start starts SSH server. Blocking
+func (srv *Server) Start() error {
+	handler := func(s ssh.Session) {
+		authorizedKey := gossh.MarshalAuthorizedKey(s.PublicKey())
+		io.WriteString(s, fmt.Sprintf("public key used by %s:\n", s.User()))
+		s.Write(authorizedKey)
+	}
+
+	publicKeyOption := ssh.PublicKeyAuth(func(ctx ssh.Context, key ssh.PublicKey) bool {
+		srv.mu.Lock()
+		defer srv.mu.Unlock()
+
+		for _, allowed := range srv.allowedKeys {
+			if ssh.KeysEqual(allowed, key) {
+				return true
+			}
+		}
+
+		return false
+	})
+
+	err := ssh.Serve(srv.listener, handler, publicKeyOption)
+	if err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func main() {
+
+	server, err := NewSSHServer()
+	if err != nil {
+		return
+	}
+
+	err = server.UpdateKeys([]string{"ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIIwAoefixS03tYDfNuFfNRMO2syYfkw/C/76m8LS8xum"})
+	if err != nil {
+		return
+	}
+
+	err = server.Start()
+	if err != nil {
+		// will throw error when Stop has been called
+	}
+}
+
+// generatePrivateKey creates RSA Private Key of specified byte size
+func generatePrivateKey(bitSize int) (*rsa.PrivateKey, error) {
+	privateKey, err := rsa.GenerateKey(rand.Reader, bitSize)
+	if err != nil {
+		return nil, err
+	}
+
+	err = privateKey.Validate()
+	if err != nil {
+		return nil, err
+	}
+
+	return privateKey, nil
+}
+
+// generatePublicKey takes a rsa.PublicKey and return bytes suitable for writing to .pub file
+// returns the key in format format "ssh-rsa ..."
+func generatePublicKey(privateKey *rsa.PublicKey) ([]byte, error) {
+	publicRsaKey, err := gossh.NewPublicKey(privateKey)
+	if err != nil {
+		return nil, err
+	}
+	return gossh.MarshalAuthorizedKey(publicRsaKey), nil
+}