[feat] allow custom public and private key

go.mod

@@ -2,14 +2,23 @@ module github.com/peanut996/CloudflareWarpSpeedTest
 
 go 1.20
 
-require github.com/cheggaaa/pb/v3 v3.1.4
+require (
+	github.com/cheggaaa/pb/v3 v3.1.4
+	golang.zx2c4.com/wireguard v0.0.0-20230704135630-469159ecf7d1
+)
 
 require (
 	github.com/VividCortex/ewma v1.2.0 // indirect
 	github.com/fatih/color v1.15.0 // indirect
+	github.com/google/btree v1.0.1 // indirect
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-isatty v0.0.19 // indirect
 	github.com/mattn/go-runewidth v0.0.14 // indirect
 	github.com/rivo/uniseg v0.2.0 // indirect
+	golang.org/x/crypto v0.6.0 // indirect
+	golang.org/x/net v0.7.0 // indirect
 	golang.org/x/sys v0.6.0 // indirect
+	golang.org/x/time v0.0.0-20191024005414-555d28b269f0 // indirect
+	golang.zx2c4.com/wintun v0.0.0-20230126152724-0fa3db229ce2 // indirect
+	gvisor.dev/gvisor v0.0.0-20221203005347-703fd9b7fbc0 // indirect
 )

go.sum

@@ -4,6 +4,8 @@ github.com/cheggaaa/pb/v3 v3.1.4 h1:DN8j4TVVdKu3WxVwcRKu0sG00IIU6FewoABZzXbRQeo=
 github.com/cheggaaa/pb/v3 v3.1.4/go.mod h1:6wVjILNBaXMs8c21qRiaUM8BR82erfgau1DQ4iUXmSA=
 github.com/fatih/color v1.15.0 h1:kOqh6YHBtK8aywxGerMG2Eq3H6Qgoqeo13Bk2Mv/nBs=
 github.com/fatih/color v1.15.0/go.mod h1:0h5ZqXfHYED7Bhv2ZJamyIOUej9KtShiJESRwBDUSsw=
+github.com/google/btree v1.0.1 h1:gK4Kx5IaGY9CD5sPJ36FHiBJ6ZXl0kilRiiCj+jdYp4=
+github.com/google/btree v1.0.1/go.mod h1:xXMiIv4Fb/0kKde4SpL7qlzvu5cMJDRkFDxJfI9uaxA=
 github.com/mattn/go-colorable v0.1.13 h1:fFA4WZxdEF4tXPZVKMLwD8oUnCTTo08duU7wxecdEvA=
 github.com/mattn/go-colorable v0.1.13/go.mod h1:7S9/ev0klgBDR4GtXTXX8a3vIGJpMovkB8vQcUbaXHg=
 github.com/mattn/go-isatty v0.0.16/go.mod h1:kYGgaQfpe5nmfYZH+SKPsOc2e4SrIfOl2e/yFXSvRLM=
@@ -13,6 +15,18 @@ github.com/mattn/go-runewidth v0.0.14 h1:+xnbZSEeDbOIg5/mE6JF0w6n9duR1l3/WmbinWV
 github.com/mattn/go-runewidth v0.0.14/go.mod h1:Jdepj2loyihRzMpdS35Xk/zdY8IAYHsh153qUoGf23w=
 github.com/rivo/uniseg v0.2.0 h1:S1pD9weZBuJdFmowNwbpi7BJ8TNftyUImj/0WQi72jY=
 github.com/rivo/uniseg v0.2.0/go.mod h1:J6wj4VEh+S6ZtnVlnTBMWIodfgj8LQOQFoIToxlJtxc=
+golang.org/x/crypto v0.6.0 h1:qfktjS5LUO+fFKeJXZ+ikTRijMmljikvG68fpMMruSc=
+golang.org/x/crypto v0.6.0/go.mod h1:OFC/31mSvZgRz0V1QTNCzfAI1aIRzbiufJtkMIlEp58=
+golang.org/x/net v0.7.0 h1:rJrUqqhjsgNp7KqAIc25s9pZnjU7TUcSY7HcVZjdn1g=
+golang.org/x/net v0.7.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs=
 golang.org/x/sys v0.0.0-20220811171246-fbc7d0a398ab/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.6.0 h1:MVltZSvRTcU2ljQOhs94SXPftV6DCNnZViHeQps87pQ=
 golang.org/x/sys v0.6.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/time v0.0.0-20191024005414-555d28b269f0 h1:/5xXl8Y5W96D+TtHSlonuFqGHIWVuyCkGJLwGh9JJFs=
+golang.org/x/time v0.0.0-20191024005414-555d28b269f0/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
+golang.zx2c4.com/wintun v0.0.0-20230126152724-0fa3db229ce2 h1:B82qJJgjvYKsXS9jeunTOisW56dUokqW/FOteYJJ/yg=
+golang.zx2c4.com/wintun v0.0.0-20230126152724-0fa3db229ce2/go.mod h1:deeaetjYA+DHMHg+sMSMI58GrEteJUUzzw7en6TJQcI=
+golang.zx2c4.com/wireguard v0.0.0-20230704135630-469159ecf7d1 h1:EY138uSo1JYlDq+97u1FtcOUwPpIU6WL1Lkt7WpYjPA=
+golang.zx2c4.com/wireguard v0.0.0-20230704135630-469159ecf7d1/go.mod h1:tqur9LnfstdR9ep2LaJT4lFUl0EjlHtge+gAjmsHUG4=
+gvisor.dev/gvisor v0.0.0-20221203005347-703fd9b7fbc0 h1:Wobr37noukisGxpKo5jAsLREcpj61RxrWYzD8uwveOY=
+gvisor.dev/gvisor v0.0.0-20221203005347-703fd9b7fbc0/go.mod h1:Dn5idtptoW1dIos9U6A2rpebLs/MtTwFacjKb8jLdQA=

main.go

@@ -3,9 +3,10 @@ package main
 import (
 	"flag"
 	"fmt"
+	"time"
+
 	"github.com/peanut996/CloudflareWarpSpeedTest/task"
 	"github.com/peanut996/CloudflareWarpSpeedTest/utils"
-	"time"
 )
 
 func init() {
@@ -38,7 +39,10 @@ CloudflareWarpSpeedTest \n` + `
         指定IP段数据；直接通过参数指定要测速的 IP 段数据，英文逗号分隔；(默认 空)
     -o result.csv
         写入结果文件；如路径含有空格请加上引号；值为空时不写入文件 [-o ""]；(默认 result.csv)
-
+    -pri 私钥
+        指定你的wireguard私钥
+    -pub 私钥
+        指定你的wireguard公钥, 默认为warp的公钥
     -full
         测速全部的端口；对 IP 段中的每个 IP 全部端口进行测速
     -h
@@ -59,6 +63,8 @@ CloudflareWarpSpeedTest \n` + `
 	flag.StringVar(&task.IPFile, "f", "", "IP段数据文件")
 	flag.StringVar(&task.IPText, "ip", "", "指定IP段数据")
 	flag.StringVar(&utils.Output, "o", "result.csv", "输出结果文件")
+	flag.StringVar(&task.PrivateKey, "pri", "", "指定private key")
+	flag.StringVar(&task.PrivateKey, "pub", "", "指定public key")
 
 	flag.BoolVar(&printVersion, "v", false, "打印程序版本")
 	flag.Usage = func() { fmt.Print(help) }
@@ -71,6 +77,7 @@ CloudflareWarpSpeedTest \n` + `
 
 func main() {
 	task.InitRandSeed()
+	task.InitHandshakePacket()
 
 	fmt.Printf("CloudflareWarpSpeedTest\n\n")
 

task/warping.go

@@ -1,15 +1,26 @@
 package task
 
 import (
+	"bytes"
+	"encoding/base64"
+	"encoding/binary"
 	"encoding/hex"
+	"errors"
 	"fmt"
-	"github.com/peanut996/CloudflareWarpSpeedTest/utils"
+	"log"
 	"math/rand"
 	"net"
+	"net/netip"
 	"sort"
 	"strconv"
 	"sync"
 	"time"
+
+	"github.com/peanut996/CloudflareWarpSpeedTest/utils"
+
+	"golang.zx2c4.com/wireguard/conn"
+	"golang.zx2c4.com/wireguard/device"
+	"golang.zx2c4.com/wireguard/tun/netstack"
 )
 
 const (
@@ -18,9 +29,14 @@ const (
 	udpConnectTimeout           = time.Millisecond * 1000
 	wireguardHandshakeRespBytes = 92
 	quickModeMaxIpNum           = 5000
+	warpPublicKey               = "bmXOC+F1FxEMF9dyiK2H5/1SUtzH0JuVo51h2wPfgyo="
 )
 
 var (
+	PrivateKey string
+
+	PublicKey string
+
 	QuickMode = false
 
 	ScanAllPort = false
@@ -229,3 +245,87 @@ func shuffleAddrs(udpAddrs *[]*UDPAddr) {
 		(*udpAddrs)[i], (*udpAddrs)[j] = (*udpAddrs)[j], (*udpAddrs)[i]
 	})
 }
+
+func InitHandshakePacket() {
+	if PrivateKey == "" && PublicKey == "" {
+		return
+	}
+
+	if PublicKey == "" {
+		PublicKey = warpPublicKey
+	}
+
+	pri, err := getNoisePrivateKeyFromBase64(PrivateKey)
+	if err != nil {
+		log.Fatalln("解析私钥失败：" + err.Error())
+	}
+
+	pub, err := getNoisePublicKeyFromBase64(PublicKey)
+	if err != nil {
+		log.Fatalln("解析公钥失败：" + err.Error())
+	}
+
+	packet := buildHandshakePacket(pri, pub)
+
+	warpHandshakePacket = packet[:]
+}
+
+func buildHandshakePacket(pri device.NoisePrivateKey, pub device.NoisePublicKey) []byte {
+	d, _, err := netstack.CreateNetTUN([]netip.Addr{}, []netip.Addr{}, 1480)
+	if err != nil {
+		log.Fatalln("构建握手包失败: " + err.Error())
+	}
+	dev := device.NewDevice(d, conn.NewDefaultBind(), device.NewLogger(0, ""))
+
+	dev.SetPrivateKey(pri)
+
+	peer, err := dev.NewPeer(pub)
+	if err != nil {
+		log.Fatalln("构建握手包失败: " + err.Error())
+	}
+	msg, err := dev.CreateMessageInitiation(peer)
+	if err != nil {
+		log.Fatalln("构建握手包失败: " + err.Error())
+	}
+
+	var buf [device.MessageInitiationSize]byte
+	writer := bytes.NewBuffer(buf[:0])
+	binary.Write(writer, binary.LittleEndian, msg)
+	packet := writer.Bytes()
+
+	generator := device.CookieGenerator{}
+	generator.Init(pub)
+	generator.AddMacs(packet)
+	return packet
+}
+
+func getNoisePrivateKeyFromBase64(b string) (device.NoisePrivateKey, error) {
+	pk := device.NoisePrivateKey{}
+	h, err := encodeBase64ToHex(b)
+	if err != nil {
+		return pk, err
+	}
+	pk.FromHex(h)
+	return pk, nil
+}
+
+func getNoisePublicKeyFromBase64(b string) (device.NoisePublicKey, error) {
+	pk := device.NoisePublicKey{}
+	h, err := encodeBase64ToHex(b)
+	if err != nil {
+		return pk, err
+	}
+	pk.FromHex(h)
+	return pk, nil
+}
+
+func encodeBase64ToHex(key string) (string, error) {
+	decoded, err := base64.StdEncoding.DecodeString(key)
+	if err != nil {
+		return "", errors.New("invalid base64 string: " + key)
+	}
+	if len(decoded) != 32 {
+		return "", errors.New("key should be 32 bytes: " + key)
+	}
+	return hex.EncodeToString(decoded), nil
+}