From dbe4b8fa17bc4ae2a86c8fbc209da673d0ce75b2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jakub=20Soko=C5=82owski?= <jakub@status.im>
Date: Tue, 21 Jul 2020 21:07:46 +0200
Subject: [PATCH] add --privkey flag to not use multiple enodes
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Signed-off-by: Jakub Sokołowski <jakub@status.im>
---
 cmd/x-check-mailserver/config.go | 12 +++++++++++-
 cmd/x-check-mailserver/flags.go  |  1 +
 cmd/x-check-mailserver/main.go   |  4 ++--
 cmd/x-check-mailserver/work.go   | 15 +++++++++++++--
 4 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/cmd/x-check-mailserver/config.go b/cmd/x-check-mailserver/config.go
index 8ac7ebf..8476e4f 100644
--- a/cmd/x-check-mailserver/config.go
+++ b/cmd/x-check-mailserver/config.go
@@ -1,10 +1,13 @@
 package main
 
 import (
+	"errors"
+	"strings"
+
 	"github.com/status-im/status-go/params"
 )
 
-func newNodeConfig(fleet string, networkID uint64) (*params.NodeConfig, error) {
+func newNodeConfig(fleet string, networkID uint64, privKey string) (*params.NodeConfig, error) {
 	c, err := params.NewNodeConfig("", networkID)
 	if err != nil {
 		return nil, err
@@ -14,6 +17,13 @@ func newNodeConfig(fleet string, networkID uint64) (*params.NodeConfig, error) {
 		return nil, err
 	}
 
+	cleanPrivKey := strings.TrimPrefix(privKey, "0x")
+	if len(cleanPrivKey) == 64 {
+		c.NodeKey = cleanPrivKey
+	} else if len(cleanPrivKey) != 64 && len(cleanPrivKey) != 0 {
+		return nil, errors.New("wrong private key length, expected 64 char hexadecimal")
+	}
+
 	c.ListenAddr = ":0"
 	c.MaxPeers = 10
 	c.IPCEnabled = true
diff --git a/cmd/x-check-mailserver/flags.go b/cmd/x-check-mailserver/flags.go
index 854406d..5c3afb6 100644
--- a/cmd/x-check-mailserver/flags.go
+++ b/cmd/x-check-mailserver/flags.go
@@ -10,6 +10,7 @@ import (
 var (
 	fleet       = pflag.StringP("fleet", "f", params.FleetProd, "cluster fleet")
 	datadir     = pflag.StringP("datadir", "d", "/tmp", "home directory for node data")
+	privkey     = pflag.StringP("privkey", "p", "", "private key for connecting to nodes, hexadecimal")
 	mailservers = pflag.StringArrayP("mailservers", "m", nil, "a list of mail servers")
 	duration    = pflag.DurationP("duration", "l", time.Hour*24, "length of time span from now")
 	channels    = pflag.StringArrayP("channels", "c", []string{"status"}, "name of one or more channels")
diff --git a/cmd/x-check-mailserver/main.go b/cmd/x-check-mailserver/main.go
index 1492591..767d483 100644
--- a/cmd/x-check-mailserver/main.go
+++ b/cmd/x-check-mailserver/main.go
@@ -34,7 +34,7 @@ func main() {
 	}()
 
 	// create config
-	config, err := newNodeConfig(*fleet, params.MainNetworkID)
+	config, err := newNodeConfig(*fleet, params.MainNetworkID, *privkey)
 	if err != nil {
 		log.Crit("failed to create a config", "err", err)
 	}
@@ -76,7 +76,7 @@ func main() {
 		nodeConfig := *config
 		log.Debug("using node config", "config", nodeConfig)
 
-		work := NewWorkUnit(msEnode, &nodeConfig)
+		work := NewWorkUnit(msEnode, &nodeConfig, *privkey)
 		go func(work *WorkUnit) {
 			if err := work.Execute(workConfig); err != nil {
 				log.Crit("failed to execute work", "err", err, "enode", work.MailServerEnode)
diff --git a/cmd/x-check-mailserver/work.go b/cmd/x-check-mailserver/work.go
index 97ea10e..6cc22da 100644
--- a/cmd/x-check-mailserver/work.go
+++ b/cmd/x-check-mailserver/work.go
@@ -5,6 +5,7 @@ import (
 	"crypto/ecdsa"
 	"fmt"
 	"go.uber.org/zap"
+	"strings"
 	"time"
 
 	"github.com/ethereum/go-ethereum/accounts"
@@ -25,6 +26,7 @@ import (
 // to a mailserver and collect received envelopes.
 type WorkUnit struct {
 	MailServerEnode string
+	PrivKey         string
 	MessageHashes   []types.HexBytes // a list of collected messages.
 
 	config    *params.NodeConfig
@@ -34,9 +36,10 @@ type WorkUnit struct {
 }
 
 // NewWorkUnit creates a new WorkUnit instance.
-func NewWorkUnit(mailEnode string, config *params.NodeConfig) *WorkUnit {
+func NewWorkUnit(mailEnode string, config *params.NodeConfig, privKey string) *WorkUnit {
 	return &WorkUnit{
 		MailServerEnode: mailEnode,
+		PrivKey:         privKey,
 		config:          config,
 	}
 }
@@ -99,6 +102,14 @@ func (u *WorkUnit) Execute(config WorkUnitConfig) error {
 	return nil
 }
 
+func (u *WorkUnit) generateKey() (*ecdsa.PrivateKey, error) {
+	if u.PrivKey != "" {
+		return crypto.HexToECDSA(strings.TrimPrefix(u.PrivKey, "0x"))
+	} else {
+		return crypto.GenerateKey()
+	}
+}
+
 func (u *WorkUnit) startNode() error {
 	u.node = node.New()
 	if err := u.node.Start(u.config, &accounts.Manager{}); err != nil {
@@ -112,7 +123,7 @@ func (u *WorkUnit) stopNode() error {
 }
 
 func (u *WorkUnit) startMessenger() error {
-	key, err := crypto.GenerateKey()
+	key, err := u.generateKey()
 	if err != nil {
 		return err
 	}