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