// Package security ports the WinterCMS plugin's SSRFGuard: the check an
// untrusted AI base URL (a user's or an organisation's) and an
// AI-returned image URL pass before anything is sent to them.
package security

import (
	"context"
	"net"
	"net/netip"
	"net/url"
	"os"
	"strings"

	"git.golem15.com/golem15/summercms/modules/compass"
	"git.golem15.com/golem15/summercms/modules/fetchguard"
)

// ConfigAllowedHosts is the config key of the host allowlist.
const ConfigAllowedHosts = "golem15.golem.ssrf.allowed_hosts"

// EnvAllowedHosts is the environment variable that replaces the
// allowlist: a comma-separated list, used when set and not empty.
const EnvAllowedHosts = "GOLEM15_SSRF_ALLOWED_HOSTS"

// DefaultAllowedHosts is the allowlist when neither the environment nor
// the config names one: OpenAI's DALL-E image host and every OpenAI host.
var DefaultAllowedHosts = []string{"oaidalleapiprodscus.blob.core.windows.net", ".openai.com"}

// UnsafeURLError is a URL the guard refused. Message is PHP's
// RuntimeException text.
type UnsafeURLError struct {
	Message string
}

func (e *UnsafeURLError) Error() string { return e.Message }

// Resolver looks a host up (net.DefaultResolver satisfies it).
type Resolver interface {
	LookupIPAddr(ctx context.Context, host string) ([]net.IPAddr, error)
}

// DefaultResolver is the resolver AssertSafeURL uses.
var DefaultResolver Resolver = net.DefaultResolver

// AllowedHosts is the allowlist: GOLEM15_SSRF_ALLOWED_HOSTS split on
// commas (entries trimmed, empty ones dropped) when it is set and not
// empty, else golem15.golem.ssrf.allowed_hosts, else DefaultAllowedHosts.
func AllowedHosts(cfg *compass.Config) []string {
	if env := os.Getenv(EnvAllowedHosts); env != "" {
		var out []string
		for _, h := range strings.Split(env, ",") {
			if h = strings.TrimSpace(h); h != "" {
				out = append(out, h)
			}
		}
		return out
	}
	if cfg != nil {
		if v, ok := cfg.Lookup(ConfigAllowedHosts); ok {
			var out []string
			switch list := v.(type) {
			case []any:
				for _, h := range list {
					if s, ok := h.(string); ok {
						out = append(out, s)
					}
				}
			case []string:
				out = append(out, list...)
			}
			return out
		}
	}
	return append([]string(nil), DefaultAllowedHosts...)
}

// AssertSafeURL is SSRFGuard::assertSafeUrl: rawURL needs a scheme and a
// host ("Invalid URL"), the https scheme ("Only https:// scheme
// allowed"), a host on the allowlist, where an entry with a leading dot
// matches any subdomain and any other entry matches exactly, both
// case-insensitively ("Host not in allowlist: <host>"), and a host that
// resolves ("Cannot resolve host: <host>") to public addresses only
// ("Resolves to private/loopback IP: <ip>", judged by
// fetchguard.IsPrivateAddr). A failure is *UnsafeURLError.
//
// The address check runs before the request, as in PHP; the guarded
// client's dial-time check runs again when it connects.
func AssertSafeURL(ctx context.Context, rawURL string, allowed []string) error {
	u, err := url.Parse(rawURL)
	if err != nil || u.Scheme == "" || u.Host == "" {
		return &UnsafeURLError{Message: "Invalid URL"}
	}
	if !strings.EqualFold(u.Scheme, "https") {
		return &UnsafeURLError{Message: "Only https:// scheme allowed"}
	}
	host := u.Hostname()
	if !hostAllowed(host, allowed) {
		return &UnsafeURLError{Message: "Host not in allowlist: " + host}
	}
	addrs, err := resolve(ctx, host)
	if err != nil || len(addrs) == 0 {
		return &UnsafeURLError{Message: "Cannot resolve host: " + host}
	}
	for _, a := range addrs {
		if fetchguard.IsPrivateAddr(a) {
			return &UnsafeURLError{Message: "Resolves to private/loopback IP: " + a.String()}
		}
	}
	return nil
}

func hostAllowed(host string, allowed []string) bool {
	h := strings.ToLower(host)
	for _, a := range allowed {
		a = strings.ToLower(a)
		if strings.HasPrefix(a, ".") {
			if strings.HasSuffix(h, a) {
				return true
			}
		} else if h == a {
			return true
		}
	}
	return false
}

func resolve(ctx context.Context, host string) ([]netip.Addr, error) {
	if a, err := netip.ParseAddr(host); err == nil {
		return []netip.Addr{a}, nil
	}
	ips, err := DefaultResolver.LookupIPAddr(ctx, host)
	if err != nil {
		return nil, err
	}
	out := make([]netip.Addr, 0, len(ips))
	for _, ip := range ips {
		if a, ok := netip.AddrFromSlice(ip.IP); ok {
			out = append(out, a.Unmap())
		}
	}
	return out, nil
}