package blacktip

import (
	"bytes"
	"crypto"
	"crypto/x509"
	"encoding/pem"
	"net"
	"net/mail"
	"regexp"
	"strings"

	"git.veillax.com/Veillax/reef/model"
	"github.com/emersion/go-msgauth/dkim"
)

func netLookupTXT(name string) ([]string, error) { return net.LookupTXT(name) }

var dkimHeaders = []string{"From", "To", "Cc", "Subject", "Date", "Message-ID", "MIME-Version", "Content-Type",
	"Content-Transfer-Encoding", "Reply-To", "In-Reply-To", "References"}

var dkimDomainRE = regexp.MustCompile(`(?i)(?:^|[;\s])d=\s*([^;\s]+)`)

// toCRLF makes every line end in CRLF. Mail travels as CRLF and DKIM is defined
// over that form; the relay hands us the message with plain LF.
func toCRLF(b []byte) []byte {
	b = bytes.ReplaceAll(b, []byte("\r\n"), []byte("\n"))
	return bytes.ReplaceAll(b, []byte("\n"), []byte("\r\n"))
}

// pickSigner finds the key for a sender domain: an exact match, else the
// longest parent domain (mail from news.example.com may be signed by example.com).
func pickSigner(signers []model.DKIMSigner, domain string) *model.DKIMSigner {
	var best *model.DKIMSigner
	for i := range signers {
		d := strings.ToLower(signers[i].Domain)
		if domain == d || strings.HasSuffix(domain, "."+d) {
			if best == nil || len(d) > len(best.Domain) {
				best = &signers[i]
			}
		}
	}
	return best
}

// sign adds a DKIM-Signature for the message's From domain if the policy has a
// key for it. It returns the bytes to send, and a short note for the log.
// Any problem results in the message being sent unsigned, never dropped.
func (s *mailSession) sign(data []byte) ([]byte, string) {
	if len(s.pol.DKIM) == 0 {
		return data, ""
	}
	msg := toCRLF(data)
	m, err := mail.ReadMessage(bytes.NewReader(msg))
	if err != nil {
		return data, "unsigned: unreadable headers"
	}
	domain := ""
	if a, err := mail.ParseAddress(m.Header.Get("From")); err == nil {
		domain = domainOf(a.Address)
	}
	if domain == "" {
		domain = domainOf(s.from)
	}
	sg := pickSigner(s.pol.DKIM, domain)
	if sg == nil {
		return data, ""
	}
	for _, existing := range m.Header["Dkim-Signature"] { // already signed for this domain by your mail server: leave it
		if mm := dkimDomainRE.FindStringSubmatch(existing); mm != nil && strings.EqualFold(mm[1], sg.Domain) {
			return data, "already signed"
		}
	}
	block, _ := pem.Decode([]byte(sg.PrivateKey))
	if block == nil {
		return data, "unsigned: bad key"
	}
	key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
	signer, ok := key.(crypto.Signer)
	if err != nil || !ok {
		return data, "unsigned: bad key"
	}
	var out bytes.Buffer
	err = dkim.Sign(&out, bytes.NewReader(msg), &dkim.SignOptions{
		Domain: sg.Domain, Selector: sg.Selector, Signer: signer, Hash: crypto.SHA256,
		HeaderCanonicalization: dkim.CanonicalizationRelaxed, BodyCanonicalization: dkim.CanonicalizationRelaxed,
		HeaderKeys: dkimHeaders,
	})
	if err != nil {
		return data, "unsigned: " + err.Error()
	}
	return out.Bytes(), "dkim " + sg.Selector + "._domainkey." + sg.Domain
}

// checkDKIM looks each key up in DNS the way a receiving server would.
func (b *Blacktip) checkDKIM(signers []model.DKIMSigner) map[string]string {
	out := map[string]string{}
	for _, sg := range signers {
		txts, err := b.lookupTXT(sg.Selector + "._domainkey." + sg.Domain)
		want := "p=" + sg.PublicKey
		status := "missing"
		if err == nil && len(txts) > 0 {
			status = "mismatch"
			for _, t := range txts {
				if strings.Contains(strings.ReplaceAll(t, " ", ""), want) {
					status = "published"
				}
			}
		}
		out[sg.KeyID] = status
	}
	return out
}