package main

import (
	"crypto/sha256"
	"encoding/pem"
	"flag"
	"fmt"
	"os"
	"strings"

	"github.com/tyler-smith/go-bip39"
	"golang.org/x/crypto/ed25519"
	"golang.org/x/crypto/ssh"
)

type KeyType string

const (
	ED25519 KeyType = "ed25519"
)

type seededRand struct {
	seed []byte
	pos  int
}

func (r *seededRand) Read(p []byte) (n int, err error) {
	for n < len(p) {
		counter := uint32(r.pos / len(r.seed))
		hash := sha256.Sum256(append(r.seed, byte(counter), byte(counter>>8), byte(counter>>16), byte(counter>>24)))
		n += copy(p[n:], hash[:])
		r.pos += len(hash)
	}
	return len(p), nil
}

func GenerateEd25519Key(seed []byte) (ed25519.PublicKey, ed25519.PrivateKey, error) {
	if len(seed) < 32 {
		return nil, nil, fmt.Errorf("[err] seed too short for ed25519; need 32 bytes, got %d", len(seed))
	}
	publicKey, privateKey, err := ed25519.GenerateKey(&seededRand{seed: seed[:32]})
	return publicKey, privateKey, err
}

func SaveKeys(privateKey ed25519.PrivateKey, privFile, pubFile string) error {
	block, err := ssh.MarshalPrivateKey(privateKey, "")
	if err != nil {
		return fmt.Errorf("[err] failed to marshal ed25519 private key: %v", err)
	}
	privBytes := pem.EncodeToMemory(block)

	pub, err := ssh.NewPublicKey(privateKey.Public())
	if err != nil {
		return fmt.Errorf("[err] failed to generate ed25519 public key: %v", err)
	}
	pubBytes := ssh.MarshalAuthorizedKey(pub)

	if err := os.WriteFile(privFile, privBytes, 0600); err != nil {
		return fmt.Errorf("[err] failed to write private key to %s: %v", privFile, err)
	}
	if err := os.WriteFile(pubFile, pubBytes, 0644); err != nil {
		return fmt.Errorf("[err] failed to write public key to %s: %v", pubFile, err)
	}
	return nil
}

func init() {
	const usageHeader = `
command-line tool that generates deterministic ed25519 ssh key pairs from a 24-word bip-39 mnemonic phrase

author: heqnx - https://heqnx.com

`
	flag.Usage = func() {
		fmt.Fprint(os.Stderr, usageHeader)
		fmt.Fprintf(os.Stderr, "usage of %s:\n", os.Args[0])
		flag.PrintDefaults()
	}
	flag.CommandLine.SetOutput(os.Stderr)
}

func main() {
	mnemonic := flag.String("mnemonic", "", "bip-39 mnemonic phrase (leave empty to generate a new 24-word mnemonic)")
	outputFile := flag.String("f", "", "output file for private key (public key will be <file>.pub)")
	flag.Parse()

	if flag.NFlag() == 0 && flag.NArg() == 0 {
		flag.Usage()
		os.Exit(1)
	}

	privFile := *outputFile
	pubFile := *outputFile + ".pub"

	var seed []byte
	if *mnemonic == "" {
		entropy, err := bip39.NewEntropy(256)
		if err != nil {
			fmt.Errorf("[err] error generating entropy: %v\n", err)
			os.Exit(1)
		}
		*mnemonic, err = bip39.NewMnemonic(entropy)
		if err != nil {
			fmt.Errorf("[err] error generating mnemonic: %v\n", err)
			os.Exit(1)
		}
		fmt.Printf("[inf] new mnemonic generated - keep it secure:\n%s\n\n", *mnemonic)
		seed = bip39.NewSeed(*mnemonic, "")
	} else {
		wordCount := len(strings.Fields(*mnemonic))
		if wordCount != 24 {
			fmt.Errorf("[err] mnemonic must contain exactly 24 words, got %d\n", wordCount)
			os.Exit(1)
		}
		if !bip39.IsMnemonicValid(*mnemonic) {
			fmt.Errorf("[err] invalid mnemonic phrase\n")
			os.Exit(1)
		}
		seed = bip39.NewSeed(*mnemonic, "")
	}

	_, priv, err := GenerateEd25519Key(seed)
	if err != nil {
		fmt.Errorf("[err] error generating ed25519 key: %v\n", err)
		os.Exit(1)
	}

	err = SaveKeys(priv, privFile, pubFile)
	if err != nil {
		fmt.Errorf("[err] error saving keys: %v\n", err)
		os.Exit(1)
	}
	fmt.Printf("[inf] ed25519 generated keys:\n  - %s (private)\n  - %s (public)\n", privFile, pubFile)
}