125 lines
2.3 KiB
Go
125 lines
2.3 KiB
Go
package handlers
|
|
|
|
import (
|
|
"archive/zip"
|
|
"encoding/json"
|
|
"fmt"
|
|
"lash"
|
|
"lash/internal/errx"
|
|
"net/http"
|
|
"os"
|
|
"strconv"
|
|
)
|
|
|
|
type FileHandler struct {
|
|
Ctx *lash.LashContext
|
|
FileData []FileData
|
|
}
|
|
|
|
type FileData struct {
|
|
Contents []byte
|
|
FileName string
|
|
}
|
|
|
|
type ValidateRequest struct {
|
|
Token string
|
|
}
|
|
|
|
type headers map[string]string
|
|
|
|
var zipHeaders headers = headers{
|
|
"Content-Disposition": "attachment; filename=lash.zip",
|
|
"Content-Type": "application/octet-stream",
|
|
}
|
|
|
|
var sent int = 0
|
|
|
|
func (h FileHandler) APIHandler(w http.ResponseWriter, r *http.Request) {
|
|
if h.hasHitMax() {
|
|
w.WriteHeader(http.StatusTooManyRequests)
|
|
os.Exit(0)
|
|
return
|
|
}
|
|
decoder := json.NewDecoder(r.Body)
|
|
|
|
var t ValidateRequest
|
|
err := decoder.Decode(&t)
|
|
if err != nil {
|
|
http.Error(w, err.Error(), http.StatusBadRequest)
|
|
return
|
|
}
|
|
|
|
if t.Token != lash.Token {
|
|
http.Error(w, "Invalid Token", http.StatusUnauthorized)
|
|
return
|
|
}
|
|
|
|
w.WriteHeader(http.StatusOK)
|
|
if len(h.FileData) == 1 {
|
|
headers := headers{
|
|
"Content-Disposition": fmt.Sprintf("attachment; filename=%s", h.FileData[0].FileName),
|
|
"Content-Type": "application/octet-stream",
|
|
"Content-Length": strconv.Itoa(len(h.FileData[0].FileName)),
|
|
}
|
|
headers.set(w)
|
|
|
|
w.Write(h.FileData[0].Contents)
|
|
} else {
|
|
zipHeaders.set(w)
|
|
h.writeZip(w)
|
|
}
|
|
sent++
|
|
}
|
|
|
|
func (h FileHandler) FileHandler(w http.ResponseWriter, r *http.Request) {
|
|
if h.hasHitMax() {
|
|
w.WriteHeader(http.StatusTooManyRequests)
|
|
os.Exit(0)
|
|
return
|
|
}
|
|
|
|
w.WriteHeader(http.StatusOK)
|
|
if len(h.FileData) == 1 {
|
|
headers := headers{
|
|
"Content-Disposition": fmt.Sprintf("attachment; filename=%s", h.FileData[0].FileName),
|
|
"Content-Type": "application/octet-stream",
|
|
"Content-Length": strconv.Itoa(len(h.FileData[0].FileName)),
|
|
}
|
|
headers.set(w)
|
|
|
|
w.Write(h.FileData[0].Contents)
|
|
|
|
} else {
|
|
zipHeaders.set(w)
|
|
h.writeZip(w)
|
|
}
|
|
sent++
|
|
}
|
|
|
|
func (h FileHandler) writeZip(hw http.ResponseWriter) {
|
|
zw := zip.NewWriter(hw)
|
|
defer zw.Close()
|
|
|
|
for _, f := range h.FileData {
|
|
w, err := zw.Create(f.FileName)
|
|
if err != nil {
|
|
errx.FatalPerror(err)
|
|
}
|
|
|
|
_, err = w.Write(f.Contents)
|
|
if err != nil {
|
|
errx.FatalPerror(err)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (h headers) set(w http.ResponseWriter) {
|
|
for k, v := range h {
|
|
w.Header().Set(k, v)
|
|
}
|
|
}
|
|
|
|
func (h FileHandler) hasHitMax() bool {
|
|
return sent >= h.Ctx.N && h.Ctx.N != -1
|
|
}
|