From 7bc7e7b7a27962976490306822aefb7ea88acb07 Mon Sep 17 00:00:00 2001
From: Pierre-Olivier Mercier <nemunaire@nemunai.re>
Date: Thu, 23 Oct 2025 15:12:19 +0700
Subject: [PATCH] Reuse domain extractes from headers

---
 pkg/analyzer/dns.go      | 30 +++------------------
 pkg/analyzer/dns_test.go | 56 ----------------------------------------
 pkg/analyzer/report.go   |  4 +--
 3 files changed, 6 insertions(+), 84 deletions(-)

diff --git a/pkg/analyzer/dns.go b/pkg/analyzer/dns.go
index 54b0d2f..ee4d7d3 100644
--- a/pkg/analyzer/dns.go
+++ b/pkg/analyzer/dns.go
@@ -52,18 +52,18 @@ func NewDNSAnalyzer(timeout time.Duration) *DNSAnalyzer {
 }
 
 // AnalyzeDNS performs DNS validation for the email's domain
-func (d *DNSAnalyzer) AnalyzeDNS(email *EmailMessage, authResults *api.AuthenticationResults) *api.DNSResults {
+func (d *DNSAnalyzer) AnalyzeDNS(email *EmailMessage, authResults *api.AuthenticationResults, headersResults *api.HeaderAnalysis) *api.DNSResults {
 	// Extract domain from From address
-	fromDomain := d.extractFromDomain(email)
-	if fromDomain == "" {
+	if headersResults.DomainAlignment.FromDomain == nil || *headersResults.DomainAlignment.FromDomain == "" {
 		return &api.DNSResults{
 			Errors: &[]string{"Unable to extract domain from email"},
 		}
 	}
+	fromDomain := *headersResults.DomainAlignment.FromDomain
 
 	results := &api.DNSResults{
 		FromDomain: fromDomain,
-		RpDomain:   d.extractRPDomain(email),
+		RpDomain:   headersResults.DomainAlignment.ReturnPathDomain,
 	}
 
 	// Determine which domain to check SPF for (Return-Path domain)
@@ -112,28 +112,6 @@ func (d *DNSAnalyzer) AnalyzeDNS(email *EmailMessage, authResults *api.Authentic
 	return results
 }
 
-// extractFromDomain extracts the domain from the email's From address
-func (d *DNSAnalyzer) extractFromDomain(email *EmailMessage) string {
-	if email.From != nil && email.From.Address != "" {
-		parts := strings.Split(email.From.Address, "@")
-		if len(parts) == 2 {
-			return strings.ToLower(strings.TrimSpace(parts[1]))
-		}
-	}
-	return ""
-}
-
-// extractRPDomain extracts the domain from the email's Return-Path address
-func (d *DNSAnalyzer) extractRPDomain(email *EmailMessage) *string {
-	if email.ReturnPath != "" {
-		parts := strings.Split(email.ReturnPath, "@")
-		if len(parts) == 2 {
-			return api.PtrTo(strings.TrimSuffix(strings.ToLower(strings.TrimSpace(parts[1])), ">"))
-		}
-	}
-	return nil
-}
-
 // checkMXRecords looks up MX records for a domain
 func (d *DNSAnalyzer) checkMXRecords(domain string) *[]api.MXRecord {
 	ctx, cancel := context.WithTimeout(context.Background(), d.Timeout)
diff --git a/pkg/analyzer/dns_test.go b/pkg/analyzer/dns_test.go
index c397726..d3deb20 100644
--- a/pkg/analyzer/dns_test.go
+++ b/pkg/analyzer/dns_test.go
@@ -22,7 +22,6 @@
 package analyzer
 
 import (
-	"net/mail"
 	"testing"
 	"time"
 )
@@ -57,61 +56,6 @@ func TestNewDNSAnalyzer(t *testing.T) {
 		})
 	}
 }
-
-func TestExtractDomain(t *testing.T) {
-	tests := []struct {
-		name           string
-		fromAddress    string
-		expectedDomain string
-	}{
-		{
-			name:           "Valid email",
-			fromAddress:    "user@example.com",
-			expectedDomain: "example.com",
-		},
-		{
-			name:           "Email with subdomain",
-			fromAddress:    "user@mail.example.com",
-			expectedDomain: "mail.example.com",
-		},
-		{
-			name:           "Email with uppercase",
-			fromAddress:    "User@Example.COM",
-			expectedDomain: "example.com",
-		},
-		{
-			name:           "Invalid email (no @)",
-			fromAddress:    "invalid-email",
-			expectedDomain: "",
-		},
-		{
-			name:           "Empty email",
-			fromAddress:    "",
-			expectedDomain: "",
-		},
-	}
-
-	analyzer := NewDNSAnalyzer(5 * time.Second)
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			email := &EmailMessage{
-				Header: make(mail.Header),
-			}
-			if tt.fromAddress != "" {
-				email.From = &mail.Address{
-					Address: tt.fromAddress,
-				}
-			}
-
-			domain := analyzer.extractFromDomain(email)
-			if domain != tt.expectedDomain {
-				t.Errorf("extractFromDomain() = %q, want %q", domain, tt.expectedDomain)
-			}
-		})
-	}
-}
-
 func TestValidateSPF(t *testing.T) {
 	tests := []struct {
 		name     string
diff --git a/pkg/analyzer/report.go b/pkg/analyzer/report.go
index b5bab60..6848a7d 100644
--- a/pkg/analyzer/report.go
+++ b/pkg/analyzer/report.go
@@ -74,11 +74,11 @@ func (r *ReportGenerator) AnalyzeEmail(email *EmailMessage) *AnalysisResults {
 
 	// Run all analyzers
 	results.Authentication = r.authAnalyzer.AnalyzeAuthentication(email)
-	results.Content = r.contentAnalyzer.AnalyzeContent(email)
-	results.DNS = r.dnsAnalyzer.AnalyzeDNS(email, results.Authentication)
 	results.Headers = r.headerAnalyzer.GenerateHeaderAnalysis(email)
+	results.DNS = r.dnsAnalyzer.AnalyzeDNS(email, results.Authentication, results.Headers)
 	results.RBL = r.rblChecker.CheckEmail(email)
 	results.SpamAssassin = r.spamAnalyzer.AnalyzeSpamAssassin(email)
+	results.Content = r.contentAnalyzer.AnalyzeContent(email)
 
 	return results
 }