package main

import (
	"bytes"
	"crypto"
	"crypto/ed25519"
	"encoding/pem"
	"fmt"
	"io"
	"os"
	"os/exec"
	"strconv"
	"strings"
	"time"

	"golang.org/x/crypto/ssh"
)

var sshHostKeys = KVSecrets[[]SSHKeyPair]{"hosts/ssh-host-keys"}

type SSHKeyPair struct {
	Type    string
	Public  string
	Private string
}

func getSSHKeyPairs(host string) (pairs []SSHKeyPair, err error) {
	pairs, _, err = sshHostKeys.Get(host)

	didGenerate := false

genLoop:
	for _, keyType := range []string{
		"rsa",
		"dsa",
		"ecdsa",
		"ed25519",
	} {
		for _, pair := range pairs {
			if pair.Type == keyType {
				continue genLoop
			}
		}

		err = func() (err error) {
			outFile, err := os.CreateTemp("/tmp", "dls-key.")
			if err != nil {
				return
			}

			outPath := outFile.Name()

			removeTemp := func() {
				os.Remove(outPath)
				os.Remove(outPath + ".pub")
			}

			removeTemp()
			defer removeTemp()

			var out, privKey, pubKey []byte

			cmd := exec.Command("ssh-keygen",
				"-N", "",
				"-C", "root@"+host,
				"-f", outPath,
				"-t", keyType)
			out, err = cmd.CombinedOutput()
			if err != nil {
				err = fmt.Errorf("ssh-keygen failed: %v: %s", err, string(out))
				return
			}

			privKey, err = os.ReadFile(outPath)
			if err != nil {
				return
			}

			pubKey, err = os.ReadFile(outPath + ".pub")
			if err != nil {
				return
			}

			pairs = append(pairs, SSHKeyPair{
				Type:    keyType,
				Public:  string(pubKey),
				Private: string(privKey),
			})
			didGenerate = true

			return
		}()

		if err != nil {
			return
		}
	}

	if didGenerate {
		err = sshHostKeys.Put(host, pairs)
		if err != nil {
			return
		}
	}

	return
}

var sshCAKeys = KVSecrets[string]{"ssh-ca-keys"}

func sshCAKey(cluster string) (caKeyPem string, err error) {
	storeKey := "clusters/" + cluster
	caKeyPem, _, err = sshCAKeys.Get(storeKey)
	if err != nil {
		return
	}

	if caKeyPem == "" {
		_, pk, err := ed25519.GenerateKey(nil)
		if err != nil {
			return "", err
		}

		pemBlock, err := ssh.MarshalPrivateKey(crypto.PrivateKey(pk), "")
		if err != nil {
			return "", err
		}

		caKeyPem = string(pem.EncodeToMemory(pemBlock))
		sshCAKeys.Put(storeKey, caKeyPem)
	}

	return
}

func sshCAPubKey(cluster string) (pubKey []byte, err error) {
	keyPem, err := sshCAKey(cluster)
	if err != nil {
		return
	}

	k, err := ssh.ParsePrivateKey([]byte(keyPem))
	if err != nil {
		return
	}

	pubKey = ssh.MarshalAuthorizedKey(k.PublicKey())
	return
}

// principal: user (login) to allow (ie: "root")
// validity: ssh-keygen validity string (ie: "+1h", "202506280811:202506281011", ""=forever)
// options:  ssh-keygen options (ie: "force-command=/bin/date +\"%F %T\"", "source-address=192.168.1.0/24,192.168.42.0/24"
func sshCASign(cluster string, userPubKey []byte, principal, validity string, options ...string) (cert []byte, err error) {
	caKey, err := sshCAKey(cluster)
	if err != nil {
		return
	}

	_, identity, _, _, err := ssh.ParseAuthorizedKey(userPubKey)
	if err != nil {
		return
	}

	userPubKeyFile, err := os.CreateTemp("/tmp", "user.pub")
	if err != nil {
		return
	}
	defer os.Remove(userPubKeyFile.Name())

	_, err = io.Copy(userPubKeyFile, bytes.NewBuffer(userPubKey))
	userPubKeyFile.Close()
	if err != nil {
		return
	}

	err = os.WriteFile(userPubKeyFile.Name(), userPubKey, 0600)
	if err != nil {
		return
	}

	serial := strconv.FormatInt(time.Now().Unix(), 10)
	cmd := exec.Command("ssh-keygen", "-q", "-s", "/dev/stdin", "-I", identity, "-z", serial, "-n", principal)

	if validity != "" {
		cmd.Args = append(cmd.Args, "-V", validity)
	}

	for _, opt := range options {
		cmd.Args = append(cmd.Args, "-O", opt)
	}

	cmd.Args = append(cmd.Args, userPubKeyFile.Name())

	stderr := new(bytes.Buffer)
	cmd.Stdin = bytes.NewBuffer([]byte(caKey))
	cmd.Stderr = stderr

	err = cmd.Run()
	if err != nil {
		err = fmt.Errorf("ssh-keygen sign failed: %s", strings.TrimSpace(stderr.String()))
		return
	}

	certFile := userPubKeyFile.Name() + "-cert.pub"
	cert, err = os.ReadFile(certFile)

	os.Remove(certFile)

	return
}