Add SSH key check + auto register of the admin account

db.go

@@ -34,12 +34,19 @@ type Host struct {
 	Comment     string
 }
 
+type UserKey struct {
+	gorm.Model
+	Key    []byte
+	UserID uint
+}
+
 type User struct {
 	// FIXME: use uuid for ID
 	gorm.Model
+	IsAdmin bool
 	Email   string // FIXME: govalidator: email
 	Name    string // FIXME: govalidator: min length 3, alphanum
-	SSHKeys []SSHKey
+	Keys    []UserKey
 	Comment string
 }
 
@@ -47,6 +54,7 @@ func dbInit(db *gorm.DB) error {
 	db.AutoMigrate(&User{})
 	db.AutoMigrate(&SSHKey{})
 	db.AutoMigrate(&Host{})
+	db.AutoMigrate(&UserKey{})
 	db.Exec(`CREATE UNIQUE INDEX uix_keys_name   ON "ssh_keys"("name") WHERE ("deleted_at" IS NULL)`)
 	db.Exec(`CREATE UNIQUE INDEX uix_hosts_name  ON "hosts"("name")    WHERE ("deleted_at" IS NULL)`)
 	db.Exec(`CREATE UNIQUE INDEX uix_users_name  ON "users"("email")   WHERE ("deleted_at" IS NULL)`)
@@ -164,7 +172,7 @@ func FindKeysByIdOrName(db *gorm.DB, queries []string) ([]*SSHKey, error) {
 
 func FindUserByIdOrEmail(db *gorm.DB, query string) (*User, error) {
 	var user User
-	if err := db.Where("id = ?", query).Or("email = ?", query).First(&user).Error; err != nil {
+	if err := db.Preload("Keys").Where("id = ?", query).Or("email = ?", query).First(&user).Error; err != nil {
 		return nil, err
 	}
 	return &user, nil

main.go

@@ -17,6 +17,10 @@ import (
 
 var version = "0.0.1"
 
+type sshportalContextKey string
+
+var userContextKey = sshportalContextKey("user")
+
 func main() {
 	app := cli.NewApp()
 	app.Name = path.Base(os.Args[0])
@@ -104,6 +108,45 @@ func server(c *cli.Context) error {
 		return errors.New("use `--demo` for now")
 	}
 
+	opts = append(opts, ssh.PublicKeyAuth(func(ctx ssh.Context, key ssh.PublicKey) bool {
+		var (
+			userKey UserKey
+			user    User
+			count   uint
+		)
+
+		// lookup user by key
+		db.Where("key = ?", key.Marshal()).First(&userKey)
+		if userKey.UserID > 0 {
+			db.Where("id = ?", userKey.UserID).First(&user)
+			ctx.SetValue(userContextKey, user)
+			return true
+		}
+
+		// check if there are users in DB
+		db.Table("users").Count(&count)
+		if count == 0 { // create an admin user
+			// if no admin, create an account for the first connection
+			user = User{
+				Name:    "Administrator",
+				Email:   "admin@sshportal",
+				Comment: "created by sshportal",
+				IsAdmin: true,
+			}
+			db.Create(&user)
+			userKey = UserKey{
+				UserID: user.ID,
+				Key:    key.Marshal(),
+			}
+			db.Create(&userKey)
+
+			ctx.SetValue(userContextKey, user)
+			return true
+		}
+
+		return false
+	}))
+
 	log.Printf("SSH Server accepting connections on %s", c.String("bind-address"))
 	return ssh.ListenAndServe(c.String("bind-address"), nil, opts...)
 }

shell.go

@@ -377,7 +377,7 @@ GLOBAL OPTIONS:
 					Usage: "Lists users",
 					Action: func(c *cli.Context) error {
 						var users []User
-						if err := db.Find(&users).Error; err != nil {
+						if err := db.Preload("Keys").Find(&users).Error; err != nil {
 							return err
 						}
 						table := tablewriter.NewWriter(s)
@@ -385,7 +385,7 @@ GLOBAL OPTIONS:
 						table.SetBorder(false)
 						table.SetCaption(true, fmt.Sprintf("Total: %d users.", len(users)))
 						for _, user := range users {
-							keys := len(user.SSHKeys)
+							keys := len(user.Keys)
 							table.Append([]string{
 								fmt.Sprintf("%d", user.ID),
 								user.Name,