From b699f1d6c335a8c451476e31f59997e6cb82e8b5 Mon Sep 17 00:00:00 2001
From: Miguel Mota <miguelmota2@gmail.com>
Date: Wed, 16 Dec 2020 22:41:46 -0800
Subject: [PATCH] ssh-server: Use hash of public key for config file path

---
 README.md         |  6 ++++++
 pkg/ssh/server.go | 51 +++++++++++++++++++++++++++++++----------------
 2 files changed, 40 insertions(+), 17 deletions(-)

diff --git a/README.md b/README.md
index e5bfff1..2fbf011 100644
--- a/README.md
+++ b/README.md
@@ -27,6 +27,12 @@ The interface is inspired by [`htop`](https://en.wikipedia.org/wiki/Htop) and sh
 
 <img src="https://user-images.githubusercontent.com/168240/41806841-043c0ca6-767a-11e8-9c51-df9fc64b3b5c.png" alt="currency convert menu" width="880" />
 
+## Demo
+
+```bash
+ssh cointop.sh
+```
+
 In action
 
 <img src="https://user-images.githubusercontent.com/168240/39569570-75b1547c-4e7a-11e8-8eac-552abaa431f0.gif" alt="screencast" width="880" />
diff --git a/pkg/ssh/server.go b/pkg/ssh/server.go
index daa415c..de35230 100644
--- a/pkg/ssh/server.go
+++ b/pkg/ssh/server.go
@@ -4,6 +4,7 @@ package ssh
 
 import (
 	"context"
+	"crypto/sha256"
 	"errors"
 	"fmt"
 	"io"
@@ -20,10 +21,10 @@ import (
 	gossh "golang.org/x/crypto/ssh"
 )
 
-// DefaultHostKeyFile ...
+// DefaultHostKeyFile is default SSH key path
 var DefaultHostKeyFile = "~/.ssh/id_rsa"
 
-// Config ...
+// Config is config struct
 type Config struct {
 	Port             uint
 	Address          string
@@ -32,7 +33,7 @@ type Config struct {
 	HostKeyFile      string
 }
 
-// Server ...
+// Server is server struct
 type Server struct {
 	port             uint
 	address          string
@@ -42,7 +43,7 @@ type Server struct {
 	hostKeyFile      string
 }
 
-// NewServer ...
+// NewServer returns a new server instance
 func NewServer(config *Config) *Server {
 	hostKeyFile := DefaultHostKeyFile
 	if config.HostKeyFile != "" {
@@ -60,7 +61,7 @@ func NewServer(config *Config) *Server {
 	}
 }
 
-// ListenAndServe ...
+// ListenAndServe starts the server
 func (s *Server) ListenAndServe() error {
 	s.sshServer = &ssh.Server{
 		Addr:        fmt.Sprintf("%s:%v", s.address, s.port),
@@ -74,13 +75,32 @@ func (s *Server) ListenAndServe() error {
 				return
 			}
 
-			tempDir, err := createTempDir()
-			if err != nil {
-				fmt.Println(err)
-				return
+			configDir := ""
+			pubKey := sshSession.PublicKey()
+			if pubKey != nil {
+				pubBytes := pubKey.Marshal()
+				if len(pubBytes) > 0 {
+					hash := sha256.Sum256(pubBytes)
+					configDir = fmt.Sprintf("/tmp/cointop_config/%x", hash)
+					err := os.MkdirAll(configDir, 0700)
+					if err != nil {
+						fmt.Println(err)
+						return
+					}
+				}
+			}
+
+			if configDir == "" {
+				tempDir, err := createTempDir()
+				if err != nil {
+					fmt.Println(err)
+					return
+				}
+				configDir = tempDir
+				defer os.RemoveAll(configDir)
 			}
 
-			configPath := fmt.Sprintf("%s/config", tempDir)
+			configPath := fmt.Sprintf("%s/config", configDir)
 			colorsDir := pathutil.NormalizePath("~/.config/cointop/colors")
 
 			cmdCtx, cancelCmd := context.WithCancel(sshSession.Context())
@@ -90,7 +110,7 @@ func (s *Server) ListenAndServe() error {
 				"--reset",
 				"--silent",
 				"--cache-dir",
-				tempDir,
+				configDir,
 				"--config",
 				configPath,
 				"--colors-dir",
@@ -128,10 +148,8 @@ func (s *Server) ListenAndServe() error {
 			io.Copy(sshSession, f)
 			f.Close()
 			cmd.Wait()
-			os.Remove(configPath)
 		},
 		PtyCallback: func(ctx ssh.Context, pty ssh.Pty) bool {
-			// TODO: check public key hash
 			return true
 		},
 		PublicKeyHandler: func(ctx ssh.Context, key ssh.PublicKey) bool {
@@ -157,19 +175,18 @@ func (s *Server) ListenAndServe() error {
 	return s.sshServer.ListenAndServe()
 }
 
-// Shutdown ...
+// Shutdown shuts down the server
 func (s *Server) Shutdown() {
 	s.sshServer.Close()
 }
 
-// setWinsize ...
+// setWinsize sets the PTY window size
 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})))
 }
 
-// createTempDir ...
-// TODO: load saved configuration based on ssh public key hash
+// createTempDir creates a temporary directory
 func createTempDir() (string, error) {
 	return ioutil.TempDir("", "")
 }