Untitled

package main

import (
    "archive/zip"
    "flag"
    "html/template"
    "io"
    "io/ioutil"
    "log"
    "net/http"
    "os"
    "path"
    "path/filepath"
    "strings"
    "sync"

    "github.com/gorilla/handlers"

    "github.com/gorilla/mux"
)

var (
    address                     = flag.String("a", "", "ip to bind to")
    port                        = flag.String("p", "80", "port server runs on")
    directory                   = flag.String("d", "public", "directory to serve files from")
    useBrowsePassword           = flag.Bool("use-browse-pass", false, "use a password for browsing")
    usePassword                 = flag.Bool("use-pass", false, "use the password or not")
    password                    = flag.String("pass", "", "password required to enter the website")
    passwordField               = flag.String("pass-field", "password", "name of the password cookie")
    browsePassword              = flag.String("browse-pass", "", "password required to browse files")
    browsePasswordIncorrectFile = flag.String("browse-pass-fail", "", "page to serve if browse password fails")
    fileNotFoundPage            = flag.String("404", "", "file to serve if a page is not found")
    passwordIncorrectFile       = flag.String("pass-fail", "", "file to serve if the password fails")
    browseTemplate              = flag.String("browse-template", "", "template file for the browse page")
    viewTemplate                = flag.String("view-template", "", "view template")
    hotReload                   = flag.Bool("hot-reload", false, "hot reload template files")
    faviconPath                 = flag.String("favicon", "", "favicon file")
    cssFile                     = flag.String("css", "", "custom css file")

    upgradeServerPort = flag.String("upgrade-server-port", "80", "port to host the upgrade server on")
    upgradeServer     = flag.Bool("host-upgrade-server", false, "host a server on port 80 that redirects to https")
    key               = flag.String("key", "", "key file for ssl")
    cert              = flag.String("cert", "", "cert file for ssl")
    useSSL            = flag.Bool("use-ssl", false, "specifies if to use ssl or not, requires key and cert to not be empty")
)

var css = `
body,
html {
    margin: 0px;
    padding: 0px;
}

a {
    color: white;
    font-family: roboto;
}

audio,
video,
img {
    width: 100%;
    max-width: 500px;
}

tr a:hover,
tr:nth-child(odd) a:hover,
a:hover {
    color: #aaaaff;
}

tr:nth-child(odd) a {
    color: #cccccc;
}

tr:nth-child(odd) {
    background-color: #323232;
}

tr:hover {
    background-color: #585858;
}

tr.selected {
    background-color: #5c633d;
}

td {
    padding-bottom: 5px;
    padding-right: 50px;
}

body {
    background-color: #232323;
}

h1 {
    color: white;
}

.dirlink {
    margin: 0px;
}

.header {
    overflow: hidden;
    width: 100%;
    margin: 0;
    padding-left: 10px;
    background-color: #1d1d1d;
    position: fixed;
    top: 0px;
    left: 0px;
}

.content {
    margin: 30px 0px 0px 10px;
    overflow: hidden;
}
`

var tmplBrowseMu sync.RWMutex
var tmplBrowse = template.Must(template.New("").Parse(
    `
<html>

<head>
    <style>
        {{.CSS}}
    </style>
</head>

<body>

    <div class="header">
        <!-- Directory list -->
        <a class="dirlink" href="/f/">/</a>
        {{- range .Dirlinks }}
        <a class="dirlink" href="/f/{{.Path}}">{{.Name}}</a>
        {{- end }}
    </div>

    <div class="content">


        <table>

            <!-- Ascend directory -->
            <tr>
                <td>
                    <a href="..">../</a>
                </td>
            </tr>

            <!-- Current directory -->
            <tr>
                <td>
                    <a href=".">.</a>
                </td>
                <td>
                    <a href="/z/{{.Rpath}}.zip">DOWNLOAD</a>
                </td>
            </tr>

            <!-- Files -->
            {{- range .Directories }}
            <tr>
                <td>
                    <a href="/f/{{$.Path}}{{.}}/">{{.}}/</a>
                </td>
                <td>
                    <a href="/z/{{$.Path}}{{.}}.zip">DOWNLOAD</a>
                </td>
            </tr>
            {{- end }} {{- range .Files }}
            <tr>
                <td>
                    <a href="/v/{{$.Path}}{{.}}">{{.}}</a>
                </td>
                <td>
                    <a href="/d/{{$.Path}}{{.}}">DOWNLOAD</a>
                </td>
                <td>
                    <a href="/f/{{$.Path}}{{.}}">OPEN</a>
                </td>
            </tr>
            {{- end }}
        </table>
    </div>
</body>

</html>
`))

const (
    ftypeImage = "image"
    ftypeVideo = "video"
    ftypeAudio = "audio"
)

var filetypes = map[string]string{
    ".mp3":  ftypeAudio,
    ".wav":  ftypeAudio,
    ".flac": ftypeAudio,

    ".png":  ftypeImage,
    ".jpeg": ftypeImage,
    ".jpg":  ftypeImage,
    ".gif":  ftypeImage,
    ".bmp":  ftypeImage,
    ".svg":  ftypeImage,

    ".mp4":  ftypeVideo,
    ".webm": ftypeVideo,
    ".ogg":  ftypeVideo,
}

var tmplFuncs = template.FuncMap{
    "Filetype": func(p string) string {
        if v, ok := filetypes[strings.ToLower(path.Ext(p))]; ok {
            return v
        }
        return ""
    },
}
var tmplViewMu sync.RWMutex
var tmplView = template.Must(template.New("").Funcs(tmplFuncs).Parse(
    `
<html>

<head>
    <style>
        {{.CSS}}
    </style>
</head>

<body>

    <!-- Directory list -->
    <div class="header">
        <a class="dirlink" href="/f/">/</a>
        {{- range .Dirlinks }}
        <a class="dirlink" href="/f/{{.Path}}">{{.Name}}</a>
        {{- end }}
    </div>

    <div class="content">
        {{- if .Path | Filetype | eq "image" }}
        <img src="/f/{{.Path}}"> {{- end }}

        {{- if .Path | Filetype | eq "video" }}
        <video src="/f/{{.Path}}"></video>
        {{- end }}

        {{- if .Path | Filetype | eq "audio" }}
        <audio autoplay controls src="/f/{{.Path}}"></audio>
        {{- end }}

        {{- if .Path | Filetype | eq "" }}
        <h1> Cannot view file </h1>
        {{- end }}

        <table>
                <!-- Ascend directory -->
                <tr>
                    <td>
                        <a href="/f/{{.Dir}}../">../</a>
                    </td>
                </tr>

                <!-- Current directory -->
                <tr>
                    <td>
                        <a href="/f/{{.Dir}}">.</a>
                    </td>
                    <td>
                        <a href="/z/{{.Dir}}.zip">DOWNLOAD</a>
                    </td>
                </tr>

                <!-- Files -->
                {{- range .Directories }}
                <tr>
                    <td>
                        <a href="/f/{{$.Dir}}{{.}}/">{{.}}/</a>
                    </td>
                    <td>
                        <a href="/z/{{$.Dir}}{{.}}.zip">DOWNLOAD</a>
                    </td>
                </tr>
                {{- end }} {{- range .Files }}
                <tr {{if eq . $.File }}class="selected"{{ end }}>
                    <td>
                        <a href="/v/{{$.Dir}}{{.}}">{{.}}</a>
                    </td>
                    <td>
                        <a href="/d/{{$.Dir}}{{.}}">DOWNLOAD</a>
                    </td>
                    <td>
                        <a href="/f/{{$.Dir}}{{.}}">OPEN</a>
                    </td>
                </tr>
                {{- end }}
            </table>
    </div>
</body>

</html>
`))

func parsePasswords(pass string) []string {
    return strings.Split(pass, ",")
}

func containsPassword(pass string, value string) bool {
    for _, v := range parsePasswords(pass) {
        if v == value {
            return true
        }
    }

    return false
}

func sendError(writer http.ResponseWriter, err string, code int) {
    log.Println(err)
    http.Error(writer, err, code)
}

type link struct {
    Name string
    Path string
}

// createLinks creates the directory links at the top of the page
func createLinks(p string) []link {
    dirlinks := []link{}
    elems := strings.Split(path.Clean(p), "/")
    for i := range elems {
        dirlinks = append(
            dirlinks,
            link{
                Name: elems[i],
                Path: strings.Join(elems[:i+1], "/") + "/",
            },
        )
    }
    return dirlinks
}

func readDir(fpath string) ([]string, []string, error) {
    directories, files := []string{}, []string{}

    dir, err := os.Open(fpath)
    if err != nil {
        return nil, nil, err
    }

    fs, err := dir.Readdir(-1)
    if err != nil {
        return nil, nil, err
    }

    for _, v := range fs {
        if v.IsDir() || v.Mode()&os.ModeSymlink != 0 {
            directories = append(directories, v.Name())
        } else {
            files = append(files, v.Name())
        }
    }

    return directories, files, nil
}

func browseHandler(w http.ResponseWriter, r *http.Request) {
    var (
        rpath       = mux.Vars(r)["path"]
        fpath       = path.Join(*directory, path.Clean(rpath))
        directories = []string{}
        files       = []string{}
    )

    f, err := os.Stat(fpath)
    if err != nil {
        if *fileNotFoundPage != "" {
            http.ServeFile(w, r, *fileNotFoundPage)
            return
        }
        sendError(w, "File not found: "+err.Error(), 404)
        return
    }

    if !f.IsDir() {
        http.ServeFile(w, r, fpath)
        return
    }

    if *useBrowsePassword && *browsePassword != "" {
        if cookie, err := r.Cookie(*passwordField); err != nil || !containsPassword(cookie.Value, *browsePassword) {
            if *browsePasswordIncorrectFile != "" {
                http.ServeFile(w, r, *browsePasswordIncorrectFile)
                return
            }
            if *passwordIncorrectFile != "" {
                http.ServeFile(w, r, *passwordIncorrectFile)
                return
            }
        }
    }

    directories, files, err = readDir(fpath)
    if err != nil {
        sendError(w, "error reading directory", http.StatusInternalServerError)
        return
    }

    if *hotReload {
        parseTemplates()
    }

    tmplBrowseMu.RLock()
    err = tmplBrowse.Execute(w, map[string]interface{}{
        "Directories": directories,
        "Files":       files,
        "Path":        path.Clean(rpath) + "/",
        "Rpath":       rpath,
        "Base":        path.Base(rpath),
        "Dirlinks":    createLinks(rpath),
        "CSS":         template.HTML("<style>" + css + "</style>"),
    })
    tmplBrowseMu.RUnlock()

    if err != nil {
        sendError(w, "Error executing template: "+err.Error(), http.StatusInternalServerError)
        return
    }
}

func downloadHandler(trimZip bool) func(w http.ResponseWriter, r *http.Request) {
    return func(w http.ResponseWriter, r *http.Request) {
        var (
            fpath = path.Join(*directory, path.Clean(mux.Vars(r)["path"]))
        )

        if trimZip {
            fpath = strings.TrimSuffix(fpath, ".zip")
        }

        stat, err := os.Stat(fpath)
        if err != nil {
            sendError(w, "Error getting stats: "+err.Error(), http.StatusInternalServerError)
            return
        }

        w.Header().Set("Content-Disposition", "attachment")

        if stat.IsDir() || stat.Mode()&os.ModeSymlink != 0 {

            if strings.HasSuffix(mux.Vars(r)["path"], "/.zip") {
                filename := "content.zip"
                w.Header().Set("Content-Disposition", `attachment; filename="`+filename+`"`)
            }

            err := archive(fpath, w)
            if err != nil {
                log.Println(err)
                return
            }
        } else {
            http.ServeFile(w, r, fpath)
        }
    }
}

func viewHandler(w http.ResponseWriter, r *http.Request) {
    var (
        rpath = mux.Vars(r)["path"]
        fpath = path.Join(*directory, path.Clean(path.Dir(rpath)))
    )

    if *hotReload {
        parseTemplates()
    }

    stat, err := os.Stat(path.Join(*directory, path.Clean(rpath)))
    if err != nil || stat.IsDir() {
        http.ServeFile(w, r, *fileNotFoundPage)
        return
    }

    directories, files, err := readDir(fpath)
    if err != nil {
        sendError(w, "error reading directory", http.StatusInternalServerError)
        return
    }

    tmplViewMu.RLock()
    err = tmplView.Execute(w, map[string]interface{}{
        "Dirlinks":    createLinks(rpath),
        "Path":        path.Clean(rpath),
        "Directories": directories,
        "Files":       files,
        "Dir":         path.Dir(rpath) + "/",
        "File":        path.Base(rpath),
        "CSS":         template.HTML("<style>" + css + "</style>"),
    })
    if err != nil {
        log.Println("Error executing view template: ", err)
    }
    tmplViewMu.RUnlock()
}

func archive(inFilePath string, writer io.Writer) error {
    zipWriter := zip.NewWriter(writer)

    basePath := filepath.Dir(inFilePath)

    err := filepath.Walk(inFilePath, func(filePath string, fileInfo os.FileInfo, err error) error {
        if err != nil || fileInfo.IsDir() {
            return err
        }

        relativeFilePath, err := filepath.Rel(basePath, filePath)
        if err != nil {
            return err
        }

        archivePath := path.Join(filepath.SplitList(relativeFilePath)...)

        file, err := os.Open(filePath)
        if err != nil {
            return err
        }
        defer func() {
            _ = file.Close()
        }()

        zipFileWriter, err := zipWriter.Create(archivePath)
        if err != nil {
            return err
        }

        _, err = io.Copy(zipFileWriter, file)
        return err
    })
    if err != nil {
        return err
    }

    return zipWriter.Close()
}

func passwordMiddleware(h http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        if c, err := r.Cookie(*passwordField); err == nil && containsPassword(c.Value, *password) {
            h.ServeHTTP(w, r)
            return
        }

        http.ServeFile(w, r, *passwordIncorrectFile)
    })
}

func parseTemplates() {
    tmplBrowseMu.Lock()
    defer tmplBrowseMu.Unlock()

    tmplViewMu.Lock()
    defer tmplViewMu.Unlock()

    if *viewTemplate != "" {
        v, err := template.New("view.html").Funcs(tmplFuncs).ParseFiles(*viewTemplate)
        if err != nil {
            log.Println(err)
            return
        }
        tmplView = v
    }
    if *browseTemplate != "" {
        t, err := template.ParseFiles(*browseTemplate)
        if err != nil {
            log.Println(err)
            return
        }
        tmplBrowse = t
    }
    if *cssFile != "" {
        if t, err := ioutil.ReadFile(*cssFile); err == nil {
            css = string(t)
        } else {
            log.Println("error loading CSS: ", err)
        }
    }

}

func fileHandler(path string) func(w http.ResponseWriter, r *http.Request) {
    return func(w http.ResponseWriter, r *http.Request) {
        http.ServeFile(w, r, path)
    }
}

func main() {
    flag.Parse()

    if *browseTemplate != "" || *viewTemplate != "" || *cssFile != "" {
        log.Println("parsing templates")
        parseTemplates()
    }

    root := mux.NewRouter()

    // Public routes
    if *faviconPath != "" {
        root.HandleFunc("/favicon", fileHandler(*faviconPath))
    }

    // Password locked routes
    func(r *mux.Router) {
        if *usePassword && *password != "" {
            r.Use(passwordMiddleware)
        }

        r.HandleFunc("/f/{path:.*}", browseHandler)
        r.HandleFunc("/d/{path:.*}", downloadHandler(false))
        r.HandleFunc("/z/{path:.*}", downloadHandler(true))
        r.HandleFunc("/v/{path:.*}", viewHandler)
        r.PathPrefix("/").Handler(http.RedirectHandler("/f/", 302))

    }(root.PathPrefix("/").Subrouter())

    log.Println("Server listening on " + *port)
    var err error
    if *key != "" && *cert != "" && *useSSL {

        // Start upgrade server. Upgrades requests from http to https via redirect
        if *upgradeServer {
            log.Println("Hosting upgrade server on port 80")
            go func() {
                err := http.ListenAndServe(*address+":"+*upgradeServerPort, handlers.LoggingHandler(os.Stdout, http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
                    http.Redirect(w, req, "https://"+req.Host+req.URL.String(), 302)
                })))
                if err != nil {
                    log.Println("Error starting upgrade server")
                    return
                }
            }()
        }

        err = http.ListenAndServeTLS(*address+":"+*port, *cert, *key, handlers.LoggingHandler(os.Stdout, root))
        if err != nil {
            log.Fatal(err)
        }

    } else {
        err = http.ListenAndServe(*address+":"+*port, handlers.LoggingHandler(os.Stdout, root))
        if err != nil {
            log.Fatal(err)
        }
    }
}