package main

import (
	"bufio"
	"fmt"
	"github.com/gliderlabs/ssh"
	gossh "golang.org/x/crypto/ssh"
	"log"
	"os"
	"strings"
)

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, error) {
	file, err := os.Open(fileName)
	if err != nil {
		return nil, fmt.Errorf("Failed to open file: '%s': %s", fileName, err)
	}
	defer file.Close()

	res := make([]ssh.PublicKey, 0)
	scanner := bufio.NewScanner(file)
	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 {
			log.Printf("Failed to parse authorized key: %v", lineText)
		} else {
			res = append(res, parsedKey)
		}
	}
	return res, nil
}

type AuthorizedPublicKeys struct {
	keys []ssh.PublicKey
}

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

func (key AuthorizedPublicKeys) authorize(ctx ssh.Context, userProvidedKey ssh.PublicKey) bool {
	for _, key := range key.keys {
		if publicKeyHandler(ctx, userProvidedKey, key) {
			return true
		}
	}
	return false
}