Security for /download and frontend style improvementsHEADmaster

1 files changed, 51 insertions, 20 deletions

diff --git a/main.go b/main.go
index 7681c86..4ece388 100644
--- a/main.go
+++ b/main.go
@@ -356,33 +356,64 @@ func handleFileServer(dir, prefix string) http.HandlerFunc {
 }
 
 func download(w http.ResponseWriter, r *http.Request) {
-	switch r.Method {
-	case http.MethodPost:
-		r.ParseForm()
-		selection := r.Form["selection"]
-		if len(selection) == 0 {
-			sendFlashError(w, r, "/", errors.New("Aucun fichier sélectionné"))
-			return
-		}
-		contentDisposition := fmt.Sprintf("attachment; filename=\"Documents.zip\"")
-		w.Header().Set("Content-Disposition", contentDisposition)
-		wr := zip.NewWriter(w)
-		defer wr.Close()
-		for _, sel := range selection {
-			wrc, err := wr.Create(filepath.Base(sel))
-			if err != nil {
-				sendError(w, r, err.Error(), http.StatusInternalServerError)
+	u, err := checkSession(w, r)
+	if u != nil && err == nil {
+		switch r.Method {
+		case http.MethodPost:
+			r.ParseForm()
+			selection := r.Form["selection"]
+			if len(selection) == 0 {
+				sendFlashError(w, r, "/", errors.New("Aucun fichier sélectionné"))
 				return
 			}
-			f, err := os.Open(sel)
+			contentDisposition := fmt.Sprintf("attachment; filename=\"Documents.zip\"")
+			w.Header().Set("Content-Disposition", contentDisposition)
+			wr := zip.NewWriter(w)
+			defer wr.Close()
+			for _, sel := range selection {
+				if filepath.Base(filepath.Dir(sel)) == u.User {
+					wrc, err := wr.Create(filepath.Base(sel))
+					if err != nil {
+						sendError(w, r, err.Error(), http.StatusInternalServerError)
+						return
+					}
+					f, err := os.Open(sel)
+					if err != nil {
+						sendError(w, r, err.Error(), http.StatusInternalServerError)
+						return
+					}
+					io.Copy(wrc, f)
+				}
+			}
+		case http.MethodGet:
+			contentDisposition := fmt.Sprintf("attachment; filename=\"Documents.zip\"")
+			w.Header().Set("Content-Disposition", contentDisposition)
+			wr := zip.NewWriter(w)
+			defer wr.Close()
+			files, err := os.ReadDir(filepath.Join(baseDocDir, u.User))
 			if err != nil {
 				sendError(w, r, err.Error(), http.StatusInternalServerError)
 				return
 			}
-			io.Copy(wrc, f)
+			for _, file := range files {
+				filePath := path.Join(baseDocDir, u.User, file.Name())
+				wrc, err := wr.Create(filepath.Base(filePath))
+				if err != nil {
+					sendError(w, r, err.Error(), http.StatusInternalServerError)
+					return
+				}
+				f, err := os.Open(filePath)
+				if err != nil {
+					sendError(w, r, err.Error(), http.StatusInternalServerError)
+					return
+				}
+				io.Copy(wrc, f)
+			}
+		default:
+			sendInvalidMethod(w, r)
 		}
-	default:
-		sendInvalidMethod(w, r)
+	} else {
+		http.Redirect(w, r, "/login", http.StatusSeeOther)
 	}
 }