package main

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

	ldap "bottin/ldapserver"

	consul "github.com/hashicorp/consul/api"
	message "github.com/lor00x/goldap/message"
	log "github.com/sirupsen/logrus"
)

// System managed attributes (cannot be changed by user, see checkRestrictedAttr)
const ATTR_MEMBEROF = "memberof"
const ATTR_ENTRYUUID = "entryuuid"
const ATTR_CREATORSNAME = "creatorsname"
const ATTR_CREATETIMESTAMP = "createtimestamp"
const ATTR_MODIFIERSNAME = "modifiersname"
const ATTR_MODIFYTIMESTAMP = "modifytimestamp"

// Attributes that we are interested in at various points
const ATTR_OBJECTCLASS = "objectclass"
const ATTR_MEMBER = "member"
const ATTR_USERPASSWORD = "userpassword"

type ConfigFile struct {
	Suffix        string   `json:"suffix"`
	Bind          string   `json:"bind"`
	BindSecure    string   `json:"bind_secure"`
	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"`
	LogLevel      string   `json:"log_level"`
}

type Config struct {
	Suffix     string
	Bind       string
	BindSecure string
	ConsulHost string
	LogLevel   log.Level

	Acl ACL

	TLSConfig *tls.Config
}

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")
var resyncFlag = flag.Bool("resync", false, "Check and re-synchronize memberOf values before launch")

func readConfig(logger *log.Logger) Config {
	config_file := ConfigFile{
		Bind:       "0.0.0.0:389",
		BindSecure: "0.0.0.0:636",
	}

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

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

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

	log_level := log.InfoLevel
	if config_file.LogLevel != "" {
		log_level, err = log.ParseLevel(config_file.LogLevel)
		if err != nil {
			logger.Fatal(err)
		}
	}

	ret := Config{
		Suffix:     config_file.Suffix,
		Bind:       config_file.Bind,
		BindSecure: config_file.BindSecure,
		ConsulHost: config_file.ConsulHost,
		Acl:        acl,
		LogLevel:   log_level,
	}

	if config_file.TLSCertFile != "" && config_file.TLSKeyFile != "" && config_file.TLSServerName != "" {
		cert_txt, err := ioutil.ReadFile(config_file.TLSCertFile)
		if err != nil {
			logger.Fatal(err)
		}
		key_txt, err := ioutil.ReadFile(config_file.TLSKeyFile)
		if err != nil {
			logger.Fatal(err)
		}
		cert, err := tls.X509KeyPair(cert_txt, key_txt)
		if err != nil {
			logger.Fatal(err)
		}
		ret.TLSConfig = &tls.Config{
			MinVersion:   tls.VersionTLS10,
			MaxVersion:   tls.VersionTLS12,
			Certificates: []tls.Certificate{cert},
			ServerName:   config_file.TLSServerName,
		}
	}

	return ret
}

func main() {
	flag.Parse()

	logger := log.New()
	logger.SetOutput(os.Stdout)
	logger.SetFormatter(&log.TextFormatter{})

	config := readConfig(logger)

	if log_level := os.Getenv("BOTTIN_LOG_LEVEL"); log_level != "" {
		level, err := log.ParseLevel(log_level)
		if err != nil {
			logger.Fatal(err)
		}
		logger.SetLevel(level)
	} else {
		logger.SetLevel(config.LogLevel)
	}

	ldap.Logger = logger

	// 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 {
		logger.Fatal(err)
	}
	kv := consul_client.KV()

	// Create bottin server
	bottin := Server{
		logger: logger,
		config: config,
		kv:     kv,
	}
	err = bottin.init()
	if err != nil {
		logger.Fatal(err)
	}

	if *resyncFlag {
		err = bottin.memberOfResync()
		if err != nil {
			logger.Fatal(err)
		}
	}

	// Create routes
	routes := ldap.NewRouteMux()

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

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

	// Create LDAP servers
	var ldapServer, ldapServerSecure *ldap.Server = nil, nil

	// Bind on standard LDAP port without TLS
	if config.Bind != "" {
		ldapServer = ldap.NewServer()
		ldapServer.Handle(routes)
		ldapServer.NewUserState = bottin.newUserState
		go func() {
			err := ldapServer.ListenAndServe(config.Bind)
			if err != nil {
				logger.Fatal(err)
			}
		}()
	}

	// Bind on LDAP secure port with TLS
	if config.BindSecure != "" {
		if config.TLSConfig != nil {
			ldapServerSecure := ldap.NewServer()
			ldapServerSecure.Handle(routes)
			ldapServerSecure.NewUserState = bottin.newUserState
			secureConn := func(s *ldap.Server) {
				s.Listener = tls.NewListener(s.Listener, config.TLSConfig)
			}
			go func() {
				err := ldapServerSecure.ListenAndServe(config.BindSecure, secureConn)
				if err != nil {
					logger.Fatal(err)
				}
			}()
		} else {
			logger.Warnf("Warning: no valid TLS configuration was provided, not binding on %s", config.BindSecure)
		}
	}

	if ldapServer == nil && ldapServerSecure == nil {
		logger.Fatal("Not doing anything.")
	}

	// 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)

	if ldapServer != nil {
		ldapServer.Stop()
	}
	if ldapServerSecure != nil {
		ldapServerSecure.Stop()
	}
}

func (server *Server) newUserState() ldap.UserState {
	return &State{
		login: Login{
			user:   "ANONYMOUS",
			groups: []string{},
		},
	}
}

func (server *Server) init() error {
	// Check that suffix is in canonical format in config file
	suffix_canonical, err := server.checkDN(server.config.Suffix, false)
	if err != nil {
		return err
	}
	if suffix_canonical != server.config.Suffix {
		return fmt.Errorf("Please write suffix in canonical format: %s", suffix_canonical)
	}

	// Check that root object exists.
	// If it does, we're done. Otherwise, we have some initialization to do.
	exists, err := server.objectExists(server.config.Suffix)
	if err != nil {
		return err
	}
	if exists {
		return nil
	}

	// We have to initialize the server.
	// Create a root object and an admin object.
	base_attributes := Entry{
		ATTR_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.putAttributes(server.config.Suffix, base_attributes)
	if err != nil {
		return err
	}

	admin_pass := make([]byte, 8)
	_, err = rand.Read(admin_pass)
	if err != nil {
		return err
	}
	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{
		ATTR_OBJECTCLASS:        []string{"simpleSecurityObject", "organizationalRole"},
		"displayname":           []string{"LDAP administrator"},
		"description":           []string{"Administrator account automatically created by Bottin"},
		"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.putAttributes(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 recommend replacing it as soon as possible.",
		admin_dn,
		admin_pass_str,
	)

	return nil
}

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

	for k, valuesNC := range attrs {
		path := prefix + "/attribute=" + k

		// Trim spaces and remove empty values
		values := []string{}
		for _, v := range valuesNC {
			vv := strings.TrimSpace(v)
			if len(vv) > 0 {
				values = append(values, vv)
			}
		}

		// If we have zero values, delete associated k/v pair
		// Otherwise, write new values
		if len(values) == 0 {
			_, err := server.kv.Delete(path, nil)
			if err != nil {
				return err
			}
		} else {
			json, err := json.MarshalIndent(values, "", "  ")
			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
	}

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

	values := []string{}
	for _, pair := range pairs {
		if strings.EqualFold(pair.Key, path+"/attribute="+attr) {
			newVals, err := parseValue(pair.Value)
			if err != nil {
				return nil, err
			}
			values = append(values, newVals...)
		}
	}

	return values, nil
}

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

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

func (server *Server) checkDN(dn string, allow_extend bool) (string, error) {
	// 1. Canonicalize: remove spaces between things and put all in lower case
	dn, err := canonicalDN(dn)
	if err != nil {
		return "", err
	}

	// 2. Check suffix (add it if allow_extend is set)
	suffix := server.config.Suffix
	if len(dn) < len(suffix) {
		if dn != suffix[len(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 {
		server.logger.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())
	}
	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, fmt.Errorf("Insufficient access rights for %#v", state.login)
	}

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

	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, fmt.Errorf("No password match")
}