package main

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net"
	"net/http"
	"net/http/httputil"
	"os"
	"regexp"
	"strings"
)

func main() {
	tr := &http.Transport{
		DialContext: func(ctx context.Context, network, addr string) (net.Conn, error) {
			return net.Dial("unix", "/var/run/docker.sock")
		},
		DisableKeepAlives: true,
	}

	proxy := &httputil.ReverseProxy{
		Director: func(req *http.Request) {
			req.URL.Scheme = "http"
			req.URL.Host = "docker" // dummy host, won't be used
		},
		Transport:      tr,
		ModifyResponse: redactSensitiveFields,
	}

	http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		log.Printf("➡️ %s %s", r.Method, r.URL.Path)

		if !isAllowed(r.Method, r.URL.Path) {
			log.Printf("⛔ Blocked: %s %s", r.Method, r.URL.Path)

			http.Error(w, "Forbidden", http.StatusForbidden)
			return
		}
		proxy.ServeHTTP(w, r)
	})

	log.Println("Unix socket proxy running on :8000")
	log.Fatal(http.ListenAndServe(":8000", nil))
}

func isAllowed(method, path string) bool {
	// White listed endpoints that are considered relatively safe for now.
	if method == "GET" {
		if strings.HasSuffix(path, "/containers/json") ||
			strings.HasSuffix(path, "/networks") ||
			strings.HasSuffix(path, "/services") ||
			strings.HasSuffix(path, "/version") ||
			strings.HasSuffix(path, "/info") ||
			strings.HasSuffix(path, "/events") ||
			strings.HasSuffix(path, "/tasks") {
			log.Printf("➡️ isAllowed %s", path)

			return true
		}
	}
	if method == "POST" {
		if strings.HasSuffix(path, "/start") ||
			strings.HasSuffix(path, "/stop") {
			return true
		}
	}
	// If it doesn't fit into any of the above, then do some more checks on it.
	return allowPath(path)

}

// Only allow specific Docker API paths (e.g., container inspect, used by Traefik)
func allowPath(path string) bool {
	pattern := `^v\d+\.\d+$`
	re := regexp.MustCompile(pattern)

	log.Printf("Inspecting %s", path)
	// Example: /v1.24/containers/<container-id>/json
	parts := strings.Split(path, "/")
	log.Printf("Parts %d parts[0] %s", len(parts), parts)
	if len(parts) == 5 && re.MatchString(parts[1]) && parts[2] == "containers" && parts[4] == "json" {
		log.Printf("Passes")
		return true
	}
	return false
}

// Remove some sensitive data from the Docker API response.
func redactSensitiveFields(resp *http.Response) error {
	data, err := io.ReadAll(resp.Body)
	if err != nil {
		log.Printf("❌ Error reading body: %v", err)
		return err
	}

	// Try to parse as a JSON object first
	var obj map[string]interface{}
	if err := json.Unmarshal(data, &obj); err == nil {
		log.Printf("➡️ Body (object): %+v", obj)

		// Redact sensitive fields from single object
		log.Printf("🧪 Item: %+v", obj)

		redactObject(obj)

		redacted, _ := json.Marshal(obj)
		resp.Body = io.NopCloser(strings.NewReader(string(redacted)))
		resp.ContentLength = int64(len(redacted))
		resp.Header.Set("Content-Length", fmt.Sprint(len(redacted)))
		return nil
	}

	// Try to parse as a JSON array
	var arr []map[string]interface{}
	if err := json.Unmarshal(data, &arr); err == nil {
		log.Printf("➡️ Body (array): %d items", len(arr))

		// Special case: filter top-level networks (from /networks)
		if strings.Contains(resp.Request.URL.Path, "/networks") {
			filtered := make([]map[string]interface{}, 0)
			for _, item := range arr {
				if name, ok := item["Name"].(string); ok && allowedNetworks[name] {
					filtered = append(filtered, item)
				} else {
					log.Printf("🗑️ Removing disallowed network: %s", name)
				}
			}
			arr = filtered
		} else {
			// Generic redaction
			for _, item := range arr {
				redactObject(item)
			}
		}

		redacted, _ := json.Marshal(arr)
		resp.Body = io.NopCloser(strings.NewReader(string(redacted)))
		resp.ContentLength = int64(len(redacted))
		resp.Header.Set("Content-Length", fmt.Sprint(len(redacted)))
		return nil
	}

	// Fallback if it's neither an object nor an array
	log.Printf("⚠️ Failed to parse response body as JSON object or array")
	resp.Body = io.NopCloser(strings.NewReader(string(data)))
	resp.ContentLength = int64(len(data))
	resp.Header.Set("Content-Length", fmt.Sprint(len(data)))
	return nil
}

var allowedNetworks = getAllowedNetworks()

func getAllowedNetworks() map[string]bool {
	env := os.Getenv("ALLOWED_NETWORKS")
	nets := strings.Split(env, ",")
	allowed := make(map[string]bool, len(nets))
	for _, n := range nets {
		n = strings.TrimSpace(n)
		if n != "" {
			allowed[n] = true
		}
	}
	return allowed
}

// This function is responsible for removing keys/properties from the response.
func redactObject(obj map[string]interface{}) {
	// Top-level redactions
	delete(obj, "HostConfig")
	delete(obj, "GraphDriver")
	delete(obj, "Mounts")

	// Config redactions
	if config, ok := obj["Config"].(map[string]interface{}); ok {
		delete(config, "Env")
		delete(config, "Volumes")
		delete(config, "User")
	}

	// NetworkSettings redactions
	netSettings, ok := obj["NetworkSettings"].(map[string]interface{})
	if !ok {
		log.Println("🔍 No NetworkSettings present")
		return
	}

	delete(netSettings, "SandboxID")
	delete(netSettings, "SandboxKey")

	networksRaw, ok := netSettings["Networks"]
	if !ok {
		log.Println("🔍 No Networks field present")
		return
	}

	// Filter out any networks that traefik is not on and therefore should not care/know about.
	switch networks := networksRaw.(type) {
	case map[string]interface{}:
		log.Printf("🌐 Found networks: %v", keysOf(networks))
		for name := range networks {
			if !allowedNetworks[name] {
				log.Printf("🗑️ Removing disallowed network: %s", name)
				delete(networks, name)
			}
		}
	case []interface{}:
		log.Printf("🌐 Found networks (array): %v", networks)
		var filtered []interface{}
		for _, v := range networks {
			if name, ok := v.(string); ok && allowedNetworks[name] {
				filtered = append(filtered, name)
			}
		}
		netSettings["Networks"] = filtered
	default:
		log.Printf("⚠️ Unknown type for Networks: %T", networksRaw)
	}
}
func keysOf(m map[string]interface{}) []string {
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	return keys
}