package main

import (
	"bufio"
	"context"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"net"
	"net/http"
	"os"
	"regexp"
	"strings"
	"time"
)

var debugMode bool

// macroRegex matches SPF macro expressions like %{s}, %{d}, %{ir}
var macroRegex = regexp.MustCompile(`%{([slodipvcrt]|ir)(\d+)?r?}`)

func Debug(format string, args ...interface{}) {
	if debugMode {
		fmt.Printf("[DEBUG] "+format+"\n", args...)
	}
}

type SPFDMARCRecord struct {
	domainName  string
	spfRecord   string
	dmarcRecord string
}

type ParsedSPFRecord struct {
	aRecord        []string
	includeRecords []ParsedIncludeRecord
	mxRecord       []string
	existsRecords  []string
	ptrRecords     []string
}

type ParsedIncludeRecord struct {
	includeRecord string
	spfText       string
	subLookup     *ParsedSPFRecord
}

type IssueScanResult struct {
	code     int
	title    string
	detail   string
	severity string
}

type IssueEngine struct {
	rolledIncludeRecords []string
}

type JSONScanResult struct {
	Domain          string                   `json:"domain"`
	SPF             string                   `json:"spf_record"`
	DMARC           string                   `json:"dmarc_record"`
	Issues          []JSONIssue              `json:"issues,omitempty"`
	Success         bool                     `json:"success"`
	Message         string                   `json:"message,omitempty"`
	TakeoverResults *SubdomainTakeoverResult `json:"subdomain_takeover,omitempty"`
}

type JSONIssue struct {
	Code     int    `json:"code"`
	Title    string `json:"title"`
	Detail   string `json:"detail"`
	Severity string `json:"severity"`
}

type TakeoverFingerprint struct {
	Service     string   `json:"service"`
	Cname       []string `json:"cname"`
	Fingerprint string   `json:"fingerprint"`
	Status      string   `json:"status"`
	Vulnerable  bool     `json:"vulnerable"`
}

type SubdomainTakeoverResult struct {
	Domain       string `json:"domain"`
	Vulnerable   bool   `json:"vulnerable"`
	Service      string `json:"service,omitempty"`
	Fingerprint  string `json:"fingerprint,omitempty"`
	ResponseBody string `json:"response_body,omitempty"`
	CnameRecord  string `json:"cname_record,omitempty"`
	ErrorMessage string `json:"error,omitempty"`
}

func isNullOrWhiteSpace(s string) bool {
	return len(strings.TrimSpace(s)) == 0
}

func preferIPv4(ips []net.IP) string {
	for _, ip := range ips {
		if ip.To4() != nil {
			return ip.String()
		}
	}
	return ips[0].String()
}

func (r *ParsedSPFRecord) mechanismCount() int {
	return len(r.aRecord) + len(r.mxRecord) + len(r.existsRecords) + len(r.ptrRecords)
}

func getSPFDMARCRecord(domain string) SPFDMARCRecord {
	var record SPFDMARCRecord
	record.domainName = domain

	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	resolver := &net.Resolver{
		PreferGo: true,
		Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
			d := net.Dialer{
				Timeout: 3 * time.Second,
			}
			return d.DialContext(ctx, network, address)
		},
	}

	spfChan := make(chan string, 1)
	go func() {
		txtRecords, err := resolver.LookupTXT(ctx, domain)
		if err == nil {
			for _, txt := range txtRecords {
				if strings.HasPrefix(txt, "v=spf1") {
					spfChan <- txt
					return
				}
			}
		}
		spfChan <- ""
	}()

	dmarcChan := make(chan string, 1)
	go func() {
		txtRecords, err := resolver.LookupTXT(ctx, "_dmarc."+domain)
		if err == nil {
			for _, txt := range txtRecords {
				if strings.HasPrefix(txt, "v=DMARC1") {
					dmarcChan <- txt
					return
				}
			}
		}
		dmarcChan <- ""
	}()

	select {
	case record.spfRecord = <-spfChan:
	case <-ctx.Done():
		// Timeout occurred
		record.spfRecord = ""
	}

	select {
	case record.dmarcRecord = <-dmarcChan:
	case <-ctx.Done():
		// Timeout occurred
		record.dmarcRecord = ""
	}

	return record
}

func reverseIP(ip string) string {
	parts := strings.Split(ip, ".")
	for i, j := 0, len(parts)-1; i < j; i, j = i+1, j-1 {
		parts[i], parts[j] = parts[j], parts[i]
	}
	return strings.Join(parts, ".")
}

func expandSPFMacro(input string, domain string, sender string, ip string) string {
	reverseIPValue := ""
	if ip != "" {
		reverseIPValue = reverseIP(ip)
	}

	localPart := ""
	domainPart := ""
	if strings.Contains(sender, "@") {
		parts := strings.SplitN(sender, "@", 2)
		localPart = parts[0]
		domainPart = parts[1]
	}

	replacements := map[string]string{
		"s":  sender,
		"l":  localPart,
		"o":  domainPart,
		"d":  domain,
		"i":  ip,
		"ir": reverseIPValue, // Reverse IP
		"v":  "in-addr",      // IPv4
		"p":  "unknown",      // Placeholder for policy domain
		"c":  ip,             // Client IP
		"r":  "unknown",      // Placeholder for receiver
		"t":  fmt.Sprintf("%d", time.Now().Unix()),
	}

	return macroRegex.ReplaceAllStringFunc(input, func(macro string) string {
		matches := macroRegex.FindStringSubmatch(macro)
		if len(matches) < 2 {
			return macro // Return as-is if no match
		}
		key := matches[1]
		if value, exists := replacements[key]; exists {
			return value
		}
		return macro // Return as-is if no replacement found
	})
}

func resolveIPFromSPF(spfRecord string, domain string) string {
	ips, err := net.LookupIP(domain)
	if err == nil && len(ips) > 0 {
		result := preferIPv4(ips)
		Debug("Resolved IP for %s: %s", domain, result)
		return result
	}

	// Return a loopback IP to prevent empty macros in SPF expansion
	Debug("Could not resolve any IP for domain %s", domain)
	return "127.0.0.1"
}

func parseSPFRecord(spfRecord string, domainName string, sender string, ip string) *ParsedSPFRecord {
	if ip == "" {
		ip = resolveIPFromSPF(spfRecord, domainName)
	}

	if isNullOrWhiteSpace(spfRecord) {
		return &ParsedSPFRecord{}
	}

	visitedDomains := make(map[string]bool)
	visitedDomains[domainName] = true

	Debug("Starting SPF parsing for domain: %s", domainName)
	return parseSPFRecordWithDepth(spfRecord, domainName, sender, ip, visitedDomains, 0)
}

func parseSPFRecordWithDepth(spfRecord string, domainName string, sender string, ip string, visitedDomains map[string]bool, depth int) *ParsedSPFRecord {
	maxDepth := 10
	if depth > maxDepth {
		Debug("Max recursion depth reached for domain: %s", domainName)
		return &ParsedSPFRecord{}
	}

	Debug("Parsing SPF record for domain: %s, depth: %d", domainName, depth)
	parsedRecord := &ParsedSPFRecord{
		aRecord:        make([]string, 0),
		includeRecords: make([]ParsedIncludeRecord, 0),
		mxRecord:       make([]string, 0),
		existsRecords:  make([]string, 0),
		ptrRecords:     make([]string, 0),
	}

	mechanisms := strings.Split(spfRecord, " ")
	for _, mechanism := range mechanisms {
		expandedMechanism := expandSPFMacro(mechanism, domainName, sender, ip)
		Debug("Expanded mechanism: %s", expandedMechanism)

		if strings.HasPrefix(expandedMechanism, "a:") || expandedMechanism == "a" {
			parsedRecord.aRecord = append(parsedRecord.aRecord, expandedMechanism)
		} else if strings.HasPrefix(expandedMechanism, "include:") {
			includeDomain := strings.TrimPrefix(expandedMechanism, "include:")
			Debug("Processing include: %s", includeDomain)

			if visitedDomains[includeDomain] {
				Debug("Loop detected for include: %s", includeDomain)
				continue
			}

			visitedDomains[includeDomain] = true
			includeSPF := getSPFDMARCRecord(includeDomain)
			if !isNullOrWhiteSpace(includeSPF.spfRecord) {
				subParsed := parseSPFRecordWithDepth(includeSPF.spfRecord, includeDomain, sender, ip, visitedDomains, depth+1)
				parsedRecord.includeRecords = append(parsedRecord.includeRecords, ParsedIncludeRecord{
					includeRecord: includeDomain,
					spfText:       includeSPF.spfRecord,
					subLookup:     subParsed,
				})
			}
		} else if strings.HasPrefix(expandedMechanism, "mx:") || expandedMechanism == "mx" {
			parsedRecord.mxRecord = append(parsedRecord.mxRecord, expandedMechanism)
		} else if strings.HasPrefix(expandedMechanism, "exists:") {
			parsedRecord.existsRecords = append(parsedRecord.existsRecords, expandedMechanism)
		} else if strings.HasPrefix(expandedMechanism, "ptr:") || expandedMechanism == "ptr" {
			parsedRecord.ptrRecords = append(parsedRecord.ptrRecords, expandedMechanism)
		} else if strings.HasPrefix(expandedMechanism, "redirect=") {
			redirectDomain := strings.TrimPrefix(expandedMechanism, "redirect=")
			Debug("Processing redirect to: %s", redirectDomain)

			if visitedDomains[redirectDomain] {
				Debug("Loop detected for redirect: %s", redirectDomain)
				continue
			}

			visitedDomains[redirectDomain] = true
			redirectSPF := getSPFDMARCRecord(redirectDomain)
			if !isNullOrWhiteSpace(redirectSPF.spfRecord) {
				redirectParsed := parseSPFRecordWithDepth(redirectSPF.spfRecord, redirectDomain, sender, ip, visitedDomains, depth+1)
				parsedRecord.aRecord = append(parsedRecord.aRecord, redirectParsed.aRecord...)
				parsedRecord.includeRecords = append(parsedRecord.includeRecords, redirectParsed.includeRecords...)
				parsedRecord.mxRecord = append(parsedRecord.mxRecord, redirectParsed.mxRecord...)
				parsedRecord.existsRecords = append(parsedRecord.existsRecords, redirectParsed.existsRecords...)
				parsedRecord.ptrRecords = append(parsedRecord.ptrRecords, redirectParsed.ptrRecords...)
			}
		}
	}

	Debug("Finished parsing SPF record for domain: %s, depth: %d", domainName, depth)
	return parsedRecord
}

func subDomainSPFNotExists(domain string) bool {
	randomSub := "random123456." + domain
	txtRecords, err := net.LookupTXT(randomSub)
	if err != nil {
		return true // Error means no wildcard record
	}

	for _, txt := range txtRecords {
		if strings.HasPrefix(txt, "v=spf1") {
			return false // Found SPF for subdomain
		}
	}
	return true
}

func (engine *IssueEngine) IssueScan(spfDMARCRecord SPFDMARCRecord, extendedSPF *ParsedSPFRecord) []IssueScanResult {
	issueResults := make([]IssueScanResult, 0)
	spfIssues := engine.SPFIssueScan(spfDMARCRecord.spfRecord, spfDMARCRecord.domainName, extendedSPF)
	issueResults = append(issueResults, spfIssues...)

	if len(issueResults) >= 1 {
		if issueResults[0].code != 0 && issueResults[0].code != 11 {
			issueResults = engine.DMARCIssueScan(spfDMARCRecord.dmarcRecord, spfDMARCRecord.domainName, issueResults)
		}
	} else {
		issueResults = engine.DMARCIssueScan(spfDMARCRecord.dmarcRecord, spfDMARCRecord.domainName, issueResults)
	}

	return issueResults
}

func (engine *IssueEngine) SPFIssueScan(spfRecord string, domainName string, parsedSPF *ParsedSPFRecord) []IssueScanResult {
	spfIssues := make([]IssueScanResult, 0)

	if subDomainSPFNotExists(domainName) {
		spfIssues = append(spfIssues, engine.issueDescriptors(6, domainName))
	}

	if isNullOrWhiteSpace(spfRecord) {
		_, err := net.LookupIP(domainName)
		if err != nil {
			// Check if it's a CNAME
			_, err = net.LookupCNAME(domainName)
			if err != nil {
				spfIssues = append(spfIssues, engine.issueDescriptors(0, domainName))
			} else {
				spfIssues = append(spfIssues, engine.issueDescriptors(1, domainName))
			}
		} else {
			spfIssues = append(spfIssues, engine.issueDescriptors(1, domainName))
		}
	} else {
		// Check for redirect and update SPF record if needed
		if strings.Contains(spfRecord, "redirect=") {
			splitSPF := strings.Split(spfRecord, " ")
			for _, part := range splitSPF {
				if strings.HasPrefix(part, "redirect=") {
					redirectDomain := strings.TrimPrefix(part, "redirect=")
					redirectRecord := getSPFDMARCRecord(redirectDomain)
					spfRecord = redirectRecord.spfRecord
					break
				}
			}
		}

		if strings.Contains(spfRecord, "+all") {
			spfIssues = append(spfIssues, engine.issueDescriptors(3, domainName))
		} else if strings.Contains(spfRecord, "~all") {
			spfIssues = append(spfIssues, engine.issueDescriptors(4, domainName))
		} else if !strings.Contains(spfRecord, "-all") || strings.Contains(spfRecord, "?all") {
			spfIssues = append(spfIssues, engine.issueDescriptors(2, domainName))
		}

		if engine.includeRollUp(parsedSPF) > 10 {
			spfIssues = append(spfIssues, engine.issueDescriptors(5, domainName))
		}

		// Check for loops
		seenIncludes := make(map[string]bool)
		for _, include := range engine.rolledIncludeRecords {
			if seenIncludes[include] {
				spfIssues = append(spfIssues, engine.issueDescriptors(12, domainName))
				break
			}
			seenIncludes[include] = true
		}
	}

	return spfIssues
}

func (engine *IssueEngine) DMARCIssueScan(dmarcRecord string, domainName string, spfIssues []IssueScanResult) []IssueScanResult {
	dmarcIssues := make([]IssueScanResult, 0)
	updatedIssues := make([]IssueScanResult, 0)

	if isNullOrWhiteSpace(dmarcRecord) {
		dmarcIssues = append(dmarcIssues, engine.issueDescriptors(7, domainName))
	} else {
		if strings.Contains(dmarcRecord, "; p=none") || strings.Contains(dmarcRecord, ";p=none") {
			dmarcIssues = append(dmarcIssues, engine.issueDescriptors(8, domainName))
		}
		if strings.Contains(dmarcRecord, "sp=none") {
			dmarcIssues = append(dmarcIssues, engine.issueDescriptors(9, domainName))
		}
		if strings.Contains(dmarcRecord, "pct=") && !strings.Contains(dmarcRecord, "pct=100") {
			dmarcIssues = append(dmarcIssues, engine.issueDescriptors(10, domainName))
		}
	}

	hasDMARCIssue7 := false
	hasDMARCIssue8 := false
	hasDMARCIssue10 := false

	for _, issue := range dmarcIssues {
		if issue.code == 7 {
			hasDMARCIssue7 = true
		} else if issue.code == 8 {
			hasDMARCIssue8 = true
		} else if issue.code == 10 {
			hasDMARCIssue10 = true
		}
	}

	if !hasDMARCIssue7 && !hasDMARCIssue8 && !hasDMARCIssue10 {
		for _, spfIssue := range spfIssues {
			if spfIssue.code != 3 {
				updatedIssues = append(updatedIssues, engine.issueDescriptorUpdate(spfIssue))
			} else {
				updatedIssues = append(updatedIssues, spfIssue)
			}
		}
	} else {
		updatedIssues = append(updatedIssues, spfIssues...)
	}

	updatedIssues = append(updatedIssues, dmarcIssues...)
	return updatedIssues
}

func (engine *IssueEngine) includeRollUp(parsedSPF *ParsedSPFRecord) int {
	if parsedSPF == nil {
		return 0
	}

	engine.rolledIncludeRecords = make([]string, 0)
	visited := make(map[string]bool)
	lookupCount := parsedSPF.mechanismCount()

	for _, includeRecord := range parsedSPF.includeRecords {
		engine.rolledIncludeRecords = append(engine.rolledIncludeRecords, includeRecord.includeRecord)
		lookupCount++

		if includeRecord.subLookup != nil {
			lookupCount += engine.countLookupsNonRecursive(includeRecord.subLookup, visited)

			if lookupCount > 10 {
				return lookupCount
			}
		}
	}

	return lookupCount
}

func (engine *IssueEngine) countLookupsNonRecursive(spf *ParsedSPFRecord, visited map[string]bool) int {
	if spf == nil {
		return 0
	}

	lookupCount := spf.mechanismCount()
	stack := []*ParsedSPFRecord{spf}
	processedIncludes := make(map[string]bool)

	for len(stack) > 0 {
		if lookupCount > 10 {
			return lookupCount
		}

		current := stack[len(stack)-1]
		stack = stack[:len(stack)-1]

		for _, include := range current.includeRecords {
			if processedIncludes[include.includeRecord] {
				continue
			}

			processedIncludes[include.includeRecord] = true
			engine.rolledIncludeRecords = append(engine.rolledIncludeRecords, include.includeRecord)
			lookupCount++

			if include.subLookup != nil {
				stack = append(stack, include.subLookup)
				lookupCount += include.subLookup.mechanismCount()
			}
		}
	}

	return lookupCount
}

func (engine *IssueEngine) issueDescriptorUpdate(issueResult IssueScanResult) IssueScanResult {
	issueResult.detail = "This issue has been mitigated through the DMARC policy 'p' qualifier being set to 'Quarantine' or 'Reject'."
	issueResult.severity = "Mitigated"
	return issueResult
}

func (engine *IssueEngine) issueDescriptors(code int, domain string) IssueScanResult {
	issueResult := IssueScanResult{}

	switch code {
	case 0:
		issueResult.code = 0
		issueResult.title = "Non-existent domain"
		issueResult.detail = "The DNS resolver raised an NXDomain error for " + domain + ". Mail receivers will be unable to resolve a DNS response for your domain and will almost certainly flag any mail as spam."
		issueResult.severity = "Low"
	case 1:
		issueResult.code = 1
		issueResult.title = "No SPF record exists"
		issueResult.detail = "There is no SPF DNS record for " + domain + ". Mail receivers have no mechanism to determine what your authorised mail servers are. Mail receivers will pass authentication checks with a \"None\" result, indicating no check could be performed. Spoofed emails are likely to be accepted."
		issueResult.severity = "High"
	case 2:
		issueResult.code = 2
		issueResult.title = "SPF \"all\" mechanism is missing or set to \"?all\""
		issueResult.detail = "The \"all\" mechanism at the end of the end of an SPF record tells receivers how to treat unauthorised (i.e. spoofed) emails - if the mechanism is missing or set to \"?all\", authentication checks will always return a \"Neutral\" result which many receivers interpret to accept all mail from " + domain + " (including spoofed emails)."
		issueResult.severity = "High"
	case 3:
		issueResult.code = 3
		issueResult.title = "SPF \"+all\" mechanism set"
		issueResult.detail = "The \"all\" mechanism at the end of the end of an SPF record tells receivers how to treat unauthorised (i.e. spoofed) emails - the \"+all\" setting tells receivers to pass/accept all mail from " + domain + " (including spoofed emails)."
		issueResult.severity = "Very High"
	case 4:
		issueResult.code = 4
		issueResult.title = "SPF \"~all\" (SoftFail) mechanism set"
		issueResult.detail = "The \"all\" mechanism at the end of the end of an SPF record tells receivers how to treat unauthorised (i.e. spoofed) emails - the \"~all\" setting tells receivers to 'SoftFail' (i.e. quarantine) emails that fail SPF authentication. In practice however, many email filters only slightly raise the total spam score and accept 'SoftFailed' (i.e. spoofed) emails."
		issueResult.severity = "Medium"
	case 5:
		issueResult.code = 5
		issueResult.title = "SPF has too many lookups for receiver validation"
		issueResult.detail = "The SPF record requires more than 10 DNS lookups for the validation process. The RFC states that maximum 10 lookups are allowed. As a result, recipients may throw a PermError instead of proceeding with SPF validation. Recipients treat these errors differently than a hard or soft SPF fail, but some will continue processing the mail (i.e. accept spoofed emails)."
		issueResult.severity = "Medium"
	case 6:
		issueResult.code = 6
		issueResult.title = "No SPF sub-domain record exists"
		issueResult.detail = "The SPF sub-domain policy is a catch-all mechanism used to prevent threat actors from maliciously spoofing sub-domains from which an explicit SPF record hasn't been set. This is typically represented through a DNS entry similar to \"* IN TXT v=spf1 -all\", effectively telling recipients to block mail if an explicit SPF entry for the sub-domain hasn't been set."
		issueResult.severity = "Medium"
	case 7:
		issueResult.code = 7
		issueResult.title = "No DMARC record exists"
		issueResult.detail = "There is no DMARC DNS Record set for the domain. Spoofed emails utilising an attack technique known as SPF-bypass are likely to be accepted."
		issueResult.severity = "High"
	case 8:
		issueResult.code = 8
		issueResult.title = "Insecure DMARC policy 'p' qualifier"
		issueResult.detail = "The DMARC policy 'p' qualifier is \"none\". If the DMARC policy is neither \"reject\" nor \"quarantine\", spoofed emails utilising an attack technique known as SPF-bypass are likely to be accepted."
		issueResult.severity = "High"
	case 9:
		issueResult.code = 9
		issueResult.title = "Insecure DMARC sub-domain 'p' qualifier"
		issueResult.detail = "The DMARC policy 'p' qualifier for sub-domains is set to \"none\". If the DMARC policy is neither \"reject\" nor \"quarantine\", spoofed emails from any " + domain + " sub-domain utilising an attack technique known as SPF-bypass are likely to be accepted."
		issueResult.severity = "High"
	case 10:
		issueResult.code = 10
		issueResult.title = "Partial DMARC coverage"
		issueResult.detail = "The DMARC \"pct\" value is set to less than '100' (i.e. 100%), meaning the DMARC policy will only be applied to a percentage of incoming mail. A threat actor can continously deliver spoofed emails (via SPF-bypass) until the DMARC policy isn't applied and mail is accepted."
		issueResult.severity = "Medium"
	case 11:
		issueResult.code = 11
		issueResult.title = "DNS Timeout during Scan"
		issueResult.detail = "There was a DNS timeout when querying " + domain + ". This will result in an SPF temperror and any mail will almost certainly be flagged as spam by mail receivers"
		issueResult.severity = "Low"
	case 12:
		issueResult.code = 12
		issueResult.title = "Trivial SPF recurse loop"
		issueResult.detail = "The SPF record is configured whereby an infinite lookup loop exists in the validation chain for " + domain + ". This will likely result in an SPF PermError. Recipients will treat these errors differently than a hard or soft SPF fail, but many will continue processing the mail (i.e. accept spoofed emails)."
		issueResult.severity = "Medium"
	}

	return issueResult
}

func getFingerprints() ([]TakeoverFingerprint, error) {
	fingerprintsURL := "https://raw.githubusercontent.com/EdOverflow/can-i-take-over-xyz/master/fingerprints.json"

	resp, err := http.Get(fingerprintsURL)
	if err != nil {
		return nil, fmt.Errorf("failed to download fingerprints: %v", err)
	}
	defer resp.Body.Close()

	body, err := io.ReadAll(resp.Body)
	if err != nil {
		return nil, fmt.Errorf("failed to read fingerprints data: %v", err)
	}

	var fingerprints []TakeoverFingerprint
	if err := json.Unmarshal(body, &fingerprints); err != nil {
		return nil, fmt.Errorf("failed to parse fingerprints JSON: %v", err)
	}

	return fingerprints, nil
}

func checkSubdomainTakeover(domain string, fingerprints []TakeoverFingerprint) SubdomainTakeoverResult {
	result := SubdomainTakeoverResult{
		Domain:     domain,
		Vulnerable: false,
	}

	cname, err := net.LookupCNAME(domain)
	if err != nil {
		result.ErrorMessage = fmt.Sprintf("Failed to lookup CNAME: %v", err)
		return result
	}

	cname = strings.TrimSuffix(cname, ".")
	result.CnameRecord = cname

	for _, fp := range fingerprints {
		for _, cnamePattern := range fp.Cname {
			if strings.Contains(cname, cnamePattern) {
				if fp.Fingerprint != "" {
					responseText := fetchDomainBody(domain)
					if responseText != "" && strings.Contains(responseText, fp.Fingerprint) && fp.Vulnerable {
						result.Vulnerable = true
						result.Service = fp.Service
						result.Fingerprint = fp.Fingerprint
						if len(responseText) > 200 {
							result.ResponseBody = responseText[:200] + "..."
						} else {
							result.ResponseBody = responseText
						}
						return result
					}
				} else if fp.Vulnerable {
					result.Vulnerable = true
					result.Service = fp.Service
					return result
				}

				break
			}
		}
	}

	return result
}

func fetchDomainBody(domain string) string {
	resp, err := http.Get("http://" + domain)
	if err != nil {
		resp, err = http.Get("https://" + domain)
		if err != nil {
			return ""
		}
	}
	body, err := io.ReadAll(resp.Body)
	resp.Body.Close()
	if err != nil {
		return ""
	}
	return string(body)
}

func main() {
	var domain string
	var interactive bool
	var jsonOutput bool
	var outputFile string
	var inputFile string
	var checkSubTakeover bool
	flag.BoolVar(&debugMode, "debug", false, "Enable debug mode")
	flag.StringVar(&domain, "domain", "", "Domain to check for spoofing vulnerabilities")
	flag.BoolVar(&interactive, "interactive", false, "Run in interactive mode")
	flag.BoolVar(&jsonOutput, "json", false, "Output results in JSON format")
	flag.StringVar(&outputFile, "output", "", "Save results to specified file")
	flag.StringVar(&inputFile, "input", "", "Read domains from specified file (one domain per line)")
	flag.BoolVar(&checkSubTakeover, "subtakeover", false, "Check for subdomain takeover vulnerabilities")
	flag.Parse()

	Debug("Debug mode enabled")
	Debug("Parsed flags: domain=%s, interactive=%v, jsonOutput=%v, outputFile=%s, inputFile=%s, checkSubTakeover=%v",
		domain, interactive, jsonOutput, outputFile, inputFile, checkSubTakeover)

	var fingerprints []TakeoverFingerprint
	if checkSubTakeover {
		var err error
		fingerprints, err = getFingerprints()
		if err != nil {
			fmt.Printf("Warning: could not load takeover fingerprints: %v\n", err)
		}
	}

	if inputFile != "" {
		processDomainFile(inputFile, jsonOutput, outputFile, fingerprints)
	} else if interactive {
		for {
			fmt.Print("Enter domain to check (or 'exit' to quit): ")
			scanner := bufio.NewScanner(os.Stdin)
			scanner.Scan()
			domain = scanner.Text()

			if strings.ToLower(domain) == "exit" {
				break
			}

			checkDomain(domain, jsonOutput, outputFile, fingerprints)
		}
	} else if domain != "" {
		checkDomain(domain, jsonOutput, outputFile, fingerprints)
	} else {
		fmt.Println("Please specify a domain with -domain, use -input to read from a file, or use -interactive")
		flag.Usage()
	}
}

func processDomainFile(inputFile string, jsonOutput bool, outputFile string, fingerprints []TakeoverFingerprint) {
	Debug("Processing input file: %s", inputFile)

	file, err := os.Open(inputFile)
	if err != nil {
		fmt.Printf("Error opening file %s: %v\n", inputFile, err)
		return
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)

	domainCount := 0
	for scanner.Scan() {
		domain := strings.TrimSpace(scanner.Text())
		if domain == "" || strings.HasPrefix(domain, "#") {
			continue
		}

		Debug("Processing domain: %s", domain)
		checkDomain(domain, jsonOutput, outputFile, fingerprints)
		domainCount++
	}

	if err := scanner.Err(); err != nil {
		fmt.Printf("Error reading file: %v\n", err)
	}

	Debug("Processed %d domains from %s", domainCount, inputFile)
	fmt.Printf("Processed %d domains from %s\n", domainCount, inputFile)
	if outputFile != "" {
		fmt.Printf("Results saved to %s\n", outputFile)
	}
}

func issuestoJSON(domain string, spfDmarcRecord SPFDMARCRecord, issues []IssueScanResult) JSONScanResult {
	result := JSONScanResult{
		Domain:  domain,
		SPF:     spfDmarcRecord.spfRecord,
		DMARC:   spfDmarcRecord.dmarcRecord,
		Issues:  make([]JSONIssue, 0),
		Success: true,
	}

	if len(issues) == 0 {
		result.Message = "No issues found. The domain is well protected against spoofing."
	} else {
		for _, issue := range issues {
			jsonIssue := JSONIssue{
				Code:     issue.code,
				Title:    issue.title,
				Detail:   issue.detail,
				Severity: issue.severity,
			}
			result.Issues = append(result.Issues, jsonIssue)
		}
	}

	return result
}

func DebugSPFRecord(prefix string, record *ParsedSPFRecord, domainName string, spfText string) {
	if record == nil {
		Debug("%sNil SPF Record", prefix)
		return
	}

	if domainName != "" && spfText != "" {
		Debug("%sDomain: %s", prefix, domainName)
		Debug("%sSPF Record: %s", prefix, spfText)
	}

	Debug("%sA Records: %v", prefix, record.aRecord)
	Debug("%sMX Records: %v", prefix, record.mxRecord)
	Debug("%sExists Records: %v", prefix, record.existsRecords)
	Debug("%sPTR Records: %v", prefix, record.ptrRecords)

	if len(record.includeRecords) > 0 {
		Debug("%sIncludes (%d):", prefix, len(record.includeRecords))
		for i, include := range record.includeRecords {
			Debug("%s  [%d] Domain: %s", prefix, i, include.includeRecord)
			if include.subLookup != nil {
				Debug("%s  [%d] SubLookup:", prefix, i)
				DebugSPFRecord(prefix+"    ", include.subLookup, include.includeRecord, include.spfText)
			} else {
				Debug("%s  [%d] SubLookup: nil", prefix, i)
			}
		}
	} else {
		Debug("%sNo includes", prefix)
	}
}

func checkDomain(domain string, jsonOutput bool, outputFile string, fingerprints []TakeoverFingerprint) {
	Debug("Checking domain: %s", domain)

	spfDmarcRecord := getSPFDMARCRecord(domain)
	Debug("Retrieved SPF record: %s", spfDmarcRecord.spfRecord)
	Debug("Retrieved DMARC record: %s", spfDmarcRecord.dmarcRecord)

	parsedSPF := parseSPFRecord(spfDmarcRecord.spfRecord, domain, "", "")

	Debug("Parsed SPF record structure:")
	DebugSPFRecord("  ", parsedSPF, domain, spfDmarcRecord.spfRecord)

	engine := &IssueEngine{}
	issues := engine.IssueScan(spfDmarcRecord, parsedSPF)
	Debug("Identified issues: %+v", issues)

	var takeoverResult *SubdomainTakeoverResult
	if fingerprints != nil {
		result := checkSubdomainTakeover(domain, fingerprints)
		takeoverResult = &result
		Debug("Subdomain takeover result: %+v", takeoverResult)
	}

	var output string
	var fileOutput string

	if jsonOutput {
		jsonResult := issuestoJSON(domain, spfDmarcRecord, issues)

		if takeoverResult != nil {
			jsonResult.TakeoverResults = takeoverResult
		}

		jsonData, err := json.MarshalIndent(jsonResult, "", "  ")
		if err != nil {
			output = fmt.Sprintf("{\"success\": false, \"message\": \"Error generating JSON: %s\"}\n", err)
			fileOutput = output
		} else {
			output = string(jsonData)

			if outputFile != "" {
				fileOutput = "[\n" + output + "\n]"
			} else {
				fileOutput = output
			}
		}
	} else {
		var builder strings.Builder
		builder.WriteString(fmt.Sprintf("\n--- Checking domain: %s ---\n", domain))
		builder.WriteString(fmt.Sprintf("SPF Record: %s\n", spfDmarcRecord.spfRecord))
		builder.WriteString(fmt.Sprintf("DMARC Record: %s\n\n", spfDmarcRecord.dmarcRecord))

		if len(issues) == 0 {
			builder.WriteString("No issues found. The domain is well protected against spoofing.\n")
		} else {
			builder.WriteString("Found the following issues:\n")
			for _, issue := range issues {
				builder.WriteString(fmt.Sprintf("\nCode %d: %s\n", issue.code, issue.title))
				builder.WriteString(fmt.Sprintf("   Severity: %s\n", issue.severity))
				builder.WriteString(fmt.Sprintf("   Detail: %s\n", issue.detail))
			}
		}

		if takeoverResult != nil {
			builder.WriteString("\n--- Subdomain Takeover Check ---\n")
			if takeoverResult.Vulnerable {
				builder.WriteString(fmt.Sprintf("VULNERABLE to subdomain takeover!\n"))
				builder.WriteString(fmt.Sprintf("Service: %s\n", takeoverResult.Service))
				if takeoverResult.CnameRecord != "" {
					builder.WriteString(fmt.Sprintf("CNAME: %s\n", takeoverResult.CnameRecord))
				}
				if takeoverResult.Fingerprint != "" {
					builder.WriteString(fmt.Sprintf("Matching Fingerprint: %s\n", takeoverResult.Fingerprint))
				}
				if takeoverResult.ResponseBody != "" {
					builder.WriteString(fmt.Sprintf("Response snippet: %s\n", takeoverResult.ResponseBody))
				}
			} else {
				builder.WriteString("Not vulnerable to subdomain takeover.\n")
				if takeoverResult.ErrorMessage != "" {
					builder.WriteString(fmt.Sprintf("Note: %s\n", takeoverResult.ErrorMessage))
				}
				if takeoverResult.CnameRecord != "" {
					builder.WriteString(fmt.Sprintf("CNAME: %s\n", takeoverResult.CnameRecord))
				}
			}
		}

		builder.WriteString("\n--- End of report ---\n")
		output = builder.String()
		fileOutput = output
	}

	fmt.Print(output)

	if outputFile != "" {
		err := os.WriteFile(outputFile, []byte(fileOutput), 0644)
		if err != nil {
			fmt.Printf("Error writing to file %s: %v\n", outputFile, err)
		} else {
			fmt.Printf("Results saved to %s\n", outputFile)
		}
	}
}