package main

import (
	"bytes"
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"os"
	"strconv"
	"strings"

	"github.com/redis/go-redis/v9"
)

var (
	ctx           = context.Background()
	redisClient   *redis.Client
	hashName      = "bl:manual"
	blocklistMode string
	masterURL     string
)

func main() {
	blocklistMode = os.Getenv("BLOCKLIST_MODE")
	hashName = os.Getenv("HASH_NAME")
	if hashName == "" {
		hashName = "bl:manual"
	}

	//DEBUG
	//blocklistMode = "master"
	//DEBUG

	if blocklistMode != "master" && blocklistMode != "slave" {
		log.Fatal("❌ BLOCKLIST_MODE has to be 'master' or 'slave'")
	} else {
		log.Println("✅ BLOCKLIST_MODE set to:", blocklistMode)
	}

	if blocklistMode == "slave" {
		masterURL = os.Getenv("MASTER_URL")
		//DEBUG
		//masterURL = "http://127.0.0.1:8081"
		//DEBUG
		if masterURL == "" {
			log.Fatal("❌ MASTER_URL has to be set in slave mode")
		}
		log.Println("✅ MASTER_URL set to:", masterURL)
	}

	if blocklistMode == "master" {
		redisAddr := os.Getenv("REDIS_ADDR")
		redisPass := os.Getenv("REDIS_PASS")
		if redisAddr == "" {
			redisAddr = "localhost:6379"
		}
		log.Println("✅ Redis Connection set to:", redisAddr)
		redisClient = redis.NewClient(&redis.Options{
			Addr:     redisAddr,
			DB:       0,
			Username: os.Getenv("REDIS_USER"),
			Password: redisPass,
		})
		if err := redisClient.Ping(ctx).Err(); err != nil {
			log.Fatalf("❌ No connection to redis-server: %v", err)
		}
	}

	portsTCP := []string{"135", "139", "445", "389", "636", "3268", "3269", "88", "3389", "3306", "27017", "5432", "25", "123", "5900", "143", "993", "1723", "21"}
	for _, port := range portsTCP {
		go startTCPListener(port)
	}

	portsUDP := []string{"135", "137", "138", "389", "3389", "88"}
	for _, port := range portsUDP {
		go startUDPListener(port)
	}

	http.HandleFunc("/", handleRoot)
	http.HandleFunc("/add", handleAdd)
	http.HandleFunc("/dashboard", handleDashboard)
	http.HandleFunc("/dashboard.json", handleDashboardJSON)
	http.HandleFunc("/metrics", handleMetrics)

	log.Println("🚀 Server listening at :8080")
	if err := http.ListenAndServe(":8080", nil); err != nil {
		log.Fatal("⚠️ ListenAndServe: ", err)
	}
}

func startTCPListener(port string) {
	ln, err := net.Listen("tcp", ":"+port)
	if err != nil {
		log.Fatalf("❌ Could not listen on TCP port %s: %v", port, err)
	}
	log.Printf("🚀 TCP Honeypod listening on port %s", port)
	for {
		conn, err := ln.Accept()
		if err != nil {
			log.Printf("⚠️ Error accepting TCP connection: %v", err)
			continue
		}
		ip, _, _ := net.SplitHostPort(conn.RemoteAddr().String())
		processIP(ip, port, "TCP", nil)
		conn.Close()
	}
}

func startUDPListener(port string) {
	addr := net.UDPAddr{
		Port: portInt(port),
		IP:   net.ParseIP("0.0.0.0"),
	}
	conn, err := net.ListenUDP("udp", &addr)
	if err != nil {
		log.Fatalf("❌ Could not listen on UDP port %s: %v", port, err)
	}
	log.Printf("🚀 UDP Honeypod listening on port %s", port)
	buf := make([]byte, 1024)
	for {
		n, remoteAddr, err := conn.ReadFromUDP(buf)
		if err != nil {
			log.Printf("⚠️ Error reading UDP packet: %v", err)
			continue
		}
		if n > 0 {
			ip := remoteAddr.IP.String()
			processIP(ip, port, "UDP", nil)
		}
	}
}

func handleRoot(w http.ResponseWriter, r *http.Request) {
	ip := getClientIP(r)
	if ip == "" {
		http.Error(w, "Wrong IP", http.StatusBadRequest)
		return
	}
	processIP(ip, "80", "HTTP", w)
}

func handleAdd(w http.ResponseWriter, r *http.Request) {
	body, err := io.ReadAll(r.Body)
	if err != nil || len(body) == 0 {
		http.Error(w, "Wrong Body", http.StatusBadRequest)
		return
	}
	parts := strings.Split(string(body), "|")
	if len(parts) != 3 {
		http.Error(w, "Wrong Body", http.StatusBadRequest)
		return
	}
	ip, port, proto := parts[0], parts[1], parts[2]
	processIP(ip, port, proto, w)
}

func processIP(ip, port, proto string, w http.ResponseWriter) {
	ipKey := ip + "/32"
	if blocklistMode == "master" {
		_ = redisClient.HSet(ctx, hashName, ipKey, 1).Err()
		_ = redisClient.SAdd(ctx, "bl:manual:"+ipKey, port+"/"+proto).Err()
		log.Printf("✅ IP %s saved with port %s/%s", ipKey, port, proto)
		if w != nil {
			w.WriteHeader(http.StatusServiceUnavailable)
			w.Write([]byte("Temporary unavailable"))
		}
	} else {
		payload := ip + "|" + port + "|" + proto
		resp, err := http.Post(masterURL+"/add", "text/plain", bytes.NewBuffer([]byte(payload)))
		if err != nil {
			log.Printf("❌ Error relaying to Master: %v", err)
			if w != nil {
				http.Error(w, "Error relaying", http.StatusBadGateway)
			}
			return
		}
		log.Printf("✅ IP %s relayed with port %s/%s", ipKey, port, proto)
		defer resp.Body.Close()
		if w != nil {
			w.WriteHeader(resp.StatusCode)
			io.Copy(w, resp.Body)
		}
	}
}

func handleMetrics(w http.ResponseWriter, r *http.Request) {
	keys, _ := redisClient.HKeys(ctx, hashName).Result()
	totalIPs := len(keys)
	w.Header().Set("Content-Type", "text/plain")
	w.Write([]byte(fmt.Sprintf("honeypod_blocked_ips_total %d\n", totalIPs)))

	// Map zum Zählen
	portHitCounter := make(map[string]int)

	for _, ip := range keys {
		ports, _ := redisClient.SMembers(ctx, "bl:manual:"+ip).Result()
		for _, p := range ports {
			portHitCounter[p]++
		}
	}

	// Jetzt pro Port+Proto Metriken ausgeben
	for p, count := range portHitCounter {
		parts := strings.Split(p, "/")
		if len(parts) == 2 {
			port, proto := parts[0], parts[1]
			metricName := fmt.Sprintf("honeypod_port_hits{port=\"%s\",protocol=\"%s\"} %d\n", port, proto, count)
			w.Write([]byte(metricName))
		}
	}
}

func handleDashboard(w http.ResponseWriter, r *http.Request) {
	keys, _ := redisClient.HKeys(ctx, hashName).Result()

	w.Header().Set("Content-Type", "text/html")
	w.Write([]byte(`
<!DOCTYPE html>
<html lang="en">
<head>
  <meta charset="UTF-8">
  <title>Honeypod Dashboard</title>
  <link href="https://cdn.jsdelivr.net/npm/bootstrap@5.3.0/dist/css/bootstrap.min.css" rel="stylesheet">
</head>
<body class="bg-dark text-light">
  <div class="container mt-4">
    <h1 class="mb-4">Honeypod Dashboard</h1>
    <table class="table table-dark table-striped">
      <thead><tr><th>IP</th><th>Ports / Protocols</th></tr></thead>
      <tbody>
`))
	for _, ip := range keys {
		ports, _ := redisClient.SMembers(ctx, "bl:manual:"+ip).Result()
		w.Write([]byte("<tr><td>" + ip + "</td><td>" + strings.Join(ports, ", ") + "</td></tr>"))
	}
	w.Write([]byte(`
      </tbody>
    </table>
  </div>
</body>
</html>
`))
}

func handleDashboardJSON(w http.ResponseWriter, r *http.Request) {
	keys, _ := redisClient.HKeys(ctx, hashName).Result()
	result := make(map[string][]string)
	for _, ip := range keys {
		ports, _ := redisClient.SMembers(ctx, "bl:manual:"+ip).Result()
		result[ip] = ports
	}
	w.Header().Set("Content-Type", "application/json")
	json.NewEncoder(w).Encode(result)
}

func getClientIP(r *http.Request) string {
	xff := r.Header.Get("X-Forwarded-For")
	xrip := r.Header.Get("X-Real-Ip")
	log.Printf("X-Forwarded-For: %s, X-Real-Ip: %s, RemoteAddr: %s", xff, xrip, r.RemoteAddr)
	if xff != "" {
		parts := strings.Split(xff, ",")
		ip := strings.TrimSpace(parts[0])
		if isValidIP(ip) {
			return ip
		}
	}

	if xrip != "" && isValidIP(xrip) {
		return xrip
	}
	host, _, err := net.SplitHostPort(r.RemoteAddr)
	if err != nil {
		return ""
	}
	if isValidIP(host) {
		return host
	}
	return ""
}

func isValidIP(ip string) bool {
	return net.ParseIP(ip) != nil
}

func portInt(port string) int {
	p, _ := strconv.Atoi(port)
	return p
}