package main

// @FIXME: Proper handling of various upper/lower case combinations
// @FIXME: Implement missing search filters (in applyFilter)
// @FIXME: Add an initial prefix to the consul key value

import (
	"crypto/tls"
	"encoding/base64"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"math/rand"
	"os"
	"os/signal"
	"syscall"

	ldap "./ldapserver"
	consul "github.com/hashicorp/consul/api"
	message "github.com/vjeantet/goldap/message"
)

const DEBUG = false

const ATTR_USERPASSWORD = "userpassword"
const ATTR_MEMBER = "member"
const ATTR_MEMBEROF = "memberof"
const ATTR_ENTRYUUID = "entryuuid"
const ATTR_CREATORSNAME = "creatorsname"
const ATTR_CREATETIMESTAMP = "createtimestamp"
const ATTR_MODIFIERSNAME = "modifiersname"
const ATTR_MODIFYTIMESTAMP = "modifytimestamp"

type ConfigFile struct {
	Suffix        string   `json:"suffix"`
	BindAddress   string   `json:"bind_address"`
	ConsulHost    string   `json:"consul_host"`
	Acl           []string `json:"acl"`
	TLSCertFile   string   `json:"tls_cert_file"`
	TLSKeyFile    string   `json:"tls_key_file"`
	TLSServerName string   `json:"tls_server_name"`
	UseStartTLS   bool     `json:"use_starttls"`
}

type Config struct {
	Suffix      string
	BindAddress string
	ConsulHost  string

	Acl ACL

	TLSConfig   *tls.Config
	UseStartTLS bool
}

type Server struct {
	logger *log.Logger
	config Config
	kv     *consul.KV
}

type State struct {
	login Login
}

type Entry map[string][]string

var configFlag = flag.String("config", "./config.json", "Configuration file path")

func readConfig() Config {
	config_file := ConfigFile{
		BindAddress: "0.0.0.0:389",
	}

	bytes, err := ioutil.ReadFile(*configFlag)
	if err != nil {
		panic(err)
	}

	err = json.Unmarshal(bytes, &config_file)
	if err != nil {
		panic(err)
	}

	acl, err := ParseACL(config_file.Acl)
	if err != nil {
		panic(err)
	}

	ret := Config{
		Suffix:      config_file.Suffix,
		BindAddress: config_file.BindAddress,
		ConsulHost:  config_file.ConsulHost,
		Acl:         acl,
		UseStartTLS: config_file.UseStartTLS,
	}

	if config_file.TLSCertFile != "" && config_file.TLSKeyFile != "" && config_file.TLSServerName != "" {
		cert_txt, err := ioutil.ReadFile(config_file.TLSCertFile)
		if err != nil {
			panic(err)
		}
		key_txt, err := ioutil.ReadFile(config_file.TLSKeyFile)
		if err != nil {
			panic(err)
		}
		cert, err := tls.X509KeyPair(cert_txt, key_txt)
		if err != nil {
			panic(err)
		}
		ret.TLSConfig = &tls.Config{
			MinVersion:   tls.VersionTLS10,
			MaxVersion:   tls.VersionTLS12,
			Certificates: []tls.Certificate{cert},
			ServerName:   config_file.TLSServerName,
		}
	} else {
		log.Printf("Warning: no TLS configuration provided, running an insecure server.")
	}

	return ret
}

func main() {
	flag.Parse()

	ldap.Logger = log.New(os.Stdout, "[ldapserver] ", log.LstdFlags)

	config := readConfig()

	// Connect to Consul
	consul_config := consul.DefaultConfig()
	if config.ConsulHost != "" {
		consul_config.Address = config.ConsulHost
	}
	consul_client, err := consul.NewClient(consul_config)
	if err != nil {
		panic(err)
	}
	kv := consul_client.KV()

	// Create gobottin server
	gobottin := Server{
		logger: log.New(os.Stdout, "[gobottin] ", log.LstdFlags),
		config: config,
		kv:     kv,
	}
	err = gobottin.init()
	if err != nil {
		panic(err)
	}

	//Create a new LDAP Server
	ldapserver := ldap.NewServer()
	ldapserver.NewUserState = func() ldap.UserState {
		return &State{
			login: Login{
				user:   "ANONYMOUS",
				groups: []string{},
			},
		}
	}

	routes := ldap.NewRouteMux()

	routes.Bind(gobottin.handleBind)
	routes.Search(gobottin.handleSearch)
	routes.Add(gobottin.handleAdd)
	routes.Compare(gobottin.handleCompare)
	routes.Delete(gobottin.handleDelete)
	routes.Modify(gobottin.handleModify)

	if config.TLSConfig != nil && config.UseStartTLS {
		routes.Extended(gobottin.handleStartTLS).
			RequestName(ldap.NoticeOfStartTLS).Label("StartTLS")
	}

	ldapserver.Handle(routes)

	go func() {
		// When CTRL+C, SIGINT and SIGTERM signal occurs
		// Then stop server gracefully
		ch := make(chan os.Signal)
		signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
		<-ch
		close(ch)
		ldapserver.Stop()
	}()

	if config.TLSConfig != nil && !config.UseStartTLS {
		secureConn := func(s *ldap.Server) {
			s.Listener = tls.NewListener(s.Listener, config.TLSConfig)
		}
		err = ldapserver.ListenAndServe(config.BindAddress, secureConn)
	} else {
		err = ldapserver.ListenAndServe(config.BindAddress)
	}
	if err != nil {
		panic(err)
	}
}

func (server *Server) init() error {
	path, err := dnToConsul(server.config.Suffix)
	if err != nil {
		return err
	}

	pair, _, err := server.kv.Get(path+"/attribute=objectClass", nil)
	if err != nil {
		return err
	}

	if pair != nil {
		return nil
	}

	base_attributes := Entry{
		"objectClass":           []string{"top", "dcObject", "organization"},
		"structuralObjectClass": []string{"Organization"},
		ATTR_CREATORSNAME:       []string{server.config.Suffix},
		ATTR_CREATETIMESTAMP:    []string{genTimestamp()},
		ATTR_ENTRYUUID:          []string{genUuid()},
	}
	suffix_dn, err := parseDN(server.config.Suffix)
	if err != nil {
		return err
	}
	base_attributes[suffix_dn[0].Type] = []string{suffix_dn[0].Value}

	err = server.addElements(server.config.Suffix, base_attributes)
	if err != nil {
		return err
	}

	admin_pass := make([]byte, 8)
	rand.Read(admin_pass)
	admin_pass_str := base64.RawURLEncoding.EncodeToString(admin_pass)
	admin_pass_hash := SSHAEncode([]byte(admin_pass_str))

	admin_dn := "cn=admin," + server.config.Suffix
	admin_attributes := Entry{
		"objectClass":           []string{"simpleSecurityObject", "organizationalRole"},
		"description":           []string{"LDAP administrator"},
		"cn":                    []string{"admin"},
		"structuralObjectClass": []string{"organizationalRole"},
		ATTR_USERPASSWORD:       []string{admin_pass_hash},
		ATTR_CREATORSNAME:       []string{server.config.Suffix},
		ATTR_CREATETIMESTAMP:    []string{genTimestamp()},
		ATTR_ENTRYUUID:          []string{genUuid()},
	}

	err = server.addElements(admin_dn, admin_attributes)
	if err != nil {
		return err
	}

	server.logger.Printf(
		"It seems to be a new installation, we created a default user for you:\n\n    dn:          %s\n    password:    %s\n\nWe didn't use true random, you should replace it as soon as possible.",
		admin_dn,
		admin_pass_str,
	)

	return nil
}

func (server *Server) addElements(dn string, attrs Entry) error {
	prefix, err := dnToConsul(dn)
	if err != nil {
		return err
	}

	for k, v := range attrs {
		path := prefix + "/attribute=" + k
		if len(v) == 0 {
			// If we have zero values, delete associated k/v pair
			_, err := server.kv.Delete(path, nil)
			if err != nil {
				return err
			}
		} else {
			json, err := json.Marshal(v)
			if err != nil {
				return err
			}
			pair := &consul.KVPair{Key: path, Value: json}
			_, err = server.kv.Put(pair, nil)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func (server *Server) getAttribute(dn string, attr string) ([]string, error) {
	path, err := dnToConsul(dn)
	if err != nil {
		return nil, err
	}

	pair, _, err := server.kv.Get(path+"/attribute="+attr, nil)
	if err != nil {
		return nil, err
	}

	if pair == nil {
		return nil, nil
	}

	return parseValue(pair.Value)
}

func (server *Server) objectExists(dn string) (bool, error) {
	prefix, err := dnToConsul(dn)
	if err != nil {
		return false, err
	}

	data, _, err := server.kv.List(prefix+"/", nil)
	if err != nil {
		return false, err
	}
	return len(data) > 0, nil
}

func (server *Server) checkSuffix(dn string, allow_extend bool) (string, error) {
	suffix := server.config.Suffix
	if len(dn) < len(suffix) {
		if dn != suffix[-len(dn):] || !allow_extend {
			return suffix, fmt.Errorf(
				"Only handling stuff under DN %s", suffix)
		}
		return suffix, nil
	} else {
		if dn[len(dn)-len(suffix):] != suffix {
			return suffix, fmt.Errorf(
				"Only handling stuff under DN %s", suffix)
		}
		return dn, nil
	}
}

func (server *Server) handleStartTLS(s ldap.UserState, w ldap.ResponseWriter, m *ldap.Message) {
	tlsConn := tls.Server(m.Client.GetConn(), server.config.TLSConfig)
	res := ldap.NewExtendedResponse(ldap.LDAPResultSuccess)
	res.SetResponseName(ldap.NoticeOfStartTLS)
	w.Write(res)

	if err := tlsConn.Handshake(); err != nil {
		log.Printf("StartTLS Handshake error %v", err)
		res.SetDiagnosticMessage(fmt.Sprintf("StartTLS Handshake error : \"%s\"", err.Error()))
		res.SetResultCode(ldap.LDAPResultOperationsError)
		w.Write(res)
		return
	}

	m.Client.SetConn(tlsConn)
}

func (server *Server) handleBind(s ldap.UserState, w ldap.ResponseWriter, m *ldap.Message) {
	state := s.(*State)
	r := m.GetBindRequest()

	result_code, err := server.handleBindInternal(state, &r)

	res := ldap.NewBindResponse(result_code)
	if err != nil {
		res.SetDiagnosticMessage(err.Error())
		server.logger.Printf("Failed bind for %s: %s", string(r.Name()), err.Error())
	}
	if result_code == ldap.LDAPResultSuccess {
		server.logger.Printf("Successfully bound to %s", string(r.Name()))
	} else {
		server.logger.Printf("Failed to bind to %s (%s)", string(r.Name()), err)
	}
	w.Write(res)
}

func (server *Server) handleBindInternal(state *State, r *message.BindRequest) (int, error) {
	// Check permissions
	if !server.config.Acl.Check(&state.login, "bind", string(r.Name()), []string{}) {
		return ldap.LDAPResultInsufficientAccessRights, nil
	}

	// Try to retrieve password and check for match
	passwd, err := server.getAttribute(string(r.Name()), ATTR_USERPASSWORD)
	if err != nil {
		return ldap.LDAPResultOperationsError, err
	}
	if passwd == nil {
		return ldap.LDAPResultNoSuchObject, nil
	}

	for _, hash := range passwd {
		valid := SSHAMatches(hash, []byte(r.AuthenticationSimple()))
		if valid {
			groups, err := server.getAttribute(string(r.Name()), ATTR_MEMBEROF)
			if err != nil {
				return ldap.LDAPResultOperationsError, err
			}
			state.login = Login{
				user:   string(r.Name()),
				groups: groups,
			}
			return ldap.LDAPResultSuccess, nil
		}
	}
	return ldap.LDAPResultInvalidCredentials, nil
}