Eric Bower
·
17 Jun 24
file_handler.go
1package pico
2
3import (
4 "bytes"
5 "errors"
6 "fmt"
7 "io"
8 "log/slog"
9 "os"
10 "path/filepath"
11 "strings"
12 "time"
13
14 "github.com/charmbracelet/ssh"
15 "github.com/charmbracelet/wish"
16 "github.com/picosh/pico/db"
17 "github.com/picosh/pico/filehandlers/util"
18 "github.com/picosh/pico/shared"
19 "github.com/picosh/send/send/utils"
20)
21
22type UploadHandler struct {
23 DBPool db.DB
24 Cfg *shared.ConfigSite
25}
26
27func NewUploadHandler(dbpool db.DB, cfg *shared.ConfigSite) *UploadHandler {
28 return &UploadHandler{
29 DBPool: dbpool,
30 Cfg: cfg,
31 }
32}
33
34func (h *UploadHandler) getAuthorizedKeyFile(user *db.User) (*utils.VirtualFile, string, error) {
35 keys, err := h.DBPool.FindKeysForUser(user)
36 text := ""
37 var modTime time.Time
38 for _, pk := range keys {
39 text += fmt.Sprintf("%s %s\n", pk.Key, pk.Name)
40 modTime = *pk.CreatedAt
41 }
42 if err != nil {
43 return nil, "", err
44 }
45 fileInfo := &utils.VirtualFile{
46 FName: "authorized_keys",
47 FIsDir: false,
48 FSize: int64(len(text)),
49 FModTime: modTime,
50 }
51 return fileInfo, text, nil
52}
53
54func (h *UploadHandler) Delete(s ssh.Session, entry *utils.FileEntry) error {
55 return errors.New("unsupported")
56}
57
58func (h *UploadHandler) Read(s ssh.Session, entry *utils.FileEntry) (os.FileInfo, utils.ReaderAtCloser, error) {
59 user, err := util.GetUser(s.Context())
60 if err != nil {
61 return nil, nil, err
62 }
63 cleanFilename := filepath.Base(entry.Filepath)
64
65 if cleanFilename == "" || cleanFilename == "." {
66 return nil, nil, os.ErrNotExist
67 }
68
69 if cleanFilename == "authorized_keys" {
70 fileInfo, text, err := h.getAuthorizedKeyFile(user)
71 if err != nil {
72 return nil, nil, err
73 }
74 reader := utils.NopReaderAtCloser(strings.NewReader(text))
75 return fileInfo, reader, nil
76 }
77
78 return nil, nil, os.ErrNotExist
79}
80
81func (h *UploadHandler) List(s ssh.Session, fpath string, isDir bool, recursive bool) ([]os.FileInfo, error) {
82 var fileList []os.FileInfo
83 user, err := util.GetUser(s.Context())
84 if err != nil {
85 return fileList, err
86 }
87 cleanFilename := filepath.Base(fpath)
88
89 if cleanFilename == "" || cleanFilename == "." || cleanFilename == "/" {
90 name := cleanFilename
91 if name == "" {
92 name = "/"
93 }
94
95 fileList = append(fileList, &utils.VirtualFile{
96 FName: name,
97 FIsDir: true,
98 })
99
100 flist, _, err := h.getAuthorizedKeyFile(user)
101 if err != nil {
102 return fileList, err
103 }
104 fileList = append(fileList, flist)
105 } else {
106 if cleanFilename == "authorized_keys" {
107 flist, _, err := h.getAuthorizedKeyFile(user)
108 if err != nil {
109 return fileList, err
110 }
111 fileList = append(fileList, flist)
112 }
113 }
114
115 return fileList, nil
116}
117
118func (h *UploadHandler) GetLogger() *slog.Logger {
119 return h.Cfg.Logger
120}
121
122func (h *UploadHandler) Validate(s ssh.Session) error {
123 var err error
124 key, err := utils.KeyText(s)
125 if err != nil {
126 return fmt.Errorf("key not found")
127 }
128
129 user, err := h.DBPool.FindUserForKey(s.User(), key)
130 if err != nil {
131 return err
132 }
133
134 if user.Name == "" {
135 return fmt.Errorf("must have username set")
136 }
137
138 util.SetUser(s.Context(), user)
139 return nil
140}
141
142type KeyWithId struct {
143 Pk ssh.PublicKey
144 ID string
145 Comment string
146}
147
148type KeyDiffResult struct {
149 Add []KeyWithId
150 Rm []string
151 Update []KeyWithId
152}
153
154func authorizedKeysDiff(keyInUse ssh.PublicKey, curKeys []KeyWithId, nextKeys []KeyWithId) KeyDiffResult {
155 update := []KeyWithId{}
156 add := []KeyWithId{}
157 for _, nk := range nextKeys {
158 found := false
159 for _, ck := range curKeys {
160 if ssh.KeysEqual(nk.Pk, ck.Pk) {
161 found = true
162
163 // update the comment field
164 if nk.Comment != ck.Comment {
165 ck.Comment = nk.Comment
166 update = append(update, ck)
167 }
168 break
169 }
170 }
171 if !found {
172 add = append(add, nk)
173 }
174 }
175
176 rm := []string{}
177 for _, ck := range curKeys {
178 // we never want to remove the key that's in the current ssh session
179 // in an effort to avoid mistakenly removing their current key
180 if ssh.KeysEqual(ck.Pk, keyInUse) {
181 continue
182 }
183
184 found := false
185 for _, nk := range nextKeys {
186 if ssh.KeysEqual(ck.Pk, nk.Pk) {
187 found = true
188 break
189 }
190 }
191 if !found {
192 rm = append(rm, ck.ID)
193 }
194 }
195
196 return KeyDiffResult{
197 Add: add,
198 Rm: rm,
199 Update: update,
200 }
201}
202
203func (h *UploadHandler) ProcessAuthorizedKeys(text []byte, logger *slog.Logger, user *db.User, s ssh.Session) error {
204 logger.Info("processing new authorized_keys")
205 dbpool := h.DBPool
206
207 curKeysStr, err := dbpool.FindKeysForUser(user)
208 if err != nil {
209 return err
210 }
211
212 splitKeys := bytes.Split(text, []byte{'\n'})
213 nextKeys := []KeyWithId{}
214 for _, pk := range splitKeys {
215 key, comment, _, _, err := ssh.ParseAuthorizedKey(bytes.TrimSpace(pk))
216 if err != nil {
217 continue
218 }
219 nextKeys = append(nextKeys, KeyWithId{Pk: key, Comment: comment})
220 }
221
222 curKeys := []KeyWithId{}
223 for _, pk := range curKeysStr {
224 key, _, _, _, err := ssh.ParseAuthorizedKey([]byte(pk.Key))
225 if err != nil {
226 continue
227 }
228 curKeys = append(curKeys, KeyWithId{Pk: key, ID: pk.ID, Comment: pk.Name})
229 }
230
231 diff := authorizedKeysDiff(s.PublicKey(), curKeys, nextKeys)
232
233 for _, pk := range diff.Add {
234 key, err := shared.KeyForKeyText(pk.Pk)
235 if err != nil {
236 continue
237 }
238
239 wish.Errorf(s, "adding pubkey (%s)\n", key)
240 logger.Info("adding pubkey", "pubkey", key)
241
242 err = dbpool.InsertPublicKey(user.ID, key, pk.Comment, nil)
243 if err != nil {
244 wish.Errorf(s, "error: could not insert pubkey: %s (%s)\n", err.Error(), key)
245 logger.Error("could not insert pubkey", "err", err.Error())
246 }
247 }
248
249 for _, pk := range diff.Update {
250 key, err := shared.KeyForKeyText(pk.Pk)
251 if err != nil {
252 continue
253 }
254
255 wish.Errorf(s, "updating pubkey with comment: %s (%s)\n", pk.Comment, key)
256 logger.Info(
257 "updating pubkey with comment",
258 "pubkey", key,
259 "comment", pk.Comment,
260 )
261
262 _, err = dbpool.UpdatePublicKey(pk.ID, pk.Comment)
263 if err != nil {
264 wish.Errorf(s, "error: could not update pubkey: %s (%s)\n", err.Error(), key)
265 logger.Error("could not update pubkey", "err", err.Error(), "key", key)
266 }
267 }
268
269 if len(diff.Rm) > 0 {
270 wish.Errorf(s, "removing pubkeys: %s\n", diff.Rm)
271 logger.Info("removing pubkeys", "pubkeys", diff.Rm)
272
273 err = dbpool.RemoveKeys(diff.Rm)
274 if err != nil {
275 wish.Errorf(s, "error: could not rm pubkeys: %s\n", err.Error())
276 logger.Error("could not remove pubkey", "err", err.Error())
277 }
278 }
279
280 return nil
281}
282
283func (h *UploadHandler) Write(s ssh.Session, entry *utils.FileEntry) (string, error) {
284 logger := h.Cfg.Logger
285 user, err := util.GetUser(s.Context())
286 if err != nil {
287 logger.Error(err.Error())
288 return "", err
289 }
290
291 filename := filepath.Base(entry.Filepath)
292 logger = logger.With(
293 "user", user.Name,
294 "filename", filename,
295 )
296
297 var text []byte
298 if b, err := io.ReadAll(entry.Reader); err == nil {
299 text = b
300 }
301
302 if filename == "authorized_keys" {
303 err := h.ProcessAuthorizedKeys(text, logger, user, s)
304 if err != nil {
305 return "", err
306 }
307 } else {
308 return "", fmt.Errorf("validation error: invalid file, received %s", entry.Filepath)
309 }
310
311 return "", nil
312}