From 3bb649f17544c81bfe928a021cd3654fce7e1dff Mon Sep 17 00:00:00 2001
From: Jeff Cody <jcody@censys.io>
Date: Tue, 28 Jan 2020 17:27:24 -0500
Subject: [PATCH] Add custom fake resolver

This adds a custom resolver, that will always resolve to the specified
ip address.  The intended usage is for when doing name-based scans,
but have a specified IP address as well.  This will provide a resolver
that can be added to a Dialer, that will cause all DNS lookups to match
the specified IP address.
---
 fake_resolver.go | 166 +++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 166 insertions(+)
 create mode 100644 fake_resolver.go

diff --git a/fake_resolver.go b/fake_resolver.go
new file mode 100644
index 0000000..a72a34e
--- /dev/null
+++ b/fake_resolver.go
@@ -0,0 +1,166 @@
+// Fake DNS lookups, so that they always returned the specified
+// IP
+//
+// Inspired by the golang net/dnsclient_unix_test.go code
+package zgrab2
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"golang.org/x/net/dns/dnsmessage"
+	"net"
+	"time"
+)
+
+// For a given IP, create a new Resolver that wraps a fake
+// DNS server. This resolver will always return an IP that
+// is represented by "ipstr", for DNS queries of the same
+// IP type.  Otherwise, it will return a DNS lookup error.
+func NewFakeResolver(ipstr string) (*net.Resolver, error) {
+	ip := net.ParseIP(ipstr)
+	if len(ip) < 4 {
+		return nil, fmt.Errorf("Fake resolver can't use non-IP '%s'", ipstr)
+	}
+	fDNS := FakeDNSServer{
+		IP: ip,
+	}
+	return &net.Resolver{
+		PreferGo: true, // Needed to force the use of the Go internal resolver
+		Dial:     fDNS.DialContext,
+	}, nil
+}
+
+type FakeDNSServer struct {
+	// Any domain name will resolve to this IP.  It can be either ipv4 or ipv6
+	IP net.IP
+}
+
+// For a given DNS query, return the hard-coded IP that is part of
+// FakeDNSServer.
+//
+// It will work with either ipv4 or ipv6 addresses; if a TypeA question
+// is received, we will only return the IP if what we have to return is
+// ipv4.  The same for TypeAAAA and ipv6.
+func (f *FakeDNSServer) fakeDNS(s string, dmsg dnsmessage.Message) (r dnsmessage.Message, err error) {
+
+	r = dnsmessage.Message{
+		Header: dnsmessage.Header{
+			ID:       dmsg.ID,
+			Response: true,
+		},
+		Questions: dmsg.Questions,
+	}
+	ipv6 := f.IP.To16()
+	ipv4 := f.IP.To4()
+	switch t := dmsg.Questions[0].Type; {
+	case t == dnsmessage.TypeA && ipv4 != nil:
+		var ip [4]byte
+		copy(ip[:], []byte(ipv4))
+		r.Answers = []dnsmessage.Resource{
+			{
+				Header: dnsmessage.ResourceHeader{
+					Name:   dmsg.Questions[0].Name,
+					Type:   dnsmessage.TypeA,
+					Class:  dnsmessage.ClassINET,
+					Length: 4,
+				},
+				Body: &dnsmessage.AResource{
+					A: ip,
+				},
+			},
+		}
+	case t == dnsmessage.TypeAAAA && ipv4 == nil:
+		var ip [16]byte
+		copy(ip[:], []byte(ipv6))
+		r.Answers = []dnsmessage.Resource{
+			{
+				Header: dnsmessage.ResourceHeader{
+					Name:   dmsg.Questions[0].Name,
+					Type:   dnsmessage.TypeAAAA,
+					Class:  dnsmessage.ClassINET,
+					Length: 16,
+				},
+				Body: &dnsmessage.AAAAResource{
+					AAAA: ip,
+				},
+			},
+		}
+	default:
+		r.Header.RCode = dnsmessage.RCodeNameError
+	}
+
+	return r, nil
+}
+
+// This merely wraps a custom net.Conn, that is only good for DNS
+// messages
+func (f *FakeDNSServer) DialContext(ctx context.Context, network,
+	address string) (net.Conn, error) {
+
+	conn := &fakeDNSPacketConn{
+		fakeDNSConn: fakeDNSConn{
+			server:  f,
+			network: network,
+			address: address,
+		},
+	}
+	return conn, nil
+}
+
+type fakeDNSConn struct {
+	net.Conn
+	server  *FakeDNSServer
+	network string
+	address string
+	dmsg    dnsmessage.Message
+}
+
+func (fc *fakeDNSConn) Read(b []byte) (int, error) {
+	resp, err := fc.server.fakeDNS(fc.address, fc.dmsg)
+	if err != nil {
+		return 0, err
+	}
+
+	bb := make([]byte, 2, 514)
+	bb, err = resp.AppendPack(bb)
+	if err != nil {
+		return 0, fmt.Errorf("cannot marshal DNS message: %v", err)
+	}
+
+	bb = bb[2:]
+	if len(b) < len(bb) {
+		return 0, errors.New("read would fragment DNS message")
+	}
+
+	copy(b, bb)
+	return len(bb), nil
+}
+
+func (fc *fakeDNSConn) Write(b []byte) (int, error) {
+	if fc.dmsg.Unpack(b) != nil {
+		return 0, fmt.Errorf("cannot unmarshal DNS message fake %s (%d)", fc.network, len(b))
+	}
+	return len(b), nil
+}
+
+func (fc *fakeDNSConn) SetDeadline(deadline time.Time) error {
+	return nil
+}
+
+func (fc *fakeDNSConn) Close() error {
+	return nil
+}
+
+type fakeDNSPacketConn struct {
+	net.PacketConn
+	fakeDNSConn
+}
+
+func (f *fakeDNSPacketConn) SetDeadline(deadline time.Time) error {
+	return nil
+}
+
+func (f *fakeDNSPacketConn) Close() error {
+	return f.fakeDNSConn.Close()
+}