package httpz

import (
	"fmt"
	"net/netip"
	"slices"
	"syscall"
)

var (
	defaultAllowedPorts = []uint16{80, 443}
)

type SafeNetControlConfig struct {
	// AllowedPorts use to explicitly allow a set of ports.
	// If not set, will default to defaultAllowedPorts.
	AllowedPorts []uint16
	// AllowedIPs use to explicitly allow a set of specific IPs.
	AllowedIPs []netip.Addr
	// AllowedIPCIDRs use to explicitly allow a set of specific IP CIDRs.
	AllowedIPCIDRs []netip.Prefix
	// BlockedIPs use to explicitly block a set of specific IPs.
	// This takes precedence over all other IP-based checks.
	BlockedIPs []netip.Addr
	// BlockedIPCIDRs use to explicitly block a set of specific IP CIDRs.
	// This takes precedence over all other IP-based checks.
	BlockedIPCIDRs []netip.Prefix
}

type ControlFunc func(network, address string, c syscall.RawConn) error

// NewSafeNetDialerControl creates a safe Control function that can be used in a net.Dialer.
// Safe means:
// - prevent access to not allowed network types. Allowed are tcp4 and tcp6 connections.
// - prevent access to not allowed ports (configurable).
// - prevent access to not allowed IP addresses:
//   - prevent access to loopback
//   - prevent access to IPv4 broadcast (non global unicast address)
//   - prevent access to multicast (non global unicast address)
//   - prevent access to link local unicast (non global unicast address)
//   - prevent access to private addresses as specified
//     by RFC 1918 (IPv4 addresses) and RFC 4193 (IPv6 addresses).
//
// This function takes owner ship of the passed SafeNetControlConfig.
// Do not modify it once passed and only use it once.
func NewSafeNetDialerControl(cfg *SafeNetControlConfig) ControlFunc {
	// shallow copy is good enough
	c := *cfg

	// apply defaults
	if len(c.AllowedPorts) == 0 {
		c.AllowedPorts = defaultAllowedPorts
	}

	return c.control
}

func (c *SafeNetControlConfig) control(network, address string, conn syscall.RawConn) error {
	if network != "tcp4" && network != "tcp6" {
		return fmt.Errorf("invalid network. Wanted tcp4 or tcp6 but got %q", network)
	}

	addrPort, err := netip.ParseAddrPort(address)
	if err != nil {
		return fmt.Errorf("invalid address. Wanted a valid host:port pair but got %q", address)
	}

	if err = c.isAllowedPort(addrPort.Port()); err != nil {
		return fmt.Errorf("invalid port. %w", err)
	}

	ip := addrPort.Addr()
	if c.isExplicitlyBlockedIPAddress(ip) {
		return fmt.Errorf("not allowed IP address. IP address %q is blocked", ip)
	}

	if c.isExplicitlyAllowedIPAddress(ip) {
		return nil
	}

	if !c.isPublicIPAddress(ip) {
		return fmt.Errorf("not allowed IP address. IP address %q is not allowed", ip)
	}

	return nil
}

func (c *SafeNetControlConfig) isAllowedPort(port uint16) error {
	if !slices.Contains(c.AllowedPorts, port) {
		return fmt.Errorf("wanted one of %v but got %d", c.AllowedPorts, port)
	}

	return nil
}

func (c *SafeNetControlConfig) isExplicitlyAllowedIPAddress(ip netip.Addr) bool {
	return containsIP(ip, c.AllowedIPs, c.AllowedIPCIDRs)
}

func (c *SafeNetControlConfig) isExplicitlyBlockedIPAddress(ip netip.Addr) bool {
	return containsIP(ip, c.BlockedIPs, c.BlockedIPCIDRs)
}

func containsIP(ip netip.Addr, ips []netip.Addr, cidrs []netip.Prefix) bool {
	if slices.Contains(ips, ip) {
		return true
	}

	if slices.ContainsFunc(cidrs, func(x netip.Prefix) bool { return x.Contains(ip) }) {
		return true
	}

	return false
}

func (c *SafeNetControlConfig) isPublicIPAddress(ip netip.Addr) bool {
	switch {
	case ip.IsLoopback():
		return false
	case !ip.IsGlobalUnicast():
		return false
	case ip.IsPrivate():
		return false
	default:
		return true
	}
}
