1package main
 2
 3import (
 4	"net/http"
 5	"os"
 6	"strconv"
 7
 8	"github.com/joho/godotenv"
 9	"github.com/sirupsen/logrus"
10	"golang.org/x/time/rate"
11)
12
13var (
14	debugSwitch = false
15	logger      = logrus.New()
16)
17
18func limit(limiter *rate.Limiter, next http.Handler) http.Handler {
19	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
20		if !limiter.Allow() {
21			status := http.StatusTooManyRequests
22			http.Error(w, http.StatusText(status), status)
23			return
24		}
25		next.ServeHTTP(w, r)
26	})
27}
28
29func getEnvDefault(key string, def string) string {
30	res := os.Getenv(key)
31	if res == "" {
32		return def
33	}
34	return res
35}
36
37func getEnvDefaultParse(key string, def string) float64 {
38	value := getEnvDefault(key, def)
39	res, err := strconv.ParseFloat(value, 64)
40	if err != nil {
41		logger.Fatal(err)
42	}
43	return res
44}
45
46func main() {
47	err := godotenv.Load()
48	if err != nil {
49		logger.Info("No .env file provided.")
50	}
51
52	logLevel, err := logrus.ParseLevel(os.Getenv("LOG_LEVEL"))
53	if err != nil {
54		logLevel = logrus.InfoLevel
55	}
56	logger.SetLevel(logLevel)
57
58	if logLevel == logrus.DebugLevel {
59		logger.SetLevel(logrus.DebugLevel)
60		logger.Debug("Debug mode enabled (rate limiting is disabled)")
61		debugSwitch = true
62	}
63
64	port := getEnvDefault("PORT", "3000")
65	burstTokens := getEnvDefaultParse("BURST_TOKENS", "3")
66	rateLimit := getEnvDefaultParse("RATE_LIMIT", "1")
67
68	adminUser = getEnvDefault("ADMIN_USER", "admin")
69	adminPass = getEnvDefault("ADMIN_PASS", "admin")
70
71	if adminUser == "admin" && adminPass == "admin" {
72		logger.Println("Admin credentials not set. Please set APP_ADMIN_USER and APP_ADMIN_PASS.")
73	}
74
75	r := http.NewServeMux()
76	r.HandleFunc("GET /", indexHandler)
77	r.HandleFunc("GET /watch", watchHandler)
78	r.HandleFunc("GET /shorts/{videoID}", shortHandler)
79	r.HandleFunc("GET /proxy/{videoID}", proxyHandler)
80	r.HandleFunc("GET /sub/{videoID}/{language}", subHandler)
81	r.HandleFunc("GET /refresh/{videoID}", refreshHandler)
82	r.HandleFunc("GET /cache", cacheHandler)
83	r.HandleFunc("GET /{videoID}", shortHandler)
84
85	r.HandleFunc("POST /download", downloadHandler)
86
87	var serveMux http.Handler
88	if debugSwitch {
89		serveMux = r
90	} else {
91		limiter := rate.NewLimiter(rate.Limit(rateLimit), int(burstTokens))
92		serveMux = limit(limiter, r)
93	}
94	logger.Info("Serving on port ", port)
95	http.ListenAndServe(":"+port, serveMux)
96}