87 lines
2.5 KiB
Go
87 lines
2.5 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
|
|
}
|
|
|
|
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:
|
|
log.Printf("Query for %s\n", q.Name[:len(q.Name)-1])
|
|
cur, err := coll.Find(context.TODO(), bson.M{"hostname":q.Name[:len(q.Name)-1]})
|
|
if err != nil {
|
|
log.Fatal(err)
|
|
}
|
|
|
|
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"
|
|
}
|
|
rr, err := dns.NewRR(fmt.Sprintf("%s %s %s",
|
|
q.Name, tp, ip))
|
|
if err == nil {
|
|
m.Answer = append(m.Answer, rr)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
w.WriteMsg(m)
|
|
}
|
|
}
|
|
|
|
func main() {
|
|
client, err := mongo.NewClient(options.Client().ApplyURI(os.Getenv("MONGO_URI")))
|
|
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 := 5354
|
|
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())
|
|
}
|
|
}
|