112 lines
2.4 KiB
Go
112 lines
2.4 KiB
Go
package main
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"log"
|
|
"os"
|
|
"strconv"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/miekg/dns"
|
|
"go.mongodb.org/mongo-driver/bson"
|
|
"go.mongodb.org/mongo-driver/mongo"
|
|
"go.mongodb.org/mongo-driver/mongo/options"
|
|
)
|
|
|
|
type InventoryItem struct {
|
|
Ip []string
|
|
}
|
|
|
|
var mongoUri string = os.Getenv("MONGO_URI")
|
|
|
|
func appendResults(etype string, name string, m *dns.Msg, cur *mongo.Cursor) int {
|
|
count := 0
|
|
for cur.Next(context.TODO()) {
|
|
var elem InventoryItem
|
|
err := cur.Decode(&elem)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
for _, ip := range elem.Ip {
|
|
tp := "A"
|
|
if strings.Contains(ip, ":") {
|
|
tp = "AAAA"
|
|
}
|
|
|
|
if etype != tp {
|
|
continue
|
|
}
|
|
|
|
log.Printf("Appending: %s %s %s\n", name, tp, ip)
|
|
rr, err := dns.NewRR(fmt.Sprintf("%s. %s %s", name, tp, ip))
|
|
if err == nil {
|
|
m.Answer = append(m.Answer, rr)
|
|
count += 1
|
|
}
|
|
}
|
|
}
|
|
return count
|
|
}
|
|
|
|
func query(tp string, name string, m *dns.Msg, coll *mongo.Collection) {
|
|
// TODO: Validate `name` against RE_FQDN
|
|
log.Printf("Query %s for %s\n", tp, name)
|
|
cur, err := coll.Find(context.TODO(), bson.M{"dns.fqdn": name})
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
if appendResults(tp, name, m, cur) == 0 {
|
|
cur, err := coll.Find(context.TODO(), bson.M{"dns.san": name})
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
appendResults(tp, name, m, cur)
|
|
}
|
|
}
|
|
|
|
func wrapper(coll *mongo.Collection) func(dns.ResponseWriter, *dns.Msg) {
|
|
return func(w dns.ResponseWriter, r *dns.Msg) {
|
|
m := new(dns.Msg)
|
|
m.SetReply(r)
|
|
m.Compress = false
|
|
switch r.Opcode {
|
|
case dns.OpcodeQuery:
|
|
for _, q := range m.Question {
|
|
switch q.Qtype {
|
|
case dns.TypeA:
|
|
query("A", q.Name[:len(q.Name)-1], m, coll)
|
|
case dns.TypeAAAA:
|
|
query("AAAA", q.Name[:len(q.Name)-1], m, coll)
|
|
}
|
|
}
|
|
}
|
|
w.WriteMsg(m)
|
|
}
|
|
}
|
|
|
|
func main() {
|
|
client, err := mongo.NewClient(options.Client().ApplyURI(mongoUri))
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
ctx, _ := context.WithTimeout(context.Background(), 10*time.Second)
|
|
err = client.Connect(ctx)
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
coll := client.Database("kspace_accounting").Collection("inventory")
|
|
defer client.Disconnect(ctx)
|
|
dns.HandleFunc(".", wrapper(coll))
|
|
port := 53
|
|
server := &dns.Server{Addr: ":" + strconv.Itoa(port), Net: "udp"}
|
|
log.Printf("Starting at %d\n", port)
|
|
err2 := server.ListenAndServe()
|
|
defer server.Shutdown()
|
|
if err2 != nil {
|
|
log.Fatalf("Failed to start server: %s\n ", err2.Error())
|
|
}
|
|
}
|