joker/main.go

package main

import (
	"net/http"
	"github.com/julienschmidt/httprouter"
	"github.com/hoisie/redis"
	"encoding/json"
	"github.com/weppos/publicsuffix-go/publicsuffix"
	"strings"
		"github.com/asaskevich/govalidator"
	"flag"
	"os"
)

var c *redis.Client

type infoResponse struct {
	RedisServer string `json:"redis"`
}

var (
	flagKey = flag.String("key", "godns:hosts", "Redis key for hash set")
	flagServer = flag.String("server", "192.168.1.9:6379", "redis host")
	filePath = flag.String("filepath", "/var/lib/joker/dist", "file path")
)

var (
	key string
	server string
)

func main() {
	flag.Parse()

	if key = os.Getenv("REDIS_KEY"); key == "" {
		key = *flagKey
	}

	if server = os.Getenv("REDIS_SERVER"); server == "" {
		server = *flagServer
	}

	c = &redis.Client{
		Addr: server,
	}

	router := httprouter.New()

	fs := http.FileServer(http.Dir(*filePath))

	serveFiles := func(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
		fs.ServeHTTP(w, r)
	}

	router.GET("/", serveFiles)
	router.GET("/js/*filepath", serveFiles)
	router.GET("/css/*filepath", serveFiles)

	router.GET("/info", getInfo)

	router.GET("/records", getRecords)
	router.GET("/records/:zone", getRecords)

	router.POST("/update", updateRecord)
	router.POST("/remove", removeRecord)

	http.ListenAndServe(":8080", router)
}

func getInfo(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	json.NewEncoder(w).Encode(&infoResponse{
		RedisServer: server,
	})
}

func getRecords(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
	zone := p.ByName("zone")

	hosts := make(map[string]string)

	if err := c.Hgetall(key, hosts); err != nil {
		return
	}

	groups := make(map[string]map[string]string)

	for host, addr := range hosts {
		domain, err := publicsuffix.Domain(host)

		if err != nil {
			domain = host
		}

		if zone != "" && !strings.HasSuffix(host, zone) {
			continue
		}

		if group, ok := groups[domain]; ok {
			group[host] = addr
		} else {
			groups[domain] = map[string]string{host:addr}
		}
	}

	json.NewEncoder(w).Encode(groups)
}

type domainRequest struct {
	Domain string `json:"domain"`
	IP string `json:"ip"`
	Group bool `json:"group"`
}

func updateRecord(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
	var req domainRequest

	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		return
	}

	req.Domain = strings.ToLower(req.Domain)

	if !govalidator.IsDNSName(req.Domain) || !govalidator.IsIP(req.IP) {
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	if _, err := c.Hset(key, req.Domain, []byte(req.IP)); err != nil {
		return
	}

	c.Publish("godns:update_record", []byte(strings.ToLower(req.Domain)))
}

func removeRecord(w http.ResponseWriter, r *http.Request, p httprouter.Params) {
	var req domainRequest

	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		return
	}

	req.Domain = strings.ToLower(req.Domain)

	if !govalidator.IsDNSName(strings.Replace(req.Domain, "*", "a", -1)) {
		w.WriteHeader(http.StatusBadRequest)
		return
	}

	if req.Group {
		hosts := make(map[string]string)

		if err := c.Hgetall(key, hosts); err != nil {
			return
		}

		for host, _ := range hosts {
			host = strings.ToLower(host)

			if !strings.HasSuffix(host, req.Domain) {
				continue
			}

			c.Hdel(key, host)

			c.Publish("godns:remove_record", []byte(host))
		}
	} else {
		if _, err := c.Hdel(key, req.Domain); err != nil {
			return
		}

		c.Publish("godns:remove_record", []byte(strings.ToLower(req.Domain)))
	}
}