From 8ff84c66dbb8b0b6580ce750b63534599c53c93f Mon Sep 17 00:00:00 2001
From: Jon Michael Aanes <jonjmaa@gmail.com>
Date: Sun, 10 Mar 2024 20:52:46 +0100
Subject: [PATCH] Argument parsing and difficulty estimation

---
 README.md                                     |  2 +-
 run.sh                                        |  1 -
 .../jmaa/pbc/PbcVanityAddressSearcher.java    | 82 +++++++++++++++----
 3 files changed, 65 insertions(+), 20 deletions(-)

diff --git a/README.md b/README.md
index 1a20380..c68c712 100644
--- a/README.md
+++ b/README.md
@@ -5,7 +5,7 @@ Tiny vanity address generator for [Partisia Blockchain](https://www.partisiabloc
 Usage:
 
 ```bash
-./run.sh "00PREFIX" "SUFFIX"
+./run.sh --prefix 00PREFIX --suffix SUFFIX
 ```
 
 Prefix must include the `00` part.
diff --git a/run.sh b/run.sh
index 4084a29..b6ce3f4 100755
--- a/run.sh
+++ b/run.sh
@@ -1,3 +1,2 @@
 #!/usr/bin/env bash
-echo $*
 mvn compile exec:java -Dexec.mainClass="jmaa.pbc.PbcVanityAddressSearcher" -Dexec.args="$*"
diff --git a/src/main/java/jmaa/pbc/PbcVanityAddressSearcher.java b/src/main/java/jmaa/pbc/PbcVanityAddressSearcher.java
index 6f32a46..15d2f50 100644
--- a/src/main/java/jmaa/pbc/PbcVanityAddressSearcher.java
+++ b/src/main/java/jmaa/pbc/PbcVanityAddressSearcher.java
@@ -1,6 +1,7 @@
 package jmaa.pbc;
 
 import java.util.List;
+import java.math.BigInteger;
 import java.util.ArrayList;
 import com.partisiablockchain.BlockchainAddress;
 import com.partisiablockchain.crypto.BlockchainPublicKey;
@@ -15,22 +16,17 @@ public final class PbcVanityAddressSearcher {
   private static final long INITIAL_PING_INTERVAL = 10_000;
 
   public static void main(String[] args) throws InterruptedException {
-    searchForPrivateKeyWithVanityAddress(args[0], args[1]);
+      final AddressRequirements requirements = parseArgumentsToRequirements(args);
+    searchForPrivateKeyWithVanityAddress(requirements);
   }
 
-  /**
-   * @param vanityPrefix Desired address prefix
-   */
-  private static void searchForPrivateKeyWithVanityAddress(String vanityPrefix, String vanitySuffix) throws InterruptedException {
-      if (!vanityPrefix.startsWith("00")) {
-          throw new IllegalArgumentException("Vanity prefix must start with 00.");
-      }
-
-      System.out.println("Searching for private key for address prefix \"%s\" and suffix \"%s\" ".formatted(vanityPrefix, vanitySuffix));
+  private static void searchForPrivateKeyWithVanityAddress(AddressRequirements addressRequirements) throws InterruptedException {
+      System.out.println("Searching for private key: %s".formatted(addressRequirements));
+      System.out.println("Difficulty: 1/%s".formatted(addressRequirements.difficulty()));
 
     List<Thread> searchers = new ArrayList<>();
     for(int threadIdx=0;threadIdx<NUM_THREADS ;threadIdx++) {
-      final Thread searcher = new Thread(new KeySearcher(threadIdx, vanityPrefix, vanitySuffix));
+      final Thread searcher = new Thread(new AddressSearcherThread(threadIdx, addressRequirements));
       searcher.start();
       searchers.add(searcher);
     }
@@ -40,7 +36,13 @@ public final class PbcVanityAddressSearcher {
     }
   }
 
-  private record KeySearcher(int threadIdx, String vanityPrefix, String vanitySuffix) implements Runnable {
+  /**
+   * Individual thread for searching for addresses.
+   *
+   * @param threadIx Id of thread
+   * @param addressRequirements Requirements for the current search.
+   */
+  private record AddressSearcherThread(int threadIdx, AddressRequirements addressRequirements) implements Runnable {
 
     @Override
     public void run() {
@@ -51,13 +53,9 @@ public final class PbcVanityAddressSearcher {
         final KeyPair keyPair = new KeyPair();
         final BlockchainAddress address = keyPair.getPublic().createAddress();
 
-        // Check vanity
-        final String addressAsString = address.writeAsString();
-        final boolean satisfiesVanity = addressAsString.startsWith(vanityPrefix) && addressAsString.endsWith(vanitySuffix);
-
         // Output if vain enough
-        if(satisfiesVanity) {
-          System.out.println("[%2d:%10d] match: addr=%s - pk=%s".formatted(threadIdx, iterationIdx, addressAsString, keyPair.getPrivateKey().toString(16)));
+        if(addressRequirements.satisfiedBy(address)) {
+          System.out.println("[%2d:%10d] match: addr=%s - pk=%s".formatted(threadIdx, iterationIdx, address.writeAsString(), keyPair.getPrivateKey().toString(16)));
         }
 
         // Ping
@@ -70,4 +68,52 @@ public final class PbcVanityAddressSearcher {
       }
     }
   }
+
+  private static AddressRequirements parseArgumentsToRequirements(String[] args)  {
+      String prefix = "00";
+      String suffix = "";
+      for (int idx = 0; idx < args.length; idx++) {
+          final String word = args[idx];
+          System.out.println("%s : %s".formatted(idx, word));
+          if ("--prefix".equals(word)) {
+              prefix = args[++idx];;
+          } else if ("--suffix".equals(word)) {
+              suffix = args[++idx];
+          } else {
+              throw new RuntimeException("Unknown argument: " + word);
+          }
+      }
+
+      return new AddressRequirements(prefix, suffix);
+  }
+
+  /**
+   * Requirements for the output {@link BlockchainAddress}.
+   *
+   * @param prefix Required prefix of address. Must include 00
+   * @param suffix Required suffix of address.
+   */
+  private record AddressRequirements(String prefix, String suffix) {
+
+      public AddressRequirements {
+          if (!prefix.startsWith("00")) {
+              throw new IllegalArgumentException("Vanity prefix must start with 00.");
+          }
+      }
+
+      /**
+       * Whether the given address satisfies the requirements.
+       */
+      public boolean satisfiedBy(BlockchainAddress address) {
+        final String addressAsString = address.writeAsString();
+        return addressAsString.startsWith(prefix) && addressAsString.endsWith(suffix);
+      }
+
+      /**
+       * Attempts to measure the "one-in-X" factor.
+       */
+      public BigInteger difficulty() {
+          return BigInteger.ONE.shiftLeft(4*(prefix.length()-2 + suffix.length()));
+      }
+  }
 }