Small refactor

config.go

@@ -1,43 +0,0 @@
-package main
-
-import (
-	"fmt"
-
-	"github.com/gliderlabs/ssh"
-	"github.com/jinzhu/gorm"
-	gossh "golang.org/x/crypto/ssh"
-)
-
-type Config struct {
-	clientConfig *gossh.ClientConfig
-	remoteAddr   string
-}
-
-func getConfig(s ssh.Session, db *gorm.DB) (*Config, error) {
-	var host Host
-	db.Where("name = ?", s.User()).Find(&host)
-	if host.Name == "" {
-		// FIXME: add available hosts
-		return nil, fmt.Errorf("No such target: %q", s.User())
-	}
-
-	config := Config{
-		remoteAddr: host.Addr,
-		clientConfig: &gossh.ClientConfig{
-			User:            host.User,
-			HostKeyCallback: gossh.InsecureIgnoreHostKey(),
-			Auth:            []gossh.AuthMethod{},
-		},
-	}
-	if host.Password != "" {
-		config.clientConfig.Auth = append(config.clientConfig.Auth, gossh.Password(host.Password))
-	}
-	if host.PrivKey != nil {
-		return nil, fmt.Errorf("auth by priv key is not yet implemented")
-	}
-	if len(config.clientConfig.Auth) == 0 {
-		return nil, fmt.Errorf("no valid authentication method for host %q", s.User())
-	}
-
-	return &config, nil
-}

db.go

@@ -1,9 +1,19 @@
 package main
 
-import "github.com/jinzhu/gorm"
+import (
+	"fmt"
 
-type Key struct {
+	"github.com/gliderlabs/ssh"
+	"github.com/jinzhu/gorm"
+	gossh "golang.org/x/crypto/ssh"
+)
+
+type SSHKey struct {
 	gorm.Model
+	Type        string
+	Fingerprint string
+	PrivKey     []byte
+	PubKey      []byte
 }
 
 type Host struct {
@@ -13,17 +23,17 @@ type Host struct {
 	User        string
 	Password    string
 	Fingerprint string
-	PrivKey     *Key
+	PrivKey     *SSHKey
 }
 
 type User struct {
 	gorm.Model
-	Keys []Key
+	SSHKeys []SSHKey
 }
 
 func dbInit(db *gorm.DB) error {
 	db.AutoMigrate(&User{})
-	db.AutoMigrate(&Key{})
+	db.AutoMigrate(&SSHKey{})
 	db.AutoMigrate(&Host{})
 	return nil
 }
@@ -35,3 +45,31 @@ func dbDemo(db *gorm.DB) error {
 	db.FirstOrCreate(&host3, &Host{Name: "ssh-chat", Addr: "chat.shazow.net:22", User: "test", Fingerprint: "MD5:e5:d5:d1:75:90:38:42:f6:c7:03:d7:d0:56:7d:6a:db"})
 	return nil
 }
+
+func RemoteHostFromSession(s ssh.Session, db *gorm.DB) (*Host, error) {
+	var host Host
+	db.Where("name = ?", s.User()).Find(&host)
+	if host.Name == "" {
+		// FIXME: add available hosts
+		return nil, fmt.Errorf("No such target: %q", s.User())
+	}
+	return &host, nil
+}
+
+func (host *Host) ClientConfig(_ ssh.Session) (*gossh.ClientConfig, error) {
+	config := gossh.ClientConfig{
+		User:            host.User,
+		HostKeyCallback: gossh.InsecureIgnoreHostKey(),
+		Auth:            []gossh.AuthMethod{},
+	}
+	if host.Password != "" {
+		config.Auth = append(config.Auth, gossh.Password(host.Password))
+	}
+	if host.PrivKey != nil {
+		return nil, fmt.Errorf("auth by priv key is not yet implemented")
+	}
+	if len(config.Auth) == 0 {
+		return nil, fmt.Errorf("no valid authentication method for host %q", host.Name)
+	}
+	return &config, nil
+}

main.go

@@ -85,13 +85,13 @@ func server(c *cli.Context) error {
 			io.WriteString(s, banner)
 			io.WriteString(s, "Configuration menu not yet implemented.\n\n")
 		default:
-			config, err := getConfig(s, db)
+			host, err := RemoteHostFromSession(s, db)
 			if err != nil {
 				io.WriteString(s, fmt.Sprintf("error: %v\n", err))
-				// FIXME: drop a menu shell?
+				// FIXME: print available hosts
 				return
 			}
-			if err := proxy(s, config); err != nil {
+			if err := proxy(s, host); err != nil {
 				io.WriteString(s, fmt.Sprintf("error: %v\n", err))
 			}
 		}

proxy.go

@@ -9,8 +9,13 @@ import (
 	gossh "golang.org/x/crypto/ssh"
 )
 
-func proxy(s ssh.Session, config *Config) error {
-	rconn, err := gossh.Dial("tcp", config.remoteAddr, config.clientConfig)
+func proxy(s ssh.Session, host *Host) error {
+	config, err := host.ClientConfig(s)
+	if err != nil {
+		return err
+	}
+
+	rconn, err := gossh.Dial("tcp", host.Addr, config)
 	if err != nil {
 		return err
 	}