From e2f831fd6a270fc85b6ccbb18f312255bf96e883 Mon Sep 17 00:00:00 2001
From: Alejandro Mery <amery@jpi.io>
Date: Wed, 23 Aug 2023 00:18:31 +0000
Subject: [PATCH] wireguard: introduce NewKeyPair, NewPrivateKey, and
 PrivateKey.Public()

Signed-off-by: Alejandro Mery <amery@jpi.io>
---
 go.mod                |  1 +
 go.sum                |  2 ++
 pkg/wireguard/keys.go | 46 +++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 49 insertions(+)

diff --git a/go.mod b/go.mod
index e1342ef..49ad042 100644
--- a/go.mod
+++ b/go.mod
@@ -10,6 +10,7 @@ require (
 	github.com/burntSushi/toml v0.3.1
 	github.com/mgechev/revive v1.3.2
 	github.com/spf13/cobra v1.7.0
+	golang.org/x/crypto v0.12.0
 	gopkg.in/gcfg.v1 v1.2.3
 )
 
diff --git a/go.sum b/go.sum
index d82aadb..8808128 100644
--- a/go.sum
+++ b/go.sum
@@ -70,6 +70,8 @@ github.com/stretchr/testify v1.7.1/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/
 github.com/stretchr/testify v1.8.0/go.mod h1:yNjHg4UonilssWZ8iaSj1OCr/vHnekPRkoO+kdMU+MU=
 github.com/stretchr/testify v1.8.2/go.mod h1:w2LPCIKwWwSfY2zedu0+kehJoqGctiVI29o6fzry7u4=
 github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
+golang.org/x/crypto v0.12.0 h1:tFM/ta59kqch6LlvYnPa0yx5a83cL2nHflFhYKvv9Yk=
+golang.org/x/crypto v0.12.0/go.mod h1:NF0Gs7EO5K4qLn+Ylc+fih8BSTeIjAP05siRnAh98yw=
 golang.org/x/exp v0.0.0-20230713183714-613f0c0eb8a1 h1:MGwJjxBy0HJshjDNfLsYO8xppfqWlA5ZT9OhtUUhTNw=
 golang.org/x/mod v0.12.0 h1:rmsUpXtvNzj340zd98LZ4KntptpfRHwpFOHG188oHXc=
 golang.org/x/mod v0.12.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
diff --git a/pkg/wireguard/keys.go b/pkg/wireguard/keys.go
index c6acea0..100eb02 100644
--- a/pkg/wireguard/keys.go
+++ b/pkg/wireguard/keys.go
@@ -2,8 +2,11 @@ package wireguard
 
 import (
 	"bytes"
+	"crypto/rand"
 	"encoding/base64"
 	"errors"
+
+	"golang.org/x/crypto/curve25519"
 )
 
 const (
@@ -89,8 +92,51 @@ func decodeKey(data string, size int) ([]byte, error) {
 	}
 }
 
+// NewPrivateKey creates a new PrivateKey
+func NewPrivateKey() (PrivateKey, error) {
+	var s [PrivateKeySize]byte
+
+	_, err := rand.Read(s[:])
+	if err != nil {
+		return []byte{}, err
+	}
+
+	// apply same clamping as wireguard-go/device/noise-helpers.go
+	s[0] &= 0xf8
+	s[31] = (s[31] & 0x7f) | 0x40
+
+	return s[:], nil
+}
+
+// Public generates the corresponding PublicKey
+func (key PrivateKey) Public() PublicKey {
+	if len(key) != PrivateKeySize {
+		return []byte{}
+	}
+
+	out := [PublicKeySize]byte{}
+	in := (*[PrivateKeySize]byte)(key)
+
+	curve25519.ScalarBaseMult(&out, in)
+	return out[:]
+}
+
 // KeyPair holds a Key pair
 type KeyPair struct {
 	PrivateKey PrivateKey
 	PublicKey  PublicKey
 }
+
+// NewKeyPair creates a new KeyPair for Wireguard
+func NewKeyPair() (*KeyPair, error) {
+	key, err := NewPrivateKey()
+	if err != nil {
+		return nil, err
+	}
+
+	out := &KeyPair{
+		PrivateKey: key,
+		PublicKey:  key.Public(),
+	}
+	return out, nil
+}