166 lines
3.7 KiB
Go
166 lines
3.7 KiB
Go
package zgrab2
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"golang.org/x/net/dns/dnsmessage"
|
|
"net"
|
|
"time"
|
|
)
|
|
|
|
// Fake DNS Resolver, to force a DNS lookup to return a pinned address
|
|
// Inspired by the golang net/dnsclient_unix_test.go code
|
|
//
|
|
// 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()
|
|
}
|