1package prose
  2
  3import (
  4	"bytes"
  5	"encoding/binary"
  6	"fmt"
  7	"io"
  8	"net/http"
  9	"os"
 10	"path/filepath"
 11	"slices"
 12	"strings"
 13
 14	exifremove "github.com/neurosnap/go-exif-remove"
 15	"github.com/picosh/pico/pkg/db"
 16	"github.com/picosh/pico/pkg/pssh"
 17	sendutils "github.com/picosh/pico/pkg/send/utils"
 18	"github.com/picosh/pico/pkg/shared"
 19	"github.com/picosh/pico/pkg/storage"
 20)
 21
 22type PostMetaData struct {
 23	Text          []byte
 24	FileSize      int
 25	TotalFileSize int
 26	Filename      string
 27	User          *db.User
 28	FeatureFlag   *db.FeatureFlag
 29	Bucket        storage.Bucket
 30}
 31
 32type UploadImgHandler struct {
 33	DBPool  DB
 34	Cfg     *Config
 35	Storage storage.StorageServe
 36}
 37
 38func NewUploadImgHandler(dbpool DB, cfg *Config, storage storage.StorageServe) *UploadImgHandler {
 39	return &UploadImgHandler{
 40		DBPool:  dbpool,
 41		Cfg:     cfg,
 42		Storage: storage,
 43	}
 44}
 45
 46func (h *UploadImgHandler) getObjectPath(fpath string) string {
 47	return filepath.Join("prose", fpath)
 48}
 49
 50func (h *UploadImgHandler) List(s *pssh.SSHServerConnSession, fpath string, isDir bool, recursive bool) ([]os.FileInfo, error) {
 51	var fileList []os.FileInfo
 52
 53	logger := pssh.GetLogger(s)
 54	user := pssh.GetUser(s)
 55
 56	if user == nil {
 57		err := fmt.Errorf("could not get user from ctx")
 58		logger.Error("error getting user from ctx", "err", err)
 59		return fileList, err
 60	}
 61
 62	cleanFilename := fpath
 63
 64	bucketName := shared.GetAssetBucketName(user.ID)
 65	bucket, err := h.Storage.GetBucket(bucketName)
 66	if err != nil {
 67		return fileList, err
 68	}
 69
 70	if cleanFilename == "" || cleanFilename == "." {
 71		name := cleanFilename
 72		if name == "" {
 73			name = "/"
 74		}
 75
 76		info := &sendutils.VirtualFile{
 77			FName:  name,
 78			FIsDir: true,
 79		}
 80
 81		fileList = append(fileList, info)
 82	} else {
 83		fp := h.getObjectPath(cleanFilename)
 84		if fp != "/" && isDir {
 85			fp += "/"
 86		}
 87
 88		foundList, err := h.Storage.ListObjects(bucket, fp, recursive)
 89		if err != nil {
 90			return fileList, err
 91		}
 92
 93		fileList = append(fileList, foundList...)
 94	}
 95
 96	return fileList, nil
 97}
 98
 99func (h *UploadImgHandler) Read(s *pssh.SSHServerConnSession, entry *sendutils.FileEntry) (os.FileInfo, sendutils.ReadAndReaderAtCloser, error) {
100	logger := pssh.GetLogger(s)
101	user := pssh.GetUser(s)
102
103	if user == nil {
104		err := fmt.Errorf("could not get user from ctx")
105		logger.Error("error getting user from ctx", "err", err)
106		return nil, nil, err
107	}
108
109	cleanFilename := filepath.Base(entry.Filepath)
110
111	if cleanFilename == "" || cleanFilename == "." {
112		return nil, nil, os.ErrNotExist
113	}
114
115	bucket, err := h.Storage.GetBucket(shared.GetAssetBucketName(user.ID))
116	if err != nil {
117		return nil, nil, err
118	}
119
120	contents, info, err := h.Storage.GetObject(bucket, h.getObjectPath(cleanFilename))
121	if err != nil {
122		return nil, nil, err
123	}
124
125	fileInfo := &sendutils.VirtualFile{
126		FName:    cleanFilename,
127		FIsDir:   false,
128		FSize:    info.Size,
129		FModTime: info.LastModified,
130	}
131
132	return fileInfo, contents, nil
133}
134
135func (h *UploadImgHandler) Write(s *pssh.SSHServerConnSession, entry *sendutils.FileEntry) (string, error) {
136	logger := pssh.GetLogger(s)
137	user := pssh.GetUser(s)
138
139	if user == nil {
140		err := fmt.Errorf("could not get user from ctx")
141		logger.Error("error getting user from ctx", "err", err)
142		return "", err
143	}
144
145	filename := filepath.Base(entry.Filepath)
146
147	var text []byte
148	if b, err := io.ReadAll(entry.Reader); err == nil {
149		text = b
150	}
151	mimeType := http.DetectContentType(text)
152	ext := filepath.Ext(filename)
153	if ext == ".svg" {
154		mimeType = "image/svg+xml"
155	}
156	// strip exif data
157	if slices.Contains([]string{"image/png", "image/jpg", "image/jpeg"}, mimeType) {
158		noExifBytes, err := exifremove.Remove(text)
159		if err == nil {
160			if len(noExifBytes) == 0 {
161				logger.Info("file silently failed to strip exif data", "filename", filename)
162			} else {
163				text = noExifBytes
164				logger.Info("stripped exif data", "filename", filename)
165			}
166		} else {
167			logger.Error("could not strip exif data", "err", err.Error())
168		}
169	}
170
171	fileSize := binary.Size(text)
172	featureFlag, _ := h.DBPool.FindFeature(user.ID, "plus")
173	// we have free tiers so users might not have a feature flag
174	// in which case we set sane defaults
175	if featureFlag == nil {
176		featureFlag = db.NewFeatureFlag(user.ID, "plus", h.Cfg.MaxSize, h.Cfg.MaxAssetSize, h.Cfg.MaxSpecialFileSize)
177	}
178	featureFlag.SetLimits(h.Cfg.MaxSize, h.Cfg.MaxAssetSize, h.Cfg.MaxSpecialFileSize)
179
180	bucket, err := h.Storage.UpsertBucket(shared.GetAssetBucketName(user.ID))
181	if err != nil {
182		return "", err
183	}
184
185	totalFileSize, err := h.Storage.GetBucketQuota(bucket)
186	if err != nil {
187		logger.Error("bucket quota", "err", err)
188		return "", err
189	}
190
191	metadata := PostMetaData{
192		Filename:      filename,
193		FileSize:      fileSize,
194		Text:          text,
195		User:          user,
196		FeatureFlag:   featureFlag,
197		Bucket:        bucket,
198		TotalFileSize: int(totalFileSize),
199	}
200
201	err = h.writeImg(s, &metadata)
202	if err != nil {
203		logger.Error("could not write img", "err", err.Error())
204		return "", err
205	}
206
207	curl := shared.NewCreateURL(h.Cfg.Protocol, h.Cfg.Domain)
208	url := curl.FullPostURL(user.Name, metadata.Filename)
209	return url + shared.StorageUsage(metadata.TotalFileSize+fileSize, int(featureFlag.Data.StorageMax)), nil
210}
211
212func (h *UploadImgHandler) Delete(s *pssh.SSHServerConnSession, entry *sendutils.FileEntry) error {
213	logger := pssh.GetLogger(s)
214	user := pssh.GetUser(s)
215
216	if user == nil {
217		err := fmt.Errorf("could not get user from ctx")
218		logger.Error("error getting user from ctx", "err", err)
219		return err
220	}
221
222	filename := filepath.Base(entry.Filepath)
223
224	logger = logger.With(
225		"filename", filename,
226	)
227
228	bucket, err := h.Storage.UpsertBucket(shared.GetAssetBucketName(user.ID))
229	if err != nil {
230		return err
231	}
232
233	logger.Info("deleting image")
234	err = h.Storage.DeleteObject(bucket, h.getObjectPath(filename))
235	if err != nil {
236		return err
237	}
238
239	return nil
240}
241
242func (h *UploadImgHandler) validateImg(data *PostMetaData) (bool, error) {
243	fileMax := data.FeatureFlag.Data.FileMax
244	if int64(data.FileSize) > fileMax {
245		return false, fmt.Errorf("ERROR: file (%s) has exceeded maximum file size (%d bytes)", data.Filename, fileMax)
246	}
247
248	storageMax := data.FeatureFlag.Data.StorageMax
249	if uint64(data.TotalFileSize+data.FileSize) > storageMax {
250		return false, fmt.Errorf("ERROR: user (%s) has exceeded (%d bytes) max (%d bytes)", data.User.Name, data.TotalFileSize, storageMax)
251	}
252
253	if !shared.IsExtAllowed(data.Filename, h.Cfg.AllowedExt) {
254		extStr := strings.Join(h.Cfg.AllowedExt, ",")
255		err := fmt.Errorf(
256			"ERROR: (%s) invalid file, format must be (%s), skipping",
257			data.Filename,
258			extStr,
259		)
260		return false, err
261	}
262
263	return true, nil
264}
265
266func (h *UploadImgHandler) metaImg(data *PostMetaData) error {
267	// if the file is empty that means we should delete it
268	// so we can skip all the meta info
269	if data.FileSize == 0 {
270		return nil
271	}
272
273	// make sure we have a bucket
274	bucket, err := h.Storage.UpsertBucket(shared.GetAssetBucketName(data.User.ID))
275	if err != nil {
276		return err
277	}
278
279	reader := bytes.NewReader([]byte(data.Text))
280	_, _, err = h.Storage.PutObject(
281		bucket,
282		h.getObjectPath(data.Filename),
283		sendutils.NopReadAndReaderAtCloser(reader),
284		&storage.ObjectInfo{},
285	)
286	if err != nil {
287		return err
288	}
289
290	return nil
291}
292
293func (h *UploadImgHandler) writeImg(s *pssh.SSHServerConnSession, data *PostMetaData) error {
294	valid, err := h.validateImg(data)
295	if !valid {
296		return err
297	}
298
299	logger := pssh.GetLogger(s)
300	logger = logger.With(
301		"filename", data.Filename,
302	)
303
304	logger.Info("uploading image")
305	err = h.metaImg(data)
306	if err != nil {
307		logger.Error("meta img", "err", err)
308		return err
309	}
310
311	return nil
312}