From e69b5ff4739b3f99a649f9d97b2b5b0e988598cf Mon Sep 17 00:00:00 2001
From: kdeme <kim.demey@gmail.com>
Date: Wed, 25 Nov 2020 00:23:28 +0100
Subject: [PATCH] Add a record create and print command

---
 beacon_chain/conf.nim               | 44 +++++++++++++++++++++++++++++
 beacon_chain/eth2_network.nim       |  2 +-
 beacon_chain/nimbus_beacon_node.nim | 27 ++++++++++++++++++
 3 files changed, 72 insertions(+), 1 deletion(-)

diff --git a/beacon_chain/conf.nim b/beacon_chain/conf.nim
index a46a8886a..a7e5111e5 100644
--- a/beacon_chain/conf.nim
+++ b/beacon_chain/conf.nim
@@ -6,6 +6,7 @@ import
   confutils, confutils/defs, confutils/std/net, stew/shims/net as stewNet,
   stew/io2, unicodedb/properties, normalize,
   eth/common/eth_types as commonEthTypes,
+  eth/p2p/discoveryv5/enr,
   json_serialization, web3/[ethtypes, confutils_defs],
   spec/[crypto, keystore, digest, datatypes, network],
   network_metadata, filepath
@@ -22,6 +23,7 @@ type
     createTestnet
     deposits
     wallets
+    record
 
   WalletsCmd* {.pure.} = enum
     create  = "Creates a new EIP-2386 wallet"
@@ -36,6 +38,10 @@ type
   VCStartUpCmd* = enum
     VCNoCommand
 
+  RecordCmd* {.pure.} = enum
+    create  = "Create a new ENR"
+    print = "Print the content of a given ENR"
+
   BeaconNodeConf* = object
     logLevel* {.
       defaultValue: "INFO"
@@ -353,6 +359,36 @@ type
       of DepositsCmd.status:
         discard
 
+    of record:
+      case recordCmd* {.command.}: RecordCmd
+      of RecordCmd.create:
+        ipExt* {.
+          desc: "External IP address"
+          name: "ip" .}: ValidIpAddress
+
+        tcpPortExt* {.
+          desc: "External TCP port"
+          name: "tcp-port" .}: Port
+
+        udpPortExt* {.
+          desc: "External UDP port"
+          name: "udp-port" .}: Port
+
+        seqNumber* {.
+          defaultValue: 1,
+          desc: "Record sequence number"
+          name: "seq-number" .}: uint
+
+        fields* {.
+          desc: "Additional record key pairs, provide as <string>:<bytes in hex>"
+          name: "field" .}: seq[(string)]
+
+      of RecordCmd.print:
+        recordPrint* {.
+          argument
+          desc: "ENR URI of the record to print"
+          name: "enr" .}: Record
+
   ValidatorClientConf* = object
     logLevel* {.
       defaultValue: "INFO"
@@ -496,6 +532,14 @@ func parseCmdArg*(T: type WalletName, input: TaintedString): T
 func completeCmdArg*(T: type WalletName, input: TaintedString): seq[string] =
   return @[]
 
+proc parseCmdArg*(T: type enr.Record, p: TaintedString): T
+    {.raises: [ConfigurationError, Defect].} =
+  if not fromURI(result, p):
+    raise newException(ConfigurationError, "Invalid ENR")
+
+proc completeCmdArg*(T: type enr.Record, val: TaintedString): seq[string] =
+  return @[]
+
 func validatorsDir*(conf: BeaconNodeConf|ValidatorClientConf): string =
   string conf.validatorsDirFlag.get(InputDir(conf.dataDir / "validators"))
 
diff --git a/beacon_chain/eth2_network.nim b/beacon_chain/eth2_network.nim
index 76afa337e..e7dc623c5 100644
--- a/beacon_chain/eth2_network.nim
+++ b/beacon_chain/eth2_network.nim
@@ -1236,7 +1236,7 @@ template tcpEndPoint(address, port): auto =
 proc getPersistentNetKeys*(rng: var BrHmacDrbgContext,
                            conf: BeaconNodeConf): KeyPair =
   case conf.cmd
-  of noCommand:
+  of noCommand, record:
     if conf.netKeyFile == "random":
       let res = PrivateKey.random(Secp256k1, rng)
       if res.isErr():
diff --git a/beacon_chain/nimbus_beacon_node.nim b/beacon_chain/nimbus_beacon_node.nim
index b34358168..cc55887aa 100644
--- a/beacon_chain/nimbus_beacon_node.nim
+++ b/beacon_chain/nimbus_beacon_node.nim
@@ -1335,3 +1335,30 @@ programMain:
 
     of WalletsCmd.restore:
       restoreWalletInteractively(rng[], config)
+
+  of record:
+    case config.recordCmd:
+    of RecordCmd.create:
+      let netKeys = getPersistentNetKeys(rng[], config)
+
+      var fieldPairs: seq[FieldPair]
+      for field in config.fields:
+        let fieldPair = field.split(":")
+        if fieldPair.len > 1:
+          fieldPairs.add(toFieldPair(fieldPair[0], hexToSeqByte(fieldPair[1])))
+        else:
+          fatal "Invalid field pair"
+          quit QuitFailure
+
+      let record = enr.Record.init(
+        config.seqNumber,
+        netKeys.seckey.asEthKey,
+        some(config.ipExt),
+        config.tcpPortExt,
+        config.udpPortExt,
+        fieldPairs).expect("Record within size limits")
+
+      echo record.toURI()
+
+    of RecordCmd.print:
+      echo $config.recordPrint