arozos-mirror/src/quota.go

package main

import (
	"encoding/json"
	"net/http"
	"os"
	"path/filepath"
	"sort"
	"strconv"
	"strings"

	fs "imuslab.com/arozos/mod/filesystem"
	"imuslab.com/arozos/mod/utils"
	//user "imuslab.com/arozos/mod/user"
)

func DiskQuotaInit() {
	//Register Endpoints
	http.HandleFunc("/system/disk/quota/setQuota", system_disk_quota_setQuota)
	http.HandleFunc("/system/disk/quota/quotaInfo", system_disk_quota_handleQuotaInfo)
	http.HandleFunc("/system/disk/quota/quotaDist", system_disk_quota_handleFileDistributionView)

	//Register Setting Interfaces
	registerSetting(settingModule{
		Name:     "Storage Quota",
		Desc:     "User Remaining Space",
		IconPath: "SystemAO/disk/quota/img/small_icon.png",
		Group:    "Disk",
		StartDir: "SystemAO/disk/quota/quota.html",
	})

	//Register the timer for running the global user quota recalculation
	nightlyManager.RegisterNightlyTask(system_disk_quota_updateAllUserQuotaEstimation)
}

//Register the handler for automatically updating all user storage quota
func system_disk_quota_updateAllUserQuotaEstimation() {
	registeredUsers := authAgent.ListUsers()
	for _, username := range registeredUsers {
		//For each user, update their current quota usage
		userinfo, _ := userHandler.GetUserInfoFromUsername(username)
		userinfo.StorageQuota.CalculateQuotaUsage()
	}
}

//Set the storage quota of the particular user
func system_disk_quota_setQuota(w http.ResponseWriter, r *http.Request) {
	userinfo, err := userHandler.GetUserInfoFromRequest(w, r)
	if err != nil {
		utils.SendErrorResponse(w, "Unknown User")
		return
	}

	//Check if admin
	if !userinfo.IsAdmin() {
		utils.SendErrorResponse(w, "Permission Denied")
		return
	}

	groupname, err := utils.PostPara(r, "groupname")
	if err != nil {
		utils.SendErrorResponse(w, "Group name not defned")
		return
	}

	quotaSizeString, err := utils.PostPara(r, "quota")
	if err != nil {
		utils.SendErrorResponse(w, "Quota not defined")
		return
	}

	quotaSize, err := utils.StringToInt64(quotaSizeString)
	if err != nil || quotaSize < 0 {
		utils.SendErrorResponse(w, "Invalid quota size given")
		return
	}
	//Qutasize unit is in MB
	quotaSize = quotaSize << 20

	systemWideLogger.PrintAndLog("Quota", "Updating "+groupname+" to "+strconv.FormatInt(quotaSize, 10)+"WIP", nil)
	utils.SendOK(w)

}

func system_disk_quota_handleQuotaInfo(w http.ResponseWriter, r *http.Request) {
	userinfo, err := userHandler.GetUserInfoFromRequest(w, r)
	if err != nil {
		utils.SendErrorResponse(w, "Unknown User")
		return
	}

	//Get quota information
	type quotaInformation struct {
		Remaining int64
		Used      int64
		Total     int64
	}

	jsonString, _ := json.Marshal(quotaInformation{
		Remaining: userinfo.StorageQuota.TotalStorageQuota - userinfo.StorageQuota.UsedStorageQuota,
		Used:      userinfo.StorageQuota.UsedStorageQuota,
		Total:     userinfo.StorageQuota.TotalStorageQuota,
	})

	utils.SendJSONResponse(w, string(jsonString))

	go func() {
		//Update this user's quota estimation in go routine
		userinfo.StorageQuota.CalculateQuotaUsage()
	}()
}

var fileExtensions = map[string][]string{
	"CAD":           {"stl", "step", "stp", "obj", "dwg", "dxf", "iges", "igs", "3dm", "sat", "sldprt", "sldasm"},
	"Music":         {"mp3", "flac", "wav", "aac", "ogg", "m4a", "wma", "opus", "alac", "aiff", "mid", "midi"},
	"Video":         {"mp4", "webm", "avi", "mov", "mkv", "flv", "wmv", "m4v", "3gp", "ts", "mts", "vob"},
	"Design":        {"psd", "ai", "xd", "sketch", "fig", "indd", "eps", "cdr", "afdesign", "afphoto"},
	"Documents":     {"pdf", "doc", "docx", "odt", "rtf", "txt", "md", "pages", "tex", "wpd", "fodt", "docm"},
	"Spreadsheets":  {"xls", "xlsx", "ods", "csv", "tsv", "numbers", "xlsm", "xlsb", "fods", "ots"},
	"Presentations": {"ppt", "pptx", "odp", "key", "pptm", "fodp", "otp"},
	"Images":        {"jpg", "jpeg", "png", "gif", "bmp", "tiff", "svg", "webp", "heic", "raw", "cr2", "nef", "arw", "avif", "jxl"},
	"Archives":      {"zip", "rar", "7z", "tar", "gz", "bz2", "xz", "zst", "lz4", "cab", "arj", "lzma"},
	"Code":          {"c", "cpp", "h", "hpp", "java", "py", "js", "ts", "go", "rs", "rb", "php", "swift", "kt", "cs", "vb", "scala", "dart", "zig", "nim", "ex", "exs"},
	"Web":           {"html", "htm", "css", "xml", "json", "yaml", "yml", "toml", "graphql", "wasm", "vue", "jsx", "tsx", "svelte"},
	"Data":          {"sql", "db", "sqlite", "mdb", "accdb", "parquet", "avro", "orc", "hdf5", "feather"},
	"Executables":   {"exe", "msi", "bat", "sh", "app", "apk", "ipa", "deb", "rpm", "run", "com"},
	"System":        {"dll", "sys", "ini", "cfg", "log", "reg", "drv", "tmp", "lock", "pid"},
	"Fonts":         {"ttf", "otf", "woff", "woff2", "eot", "fon", "psf", "pfb", "pfm"},
	"Ebooks":        {"epub", "mobi", "azw", "djvu", "azw3", "fb2", "lit", "cbz", "cbr"},
	"DiskImages":    {"iso", "img", "dmg", "vhd", "vhdx", "vmdk", "qcow2", "vdi", "ova", "ovf"},
	"3DModels":      {"fbx", "blend", "dae", "3ds", "gltf", "glb", "abc", "usd", "usda", "x3d"},
	"Scripts":       {"pl", "lua", "r", "matlab", "ps1", "psm1", "bash", "zsh", "fish", "awk", "sed", "tcl"},
}

//Get all the users files and return size + count per category
func system_disk_quota_handleFileDistributionView(w http.ResponseWriter, r *http.Request) {
	userinfo, err := userHandler.GetUserInfoFromRequest(w, r)
	if err != nil {
		utils.SendErrorResponse(w, "Unknown User")
		return
	}

	// Build reverse lookup: extension → category
	extToCategory := make(map[string]string, 300)
	for category, exts := range fileExtensions {
		for _, ext := range exts {
			extToCategory[ext] = category
		}
	}

	type distEntry struct {
		Count int64
		Size  int64
	}
	fileDist := map[string]*distEntry{}

	userFileSystemHandlers := userinfo.GetAllFileSystemHandler()
	for _, thisHandler := range userFileSystemHandlers {
		if thisHandler.Hierarchy == "user" && !thisHandler.IsNetworkDrive() {
			// We only calculate the file distribution for the user's local storage, not network drives, and we only calculate for the user level handler to avoid duplication
			thispath := filepath.ToSlash(filepath.Clean(thisHandler.Path)) + "/users/" + userinfo.Username + "/"
			filepath.Walk(thispath, func(fpath string, info os.FileInfo, err error) error {
				if err != nil {
					return err
				}
				if !info.IsDir() {
					ext := strings.ToLower(strings.TrimPrefix(filepath.Ext(fpath), "."))
					category, ok := extToCategory[ext]
					if !ok {
						mime, _, merr := fs.GetMime(fpath)
						if merr != nil || mime == "" {
							category = "Others"
						} else {
							mediaType := strings.SplitN(mime, "/", 2)[0]
							category = strings.Title(mediaType)
						}
					}
					if fileDist[category] == nil {
						fileDist[category] = &distEntry{}
					}
					fileDist[category].Count++
					fileDist[category].Size += info.Size()
				}
				return nil
			})
		}
	}

	type kv struct {
		Mime  string
		Count int64
		Size  int64
	}

	var ss []kv
	for k, v := range fileDist {
		ss = append(ss, kv{k, v.Count, v.Size})
	}

	sort.Slice(ss, func(i, j int) bool {
		return ss[i].Size > ss[j].Size
	})

	jsonString, _ := json.Marshal(ss)
	utils.SendJSONResponse(w, string(jsonString))
}