package main

import (
	"crypto/sha256"
	"crypto/tls"
	"encoding/base64"
	"fmt"
	"gopkg.in/ldap.v2"
	"log"
	"strings"
)

// Search ldap for keys...
type LdapInfo struct {
	Server        string
	Port          int
	User          string
	Password      string
	SSLSkipVerify bool
	UserDNFmt     string
}

func (i *LdapInfo) LdapConnect() (*ldap.Conn, error) {
	var tlsConf *tls.Config
	if i.SSLSkipVerify {
		tlsConf = &tls.Config{InsecureSkipVerify: true}
	} else {
		tlsConf = &tls.Config{ServerName: i.Server}
	}
	l, err := ldap.DialTLS("tcp", fmt.Sprintf("%s:%d", i.Server, i.Port), tlsConf)
	if err != nil {
		return nil, fmt.Errorf("LDAP Unable to connect to %s on port %d. Got error: %s", i.Server, i.Port, err)
	}

	return l, nil
}

func (i *LdapInfo) LdapConnectBind() (*ldap.Conn, error) {
	// check if we have LDAP credentials
	if i.User == "" || i.Password == "" {
		return nil, fmt.Errorf("LDAP Bind user and/or password missing")
	}

	l, err := i.LdapConnect()
	if err != nil {
		return nil, err
	}

	err = l.Bind(i.User, i.Password)
	if err != nil {
		l.Close()
		return nil, err
	}
	return l, nil
}

func (i *LdapInfo) UserDN(username string) string {
	if i.UserDNFmt != "" {
		return fmt.Sprintf(i.UserDNFmt, username)
	} else {
		return fmt.Sprintf("uid=%s,ou=People,dc=nordu,dc=net", username)
	}
}

type SSHPubKey struct {
	Format      string
	Key         string
	Comment     string
	Fingerprint string
}

func NewSSHPubKey(ssh_key string) SSHPubKey {
	key_parts := strings.SplitN(ssh_key, " ", 3)
	comment := ""
	if len(key_parts) > 2 {
		comment = key_parts[2]
	}
	return SSHPubKey{key_parts[0], key_parts[1], comment, calculateFingerprint(key_parts[1])}
}

func (k SSHPubKey) String() string {
	return fmt.Sprintf("%s %s %s", k.Format, k.Key, k.Comment)
}

func (k SSHPubKey) KeyEnd() string {
	i := len(k.Key) - 8
	if i < 0 {
		i = 0
	}
	return k.Key[i:]
}

// Get SSH keys
//    Preferably keyformat, fingerprint, comment, but full key is how it is now
func (i *LdapInfo) GetSSHKeys(username string) ([]SSHPubKey, error) {
	l, err := i.LdapConnect()
	if err != nil {
		return nil, err
	}
	defer l.Close()

	searchRequest := ldap.NewSearchRequest(
		i.UserDN(username),
		ldap.ScopeWholeSubtree, ldap.NeverDerefAliases, 0, 0, false,
		"(objectClass=person)",
		[]string{"sshPublicKey"},
		nil,
	)

	sr, err := l.Search(searchRequest)
	if err != nil {
		return nil, fmt.Errorf("LDAP Search for user '%s' failed: %s", username, err)
	}

	if len(sr.Entries) < 1 {
		return nil, fmt.Errorf("LDAP User %s does not exist", username)
	} else if len(sr.Entries) > 1 {
		return nil, fmt.Errorf("LDAP User %s returned more than one enty (Results: %d)", username, len(sr.Entries))
	}

	pubKeys := make([]SSHPubKey, len(sr.Entries[0].GetAttributeValues("sshPublicKey")))
	for i, key := range sr.Entries[0].GetAttributeValues("sshPublicKey") {
		pubKeys[i] = NewSSHPubKey(key)
	}
	return pubKeys, nil

}

// Add ssh key
func (i *LdapInfo) AddSSHKey(username string, ssh_keys []string) error {
	l, err := i.LdapConnectBind()
	if err != nil {
		return err
	}
	defer l.Close()

	dn := i.UserDN(username)
	// Add objectClass ldapPublicKey if missing
	mod := ldap.NewModifyRequest(dn)
	mod.Add("objectClass", []string{"ldapPublicKey"})
	err = l.Modify(mod)
	if err != nil {
		// check err if type or value exist
	}

	// Add keys
	// One mod req per key, to handle errors for existing keys
	for _, key := range ssh_keys {
		if err = validateSSHkey(key); err != nil {
			log.Println(err)
		} else {
			mod = ldap.NewModifyRequest(dn)
			mod.Add("sshPublicKey", []string{key})
			err = l.Modify(mod)
			if err != nil {
				// Ignore Attribute or value exists errors
				if !strings.Contains(err.Error(), "Code 20") {
					return err
				}
			}
		}
	}

	return nil
}

// Delete ssh key
// 	Use fingerprint, or full key to delete
func (i *LdapInfo) DeleteSSHKey(username, ssh_key string) error {
	l, err := i.LdapConnectBind()
	if err != nil {
		return err
	}
	defer l.Close()

	del := ldap.NewModifyRequest(i.UserDN(username))
	del.Delete("sshPublicKey", []string{ssh_key})
	err = l.Modify(del)
	if err != nil {
		// Ignore error about No such attribute
		if !strings.Contains(err.Error(), "Code 16") {
			return err
		}
	}
	return nil
}

// Sanity checks on a ssh key.
// Checks if key has 2-3 parts (key_format, key, comment)
// Checks if key_format is: ssh-rsa or ssh-ed25519
// 		If ssh-rsa check that key is at least 2048 bit
// Check that key can be base64 decoded
func validateSSHkey(ssh_key string) error {
	key_parts := strings.SplitN(ssh_key, " ", 3)
	if len(key_parts) < 2 || len(key_parts) > 3 {
		return fmt.Errorf("SSH key is invalid. Expected 2-3 parts, got %d. Key was: '%s'", len(key_parts), ssh_key)
	}
	//  Check base64
	decoded_key, err := base64.StdEncoding.DecodeString(key_parts[1])
	if err != nil {
		return fmt.Errorf("SSH key is not properly base64 encoded. Key was: %s", ssh_key)
	}
	// 	Check keyformat: ssh-rsa, ssh-ed25519
	switch key_parts[0] {
	case "ssh-rsa":
		padding := strings.Count(key_parts[1], "=")
		key_length := (len(decoded_key) - padding) * 8
		if key_length < 2048 {
			return fmt.Errorf("SSH rsa key should at least be a 2048 bit key. Was: %d", key_length)
		}
	case "ssh-ed25519":
		// nothing to check
	default:
		return fmt.Errorf("SSH key is not an acceptable format (ssh-rsa or ssh-ed25519). Key was: %s", ssh_key)
	}
	// Key looks ok
	return nil
}

func calculateFingerprint(ssh_key string) string {
	key, _ := base64.StdEncoding.DecodeString(ssh_key)
	fingerprint := sha256.Sum256(key)
	return fmt.Sprintf("SHA256:%s", base64.StdEncoding.EncodeToString(fingerprint[:]))
	//return fmt.Sprintf("SHA256:%x", fingerprint)
}

//// set_nordunet_ldap_pw_sasl used on sso pw set if change pw fail?