From 673f70c5d06cc686a91840f5f9054a3eef455130 Mon Sep 17 00:00:00 2001
From: Milo Davis <davis.mil@husky.neu.edu>
Date: Sun, 15 Oct 2017 12:42:58 +0200
Subject: [PATCH] Client Keys: Vanity keys

---
 src/client/client_keys.ml | 66 +++++++++++++++++++++++++++++++++++++--
 src/utils/base58.ml       | 11 +++++++
 src/utils/base58.mli      |  2 ++
 3 files changed, 77 insertions(+), 2 deletions(-)

diff --git a/src/client/client_keys.ml b/src/client/client_keys.ml
index fe035e3c8..814972d6c 100644
--- a/src/client/client_keys.ml
+++ b/src/client/client_keys.ml
@@ -68,6 +68,59 @@ let gen_keys ?seed cctxt name =
     "I generated a brand new pair of keys under the name '%s'." name >>= fun () ->
   return ()
 
+let gen_keys_containing ?(prefix=false) ~containing ~name (cctxt : Client_commands.context) =
+  let unrepresentable =
+    List.filter (fun s -> not @@ Base58.Alphabet.all_in_alphabet Base58.Alphabet.bitcoin s) containing in
+  match unrepresentable with
+  | _ :: _ ->
+      cctxt.warning
+        "The following can't be written in the key alphabet (%a): %a"
+        Base58.Alphabet.pp Base58.Alphabet.bitcoin
+        (Format.pp_print_list
+           ~pp_sep:(fun ppf () -> Format.fprintf ppf ", ")
+           (fun ppf s -> Format.fprintf ppf "'%s'" s))
+        unrepresentable >>= return
+  | [] ->
+      Public_key_hash.mem cctxt name >>=? fun name_exists ->
+      if name_exists && not cctxt.config.force
+      then
+        cctxt.warning
+          "Key for name '%s' already exists. Use -force to update." name >>= return
+      else
+        begin
+          cctxt.message "This process uses a brute force search and \
+                         may take a long time to find a key." >>= fun () ->
+          let matches =
+            if prefix then
+              let containing_tz1 = List.map ((^) "tz1") containing in
+              (fun key -> List.exists
+                  (fun containing ->
+                     String.sub key 0 (String.length containing) = containing)
+                  containing_tz1)
+            else
+              let re = Str.regexp (String.concat "\\|" containing) in
+              (fun key -> try ignore (Str.search_forward re key 0); true
+                with Not_found -> false) in
+          let rec loop attempts =
+            let seed = Seed.generate () in
+            let secret_key, public_key = Sodium.Sign.seed_keypair seed in
+            let hash = Ed25519.Public_key_hash.to_b58check @@ Ed25519.Public_key.hash public_key in
+            if matches hash
+            then
+              Secret_key.add cctxt name secret_key >>=? fun () ->
+              Public_key.add cctxt name public_key >>=? fun () ->
+              Public_key_hash.add cctxt name (Ed25519.Public_key.hash public_key) >>=? fun () ->
+              return hash
+            else begin if attempts mod 25_000 = 0
+              then cctxt.message "Tried %d keys without finding a match" attempts
+              else Lwt.return () end >>= fun () ->
+              loop (attempts + 1) in
+          loop 1 >>=? fun key_hash ->
+          cctxt.message
+            "Generated '%s' under the name '%s'." key_hash name >>= fun () ->
+          return ()
+        end
+
 let check_keys_consistency pk sk =
   let message = MBytes.of_string "Voulez-vous coucher avec moi, ce soir ?" in
   let signature = Ed25519.sign sk message in
@@ -137,6 +190,15 @@ let commands () =
        @@ stop)
       (fun () name cctxt -> gen_keys cctxt name) ;
 
+    command ~group ~desc: "Generate keys including the given string"
+      (args1 (switch ~doc:"The key must begin with tz1[containing]" ~parameter:"-prefix"))
+      (prefixes [ "gen" ; "vanity" ; "keys" ]
+       @@ Public_key_hash.fresh_alias_param
+       @@ prefix "matching"
+       @@ (seq_of_param @@ string ~name:"strs" ~desc:"String key must contain"))
+      (fun prefix name containing cctxt ->
+         gen_keys_containing ~prefix ~containing ~name cctxt) ;
+
     command ~group ~desc: "add a secret key to the wallet"
       no_options
       (prefixes [ "add" ; "secret" ; "key" ]
@@ -156,7 +218,7 @@ let commands () =
                (check_keys_consistency pk sk || cctxt.config.force)
                (failure
                   "public and secret keys '%s' don't correspond, \
-                   please don't use -force true" name) >>=? fun () ->
+                   please don't use -force" name) >>=? fun () ->
              Secret_key.add cctxt name sk) ;
 
     command ~group ~desc: "add a public key to the wallet"
@@ -223,7 +285,7 @@ let commands () =
       (fixed [ "forget" ; "all" ; "keys" ])
       (fun () cctxt ->
          fail_unless cctxt.config.force
-           (failure "this can only used with option -force true") >>=? fun () ->
+           (failure "this can only used with option -force") >>=? fun () ->
          Public_key.save cctxt [] >>=? fun () ->
          Secret_key.save cctxt [] >>=? fun () ->
          Public_key_hash.save cctxt []) ;
diff --git a/src/utils/base58.ml b/src/utils/base58.ml
index e3c349b95..1bf43f6a1 100644
--- a/src/utils/base58.ml
+++ b/src/utils/base58.ml
@@ -44,6 +44,17 @@ module Alphabet = struct
 
   let default = bitcoin
 
+  let all_in_alphabet alphabet string =
+    let ok = Array.make 256 false in
+    String.iter (fun x -> ok.(Char.code x) <- true) alphabet.encode ;
+    let res = ref true in
+    for i = 0 to (String.length string) - 1 do
+      res := !res && ok.(Char.code string.[i])
+    done;
+    !res
+
+  let pp ppf { encode } = Format.fprintf ppf "%s" encode
+
 end
 
 let count_trailing_char s c =
diff --git a/src/utils/base58.mli b/src/utils/base58.mli
index 12afa6aff..3cb7cae4b 100644
--- a/src/utils/base58.mli
+++ b/src/utils/base58.mli
@@ -64,6 +64,8 @@ module Alphabet : sig
   val ripple: t
   val flickr: t
   val make: string -> t
+  val all_in_alphabet :  t -> string -> bool
+  val pp : Format.formatter -> t -> unit
 end
 
 (** Encoder for a given kind of data. *)