package crypto

import (
	"crypto/rand"
	"encoding/base64"
	"fmt"

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

// Expected key length for a WireGuard key.
const wireguardKeyLen = 32

type wireguardKeyPair struct {
	publicKey  string
	privateKey string
}

// Generate an asymmetric key pair using Curve25519, as per Wireguard protocol: https://wireguard.com/protocol.
func GenerateWireguardKeyPair() (AsymmetricKeyPair, error) {
	privateKey, err := generateWireguardPrivateKey()
	if err != nil {
		return nil, err
	}
	publicKey := privateKey.PublicKey()

	return &wireguardKeyPair{
		publicKey:  publicKey.String(),
		privateKey: privateKey.String(),
	}, nil
}

func (wgKeyPair *wireguardKeyPair) PublicKey() string {
	return wgKeyPair.publicKey
}

func (wgKeyPair *wireguardKeyPair) PrivateKey() string {
	return wgKeyPair.privateKey
}

// A wireguardKey is a public, private, or pre-shared secret key. The wireguardKey constructor
// functions in this package can be used to create Keys suitable for each of
// these applications.
type wireguardKey [wireguardKeyLen]byte

// newWireguardKey creates a wireguardKey from an existing byte slice. The byte slice must be
// exactly 32 bytes in length.
func newWireguardKey(bytes []byte) (wireguardKey, error) {
	if len(bytes) != wireguardKeyLen {
		return wireguardKey{}, fmt.Errorf("incorrect key size: %d", len(bytes))
	}

	var wgKey wireguardKey
	copy(wgKey[:], bytes)

	return wgKey, nil
}

// PublicKey computes a public key from the private key k.
// PublicKey should only be called when k is a private key.
func (wgKey wireguardKey) PublicKey() wireguardKey {
	var (
		publicKey  [wireguardKeyLen]byte
		privateKey = [wireguardKeyLen]byte(wgKey)
	)

	// ScalarBaseMult uses the correct base value per https://cr.yp.to/ecdh.html,
	// so no need to specify it.
	curve25519.ScalarBaseMult(&publicKey, &privateKey)

	return wireguardKey(publicKey)
}

// String returns the base64-encoded string representation of a Key.
func (wgKey wireguardKey) String() string {
	return base64.StdEncoding.EncodeToString(wgKey[:])
}

// generateWireguardPrivateKey generates a wireguardKey suitable for use as a private key from a
// cryptographically safe source.
func generateWireguardPrivateKey() (wireguardKey, error) {
	wgKey, err := generateWireguardKey()
	if err != nil {
		return wireguardKey{}, err
	}

	// Modify random bytes using algorithm described at:
	// https://cr.yp.to/ecdh.html.
	wgKey[0] &= 248
	wgKey[31] &= 127
	wgKey[31] |= 64

	return wgKey, nil
}

// generateWireguardKey generates a wireguardKey suitable for use as a pre-shared secret key from
// a cryptographically safe source.
//
// The output wireguardKey should not be used as a private key; use generateWireguardPrivateKey
// instead.
func generateWireguardKey() (wireguardKey, error) {
	bytes := make([]byte, wireguardKeyLen)
	if _, err := rand.Read(bytes); err != nil {
		return wireguardKey{}, fmt.Errorf("failed to read random bytes: %v", err)
	}

	return newWireguardKey(bytes)
}