package egress

import (
	"context"
	"errors"
	"fmt"
	"net"
	"strings"
	"time"
)

const connectTimeout = 5 * time.Second

var ErrDestinationDenied = errors.New("extension egress destination is denied")

type ipResolver interface {
	LookupIP(ctx context.Context, network, host string) ([]net.IP, error)
}

type dialContextFunc func(context.Context, string, string) (net.Conn, error)

type guardedDialer struct {
	resolver     ipResolver
	dial         dialContextFunc
	allowedCIDRs []*net.IPNet
	allowPrivate bool
}

func newGuardedDialer(config Config) *guardedDialer {
	resolver := config.Resolver
	if resolver == nil {
		resolver = net.DefaultResolver
	}
	dial := config.DialContext
	if dial == nil {
		netDialer := &net.Dialer{Timeout: connectTimeout, KeepAlive: 30 * time.Second}
		dial = netDialer.DialContext
	}
	return &guardedDialer{
		resolver:     resolver,
		dial:         dial,
		allowedCIDRs: cloneIPNetworks(config.TargetNetworks),
		allowPrivate: config.AllowPrivate,
	}
}

func (dialer *guardedDialer) DialContext(
	ctx context.Context,
	network string,
	address string,
) (net.Conn, error) {
	host, port, err := net.SplitHostPort(address)
	if err != nil {
		return nil, fmt.Errorf("%w: invalid address", ErrDestinationDenied)
	}
	var ips []net.IP
	if literal := net.ParseIP(strings.Trim(host, "[]")); literal != nil {
		ips = []net.IP{literal}
	} else {
		ips, err = dialer.resolver.LookupIP(ctx, "ip", host)
		if err != nil {
			return nil, fmt.Errorf("resolve extension egress destination: %w", err)
		}
	}
	if len(ips) == 0 {
		return nil, fmt.Errorf("%w: no destination addresses", ErrDestinationDenied)
	}
	for _, ip := range ips {
		if !dialer.allowedIP(ip) {
			return nil, fmt.Errorf("%w: resolved address is not public", ErrDestinationDenied)
		}
	}

	connectCtx, cancel := context.WithTimeout(ctx, connectTimeout)
	defer cancel()
	var lastErr error
	for _, ip := range ips {
		target := net.JoinHostPort(ip.String(), port)
		connection, err := dialer.dial(connectCtx, network, target)
		if err == nil {
			return connection, nil
		}
		lastErr = err
		if connectCtx.Err() != nil {
			break
		}
	}
	if lastErr == nil {
		lastErr = errors.New("no address connected")
	}
	return nil, fmt.Errorf("dial extension egress destination: %w", lastErr)
}

func (dialer *guardedDialer) allowedIP(ip net.IP) bool {
	if ip == nil {
		return false
	}
	if ipv4 := ip.To4(); ipv4 != nil {
		ip = ipv4
	}
	if hardDeniedIP(ip) {
		return false
	}
	if dialer.allowPrivate {
		return true
	}
	if len(dialer.allowedCIDRs) != 0 {
		for _, network := range dialer.allowedCIDRs {
			if network.Contains(ip) {
				return true
			}
		}
		return false
	}
	return allowedPublicIP(ip)
}

func allowedPublicIP(ip net.IP) bool {
	if ip.IsPrivate() {
		return false
	}
	for _, network := range blockedPublicIPNets {
		if network.Contains(ip) {
			return false
		}
	}
	return true
}

func hardDeniedIP(ip net.IP) bool {
	if !ip.IsGlobalUnicast() ||
		ip.IsLoopback() ||
		ip.IsLinkLocalUnicast() ||
		ip.IsLinkLocalMulticast() ||
		ip.IsMulticast() ||
		ip.IsUnspecified() {
		return true
	}
	for _, address := range hardDeniedMetadataIPs {
		if ip.Equal(address) {
			return true
		}
	}
	for _, network := range hardDeniedIPNets {
		if network.Contains(ip) {
			return true
		}
	}
	return false
}

// net.IP.IsGlobalUnicast reports true for several IANA special-purpose ranges
// that are not publicly routable (for example CGNAT, benchmarking, and
// documentation prefixes). Those ranges can still route inside an operator's
// network, so treating IsGlobalUnicast as "public" would leave a DNS-rebinding
// path into internal infrastructure.
var blockedPublicIPNets = mustParseBlockedPublicIPNets([]string{
	"0.0.0.0/8",
	"10.0.0.0/8",
	"100.64.0.0/10",
	"127.0.0.0/8",
	"169.254.0.0/16",
	"172.16.0.0/12",
	"192.0.0.0/24",
	"192.0.2.0/24",
	"192.168.0.0/16",
	"198.18.0.0/15",
	"198.51.100.0/24",
	"203.0.113.0/24",
	"224.0.0.0/4",
	"240.0.0.0/4",
	"255.255.255.255/32",
	"::/128",
	"::1/128",
	"64:ff9b::/96",
	"100::/64",
	"2001:db8::/32",
	"fc00::/7",
	"fe80::/10",
	"ff00::/8",
})

var hardDeniedMetadataIPs = []net.IP{
	net.ParseIP("100.100.100.200"),
	net.ParseIP("168.63.129.16"),
	net.ParseIP("fd00:ec2::254"),
}

var hardDeniedIPNets = mustParseBlockedPublicIPNets([]string{
	"100.64.0.0/10",
})

func mustParseBlockedPublicIPNets(values []string) []*net.IPNet {
	result := make([]*net.IPNet, 0, len(values))
	for _, value := range values {
		_, network, err := net.ParseCIDR(value)
		if err != nil {
			panic(err)
		}
		result = append(result, network)
	}
	return result
}

func cloneIPNetworks(networks []*net.IPNet) []*net.IPNet {
	cloned := make([]*net.IPNet, 0, len(networks))
	for _, network := range networks {
		if network == nil {
			continue
		}
		cloned = append(cloned, &net.IPNet{
			IP:   append(net.IP(nil), network.IP...),
			Mask: append(net.IPMask(nil), network.Mask...),
		})
	}
	return cloned
}