probo/pkg/agent/tools/security/dns_records.go

// Copyright (c) 2026 Probo Inc <hello@probo.com>.
//
// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
// SOFTWARE.

package security

import (
	"context"
	"fmt"
	"strings"

	"codeberg.org/miekg/dns"
	"go.probo.inc/probo/pkg/agent"
)

type (
	dnsRecordsParams struct {
		Domain string `json:"domain" jsonschema:"The domain to query DNS records for (e.g. example.com)"`
	}

	dnsRecordsResult struct {
		A           []string `json:"a_records,omitempty"`
		AAAA        []string `json:"aaaa_records,omitempty"`
		MX          []string `json:"mx_records,omitempty"`
		CNAME       []string `json:"cname_records,omitempty"`
		TXT         []string `json:"txt_records,omitempty"`
		NS          []string `json:"ns_records,omitempty"`
		ErrorDetail string   `json:"error_detail,omitempty"`
	}

	queryOption func(*dns.MsgHeader)
)

func CheckDNSRecordsTool() agent.Tool {
	return agent.FunctionTool(
		"check_dns_records",
		"Query DNS records for a domain (A, AAAA, MX, CNAME, TXT, NS). Reveals hosting provider, email provider, and additional security signals.",
		func(ctx context.Context, p dnsRecordsParams) (agent.ToolResult, error) {
			fqdn := p.Domain
			if !strings.HasSuffix(fqdn, ".") {
				fqdn = fqdn + "."
			}

			hdr := dns.Header{Name: fqdn, Class: dns.ClassINET}
			client := dns.NewClient()

			var (
				result dnsRecordsResult
				errs   []string
			)

			// A records.

			if answers, err := queryDNS(ctx, client, &dns.A{Hdr: hdr}); err != nil {
				errs = append(errs, fmt.Sprintf("A query failed: %s", err))
			} else {
				for _, rr := range answers {
					if a, ok := rr.(*dns.A); ok {
						result.A = append(result.A, a.A.String())
					}
				}
			}

			// AAAA records.
			if answers, err := queryDNS(ctx, client, &dns.AAAA{Hdr: hdr}); err != nil {
				errs = append(errs, fmt.Sprintf("AAAA query failed: %s", err))
			} else {
				for _, rr := range answers {
					if aaaa, ok := rr.(*dns.AAAA); ok {
						result.AAAA = append(result.AAAA, aaaa.AAAA.String())
					}
				}
			}

			// MX records.
			if answers, err := queryDNS(ctx, client, &dns.MX{Hdr: hdr}); err != nil {
				errs = append(errs, fmt.Sprintf("MX query failed: %s", err))
			} else {
				for _, rr := range answers {
					if mx, ok := rr.(*dns.MX); ok {
						result.MX = append(result.MX, fmt.Sprintf("%d %s", mx.Preference, strings.TrimSuffix(mx.Mx, ".")))
					}
				}
			}

			// CNAME records.
			if answers, err := queryDNS(ctx, client, &dns.CNAME{Hdr: hdr}); err != nil {
				errs = append(errs, fmt.Sprintf("CNAME query failed: %s", err))
			} else {
				for _, rr := range answers {
					if cname, ok := rr.(*dns.CNAME); ok {
						result.CNAME = append(result.CNAME, strings.TrimSuffix(cname.Target, "."))
					}
				}
			}

			// TXT records.
			if answers, err := queryDNS(ctx, client, &dns.TXT{Hdr: hdr}); err != nil {
				errs = append(errs, fmt.Sprintf("TXT query failed: %s", err))
			} else {
				for _, rr := range answers {
					if txt, ok := rr.(*dns.TXT); ok {
						result.TXT = append(result.TXT, strings.Join(txt.Txt, ""))
					}
				}
			}

			// NS records.
			if answers, err := queryDNS(ctx, client, &dns.NS{Hdr: hdr}); err != nil {
				errs = append(errs, fmt.Sprintf("NS query failed: %s", err))
			} else {
				for _, rr := range answers {
					if ns, ok := rr.(*dns.NS); ok {
						result.NS = append(result.NS, strings.TrimSuffix(ns.Ns, "."))
					}
				}
			}

			if len(errs) > 0 {
				result.ErrorDetail = strings.Join(errs, "; ")
			}

			return agent.ResultJSON(result), nil
		},
	)
}

func withDNSSEC() queryOption {
	return func(h *dns.MsgHeader) {
		h.UDPSize = 4096
		h.Security = true
	}
}

func queryDNS(ctx context.Context, client *dns.Client, question dns.RR, opts ...queryOption) ([]dns.RR, error) {
	msg := &dns.Msg{
		MsgHeader: dns.MsgHeader{
			ID:               dns.ID(),
			RecursionDesired: true,
		},
	}
	for _, opt := range opts {
		opt(&msg.MsgHeader)
	}

	msg.Question = []dns.RR{question}

	resp, _, err := client.Exchange(ctx, msg, "udp", defaultResolverAddr)
	if err == nil && resp.Truncated {
		resp, _, err = client.Exchange(ctx, msg, "tcp", defaultResolverAddr)
	}

	if err != nil {
		return nil, err
	}

	if resp.Rcode != dns.RcodeSuccess {
		return nil, fmt.Errorf("cannot execute DNS query: %s", dns.RcodeToString[resp.Rcode])
	}

	return resp.Answer, nil
}