CCTV/internal/server.go

package internal

import (
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"text/template"
	"time"

	"github.com/golang-jwt/jwt/v5"
)

var (
	MaxFailedAttempts = 3
	BanDuration       = 5 * time.Minute
	UploadDir         = ""
	failedAttempts    = make(map[string]int)
	bannedIPs         = make(map[string]time.Time)
	JWT_Secret        = ""
)

type Claims struct {
	Username string `json:"username"`
	jwt.RegisteredClaims
}

type picture struct {
	picture string
}

func SetShit() {
	if maxAttempts, err := strconv.Atoi(os.Getenv("CCTV_MaxFailedAttempts")); err == nil {
		MaxFailedAttempts = maxAttempts
	}
	if banDuration, err := time.ParseDuration(os.Getenv("CCTV_BanDuration")); err == nil {
		BanDuration = banDuration * time.Minute
	}

	JWT_Secret = os.Getenv("CCTV_JWT_Secret")
	UploadDir = os.Getenv("CCTV_UploadDir")
}

func LandingHandler(w http.ResponseWriter, r *http.Request) {
	content, err := os.ReadFile("views/index.html")
	if err != nil {
		http.Error(w, "Error reading file", http.StatusInternalServerError)
		return
	}
	fmt.Fprint(w, string(content))
}

func LoginHandler(w http.ResponseWriter, r *http.Request) {
	ip := getClientIP(r)

	mu.Lock()

	if banTime, banned := bannedIPs[ip]; banned {
		fmt.Print("Banned IP: ", ip, " until ", banTime, "\n")
		if time.Now().Before(banTime) {
			bannedIPs[ip] = time.Now().Add(BanDuration)
			http.Error(w, "You are banned", http.StatusForbidden)
			mu.Unlock()
			return
		}
		delete(bannedIPs, ip)
	}

	mu.Unlock()

	username := r.FormValue("username")
	password := r.FormValue("password")

	if username == "admin" && password == "admin" {
		exTime := time.Now().Add(5 * time.Minute)
		claims := &Claims{
			Username: username,
			RegisteredClaims: jwt.RegisteredClaims{
				ExpiresAt: jwt.NewNumericDate(exTime),
			},
		}

		token := jwt.NewWithClaims(jwt.SigningMethodHS256, claims)
		tokenString, err := token.SignedString([]byte(JWT_Secret))
		if err != nil {
			http.Error(w, "Error signing token", http.StatusInternalServerError)
			return
		}

		http.SetCookie(w, &http.Cookie{
			Name:     "token",
			Value:    tokenString,
			Expires:  exTime,
			SameSite: http.SameSiteStrictMode,
			Secure:   true,
		})
		http.Redirect(w, r, "/dash", http.StatusSeeOther)
		return
	} else {
		LogFailedAttempt(ip)

		mu.Lock()
		failedAttempts[ip]++
		if failedAttempts[ip] >= MaxFailedAttempts {
			bannedIPs[ip] = time.Now().Add(BanDuration)
			delete(failedAttempts, ip)

			LogBan(ip)
		}
		mu.Unlock()

		http.Error(w, "Unauthorized", http.StatusUnauthorized)
	}

}

func AuthClientMiddleware(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		authCookie, err := r.Cookie("token")
		if err != nil {
			http.Error(w, "Forbidden", http.StatusForbidden)
			return
		}

		claims := &Claims{}
		token, err := jwt.ParseWithClaims(authCookie.Value, claims, func(token *jwt.Token) (interface{}, error) {
			return []byte(JWT_Secret), nil
		})

		if err != nil || !token.Valid {
			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}

		exTime := time.Now().Add(5 * time.Minute)
		claims.ExpiresAt = jwt.NewNumericDate(exTime)
		tokenString, err := jwt.NewWithClaims(jwt.SigningMethodHS256, claims).SignedString([]byte(JWT_Secret))
		if err == nil {
			http.SetCookie(w, &http.Cookie{
				Name:     "token",
				Value:    tokenString,
				Expires:  exTime,
				SameSite: http.SameSiteStrictMode,
				Secure:   true,
			})
		}

		next(w, r)
	}
}

func DashboardHandler(w http.ResponseWriter, r *http.Request) {
	// get the newest picture name
	pictures, err := os.ReadDir(UploadDir)
	if err != nil {
		http.Error(w, "Error reading directory", http.StatusInternalServerError)
		return
	}

	if len(pictures) == 0 {
		http.Error(w, "No pictures found", http.StatusNotFound)
		return
	}

	// get the newest picture
	newestPicture := pictures[0]
	for _, picture := range pictures {
		pictureInfo, err := picture.Info()
		if err != nil {
			http.Error(w, "Error getting file info", http.StatusInternalServerError)
			return
		}
		newestPictureInfo, err := newestPicture.Info()
		if err != nil {
			http.Error(w, "Error getting file info", http.StatusInternalServerError)
			return
		}
		if pictureInfo.ModTime().After(newestPictureInfo.ModTime()) {
			newestPicture = picture
		}
	}

	p := picture{picture: newestPicture.Name()}

	tmpl := template.Must(template.ParseFiles("views/dash.html"))
	tmpl.Execute(w, p)
}

func UploadHandler(w http.ResponseWriter, r *http.Request) {
	err := r.ParseMultipartForm(10 << 20)
	if err != nil {
		http.Error(w, "Error parsing form", http.StatusBadRequest)
		return
	}

	file, handler, err := r.FormFile("file")
	if err != nil {
		http.Error(w, "Error retrieving file", http.StatusBadRequest)
		return
	}

	defer file.Close()

	filePath := filepath.Join(UploadDir, handler.Filename)

	if _, err := os.Stat(filePath); err == nil {
		http.Error(w, "File already exists", http.StatusBadRequest)
		return
	}

	out, err := os.Create(filePath)
	if err != nil {
		http.Error(w, "Error creating file", http.StatusInternalServerError)
		return
	}

	defer out.Close()

	_, err = io.Copy(out, file)
	if err != nil {
		http.Error(w, "Error saving file", http.StatusInternalServerError)
		return
	}

	LogSuccessUpload(getClientIP(r), handler.Filename)
}

func AuthMiddleware(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		ip := getClientIP(r)

		mu.Lock()

		// Check if IP is banned
		if banTime, banned := bannedIPs[ip]; banned {
			fmt.Print("Banned IP: ", ip, " until ", banTime, "\n")
			if time.Now().Before(banTime) {
				bannedIPs[ip] = time.Now().Add(BanDuration)

				http.Error(w, "You are banned", http.StatusForbidden)
				mu.Unlock()
				return
			}
			delete(bannedIPs, ip)
		}

		mu.Unlock()

		authHeader := r.Header.Get("Authorization")
		fmt.Println(authHeader,"-",  os.Getenv("CCTV_BearerToken"))
		if !strings.HasPrefix(authHeader, "Bearer ") || strings.TrimPrefix(authHeader, "Bearer ") != os.Getenv("CCTV_BearerToken") {
			LogFailedAttempt(ip)

			mu.Lock()
			failedAttempts[ip]++
			if failedAttempts[ip] >= MaxFailedAttempts {
				bannedIPs[ip] = time.Now().Add(BanDuration)
				delete(failedAttempts, ip)

				LogBan(ip)
			}
			mu.Unlock()

			http.Error(w, "Unauthorized", http.StatusUnauthorized)
			return
		}

		mu.Lock()
		delete(failedAttempts, ip)
		mu.Unlock()

		next(w, r)
	}

}

func HandleFileServer(dir, prefix string) http.HandlerFunc {
	fs := http.FileServer(http.Dir(dir))
	realHandler := http.StripPrefix(prefix, fs).ServeHTTP
	return func(w http.ResponseWriter, r *http.Request) {
		LogFileServerRequest(getClientIP(r), r.URL.Path)
		realHandler(w, r)
	}
}

func HealthHandler(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusOK)
}

// Get client IP address from request
func getClientIP(r *http.Request) string {
	ip, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return r.RemoteAddr // Fallback
	}
	return ip
}