Add vendor assessment agent

Signed-off-by: Aurélien Sibiril <81782+aureliensibiril@users.noreply.github.com>

pkg/agent/tools/security/dns_records.go

@@ -0,0 +1,166 @@
+// Copyright (c) 2026 Probo Inc <hello@getprobo.com>.
+//
+// Permission to use, copy, modify, and/or distribute this software for any
+// purpose with or without fee is hereby granted, provided that the above
+// copyright notice and this permission notice appear in all copies.
+//
+// THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES WITH
+// REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF MERCHANTABILITY
+// AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY SPECIAL, DIRECT,
+// INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES WHATSOEVER RESULTING FROM
+// LOSS OF USE, DATA OR PROFITS, WHETHER IN AN ACTION OF CONTRACT, NEGLIGENCE OR
+// OTHER TORTIOUS ACTION, ARISING OUT OF OR IN CONNECTION WITH THE USE OR
+// PERFORMANCE OF THIS 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
+			var 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
+}