package main

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

var urlDB *bolt.DB
var wordList []string
var virustotalKey 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")

	b, err = os.ReadFile("virustotal.key")
	fec(err)
	if len(b) == 0 {
		println("virustotal.key is empty")
		os.Exit(1)
	}
	virustotalKey = string(b)

	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 apiRequest(method string, url string, body io.Reader, headers map[string]string) []byte {
	req, _ := http.NewRequest(method, url, body)
	for k, v := range headers {
		req.Header.Add(k, v)
	}
	resp, _ := http.DefaultClient.Do(req)
	jay_son, _ := io.ReadAll(resp.Body)
	resp.Body.Close()
	return jay_son
}

func testUrlSafe(inspectionUrl string) bool {
	var jsonDec fastjson.Parser
	url := "https://www.virustotal.com/api/v3/urls"

	headers := map[string]string{
		"content-type": "application/x-www-form-urlencoded",
		"accept":       "application/json",
		"x-apikey":     virustotalKey,
	}

	jsonB := apiRequest("POST", url, strings.NewReader("url="+inspectionUrl), headers)

	jay_son, _ := jsonDec.ParseBytes(jsonB)
	url = string(jay_son.GetStringBytes("data", "links", "self"))

	time.Sleep(30 * time.Second)

	delete(headers, "content-type")
	jsonB = apiRequest("GET", url, nil, headers)

	jay_son, _ = jsonDec.ParseBytes(jsonB)
	harmless := jay_son.GetFloat64("data", "attributes", "stats", "harmless")
	undetected := jay_son.GetFloat64("data", "attributes", "stats", "undetected")
	malicious := jay_son.GetFloat64("data", "attributes", "stats", "malicious")
	suspicious := jay_son.GetFloat64("data", "attributes", "stats", "suspicious")

	return (malicious+suspicious)/(harmless+undetected+malicious+suspicious) > 0.25
}

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}}", resp)
	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 deleteDangerUrl(url urlStruc) {
	if !testUrlSafe(url.Url) {
		return
	}
	url.Url = "https://www.youtube.com/watch?v=dQw4w9WgXcQ"
	writeURL(url)
}

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
	}
	if !strings.Contains(url.Url, ".") || !(strings.HasPrefix(url.Url, "http://") || strings.HasPrefix(url.Url, "https://")) {
		createPage(w, api, 400, "Error: Invalid URL")
		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
		}
	}
	go deleteDangerUrl(url)
	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, "/")
	path = strings.TrimPrefix(path, r.URL.Host)
	path = strings.TrimPrefix(path, "/")
	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())
	print("\nServer stopped\n\n")
}