feat: add category detection for file attachments and update related interfaces

4 files changed, 102 insertions, 0 deletions

diff --git a/shrine/models/letter.go b/shrine/models/letter.go
index f74de66..7668758 100644
--- a/shrine/models/letter.go
+++ b/shrine/models/letter.go
@@ -50,6 +50,7 @@ type LetterAttachment struct {
 	FilePath    string `gorm:"size:512;not null"`
 	FileSize    int64  `gorm:"not null"`
 	ContentType string `gorm:"size:100;not null"`
+	Category    string `gorm:"size:20;not null;default:other"`
 }
 
 func (self *Letter) BeforeCreate(tx *gorm.DB) error {
@@ -89,6 +90,7 @@ func (self *LetterAttachment) ToResponse() letter.AttachmentResponse {
 		URL:         storage.ResolveCDN(self.FilePath),
 		FileSize:    self.FileSize,
 		ContentType: self.ContentType,
+		Category:    self.Category,
 	}
 }
 
diff --git a/shrine/services/letter.go b/shrine/services/letter.go
index e9bd825..940e9c7 100644
--- a/shrine/services/letter.go
+++ b/shrine/services/letter.go
@@ -11,6 +11,7 @@ import (
 	"shrine/types/common"
 	"shrine/types/hypertext"
 	"shrine/types/letter"
+	"shrine/utils/files"
 	"shrine/utils/meta"
 	"shrine/utils/storage"
 	"strings"
@@ -256,6 +257,7 @@ func UploadLetterAttachment(userID uint, fileName string, fileSize int64, conten
 		FileName:    fileName,
 		FileSize:    fileSize,
 		ContentType: contentType,
+		Category:    files.DetectCategory(contentType),
 	}
 
 	if err := repositories.UploadAttachment(userID, &attachment); err != nil {
diff --git a/shrine/types/letter/response.go b/shrine/types/letter/response.go
index 6cbd6ed..fe1c0d7 100644
--- a/shrine/types/letter/response.go
+++ b/shrine/types/letter/response.go
@@ -18,6 +18,7 @@ type AttachmentResponse struct {
 	URL         string `json:"url"`
 	FileSize    int64  `json:"file_size"`
 	ContentType string `json:"content_type"`
+	Category    string `json:"category"`
 }
 
 type MessageResponse struct {
diff --git a/shrine/utils/files/category.go b/shrine/utils/files/category.go
new file mode 100644
index 0000000..d9d87e4
--- /dev/null
+++ b/shrine/utils/files/category.go
@@ -0,0 +1,97 @@
+package files
+
+import (
+	"shrine/utils/collections"
+	"strings"
+)
+
+var archiveMimes = collections.SetOf(
+	"application/zip",
+	"application/x-rar-compressed",
+	"application/gzip",
+	"application/x-7z-compressed",
+	"application/x-tar",
+	"application/x-bzip2",
+	"application/x-xz",
+	"application/x-compress",
+)
+
+var documentMimes = collections.SetOf(
+	"application/msword",
+	"application/vnd.openxmlformats-officedocument.wordprocessingml.document",
+	"application/vnd.oasis.opendocument.text",
+	"application/rtf",
+	"application/epub+zip",
+)
+
+var spreadsheetMimes = collections.SetOf(
+	"application/vnd.ms-excel",
+	"application/vnd.openxmlformats-officedocument.spreadsheetml.sheet",
+	"application/vnd.oasis.opendocument.spreadsheet",
+	"text/csv",
+)
+
+var presentationMimes = collections.SetOf(
+	"application/vnd.ms-powerpoint",
+	"application/vnd.openxmlformats-officedocument.presentationml.presentation",
+	"application/vnd.oasis.opendocument.presentation",
+)
+
+var codeMimes = collections.SetOf(
+	"application/javascript",
+	"application/json",
+	"application/xml",
+	"application/x-httpd-php",
+	"application/x-sh",
+	"application/x-python",
+	"application/typescript",
+)
+
+var databaseMimes = collections.SetOf(
+	"application/x-sqlite3",
+	"application/vnd.ms-access",
+)
+
+
+func DetectCategory(contentType string) string {
+	if strings.HasPrefix(contentType, "image/") {
+		return "image"
+	}
+	if strings.HasPrefix(contentType, "video/") {
+		return "video"
+	}
+	if strings.HasPrefix(contentType, "audio/") {
+		return "audio"
+	}
+	if strings.Contains(contentType, "font") {
+		return "font"
+	}
+	if contentType == "application/pdf" {
+		return "document"
+	}
+	if archiveMimes.Has(contentType) {
+		return "archive"
+	}
+	if documentMimes.Has(contentType) {
+		return "document"
+	}
+	if spreadsheetMimes.Has(contentType) {
+		return "spreadsheet"
+	}
+	if presentationMimes.Has(contentType) {
+		return "presentation"
+	}
+	if codeMimes.Has(contentType) {
+		return "code"
+	}
+	if databaseMimes.Has(contentType) {
+		return "database"
+	}
+	if strings.HasPrefix(contentType, "text/x-") {
+		return "code"
+	}
+	if strings.HasPrefix(contentType, "text/") {
+		return "document"
+	}
+	return "other"
+}
\ No newline at end of file