Merge pull request #248 from jwessel/master

feat: Allow removal by user for 'userkey rm'
moul · Mar 26, 2021 · 40bbea5 · 40bbea5
2 parents e455d50 + c1c4c55
commit 40bbea5
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 1 deletion.
diff --git a/pkg/bastion/shell.go b/pkg/bastion/shell.go
@@ -2141,7 +2141,16 @@ GLOBAL OPTIONS:
 						if err := myself.CheckRoles([]string{"admin"}); err != nil {
 							return err
 						}
-
+						if err := dbmodels.UserKeysByIdentifiers(db, c.Args()).Find(&dbmodels.UserKey{}).Error; err != nil {
+							var user dbmodels.User
+							if err := dbmodels.UsersByIdentifiers(db, c.Args()).First(&user).Error; err != nil {
+								return err
+							}
+							if err := dbmodels.UserKeysByUserID(db, []string{fmt.Sprint(user.ID)}).Find(&dbmodels.UserKey{}).Error; err != nil {
+								return err
+							}
+							return dbmodels.UserKeysByUserID(db, []string{fmt.Sprint(user.ID)}).Delete(&dbmodels.UserKey{}).Error
+						}
 						return dbmodels.UserKeysByIdentifiers(db, c.Args()).Delete(&dbmodels.UserKey{}).Error
 					},
 				},

diff --git a/pkg/dbmodels/dbmodels.go b/pkg/dbmodels/dbmodels.go
@@ -371,6 +371,9 @@ func UserKeysPreload(db *gorm.DB) *gorm.DB {
 func UserKeysByIdentifiers(db *gorm.DB, identifiers []string) *gorm.DB {
 	return db.Where("id IN (?)", identifiers)
 }
+func UserKeysByUserID(db *gorm.DB, identifiers []string) *gorm.DB {
+	return db.Where("user_id IN (?)", identifiers)
+}
 
 // UserRole helpers