pcdns.go

package main

import (
	"errors"
	"flag"
	"github.com/boltdb/bolt"
	"github.com/miekg/dns"
	"log"
	"math"
	"net"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"
)

var (
	tsig     *string
	db_path  *string
	port     *int
	bdb      *bolt.DB
	logfile  *string
	pid_file *string
)

const rr_bucket = "rr"
const ip_notice = "192.168.200.1"

func isMalicious(name string) bool {
	if name == "www.gfxworld.ws." {
		return true
	}
	return false
}

func getGatewayPolicy(mac string) map[string]map[string]bool {
	policyMap := map[string]map[string]map[string]bool {
		"08:00:27:9f:e2:a1": map[string]map[string]bool {
			"56:84:7a:fe:97:99": map[string]bool {
				"www.google.com.": true,
				"www.yahoo.com.": true,
			},
			//"12:34:56:78:aa:bb": false,
		},
	}

	return policyMap[mac]
}

func getKey(domain string, rtype uint16) (r string, e error) {
	if n, ok := dns.IsDomainName(domain); ok {
		labels := dns.SplitDomainName(domain)

		// Reverse domain, starting from top-level domain
		// eg. ".com.mkaczanowski.test"
		var tmp string
		for i := 0; i < int(math.Floor(float64(n / 2))); i++ {
			tmp = labels[i]
			labels[i] = labels[n - 1]
			labels[n - 1] = tmp
		}

		reverse_domain := strings.Join(labels, ".")
		r = strings.Join([]string{reverse_domain, strconv.Itoa(int(rtype))}, "_")
	} else {
		e = errors.New("Invailid domain: " + domain)
		log.Println(e.Error())
	}

	return r, e
}

func createBucket(bucket string) (err error) {
	err = bdb.Update(func(tx *bolt.Tx) error {
		_, err := tx.CreateBucketIfNotExists([]byte(bucket))
		if err != nil {
			e := errors.New("Create bucket: " + bucket)
			log.Println(e.Error())

			return e
		}

		return nil
	})

	return err
}

func deleteRecord(domain string, rtype uint16) (err error) {
	key, _ := getKey(domain, rtype)
	err = bdb.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(rr_bucket))
		err := b.Delete([]byte(key))

		if err != nil {
			e := errors.New("Delete record failed for domain: " + domain)
			log.Println(e.Error())

			return e
		}

		return nil
	})

	return err
}

func storeRecord(rr dns.RR) (err error) {
	key, _ := getKey(rr.Header().Name, rr.Header().Rrtype)
	err = bdb.Update(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(rr_bucket))
		err := b.Put([]byte(key), []byte(rr.String()))

		if err != nil {
			e := errors.New("Store record failed: " + rr.String())
			log.Println(e.Error())

			return e
		}

		return nil
	})

	return err
}

func getRecord(domain string, rtype uint16) (rr dns.RR, err error) {
	key, _ := getKey(domain, rtype)
	var v []byte

	err = bdb.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(rr_bucket))
		v = b.Get([]byte(key))

		if string(v) == "" {
			e := errors.New("Record not found, key: " + key)
			log.Println(e.Error())

			return e
		}

		return nil
	})

	if err == nil {
		rr, err = dns.NewRR(string(v))
	}

	return rr, err
}

func updateRecord(r dns.RR, q *dns.Question) {
	var (
		rr dns.RR
		name string
		rtype uint16
		ttl uint32
		ip net.IP
	)

	header := r.Header()
	name = header.Name
	rtype = header.Rrtype
	ttl = header.Ttl

	if _, ok := dns.IsDomainName(name); ok {
		if header.Class == dns.ClassANY && header.Rdlength == 0 {
			// Delete record
			deleteRecord(name, rtype)
		} else {
			// Add record
			rheader := dns.RR_Header{
				Name:   name,
				Rrtype: rtype,
				Class:  dns.ClassINET,
				Ttl:    ttl,
			}

			if a, ok := r.(*dns.A); ok {
				rrr, err := getRecord(name, rtype)
				if err == nil {
					rr = rrr.(*dns.A)
				} else {
					rr = new(dns.A)
				}

				ip = a.A
				rr.(*dns.A).Hdr = rheader
				rr.(*dns.A).A = ip
			} else if a, ok := r.(*dns.AAAA); ok {
				rrr, err := getRecord(name, rtype)
				if err == nil {
					rr = rrr.(*dns.AAAA)
				} else {
					rr = new(dns.AAAA)
				}

				ip = a.AAAA
				rr.(*dns.AAAA).Hdr = rheader
				rr.(*dns.AAAA).AAAA = ip
			}

			storeRecord(rr)
		}
	}
}

func parseQuery(m *dns.Msg, req *dns.Msg) {
	var rr dns.RR
	var rgMac, clientMac string

	opt := req.IsEdns0();
	if opt != nil {
		for _, o := range opt.Option {
			lo := o.(*dns.EDNS0_LOCAL);
			if lo.Option() == 65001 {
				rgMac = string(lo.Data)
			} else if lo.Option() == 65002 {
				clientMac = string(lo.Data)
			}
		}
		if rgMac != "" {
			log.Printf("Gateway MAC address >%s", rgMac)
		}
		if clientMac != "" {
			log.Printf("Client MAC address >%s", clientMac)
		}
	}

	policy := getGatewayPolicy(rgMac)

	for _, q := range m.Question {
		malicious := isMalicious(q.Name)
		if (malicious) {
			log.Printf("Domain %s is detected malicious and blocked", q.Name)
			mr := new(dns.A)
			mr.A = net.ParseIP(ip_notice)
			mr.Hdr = dns.RR_Header{Name: q.Name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 86400}

			m.Answer = append(m.Answer, mr)
			return
		}

		blocked := policy[clientMac][q.Name]
		if blocked {
			log.Printf("Domain %s is blocked under Gateway and Client's policy", q.Name)
			br := new(dns.A)
			br.A = net.ParseIP(ip_notice)
			br.Hdr = dns.RR_Header{Name: q.Name, Rrtype: dns.TypeA, Class: dns.ClassINET, Ttl: 3600}

			m.Answer = append(m.Answer, br)
			return
		}

		if read_rr, e := getRecord(q.Name, q.Qtype); e == nil {
			rr = read_rr.(dns.RR)
			if rr.Header().Name == q.Name {
				m.Answer = append(m.Answer, rr)
			}
		}
	}
}

func handleDnsRequest(w dns.ResponseWriter, r *dns.Msg) {
	m := new(dns.Msg)
	m.SetReply(r)
	m.Compress = false

	log.Printf("DNS Request >\n%s", r.String())
	switch r.Opcode {
	case dns.OpcodeQuery:
		parseQuery(m, r)

	case dns.OpcodeUpdate:
		for _, question := range r.Question {
			for _, rr := range r.Ns {
				updateRecord(rr, &question)
			}
		}
	}

	if r.IsTsig() != nil {
		if w.TsigStatus() == nil {
			m.SetTsig(r.Extra[len(r.Extra) - 1].(*dns.TSIG).Hdr.Name,
				dns.HmacMD5, 300, time.Now().Unix())
		} else {
			log.Println("Status", w.TsigStatus().Error())
		}
	}

	w.WriteMsg(m)
}

func serve(name, secret string, port int) {
	server := &dns.Server{Addr: ":" + strconv.Itoa(port), Net: "udp"}

	if name != "" {
		server.TsigSecret = map[string]string{name: secret}
	}

	err := server.ListenAndServe()
	defer server.Shutdown()

	if err != nil {
		log.Fatalf("Failed to setup the udp server: %sn", err.Error())
	}
}

func main() {
	var (
		name string   // tsig keyname
		secret string   // tsig base64
		fh     *os.File // logfile handle
	)

	// Parse flags
	logfile = flag.String("logfile", "", "path to log file")
	port = flag.Int("port", 53, "server port")
	tsig = flag.String("tsig", "", "use MD5 hmac tsig: keyname:base64")
	db_path = flag.String("db_path", "./dyndns.db", "location where db will be stored")
	pid_file = flag.String("pid", "./go-dyndns.pid", "pid file location")

	flag.Parse()

	// Open db
	db, err := bolt.Open(*db_path, 0600,
		&bolt.Options{Timeout: 10 * time.Second})

	if err != nil {
		log.Fatal(err)
	}
	defer db.Close()
	bdb = db

	// Create dns bucket if doesn't exist
	createBucket(rr_bucket)

	// Attach request handler func
	dns.HandleFunc(".", handleDnsRequest)

	// Tsig extract
	if *tsig != "" {
		a := strings.SplitN(*tsig, ":", 2)
		name, secret = dns.Fqdn(a[0]), a[1]
	}

	// Logger setup
	if *logfile != "" {
		if _, err := os.Stat(*logfile); os.IsNotExist(err) {
			if file, err := os.Create(*logfile); err != nil {
				if err != nil {
					log.Panic("Couldn't create log file: ", err)
				}

				fh = file
			}
		} else {
			fh, _ = os.OpenFile(*logfile, os.O_RDWR | os.O_CREATE | os.O_APPEND, 0666)
		}
		defer fh.Close()
		log.SetOutput(fh)
	}

	// Pidfile
	file, err := os.OpenFile(*pid_file, os.O_RDWR | os.O_CREATE, 0666)
	if err != nil {
		log.Panic("Couldn't create pid file: ", err)
	} else {
		file.Write([]byte(strconv.Itoa(syscall.Getpid())))
		defer file.Close()
	}

	// Start server
	go serve(name, secret, *port)

	sig := make(chan os.Signal)
	signal.Notify(sig, syscall.SIGINT, syscall.SIGTERM)

endless:
	for {
		select {
		case s := <-sig:
			log.Printf("Signal (%d) received, stoppingn", s)
			break endless
		}
	}
}