package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"strconv"
	"strings"
	"syscall"
	"text/template"

	"gopkg.in/yaml.v2"
)

var zoneTemplate = template.Must(template.New("bind").Parse(`$ORIGIN .
$TTL 86400      ; 1 day
{{ .Zone.Name }}               IN SOA  {{ .Zone.PrimaryNS }} {{ .Zone.HostmasterEmail }} (
                                1          ; serial
                                604800     ; refresh (1 week)
                                86400      ; retry (1 day)
                                2419200    ; expire (4 weeks)
                                86400      ; minimum (1 day)
                                )
{{ range .Zone.Nameservers }}
                        NS      {{ . }}
{{ end }}
`))

var cfgTemplate = template.Must(template.New("bind").Parse(`// vim:ft=named
{{ range .Keys -}}
key "{{ .Name }}" {
	algorithm {{ .Algorithm }};
	secret "{{ .KeyData }}";
};
{{ end }}

{{ range $name, $value := .Acls -}}
acl {{ $name }} {
	{{ range $value -}}
	{{ . }};
	{{ end }}
};
{{ end }}

{{ range .Views -}}
view {{ .Name }} {
	match-clients {
		{{ range .MatchClients -}}
			{{ . }};
		{{ end }}
	};

	{{ .RawInclude }}

	{{- if .Servers }}
	{{ range .Servers -}}
	server {{ .Address }} {
		keys {{ .Key }};
	};
	{{- end }}
	{{ end -}}

	{{ if .NotifySource -}}notify-source {{ .NotifySource }};{{- end }}
	{{ if .AlsoNotify -}}
	also-notify {
	{{- range .AlsoNotify }}
		{{ .Address }}{{ if .Key }} key {{ .Key }}{{ end }};
	{{- end }}
	};
	{{- end }}

	include "/etc/bind/named.conf.default-zones";

	{{ range .Zones }}
	zone "{{ .Name }}" {
		{{ if .InView -}}
		in-view {{ .InView }};
		{{- else -}}
		type {{ .Type }};
		{{ if .ForwardOnly }}forward only;{{ end }}
		{{- if .File }}file "{{ .File }}";{{ end }}
		{{ if .UpdateGrants -}}
		update-policy {
		{{- range .UpdateGrants }}
			grant {{ . }} zonesub ANY;
		{{- end }}
		};
		{{- end }}
		{{- if .AllowQuery -}}
		allow-query {
		{{- range .AllowQuery }}
			{{ . }};
		{{- end }}
		};
		{{- end -}}
		{{ if .Forwarders -}}
		forwarders {
		{{- range .Forwarders }}
			{{ . }};
		{{- end }}
		};
		{{- end }}
		{{ if .Primaries -}}
		masters {
		{{- range .Primaries }}
			{{ . }};
		{{- end }}
		};
		{{- end -}}
		{{- end }}
	};
	{{ end }}
};
{{ end }}
`))

type Server struct {
	Address    string
	IPs        []string            `yaml:"ips"`
	Type       string              `yaml:"type"`
	Key        *string             `yaml:"key"`
	Forwarders map[string][]string `yaml:"forwarders"`
}

type ServerMap map[string]*Server

func (l *ServerMap) UnmarshalYAML(unmarshal func(v interface{}) error) error {
	var rawData map[string]*Server
	if err := unmarshal(&rawData); err != nil {
		return err
	}

	for k, v := range rawData {
		v.Address = k
	}

	*l = ServerMap(rawData)

	return nil
}

type Config struct {
	Keys        map[string]map[string]map[string]string `yaml:"keys"`
	StaticAcls  map[string][]string                     `yaml:"static-acls"`
	DynamicAcls map[string]struct {
		Generator string `yaml:"generator"`
		Filter    string `yaml:"filter"`
	} `yaml:"dynamic-acls"`
	Servers ServerMap `yaml:"servers"`
	Views   map[string]struct {
		MatchClients []string `yaml:"match-clients"`
		RawInclude   string   `yaml:"raw-include"`
	} `yaml:"views"`
	Zones []struct {
		Name            string   `yaml:"name"`
		Type            string   `yaml:"type"` // forward-only is only valid type
		MasterViews     []string `yaml:"master-views"`
		InViews         []string `yaml:"in-views"`
		AllowUpdateKeys []string `yaml:"allow-update-keys"`
		AllowQuery      []string `yaml:"allow-query"`
	} `yaml:"zones"`
}

type TemplateData struct {
	Keys  []MaterializedKey
	Acls  map[string][]string
	Views []MaterializedView
}

type MaterializedKey struct {
	Name      string
	Algorithm string
	KeyData   string
}

type MaterializedView struct {
	Name         string
	MatchClients []string
	RawInclude   string
	Zones        []MaterializedZone
	NotifySource *string
	AlsoNotify   []Server
	Servers      []Server
}

type MaterializedZone struct {
	Name         string
	Type         string
	File         string
	InView       *string
	ForwardOnly  bool
	UpdateGrants []string
	Forwarders   []string
	AllowQuery   []string
	Primaries    []string
}

func inSlice(needle string, haystack []string) bool {
	for _, v := range haystack {
		if v == needle {
			return true
		}
	}
	return false
}

func materializeZones(cfg *Config, forView string, forServer Server) []MaterializedZone {
	out := []MaterializedZone{}

	for _, zone := range cfg.Zones {
		var forwarders []string
		forwardOnly := false
		zoneType := forServer.Type
		if zone.Type == "forward-only" {
			zoneType = "forward"
			forwardOnly = true

			if fw, ok := forServer.Forwarders[zone.Name]; ok {
				forwarders = fw
			} else {
				// Some forward-only zones don't have forwarders in a given
				// region, so suppress them if there's nothing to forward to
				continue
			}
		}

		masteredHere := inSlice(forView, zone.MasterViews)
		exclusiveZone := len(zone.InViews) > 0

		// Zones mastered in a view are always materialized there.
		//
		// Zones that specify an in-views clause are only materialized in a
		// view if that view is listed in the in-views list.
		//
		// Zones not mastered in a view and without an in-views list are always
		// materialized.
		if masteredHere || !exclusiveZone || (exclusiveZone && inSlice(forView, zone.InViews)) {
			keys := []string{}
			primaries := []string{}

			if zoneType == "primary" {
				// Only primary zones may be updated
				for _, k := range zone.AllowUpdateKeys {
					keys = append(keys, fmt.Sprintf("%s-%s", k, forView))
				}
			} else if len(forwarders) == 0 { // forward-only zones have forwarders, not primaries
				// Secondary zones use all primary servers as their primaries
				for _, v := range cfg.Servers {
					if v.Type == "primary" {
						primaries = append(primaries, v.Address)
					}
				}
			}

			// Zones not mastered here are considered to be mastered in their
			// first master view. This may not be correct for views nested in
			// views.
			var masterView *string
			if !masteredHere {
				masterView = &zone.MasterViews[0]
			}

			// forward-only zones are memory resident and have no file
			fileName := fmt.Sprintf("%s/db.%s", forView, zone.Name)
			if zoneType == "forward" {
				fileName = ""
			}

			out = append(out, MaterializedZone{
				Name:         zone.Name,
				Type:         zoneType,
				File:         fileName,
				InView:       masterView,
				ForwardOnly:  forwardOnly,
				UpdateGrants: keys,
				Forwarders:   forwarders,
				AllowQuery:   zone.AllowQuery,
				Primaries:    primaries,
			})
		}
	}

	return out
}

func getSecondaryServers(cfg *Config, viewName string) []Server {
	out := []Server{}
	for _, v := range cfg.Servers {
		if v.Type == "secondary" {
			kv := fmt.Sprintf("%s-%s", *v.Key, viewName)
			out = append(out, Server{
				Address: v.Address,
				Key:     &kv,
			})
		}
	}
	return out
}

func materializeViews(cfg *Config, forServer Server) []MaterializedView {
	out := []MaterializedView{}

	for viewName, configView := range cfg.Views {
		var notifySource *string
		var servers []Server
		var alsoNotify []Server

		if forServer.Type == "primary" {
			// Only primary zones do notifies
			notifySource = &forServer.Address
			alsoNotify = getSecondaryServers(cfg, viewName)
		} else {
			myKey := fmt.Sprintf("%s-%s", *forServer.Key, viewName)

			for _, s := range cfg.Servers {
				if s.Type == "primary" {
					servers = append(servers, Server{
						Address: s.Address,
						Key:     &myKey,
					})
				}
			}
		}

		out = append(out, MaterializedView{
			Name:         viewName,
			MatchClients: configView.MatchClients,
			RawInclude:   configView.RawInclude,
			Zones:        materializeZones(cfg, viewName, forServer),
			AlsoNotify:   alsoNotify,
			NotifySource: notifySource,
			Servers:      servers,
		})
	}
	return out
}

func materializeKeys(cfg *Config) []MaterializedKey {
	out := []MaterializedKey{}
	// Flatten the key YAML structure to <key_name>-<view>
	for name, subKeys := range cfg.Keys {
		for view, key := range subKeys {
			for keyAlgo, keyData := range key {
				out = append(out, MaterializedKey{
					Name:      fmt.Sprintf("%s-%s", name, view),
					Algorithm: keyAlgo,
					KeyData:   keyData,
				})
			}
		}
	}
	return out
}

func aclGenerateServers(cfg *Config) []string {
	lines := []string{}

	for _, server := range cfg.Servers {
		lines = append(lines, server.Address)
		lines = append(lines, server.IPs...)
	}

	return lines
}

// Filters use the https://golang.org/pkg/path/filepath/#Match syntax. If no
// filter is specified then the default is "*".
func aclGenerateKeys(cfg *Config, keys []MaterializedKey, filter string) []string {
	lines := []string{}

	if filter == "" {
		filter = "*"
	}

	for _, k := range keys {
		if ok, _ := filepath.Match(filter, k.Name); ok {
			lines = append(lines, fmt.Sprintf("key %s", k.Name))
		}
	}

	return lines
}

func materializeAcls(cfg *Config, keys []MaterializedKey) map[string][]string {
	out := map[string][]string{}

	keyIndex := map[string]bool{}
	for _, k := range keys {
		keyIndex[k.Name] = true
	}

	// Generate dynamic ACLs
	for aclName, acl := range cfg.DynamicAcls {
		switch acl.Generator {
		case "servers":
			out[aclName] = aclGenerateServers(cfg)
		case "keys":
			out[aclName] = aclGenerateKeys(cfg, keys, acl.Filter)
		}
	}

	// For static key-based ACLs only emit key lines for keys we have
	//
	// TODO: This was for an old iteration of the config format. All key-based
	// ACLs should be generated now. If that's the case then this loop can just
	// be removed.
	for aclName, acl := range cfg.StaticAcls {
		validLines := []string{}
		for _, line := range acl {
			parts := strings.Split(line, " ")
			if parts[0] == "key" {
				if _, ok := keyIndex[parts[1]]; !ok {
					continue
				}
			}

			validLines = append(validLines, line)
		}
		out[aclName] = validLines
	}

	return out
}

func loadYaml(filename string, into interface{}) error {
	rawData, err := ioutil.ReadFile(filename)
	if err != nil {
		return err
	}

	if err = yaml.Unmarshal(rawData, into); err != nil {
		return err
	}

	return nil
}

func lookupUserGroup(un, gn string) (int, int, error) {
	u, err := user.Lookup(un)
	if err != nil {
		return 0, 0, err
	}

	g, err := user.LookupGroup(gn)
	if err != nil {
		return 0, 0, nil
	}

	uid, err := strconv.Atoi(u.Uid)
	if err != nil {
		return 0, 0, err
	}

	gid, err := strconv.Atoi(g.Gid)
	if err != nil {
		return 0, 0, err
	}

	return uid, gid, nil
}

func generateNamedConf(forServer string, cfgFile, keysFile string) []byte {
	var out bytes.Buffer

	cfg := Config{}
	if err := loadYaml(cfgFile, &cfg); err != nil {
		log.Fatalf("error: %s", err)
	}

	if err := loadYaml(keysFile, &cfg); err != nil {
		log.Fatalf("error: %s", err)
	}

	server, ok := cfg.Servers[forServer]
	if !ok {
		log.Fatalf("No server %s\n", forServer)
	}

	keys := materializeKeys(&cfg)

	cfgTemplate.Execute(&out, &TemplateData{
		Keys:  keys,
		Acls:  materializeAcls(&cfg, keys),
		Views: materializeViews(&cfg, *server),
	})

	return out.Bytes()
}

// This will replace the current process
func execToProcess(args []string) {
	log.Printf("running: %s %+v", args[1], args[1:])
	if err := syscall.Exec(args[1], args[1:], []string{}); err != nil {
		log.Fatalf("error: exec failed %s", err)
	}
}

func writeFile(path string, contents []byte, mode os.FileMode, uid, gid int) error {
	if err := ioutil.WriteFile(path, contents, mode); err != nil {
		return err
	}

	if err := os.Chown(path, uid, gid); err != nil {
		return err
	}

	return nil
}

// This is used for RNDC and only from the current host, so just generate
// it fresh each container start.
func generateRNDCKey(path string, uid, gid int) error {
	var out bytes.Buffer
	cmd := exec.Command("ddns-confgen", "-q", "-k", "rndc-key")
	cmd.Stdout = &out
	if err := cmd.Run(); err != nil {
		return err
	}

	if err := writeFile(path, out.Bytes(), 0440, uid, gid); err != nil {
		return err
	}

	return nil
}

func ensureDirExists(path string, mode os.FileMode, uid, gid int) {
	_ = os.Mkdir(path, mode)
	os.Chown(path, uid, gid)
}

func realMain() {
	forServer := os.Getenv("SERVER_ADDRESS")
	if forServer == "" {
		log.Fatalf("error: SERVER_ADDRESS is not in environment")
	}

	namedUser, namedGroup, err := lookupUserGroup("named", "named")
	if err != nil {
		log.Fatalf("error: %s", err)
	}

	cfg := generateNamedConf(forServer, "/etc/bind/zones.yaml", "keys.yaml")
	if err := writeFile("/etc/bind/named_local.conf", cfg, 0644, namedUser, namedGroup); err != nil {
		log.Fatalf("error: %s", err)
	}

	// Ensure these directories exist
	ensureDirExists("/etc/bind/local/cache", 0755, namedUser, namedGroup)
	ensureDirExists("/etc/bind/local/cache/internal", 0755, namedUser, namedGroup)
	ensureDirExists("/etc/bind/local/cache/external", 0755, namedUser, namedGroup)

	// TODO: Write missing zone files for master

	if err := generateRNDCKey("/etc/bind/rndc.key", namedUser, namedGroup); err != nil {
		log.Fatalf("error: %s", err)
	}

	execToProcess(os.Args)
}

func testMain() {
	forServer := os.Getenv("SERVER_ADDRESS")
	if forServer == "" {
		log.Fatalf("error: SERVER_ADDRESS is not in environment")
	}

	namedUser, namedGroup, err := lookupUserGroup("mcrute", "mcrute")
	if err != nil {
		log.Fatalf("error: %s", err)
	}

	cfg := generateNamedConf(forServer, "zones.yaml", "keys.yaml")
	if err := writeFile("named_local.conf", cfg, 0644, namedUser, namedGroup); err != nil {
		log.Fatalf("error: %s", err)
	}
}

func main() {
	testMain()
}