M .gitignore → .gitignore

@@ -1,3 +1,4 @@ 
 .idea
 urlshort
-urls.db+urls.db
+virustotal.key

M go.mod → go.mod

@@ -3,6 +3,7 @@ 
 go 1.21.0
 
 require (
+	github.com/valyala/fastjson v1.6.4 // indirect
 	go.etcd.io/bbolt v1.3.7 // indirect
 	golang.org/x/sys v0.4.0 // indirect
 )

M go.sum → go.sum

@@ -1,3 +1,5 @@ 
+github.com/valyala/fastjson v1.6.4 h1:uAUNq9Z6ymTgGhcm0UynUAB6tlbakBrz6CQFax3BXVQ=
+github.com/valyala/fastjson v1.6.4/go.mod h1:CLCAqky6SMuOcxStkYQvblddUtoRxhYMGLrsQns1aXY=
 go.etcd.io/bbolt v1.3.7 h1:j+zJOnnEjF/kyHlDDgGnVL/AIqIJPq8UoB2GSNfkUfQ=
 go.etcd.io/bbolt v1.3.7/go.mod h1:N9Mkw9X8x5fupy0IKsmuqVtoGDyxsaDlbk4Rd05IAQw=
 golang.org/x/sys v0.4.0 h1:Zr2JFtRQNX3BCZ8YtxRE9hNJYC8J6I1MVbMg6owUp18=

M main.go → main.go

@@ -4,7 +4,9 @@ import (
 	"context"
 	"encoding/json"
 	"errors"
+	"github.com/valyala/fastjson"
 	bolt "go.etcd.io/bbolt"
+	"io"
 	"math/rand"
 	"net/http"
 	"os"
@@ -15,6 +17,7 @@ )
 
 var urlDB *bolt.DB
 var wordList []string
+var virustotalKey string
 
 type urlStruc struct {
 	Date time.Time
@@ -40,6 +43,14 @@ 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)
@@ -86,6 +97,46 @@ 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")
@@ -178,6 +229,14 @@ 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
@@ -196,6 +255,10 @@ 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() + "-"
@@ -212,6 +275,7 @@ createPage(w, api, 400, "Error: Custom URL already exists")
 			return
 		}
 	}
+	go deleteDangerUrl(url)
 	err = writeURL(url)
 	if hec(w, err, 500) {
 		return