296
            

             package main

import (
	"context"
	"encoding/json"
	"errors"
	bolt "go.etcd.io/bbolt"
	"math/rand"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"time"
)

var urlDB *bolt.DB
var wordList []string

type urlStruc struct {
	date time.Time
	path string
	url  string
	uses int
}

func fec(err error) { // fatal error check
	if err != nil {
		println(err.Error())
		urlDB.Close()
		os.Exit(1)
	}
}

func openFiles() {
	b, err := os.ReadFile("wordlist")
	fec(err)
	if len(b) == 0 {
		println("wordlist is empty")
		os.Exit(1)
	}
	wordList = strings.Split(string(b), "\n")

	var db *bolt.DB
	db, err = bolt.Open("urls.db", 0600, nil)
	fec(err)
	fec(db.Update(func(tx *bolt.Tx) error {
		_, err = tx.CreateBucketIfNotExists([]byte("urls"))
		return err
	}))
	urlDB = db
}

func getWord() string {
	return wordList[rand.Intn(len(wordList))]
}

func readURL(words string) (urlStruc, error) {
	if words == "" {
		return urlStruc{}, errors.New("No words provided")
	}
	var url urlStruc
	var err error
	err = urlDB.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("urls"))
		return json.Unmarshal(b.Get([]byte(words)), &url)
	})
	return url, err
}

func checkURL(words string) bool {
	url, err := readURL(words)
	if err == nil && url.url != "" {
		return true
	}
	return false
}

func writeURL(url urlStruc) error {
	jsonUrl, err := json.Marshal(url)
	if err != nil {
		return err
	}
	return urlDB.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte("urls"))
		return b.Put([]byte(url.path), jsonUrl)
	})
}

func http404(w http.ResponseWriter) {
	w.WriteHeader(404)
	b, err := os.ReadFile("html/404.html")
	if err != nil {
		b = []byte("404 not found")
		w.Header().Set("Content-Type", "text/plain")
	} else {
		w.Header().Set("Content-Type", "text/html")
	}
	w.Write(b)
}

func http500(w http.ResponseWriter) {
	w.WriteHeader(500)
	b, err := os.ReadFile("html/500.html")
	if err != nil {
		b = []byte("500 internal server error")
		w.Header().Set("Content-Type", "text/plain")
	} else {
		w.Header().Set("Content-Type", "text/html")
	}
	w.Write(b)
}

func hec(w http.ResponseWriter, err error, code int) bool {
	if err == nil {
		return false
	}
	switch code {
	case 404:
		http404(w)
		break
	case 500:
		http500(w)
		break
	}
	return true
}

func intStr(n int) string {
	if n == 0 {
		return "0"
	}
	var s, sign string
	if n < 0 {
		sign = "-"
		n = -n
	}
	for ; n > 0; n /= 10 {
		s = string(n%10+int('0')) + s
	}
	return sign + s
}

func infoPage(w http.ResponseWriter, url urlStruc) {
	w.WriteHeader(200)
	w.Header().Set("Content-Type", "text/html")
	b, err := os.ReadFile("html/info.html")
	if hec(w, err, 500) {
		return
	}
	s := strings.ReplaceAll(string(b), "{{url}}", url.url)
	s = strings.ReplaceAll(s, "{{path}}", url.path)
	s = strings.ReplaceAll(s, "{{uses}}", intStr(url.uses))
	s = strings.ReplaceAll(s, "{{date}}", url.date.Format("2006-01-02 15:04:05"))
	w.Write([]byte(s))
}

func createPage(w http.ResponseWriter, api bool, code int, resp string) {
	w.WriteHeader(code)
	if api {
		w.Header().Set("Content-Type", "text/plain")
		w.Write([]byte(resp))
		return
	}
	s := "Error loading HTML\nYour request was still processed.\nRequest response: " + resp + "\n"
	b, err := os.ReadFile("html/create.html")
	if err != nil {
		w.Header().Set("Content-Type", "text/plain")
		w.Write([]byte(s))
		return
	}
	s = strings.ReplaceAll(string(b), "{{response}}", s)
	if code != 200 {
		s = strings.ReplaceAll(s, "{{title}}", "Your request failed:")
	} else {
		s = strings.ReplaceAll(s, "{{title}}", "Your short URL:")
	}
	w.Header().Set("Content-Type", "text/html")
	w.Write([]byte(s))
}

func createHandler(w http.ResponseWriter, r *http.Request) {
	var err error
	var url urlStruc
	var api bool = false
	url.date = time.Now()
	url.uses = 0
	err = r.ParseForm()
	if hec(w, err, 500) {
		return
	}
	if r.Form.Get("mode") == "api" {
		api = true
	}
	url.url = r.Form.Get("url")
	if url.url == "" {
		createPage(w, api, 400, "Error: No URL provided")
		return
	}
	url.path = r.Form.Get("custom")
	if url.path == "" {
		url.path = getWord() + "-" + getWord() + "-"
		for {
			url.path += getWord()
			if !checkURL(url.path) {
				break
			}
			url.path += "-"
		}
	} else {
		if checkURL(url.path) {
			createPage(w, api, 400, "Error: Custom URL already exists")
			return
		}
	}
	err = writeURL(url)
	if hec(w, err, 500) {
		return
	}
	createPage(w, api, 200, r.URL.Host+"/"+url.path)
}

func httpHandler(w http.ResponseWriter, r *http.Request) {
	if len(r.URL.Path) < 2 {
		if r.Method == "POST" {
			createHandler(w, r)
			return
		}
		http.ServeFile(w, r, "html/index.html")
		return
	}
	path := strings.TrimSuffix(r.URL.Path[1:], "/")
	var err error
	if path == "favicon.ico" {
		_, err = os.Stat("favicon.ico")
		if hec(w, err, 404) {
			return
		}
		http.ServeFile(w, r, "favicon.ico")
		return
	}
	var info bool = false
	if strings.HasSuffix(path, "/info") {
		info = true
		path = strings.TrimSuffix(path, "/info")
	}
	var url urlStruc
	url, err = readURL(path)
	if err != nil || url.url == "" {
		http404(w)
		return
	}
	if info {
		infoPage(w, url)
		return
	}
	http.Redirect(w, r, url.url, http.StatusFound)
	url.uses++
	_ = writeURL(url)
}

func main() {
	var (
		port   string = ":8088"
		err    error
		server http.Server
	)

	if len(os.Args) == 2 {
		port = ":" + os.Args[1]
	}

	println("Starting server on port " + port)

	openFiles()

	http.HandleFunc("/", httpHandler)

	server = http.Server{
		Addr:    port,
		Handler: nil,
	}

	go func() {
		err = server.ListenAndServe()
		if err != http.ErrServerClosed && err != nil {
			urlDB.Close()
			panic(err)
		}
	}()
	quit := make(chan os.Signal, 1)
	signal.Notify(quit, os.Interrupt)
	<-quit

	urlDB.Close()
	server.Shutdown(context.Background())
}