package main

import (
	"bufio"
	"converge/pkg/agent/session"
	"fmt"
	"github.com/fsnotify/fsnotify"
	"github.com/gliderlabs/ssh"
	gossh "golang.org/x/crypto/ssh"
	"log"
	"os"
	"path/filepath"
	"strings"
	"sync"
)

func publicKeyHandler(ctx ssh.Context, key gossh.PublicKey, authorizedKey gossh.PublicKey) bool {
	providedKey := gossh.MarshalAuthorizedKey(key)

	if ssh.KeysEqual(key, authorizedKey) {
		log.Printf("Successful login from %s", ctx.RemoteAddr())
		return true
	}

	log.Printf("Failed login attempt from %s with key: %s", ctx.RemoteAddr(), strings.TrimSpace(string(providedKey)))
	return false
}

func readSshPublicKeys(fileName string) ([]ssh.PublicKey, []string, error) {
	file, err := os.Open(fileName)
	if err != nil {
		return nil, nil, fmt.Errorf("Failed to open file: '%s': %s", fileName, err)
	}
	defer file.Close()

	res := make([]ssh.PublicKey, 0)
	scanner := bufio.NewScanner(file)
	var errorKeys []string
	for scanner.Scan() {
		lineText := scanner.Text()
		ind := strings.Index(lineText, "#")
		if ind >= 0 {
			lineText = lineText[:ind]
		}
		log.Println("Reading public key " + lineText)
		lineText = strings.Trim(lineText, "")
		if lineText == "" {
			continue
		}
		line := []byte(lineText)
		parsedKey, _, _, _, err := ssh.ParseAuthorizedKey(line)
		if err != nil {
			errorKeys = append(errorKeys, lineText)
			log.Printf("Failed to parse authorized key: %v", lineText)
		} else {
			res = append(res, parsedKey)
		}
	}
	return res, errorKeys, nil
}

type AuthorizedPublicKeys struct {
	authorizedKeysFile string
	mutex              sync.Mutex
	publicKeys         []ssh.PublicKey
}

func NewAuthorizedPublicKeys(authorizedKeysFile string) (*AuthorizedPublicKeys, error) {
	pubkeys := AuthorizedPublicKeys{
		authorizedKeysFile: authorizedKeysFile,
		mutex:              sync.Mutex{},
		publicKeys:         nil,
	}
	pubkeys.publicKeys, _ = pubkeys.Parse()
	if len(pubkeys.publicKeys) == 0 {
		return nil, fmt.Errorf("No valid public keys found, login is impossible, exiting")
	}
	go pubkeys.monitorAuthorizedKeysFile(authorizedKeysFile)
	return &pubkeys, nil
}

func (pubkeys *AuthorizedPublicKeys) setPubKeys(keys []ssh.PublicKey) {
	pubkeys.mutex.Lock()
	defer pubkeys.mutex.Unlock()
	pubkeys.publicKeys = keys
}

func (pubkeys *AuthorizedPublicKeys) getPubKeys() []ssh.PublicKey {
	pubkeys.mutex.Lock()
	defer pubkeys.mutex.Unlock()
	return pubkeys.publicKeys
}

func (pubkeys *AuthorizedPublicKeys) monitorAuthorizedKeysFile(authorizedPublicKeysFile string) {
	dir := filepath.Dir(authorizedPublicKeysFile)
	watcher, err := fsnotify.NewWatcher()
	if err != nil {
		log.Printf("Cannot watch authorized keys file '%s', user notifications about changes will not work: %v",
			authorizedPublicKeysFile, err)
	}
	defer watcher.Close()
	log.Println("XXX: monitor " + dir)
	err = watcher.Add(dir)
	if err != nil {
		log.Printf("Cannot watch hold file %s, user notifications for change in expiry time will be unavailable: %v", authorizedPublicKeysFile, err)
	}
	for {
		select {
		case event, ok := <-watcher.Events:
			if !ok {
				session.MessageUsers(
					fmt.Sprintf("Watching authorized keys file '%s' stopped", authorizedPublicKeysFile))
				return
			}
			base := filepath.Base(event.Name)
			log.Println("CHANGE " + base + " " + authorizedPublicKeysFile + " " + filepath.Base(authorizedPublicKeysFile))
			if base == filepath.Base(authorizedPublicKeysFile) {
				keys, errorKeys := pubkeys.Parse()
				for _, errorKey := range errorKeys {
					session.MessageUsers(fmt.Sprintf("Public key '%s' is invalid", errorKey))
				}

				if len(keys) == 0 {
					session.MessageUsers(
						fmt.Sprintf("Authorized keys file '%s' does not exist or does not contain any valid keys, using last known configuration", authorizedPublicKeysFile))
				} else {
					session.MessageUsers(fmt.Sprintf("Updated authorized keys, now %d valid key(s)", len(keys)))
					pubkeys.setPubKeys(keys)
				}
			}

		case err, ok := <-watcher.Errors:
			if ok {
				session.MessageUsers(fmt.Sprintf(
					"Watching authorized keys file '%s' stopped",
					pubkeys.authorizedKeysFile))
			}
			session.MessageUsers(fmt.Sprintf(
				"Watching authorized keys file '%s' stopped: %v",
				pubkeys.authorizedKeysFile, err))
			return
		}
	}
}

func (pubkeys *AuthorizedPublicKeys) Parse() ([]ssh.PublicKey, []string) {
	if pubkeys.authorizedKeysFile == "" {
		return nil, nil
	}
	keys, errorKeys, err := readSshPublicKeys(pubkeys.authorizedKeysFile)
	if os.IsNotExist(err) {
		log.Printf("Authorized keys file '%s' not found.", pubkeys.authorizedKeysFile)
		return nil, nil
	}
	if err != nil {
		log.Println("Public key authentication will not work since no public keys were found.")
	}
	return keys, errorKeys
}

func (key *AuthorizedPublicKeys) authorize(ctx ssh.Context, userProvidedKey ssh.PublicKey) bool {
	//
	keys, _ := key.Parse()
	if len(keys) == 0 {
		keys = key.getPubKeys()
	}
	if len(keys) > 0 {
		log.Printf("A total of %d authorized ssh keys were found", len(keys))
	} else {
		log.Printf("No valid public keys were found, login is impossible")
		// keep agent running, a user may still be logged in and could change the
		// authorizedk eys file
		// TODO: if no users are logged in, the agent should exit.
	}
	for _, key := range keys {
		if publicKeyHandler(ctx, userProvidedKey, key) {
			return true
		}
	}
	return false
}