From 2da476eef426e7bbb42915aa2178ebb3cc5a553c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E3=83=81=E3=82=BB?=
 <123655015+chise0713@users.noreply.github.com>
Date: Tue, 12 Dec 2023 09:52:10 +0800
Subject: [PATCH] Command: Add `wg` for wireguard key generation (#2794)

* Command: Add `wg` for wireguard key generation

* Command: Merge `x25519` and `wg`
---
 main/commands/all/commands.go   |  1 +
 main/commands/all/curve25519.go | 57 +++++++++++++++++++++++++++++++++
 main/commands/all/wg.go         | 27 ++++++++++++++++
 main/commands/all/x25519.go     | 55 ++-----------------------------
 4 files changed, 87 insertions(+), 53 deletions(-)
 create mode 100644 main/commands/all/curve25519.go
 create mode 100644 main/commands/all/wg.go

diff --git a/main/commands/all/commands.go b/main/commands/all/commands.go
index 9b8b49e0..41d0e0f1 100644
--- a/main/commands/all/commands.go
+++ b/main/commands/all/commands.go
@@ -16,5 +16,6 @@ func init() {
 		tls.CmdTLS,
 		cmdUUID,
 		cmdX25519,
+		cmdWG,
 	)
 }
diff --git a/main/commands/all/curve25519.go b/main/commands/all/curve25519.go
new file mode 100644
index 00000000..25cc812e
--- /dev/null
+++ b/main/commands/all/curve25519.go
@@ -0,0 +1,57 @@
+package all
+
+import (
+	"crypto/rand"
+	"encoding/base64"
+	"fmt"
+
+	"golang.org/x/crypto/curve25519"
+)
+
+func Curve25519Genkey(StdEncoding bool, input_base64 string) {
+	var output string
+	var err error
+	var privateKey, publicKey []byte
+	var encoding *base64.Encoding
+	if *input_stdEncoding || StdEncoding {
+		encoding = base64.StdEncoding
+	} else {
+		encoding = base64.RawURLEncoding
+	}
+
+	if len(input_base64) > 0 {
+		privateKey, err = encoding.DecodeString(input_base64)
+		if err != nil {
+			output = err.Error()
+			goto out
+		}
+		if len(privateKey) != curve25519.ScalarSize {
+			output = "Invalid length of private key."
+			goto out
+		}
+	}
+
+	if privateKey == nil {
+		privateKey = make([]byte, curve25519.ScalarSize)
+		if _, err = rand.Read(privateKey); err != nil {
+			output = err.Error()
+			goto out
+		}
+	}
+
+	// Modify random bytes using algorithm described at:
+	// https://cr.yp.to/ecdh.html.
+	privateKey[0] &= 248
+	privateKey[31] &= 127 | 64
+
+	if publicKey, err = curve25519.X25519(privateKey, curve25519.Basepoint); err != nil {
+		output = err.Error()
+		goto out
+	}
+
+	output = fmt.Sprintf("Private key: %v\nPublic key: %v",
+		encoding.EncodeToString(privateKey),
+		encoding.EncodeToString(publicKey))
+out:
+	fmt.Println(output)
+}
diff --git a/main/commands/all/wg.go b/main/commands/all/wg.go
new file mode 100644
index 00000000..70da4668
--- /dev/null
+++ b/main/commands/all/wg.go
@@ -0,0 +1,27 @@
+package all
+
+import (
+	"github.com/xtls/xray-core/main/commands/base"
+)
+
+var cmdWG = &base.Command{
+	UsageLine: `{{.Exec}} wg [-i "private key (base64.StdEncoding)"]`,
+	Short:     `Generate key pair for wireguard key exchange`,
+	Long: `
+Generate key pair for wireguard key exchange.
+
+Random: {{.Exec}} wg
+
+From private key: {{.Exec}} wg -i "private key (base64.StdEncoding)"
+`,
+}
+
+func init() {
+	cmdWG.Run = executeWG // break init loop
+}
+
+var input_wireguard = cmdWG.Flag.String("i", "", "")
+
+func executeWG(cmd *base.Command, args []string) {
+	Curve25519Genkey(true, *input_wireguard)
+}
diff --git a/main/commands/all/x25519.go b/main/commands/all/x25519.go
index 814cca72..73f669b2 100644
--- a/main/commands/all/x25519.go
+++ b/main/commands/all/x25519.go
@@ -1,12 +1,7 @@
 package all
 
 import (
-	"crypto/rand"
-	"encoding/base64"
-	"fmt"
-
 	"github.com/xtls/xray-core/main/commands/base"
-	"golang.org/x/crypto/curve25519"
 )
 
 var cmdX25519 = &base.Command{
@@ -26,55 +21,9 @@ func init() {
 	cmdX25519.Run = executeX25519 // break init loop
 }
 
-var input_base64 = cmdX25519.Flag.String("i", "", "")
 var input_stdEncoding = cmdX25519.Flag.Bool("std-encoding", false, "")
+var input_x25519 = cmdX25519.Flag.String("i", "", "")
 
 func executeX25519(cmd *base.Command, args []string) {
-	var output string
-	var err error
-	var privateKey []byte
-	var publicKey []byte
-	var encoding *base64.Encoding
-	if len(*input_base64) > 0 {
-		privateKey, err = base64.RawURLEncoding.DecodeString(*input_base64)
-		if err != nil {
-			output = err.Error()
-			goto out
-		}
-		if len(privateKey) != curve25519.ScalarSize {
-			output = "Invalid length of private key."
-			goto out
-		}
-	}
-
-	if privateKey == nil {
-		privateKey = make([]byte, curve25519.ScalarSize)
-		if _, err = rand.Read(privateKey); err != nil {
-			output = err.Error()
-			goto out
-		}
-	}
-
-	// Modify random bytes using algorithm described at:
-	// https://cr.yp.to/ecdh.html.
-	privateKey[0] &= 248
-	privateKey[31] &= 127
-	privateKey[31] |= 64
-
-	if publicKey, err = curve25519.X25519(privateKey, curve25519.Basepoint); err != nil {
-		output = err.Error()
-		goto out
-	}
-
-	if *input_stdEncoding {
-		encoding = base64.StdEncoding
-	} else {
-		encoding = base64.RawURLEncoding
-	}
-
-	output = fmt.Sprintf("Private key: %v\nPublic key: %v",
-		encoding.EncodeToString(privateKey),
-		encoding.EncodeToString(publicKey))
-out:
-	fmt.Println(output)
+	Curve25519Genkey(false, *input_x25519)
 }