repos / pico

pico services mono repo
git clone https://github.com/picosh/pico.git

pico / pkg / shared
Antonio Mika  ·  2025-04-30

ssh.go

 1package shared
 2
 3import (
 4	"fmt"
 5	"log/slog"
 6	"strings"
 7
 8	"github.com/picosh/pico/pkg/db"
 9	"github.com/picosh/utils"
10	"golang.org/x/crypto/ssh"
11)
12
13const adminPrefix = "admin__"
14
15type SshAuthHandler struct {
16	DB     AuthFindUser
17	Logger *slog.Logger
18}
19
20type AuthFindUser interface {
21	FindUserByPubkey(key string) (*db.User, error)
22	FindUserByName(name string) (*db.User, error)
23	FindFeature(userID, name string) (*db.FeatureFlag, error)
24}
25
26func NewSshAuthHandler(dbh AuthFindUser, logger *slog.Logger) *SshAuthHandler {
27	return &SshAuthHandler{
28		DB:     dbh,
29		Logger: logger,
30	}
31}
32
33func (r *SshAuthHandler) PubkeyAuthHandler(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
34	pubkey := utils.KeyForKeyText(key)
35	user, err := r.DB.FindUserByPubkey(pubkey)
36	if err != nil {
37		r.Logger.Error(
38			"could not find user for key",
39			"keyType", key.Type(),
40			"key", string(key.Marshal()),
41			"err", err,
42		)
43		return nil, err
44	}
45
46	if user.Name == "" {
47		r.Logger.Error("username is not set")
48		return nil, fmt.Errorf("username is not set")
49	}
50
51	// impersonation
52	var impID string
53	usr := conn.User()
54	if strings.HasPrefix(usr, adminPrefix) {
55		ff, err := r.DB.FindFeature(user.ID, "admin")
56		if err == nil && ff.IsValid() {
57			impersonate := strings.TrimPrefix(usr, adminPrefix)
58			impersonatedUser, err := r.DB.FindUserByName(impersonate)
59			if err == nil {
60				impID = user.ID
61				user = impersonatedUser
62			}
63		}
64	}
65
66	perms := &ssh.Permissions{
67		Extensions: map[string]string{
68			"user_id": user.ID,
69			"pubkey":  pubkey,
70		},
71	}
72
73	if impID != "" {
74		perms.Extensions["imp_id"] = impID
75	}
76
77	return perms, nil
78}
79
80func FindPlusFF(dbpool db.DB, cfg *ConfigSite, userID string) *db.FeatureFlag {
81	ff, _ := dbpool.FindFeature(userID, "plus")
82	// we have free tiers so users might not have a feature flag
83	// in which case we set sane defaults
84	if ff == nil {
85		ff = db.NewFeatureFlag(
86			userID,
87			"plus",
88			cfg.MaxSize,
89			cfg.MaxAssetSize,
90			cfg.MaxSpecialFileSize,
91		)
92	}
93	// this is jank
94	ff.Data.StorageMax = ff.FindStorageMax(cfg.MaxSize)
95	ff.Data.FileMax = ff.FindFileMax(cfg.MaxAssetSize)
96	ff.Data.SpecialFileMax = ff.FindSpecialFileMax(cfg.MaxSpecialFileSize)
97	return ff
98}