package blacktip

import (
	"errors"
	"fmt"
	"io"
	"net/http"
	"strconv"
	"strings"
	"sync"
	"time"

	"git.veillax.com/Veillax/reef/netguard"
)

// MTA-STS (RFC 8461), sending side. A recipient domain can publish a policy that
// says "mail to us must go to these servers, over verified TLS". Without
// checking it, an attacker who can tamper with DNS or the network can redirect
// or downgrade your outgoing mail. When a policy says "enforce" we only deliver
// to the listed servers and only over TLS with a valid certificate.

// STSPolicy is a recipient domain's MTA-STS policy.
type STSPolicy struct {
	Mode   string   // "enforce", "testing" or "none"
	MX     []string // exact host names or "*.example.com" patterns
	MaxAge time.Duration
}

// Allows reports whether host is a mail server the policy permits.
func (p *STSPolicy) Allows(host string) bool {
	host = strings.ToLower(strings.TrimSuffix(host, "."))
	for _, pat := range p.MX {
		pat = strings.ToLower(strings.TrimSuffix(pat, "."))
		if strings.HasPrefix(pat, "*.") {
			// a wildcard stands for exactly one label
			if rest, ok := strings.CutSuffix(host, pat[1:]); ok && rest != "" && !strings.Contains(rest, ".") {
				return true
			}
		} else if pat == host {
			return true
		}
	}
	return false
}

func parseSTS(text string) (*STSPolicy, error) {
	p := &STSPolicy{}
	version := ""
	for _, line := range strings.Split(strings.ReplaceAll(text, "\r\n", "\n"), "\n") {
		k, v, ok := strings.Cut(line, ":")
		if !ok {
			continue
		}
		k, v = strings.ToLower(strings.TrimSpace(k)), strings.TrimSpace(v)
		switch k {
		case "version":
			version = v
		case "mode":
			p.Mode = strings.ToLower(v)
		case "mx":
			p.MX = append(p.MX, v)
		case "max_age":
			n, err := strconv.ParseInt(v, 10, 64)
			if err != nil || n < 0 {
				return nil, errors.New("bad max_age")
			}
			if n > 31557600 {
				n = 31557600 // a year is the most the RFC allows
			}
			p.MaxAge = time.Duration(n) * time.Second
		}
	}
	switch {
	case version != "STSv1":
		return nil, errors.New("not an STSv1 policy")
	case p.Mode != "enforce" && p.Mode != "testing" && p.Mode != "none":
		return nil, fmt.Errorf("unknown mode %q", p.Mode)
	case p.Mode != "none" && len(p.MX) == 0:
		return nil, errors.New("policy lists no mx")
	case p.MaxAge <= 0:
		return nil, errors.New("policy has no max_age")
	}
	return p, nil
}

type stsEntry struct {
	id      string
	pol     *STSPolicy
	expires time.Time
}

type stsCache struct {
	mu sync.Mutex
	m  map[string]*stsEntry
}

func (b *Blacktip) lookupTXT(name string) ([]string, error) {
	if b.cfg.MailLookupTXT != nil {
		return b.cfg.MailLookupTXT(name)
	}
	return netLookupTXT(name)
}

// fetchSTS downloads and parses the policy file. The recipient domain is
// attacker-controlled, so the request goes through the guarded client: no
// internal addresses, no redirects.
func fetchSTS(client *http.Client, domain string) (*STSPolicy, error) {
	resp, err := client.Get("https://mta-sts." + domain + "/.well-known/mta-sts.txt")
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return nil, fmt.Errorf("policy server answered %s", resp.Status)
	}
	body, err := io.ReadAll(io.LimitReader(resp.Body, 64<<10))
	if err != nil {
		return nil, err
	}
	return parseSTS(string(body))
}

// stsFor returns the domain's current MTA-STS policy, or nil if it has none.
func (b *Blacktip) stsFor(domain string) *STSPolicy {
	if b.cfg.MailSTS != nil {
		p, _ := b.cfg.MailSTS(domain)
		return p
	}
	id := ""
	if txts, err := b.lookupTXT("_mta-sts." + domain); err == nil {
		for _, t := range txts {
			if strings.HasPrefix(strings.TrimSpace(t), "v=STSv1") {
				for _, f := range strings.Split(t, ";") {
					if k, v, ok := strings.Cut(strings.TrimSpace(f), "="); ok && k == "id" {
						id = v
					}
				}
			}
		}
	}
	now := time.Now()
	b.sts.mu.Lock()
	if b.sts.m == nil {
		b.sts.m = map[string]*stsEntry{}
	}
	ent := b.sts.m[domain]
	b.sts.mu.Unlock()
	valid := ent != nil && now.Before(ent.expires)
	if id == "" {
		if valid {
			return ent.pol // a cached policy stays in force until it expires
		}
		return nil
	}
	if valid && ent.id == id {
		return ent.pol
	}
	pol, err := fetchSTS(netguard.Client(b.cfg.MailAllowPrivate), domain)
	if err != nil {
		if valid {
			return ent.pol
		}
		return nil
	}
	b.sts.mu.Lock()
	b.sts.m[domain] = &stsEntry{id: id, pol: pol, expires: now.Add(pol.MaxAge)}
	b.sts.mu.Unlock()
	return pol
}