filehole/main.go

package main

import (
	"crypto/rand"
	"io"
	"net/http"
	"os"
	"strconv"
	"time"

	"github.com/gabriel-vasile/mimetype"
	"github.com/gorilla/mux"

	"github.com/rs/zerolog"
	"github.com/rs/zerolog/log"

	"github.com/boltdb/bolt"
)

func shortID(length int64) string {
	const chars = "abcdefghjkmnpqrstuvwxyzABCDEFGHJKLMNPQRSTUVWXYZ123456789"
	ll := len(chars)
	b := make([]byte, length)
	rand.Read(b) // generates len(b) random bytes
	for i := int64(0); i < length; i++ {
		b[i] = chars[int(b[i])%ll]
	}
	return string(b)
}

var db *bolt.DB

func UploadHandler(w http.ResponseWriter, r *http.Request) {
	// url_len sanitize
	inpUrlLen := r.FormValue("url_len")
	sanUrlLen, err := strconv.ParseInt(inpUrlLen, 10, 64)
	if err != nil {
		sanUrlLen = 24
	}
	if sanUrlLen < 5 || sanUrlLen > 169 {
		w.Write([]byte("url_len needs to be between 5 and 169\n"))
		return
	}

	// expiry sanitize
	inpExpiry := r.FormValue("expiry")
	sanExpiry, err := strconv.ParseInt(inpExpiry, 10, 64)
	if err != nil {
		sanExpiry = 86400
	}
	if sanExpiry < 5 || sanExpiry > 432000 {
		w.Write([]byte("expiry needs to be between 5 and 432000\n"))
		return
	}

	// mimetype check
	file, _, err := r.FormFile("file")
	if err != nil {
		w.Write([]byte("error reading your file parameter\n"))
		return
	}
	defer file.Close()

	mtype, err := mimetype.DetectReader(file)
	if err != nil {
		w.Write([]byte("error detecting the mime type of your file\n"))
		return
	}
	file.Seek(0, 0)

	// ready for upload
	name := shortID(sanUrlLen) + mtype.Extension()

	err = db.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("expiry"))
		err := b.Put([]byte(name), []byte(strconv.FormatInt(time.Now().Unix()+sanExpiry, 10)))
		return err
	})
	if err != nil {
		log.Error().Err(err).Msg("Failed to put expiry")
	}

	log.Info().Str("mtype", mtype.String()).Str("ext", mtype.Extension()).Int64("expiry", sanExpiry).Int64("url_len", sanUrlLen).Msg("Writing new file")

	f, err := os.OpenFile("data/"+name, os.O_WRONLY|os.O_CREATE, 0644)
	if err != nil {
		log.Error().Err(err).Msg("Error opening a file for write")
		w.Write([]byte("internal error\n"))
		return
	}
	defer f.Close()

	io.Copy(f, file)

	w.Write([]byte("https://u.filehole.org/" + name + "\n"))
}

func ExpiryDoer() {
	for {
		removed := 0
		db.Update(func(tx *bolt.Tx) error {
			b := tx.Bucket([]byte("expiry"))
			c := b.Cursor()
			for k, v := c.First(); k != nil; k, v = c.Next() {
				expiryTime, err := strconv.ParseInt(string(v), 10, 64)
				if err != nil {
					log.Error().Err(err).Bytes("k", k).Bytes("v", v).Msg("Expiry time could not be parsed")
					continue
				}
				if time.Now().After(time.Unix(expiryTime, 0)) {
					os.Remove("data/" + string(k))
					removed += 1
					c.Delete()
				}
			}
			return nil
		})
		if removed >= 1 {
			log.Info().Int("amt", removed).Msg("Purged based on expiry")
		}
		time.Sleep(5 * time.Second)
	}
}

func main() {
	log.Logger = log.Output(zerolog.ConsoleWriter{Out: os.Stderr})

	var err error
	db, err = bolt.Open("filehole.db", 0600, nil)
	if err != nil {
		log.Fatal().Err(err).Msg("dangerous database activity")
	}
	db.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists([]byte("expiry"))
		if err != nil {
			log.Fatal().Err(err).Msg("Error creating expiry bucket")
			return err
		}
		return nil
	})

	r := mux.NewRouter()
	r.HandleFunc("/", UploadHandler).Methods("POST")
	r.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		http.ServeFile(w, r, "index.html")
	}).Methods("GET")
	http.Handle("/", r)

	go ExpiryDoer()
	http.ListenAndServe("127.0.0.1:8000", r)

	db.Close()
}