1package postgres
  2
  3import (
  4	"fmt"
  5	"strings"
  6
  7	"github.com/jmoiron/sqlx"
  8	"github.com/picosh/pico/pkg/db"
  9)
 10
 11func (me *PsqlDB) insertPublicKeyWithTx(userID, key, name string, tx *sqlx.Tx) error {
 12	pk, _ := me.findPublicKeyByKey(key)
 13	if pk != nil {
 14		return db.ErrPublicKeyTaken
 15	}
 16	query := `INSERT INTO public_keys (user_id, public_key, name) VALUES ($1, $2, $3)`
 17	_, err := tx.Exec(query, userID, key, name)
 18	return err
 19}
 20
 21func (me *PsqlDB) InsertPublicKey(userID, key, name string) error {
 22	pk, _ := me.findPublicKeyByKey(key)
 23	if pk != nil {
 24		return db.ErrPublicKeyTaken
 25	}
 26	query := `INSERT INTO public_keys (user_id, public_key, name) VALUES ($1, $2, $3)`
 27	_, err := me.Db.Exec(query, userID, key, name)
 28	return err
 29}
 30
 31func (me *PsqlDB) UpdatePublicKey(pubkeyID, name string) (*db.PublicKey, error) {
 32	pk, err := me.findPublicKey(pubkeyID)
 33	if err != nil {
 34		return nil, err
 35	}
 36
 37	query := `UPDATE public_keys SET name=$1 WHERE id=$2;`
 38	_, err = me.Db.Exec(query, name, pk.ID)
 39	if err != nil {
 40		return nil, err
 41	}
 42
 43	pk, err = me.findPublicKey(pubkeyID)
 44	if err != nil {
 45		return nil, err
 46	}
 47	return pk, nil
 48}
 49
 50func (me *PsqlDB) findPublicKeyByKey(key string) (*db.PublicKey, error) {
 51	var keys []*db.PublicKey
 52	rs, err := me.Db.Queryx(`SELECT id, user_id, name, public_key, created_at FROM public_keys WHERE public_key = $1`, key)
 53	if err != nil {
 54		return nil, err
 55	}
 56	defer func() { _ = rs.Close() }()
 57
 58	for rs.Next() {
 59		pk := &db.PublicKey{}
 60		err := rs.Scan(&pk.ID, &pk.UserID, &pk.Name, &pk.Key, &pk.CreatedAt)
 61		if err != nil {
 62			return nil, err
 63		}
 64
 65		keys = append(keys, pk)
 66	}
 67
 68	if rs.Err() != nil {
 69		return nil, rs.Err()
 70	}
 71
 72	if len(keys) == 0 {
 73		return nil, fmt.Errorf("pubkey not found in our database: [%s]", key)
 74	}
 75
 76	// When we run PublicKeyByKey and there are multiple public keys returned from the database
 77	// that should mean that we don't have the correct username for this public key.
 78	// When that happens we need to reject the authentication and ask the user to provide the correct
 79	// username when using ssh.  So instead of `ssh <domain>` it should be `ssh user@<domain>`
 80	if len(keys) > 1 {
 81		return nil, &db.ErrMultiplePublicKeys{}
 82	}
 83
 84	return keys[0], nil
 85}
 86
 87func (me *PsqlDB) findPublicKey(pubkeyID string) (*db.PublicKey, error) {
 88	pk := &db.PublicKey{}
 89	err := me.Db.Get(pk, `SELECT * FROM public_keys WHERE id = $1`, pubkeyID)
 90	if err != nil {
 91		return nil, err
 92	}
 93	return pk, nil
 94}
 95
 96func (me *PsqlDB) FindKeysByUser(user *db.User) ([]*db.PublicKey, error) {
 97	var keys []*db.PublicKey
 98	err := me.Db.Select(&keys, `SELECT * FROM public_keys WHERE user_id = $1 ORDER BY created_at ASC`, user.ID)
 99	if err != nil {
100		return nil, err
101	}
102	return keys, nil
103}
104
105func (me *PsqlDB) RemoveKeys(keyIDs []string) error {
106	param := "{" + strings.Join(keyIDs, ",") + "}"
107	_, err := me.Db.Exec(`DELETE FROM public_keys WHERE id = ANY($1::uuid[])`, param)
108	return err
109}