nirvash/archetype/fileManager.go

package archetype

import (
	"errors"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"strings"
)

type SimpleFileManager struct {
	Root        string
	ShowHTML    bool
	ShowHidden  bool
	maxUploadMB int64
}

type FileData struct {
	Error  string
	Path   string
	Name   string
	Parent string
	IsDir  bool
}

type FileListing struct {
	Error   string
	Root    string
	Up      string
	SubDirs []string
	Files   []string
}

type FileManager interface {
	Init(cfg *Config) error
	ListSubTree(root string) FileListing
	GetFileData(slug string) FileData
	AddFile(path string, req *http.Request) error
	MkDir(path, newDir string) error
	Remove(path string) error
	Rename(oldFullPath, newPath, newName string) error
	MaxUploadMB() int64
}

func (self *SimpleFileManager) Init(cfg *Config) error {
	self.Root = filepath.Clean(cfg.StaticRoot)
	self.ShowHTML = cfg.StaticShowHTML
	self.ShowHidden = cfg.StaticShowHidden
	self.maxUploadMB = cfg.StaticMaxUploadMB
	return nil
}

func (self *SimpleFileManager) ListSubTree(root string) FileListing {
	list := FileListing{}

	fullPath := filepath.Join(self.Root, root)

	if !strings.HasPrefix(filepath.Clean(fullPath), self.Root) {
		list.Error = "You cannot escape!"
		return list
	}

	files, err := ioutil.ReadDir(fullPath)

	if err != nil {
		list.Error = err.Error()
		return list
	}

	list.Root = root
	if !strings.HasSuffix(list.Root, "/") {
		list.Root += "/"
	}
	if !strings.HasPrefix(list.Root, "/") {
		list.Root = "/" + list.Root
	}

	levels := strings.Split(root, "/")

	if list.Root != "/" {
		list.Up = "/"
	}
	if len(levels) >= 2 && list.Root != "/" {
		list.Up = strings.Join(levels[:len(levels)-1], "/")
		if !strings.HasPrefix(list.Up, "/") {
			list.Up = "/" + list.Up
		}
	}

	for _, file := range files {
		if !self.ShowHidden && strings.HasPrefix(file.Name(), ".") {
			continue
		}
		if file.IsDir() {
			list.SubDirs = append(list.SubDirs, file.Name())
		} else {
			if !self.ShowHTML && strings.HasSuffix(file.Name(), ".html") {
				continue
			}
			list.Files = append(list.Files, file.Name())
		}
	}

	return list
}

func (self *SimpleFileManager) GetFileData(slug string) FileData {
	fullPath := filepath.Join(self.Root, slug)
	fileInfo, err := os.Stat(fullPath)

	if err != nil {
		return FileData{
			Error: err.Error(),
		}
	}
	if !strings.HasPrefix(filepath.Clean(fullPath), self.Root) {
		return FileData{
			Error: "You cannot escape!",
		}
	}

	cleanedSlug := filepath.Clean(slug)
	fileBase := filepath.Base(cleanedSlug)
	parent := strings.TrimSuffix("/"+cleanedSlug, "/"+fileBase)
	return FileData{
		Path:   filepath.Clean(slug),
		Name:   fileBase,
		Parent: parent,
		IsDir:  fileInfo.IsDir(),
	}
}

func (self *SimpleFileManager) Remove(slug string) error {
	fullPath := filepath.Join(self.Root, slug)

	if !strings.HasPrefix(filepath.Clean(fullPath), self.Root) {
		return errors.New("You cannot escape!")
	}

	_, err := os.Stat(fullPath)
	if err != nil {
		return err
	}

	return os.RemoveAll(fullPath)
}

func (self *SimpleFileManager) AddFile(path string, req *http.Request) error {
	fullPath := filepath.Join(self.Root, path)

	if !strings.HasPrefix(filepath.Clean(fullPath), filepath.Clean(self.Root)) {
		return errors.New("You cannot escape!")
	}

	_, err := os.Stat(fullPath)
	if err != nil {
		return err
	}

	req.ParseMultipartForm(self.maxUploadMB << 20)
	file, header, err := req.FormFile("file")
	if err != nil {
		return err
	}

	fileData, err := ioutil.ReadAll(file)
	if err != nil {
		return err
	}

	destPath := filepath.Join(fullPath, header.Filename)
	dest, err := os.Create(destPath)
	if err != nil {
		return err
	}
	defer dest.Close()

	dest.Write(fileData)
	return nil
}

func (self *SimpleFileManager) MkDir(path, newDir string) error {
	fullPath := filepath.Join(self.Root, path)
	if !strings.HasPrefix(filepath.Clean(fullPath), self.Root) {
		return errors.New("You cannot escape!")
	}

	_, err := os.Stat(fullPath)
	if err != nil {
		return err
	}

	if strings.Contains(newDir, "/") || strings.Contains(newDir, "\\") {
		return errors.New("Cannot create nested directories at once")
	}

	newDirPath := filepath.Join(fullPath, newDir)
	_, err = os.Stat(newDirPath)
	if !os.IsNotExist(err) {
		return errors.New("Directory exists")
	}

	return os.Mkdir(newDirPath, 0755)
}

func (self *SimpleFileManager) Rename(oldFullPath, newPath, newName string) error {
	fullPath := filepath.Join(self.Root, oldFullPath)

	if !strings.HasPrefix(filepath.Clean(fullPath), self.Root) {
		return errors.New("You cannot escape!")
	}

	_, err := os.Stat(fullPath)
	if err != nil {
		return err
	}

	newParent := filepath.Clean(filepath.Join(self.Root, newPath))
	if !strings.HasPrefix(newParent, self.Root) {
		return errors.New("You cannot escape!")
	}

	_, err = os.Stat(newParent)
	if err != nil {
		return err
	}

	if newName == "" {
		_, oldName := filepath.Split(oldFullPath)
		newName = oldName
	}

	newFullPath := filepath.Join(newParent, newName)
	if !strings.HasPrefix(filepath.Clean(newFullPath), newParent) {
		return errors.New("You cannot escape!")
	}

	return os.Rename(fullPath, filepath.Join(newParent, newName))
}

func (self *SimpleFileManager) MaxUploadMB() int64 {
	return self.maxUploadMB
}