Antonio Mika
·
2025-03-12
handler.go
1package pobj
2
3import (
4 "bytes"
5 "encoding/binary"
6 "fmt"
7 "io"
8 "log/slog"
9 "os"
10 "path/filepath"
11 "time"
12
13 "github.com/picosh/pico/pkg/pobj/storage"
14 "github.com/picosh/pico/pkg/pssh"
15 "github.com/picosh/pico/pkg/send/utils"
16)
17
18type ctxBucketKey struct{}
19
20func getBucket(ctx *pssh.SSHServerConnSession) (storage.Bucket, error) {
21 bucket, ok := ctx.Value(ctxBucketKey{}).(storage.Bucket)
22 if !ok {
23 return bucket, fmt.Errorf("bucket not set on `ssh.Context()` for connection")
24 }
25 if bucket.Name == "" {
26 return bucket, fmt.Errorf("bucket not set on `ssh.Context()` for connection")
27 }
28 return bucket, nil
29}
30func setBucket(ctx *pssh.SSHServerConnSession, bucket storage.Bucket) {
31 ctx.SetValue(ctxBucketKey{}, bucket)
32}
33
34type FileData struct {
35 *utils.FileEntry
36 Text []byte
37 User string
38 Bucket storage.Bucket
39}
40
41type Config struct {
42 Logger *slog.Logger
43 Storage storage.ObjectStorage
44 AssetNames AssetNames
45}
46
47type UploadAssetHandler struct {
48 Cfg *Config
49}
50
51var _ utils.CopyFromClientHandler = &UploadAssetHandler{}
52var _ utils.CopyFromClientHandler = (*UploadAssetHandler)(nil)
53
54func NewUploadAssetHandler(cfg *Config) *UploadAssetHandler {
55 if cfg.AssetNames == nil {
56 cfg.AssetNames = &AssetNamesBasic{}
57 }
58
59 return &UploadAssetHandler{
60 Cfg: cfg,
61 }
62}
63
64func (h *UploadAssetHandler) GetLogger(s *pssh.SSHServerConnSession) *slog.Logger {
65 return h.Cfg.Logger
66}
67
68func (h *UploadAssetHandler) Delete(s *pssh.SSHServerConnSession, entry *utils.FileEntry) error {
69 h.Cfg.Logger.Info("deleting file", "file", entry.Filepath)
70 bucket, err := getBucket(s)
71 if err != nil {
72 h.Cfg.Logger.Error(err.Error())
73 return err
74 }
75
76 objectFileName, err := h.Cfg.AssetNames.ObjectName(s, entry)
77 if err != nil {
78 return err
79 }
80 return h.Cfg.Storage.DeleteObject(bucket, objectFileName)
81}
82
83func (h *UploadAssetHandler) Read(s *pssh.SSHServerConnSession, entry *utils.FileEntry) (os.FileInfo, utils.ReadAndReaderAtCloser, error) {
84 fileInfo := &utils.VirtualFile{
85 FName: filepath.Base(entry.Filepath),
86 FIsDir: false,
87 FSize: entry.Size,
88 FModTime: time.Unix(entry.Mtime, 0),
89 }
90 h.Cfg.Logger.Info("reading file", "file", fileInfo)
91
92 bucketName, err := h.Cfg.AssetNames.BucketName(s)
93 if err != nil {
94 return nil, nil, err
95 }
96 bucket, err := h.Cfg.Storage.GetBucket(bucketName)
97 if err != nil {
98 return nil, nil, err
99 }
100
101 fname, err := h.Cfg.AssetNames.ObjectName(s, entry)
102 if err != nil {
103 return nil, nil, err
104 }
105 contents, info, err := h.Cfg.Storage.GetObject(bucket, fname)
106 if err != nil {
107 return nil, nil, err
108 }
109
110 fileInfo.FSize = info.Size
111 fileInfo.FModTime = info.LastModified
112
113 reader := NewAllReaderAt(contents)
114
115 return fileInfo, reader, nil
116}
117
118func (h *UploadAssetHandler) List(s *pssh.SSHServerConnSession, fpath string, isDir bool, recursive bool) ([]os.FileInfo, error) {
119 h.Cfg.Logger.Info(
120 "listing path",
121 "dir", fpath,
122 "isDir", isDir,
123 "recursive", recursive,
124 )
125 var fileList []os.FileInfo
126
127 cleanFilename := fpath
128
129 bucketName, err := h.Cfg.AssetNames.BucketName(s)
130 if err != nil {
131 return fileList, err
132 }
133 bucket, err := h.Cfg.Storage.GetBucket(bucketName)
134 if err != nil {
135 return fileList, err
136 }
137
138 fname, err := h.Cfg.AssetNames.ObjectName(s, &utils.FileEntry{Filepath: cleanFilename})
139 if err != nil {
140 return fileList, err
141 }
142
143 if fname == "" || fname == "." {
144 name := fname
145 if name == "" {
146 name = "/"
147 }
148
149 info := &utils.VirtualFile{
150 FName: name,
151 FIsDir: true,
152 }
153
154 fileList = append(fileList, info)
155 } else {
156 name := fname
157 if name != "/" && isDir {
158 name += "/"
159 }
160
161 foundList, err := h.Cfg.Storage.ListObjects(bucket, name, recursive)
162 if err != nil {
163 return fileList, err
164 }
165
166 fileList = append(fileList, foundList...)
167 }
168
169 return fileList, nil
170}
171
172func (h *UploadAssetHandler) Validate(s *pssh.SSHServerConnSession) error {
173 var err error
174 userName := s.User()
175
176 assetBucket, err := h.Cfg.AssetNames.BucketName(s)
177 if err != nil {
178 return err
179 }
180 bucket, err := h.Cfg.Storage.UpsertBucket(assetBucket)
181 if err != nil {
182 return err
183 }
184 setBucket(s, bucket)
185
186 pk, _ := utils.KeyText(s)
187 h.Cfg.Logger.Info(
188 "attempting to upload files",
189 "user", userName,
190 "bucket", bucket.Name,
191 "publicKey", pk,
192 )
193 return nil
194}
195
196func (h *UploadAssetHandler) Write(s *pssh.SSHServerConnSession, entry *utils.FileEntry) (string, error) {
197 var origText []byte
198 if b, err := io.ReadAll(entry.Reader); err == nil {
199 origText = b
200 }
201 fileSize := binary.Size(origText)
202 // TODO: hack for now until I figure out how to get correct
203 // filesize from sftp,scp,rsync
204 entry.Size = int64(fileSize)
205 userName := s.User()
206
207 bucket, err := getBucket(s)
208 if err != nil {
209 h.Cfg.Logger.Error(err.Error())
210 return "", err
211 }
212
213 data := &FileData{
214 FileEntry: entry,
215 User: userName,
216 Text: origText,
217 Bucket: bucket,
218 }
219 err = h.writeAsset(s, data)
220 if err != nil {
221 h.Cfg.Logger.Error(err.Error())
222 return "", err
223 }
224
225 url, err := h.Cfg.AssetNames.PrintObjectName(s, entry, bucket.Name)
226 if err != nil {
227 return "", err
228 }
229 return url, nil
230}
231
232func (h *UploadAssetHandler) validateAsset(_ *FileData) (bool, error) {
233 return true, nil
234}
235
236func (h *UploadAssetHandler) writeAsset(s *pssh.SSHServerConnSession, data *FileData) error {
237 valid, err := h.validateAsset(data)
238 if !valid {
239 return err
240 }
241
242 objectFileName, err := h.Cfg.AssetNames.ObjectName(s, data.FileEntry)
243 if err != nil {
244 return err
245 }
246 reader := bytes.NewReader(data.Text)
247
248 h.Cfg.Logger.Info(
249 "uploading file to bucket",
250 "user",
251 data.User,
252 "bucket",
253 data.Bucket.Name,
254 "object",
255 objectFileName,
256 )
257
258 _, _, err = h.Cfg.Storage.PutObject(
259 data.Bucket,
260 objectFileName,
261 utils.NopReadAndReaderAtCloser(reader),
262 data.FileEntry,
263 )
264 if err != nil {
265 return err
266 }
267
268 return nil
269}