Add --format option for RNG cmdlet

5 files changed, 73 insertions, 43 deletions

diff --git a/src/cli/cli.cpp b/src/cli/cli.cpp
index 1918a4fbd..25edd5df2 100644
--- a/src/cli/cli.cpp
+++ b/src/cli/cli.cpp
@@ -12,6 +12,18 @@
 #include <iostream>
 #include <fstream>
 
+#if defined(BOTAN_HAS_HEX_CODEC)
+   #include <botan/hex.h>
+#endif
+
+#if defined(BOTAN_HAS_BASE64_CODEC)
+   #include <botan/base64.h>
+#endif
+
+#if defined(BOTAN_HAS_BASE58_CODEC)
+   #include <botan/base58.h>
+#endif
+
 namespace Botan_CLI {
 
 Command::Command(const std::string& cmd_spec) : m_spec(cmd_spec)
@@ -240,6 +252,39 @@ std::string Command::get_passphrase(const std::string& prompt)
    return pass;
    }
 
+//static
+std::string Command::format_blob(const std::string& format,
+                                 const uint8_t bits[], size_t len)
+   {
+#if defined(BOTAN_HAS_HEX_CODEC)
+   if(format == "hex")
+      {
+      return Botan::hex_encode(bits, len);
+      }
+#endif
+
+#if defined(BOTAN_HAS_BASE64_CODEC)
+   if(format == "base64")
+      {
+      return Botan::base64_encode(bits, len);
+      }
+#endif
+
+#if defined(BOTAN_HAS_BASE58_CODEC)
+   if(format == "base58")
+      {
+      return Botan::base58_encode(bits, len);
+      }
+   if(format == "base58check")
+      {
+      return Botan::base58_check_encode(bits, len);
+      }
+#endif
+
+   // If we supported format, we would have already returned
+   throw CLI_Usage_Error("Unknown or unsupported format type");
+   }
+
 // Registration code
 
 Command::Registration::Registration(const std::string& name, Command::cmd_maker_fn maker_fn)
diff --git a/src/cli/cli.h b/src/cli/cli.h
index b6f6a3076..838fb9eb3 100644
--- a/src/cli/cli.h
+++ b/src/cli/cli.h
@@ -122,6 +122,15 @@ class Command
 
       bool flag_set(const std::string& flag_name) const;
 
+      static std::string format_blob(const std::string& format, const uint8_t bits[], size_t len);
+
+      template<typename Alloc>
+      static std::string format_blob(const std::string& format,
+                                     const std::vector<uint8_t, Alloc>& vec)
+         {
+         return format_blob(format, vec.data(), vec.size());
+         }
+
       std::string get_arg(const std::string& opt_name) const;
 
       /**
diff --git a/src/cli/cli_rng.cpp b/src/cli/cli_rng.cpp
index 90b30a820..a7210f27d 100644
--- a/src/cli/cli_rng.cpp
+++ b/src/cli/cli_rng.cpp
@@ -90,7 +90,7 @@ cli_make_rng(const std::string& rng_type, const std::string& hex_drbg_seed)
 class RNG final : public Command
    {
    public:
-      RNG() : Command("rng --system --rdrand --auto --entropy --drbg --drbg-seed= *bytes") {}
+      RNG() : Command("rng --format=hex --system --rdrand --auto --entropy --drbg --drbg-seed= *bytes") {}
 
       std::string group() const override
          {
@@ -104,6 +104,7 @@ class RNG final : public Command
 
       void go() override
          {
+         const std::string format = get_arg("format");
          std::string type = get_arg("rng-type");
 
          if(type.empty())
@@ -123,7 +124,17 @@ class RNG final : public Command
 
          for(const std::string& req : get_arg_list("bytes"))
             {
-            output() << Botan::hex_encode(rng->random_vec(Botan::to_u32bit(req))) << "\n";
+            const size_t req_len = Botan::to_u32bit(req);
+            const auto blob = rng->random_vec(req_len);
+
+            if(format == "binary" || format == "raw")
+               {
+               output().write(reinterpret_cast<const char*>(blob.data()), blob.size());
+               }
+            else
+               {
+               output() << format_blob(format, blob) << "\n";
+               }
             }
          }
    };
diff --git a/src/cli/hash.cpp b/src/cli/hash.cpp
index a7481ba5e..8e59b2ab5 100644
--- a/src/cli/hash.cpp
+++ b/src/cli/hash.cpp
@@ -6,17 +6,8 @@
 
 #include "cli.h"
 
-#if defined(BOTAN_HAS_HASH) && defined(BOTAN_HAS_HEX_CODEC)
+#if defined(BOTAN_HAS_HASH)
    #include <botan/hash.h>
-   #include <botan/hex.h>
-#endif
-
-#if defined(BOTAN_HAS_BASE64_CODEC)
-   #include <botan/base64.h>
-#endif
-
-#if defined(BOTAN_HAS_BASE58_CODEC)
-   #include <botan/base58.h>
 #endif
 
 namespace Botan_CLI {
@@ -38,36 +29,6 @@ class Hash final : public Command
          return "Compute the message digest of given file(s)";
          }
 
-      template<typename Alloc>
-      static std::string format_digest(const std::string& format,
-                                       const std::vector<uint8_t, Alloc>& vec)
-         {
-         if(format == "hex")
-            {
-            return Botan::hex_encode(vec);
-            }
-#if defined(BOTAN_HAS_BASE64_CODEC)
-         else if(format == "base64")
-            {
-            return Botan::base64_encode(vec);
-            }
-#endif
-#if defined(BOTAN_HAS_BASE58_CODEC)
-         else if(format == "base58")
-            {
-            return Botan::base58_encode(vec);
-            }
-         else if(format == "base58check")
-            {
-            return Botan::base58_check_encode(vec);
-            }
-#endif
-         else
-            {
-            throw CLI_Usage_Error("Unknown format for encoding digests");
-            }
-         }
-
       void go() override
          {
          const std::string hash_algo = get_arg("algo");
@@ -95,7 +56,7 @@ class Hash final : public Command
                auto update_hash = [&](const uint8_t b[], size_t l) { hash_fn->update(b, l); };
                read_file(fsname, update_hash, buf_size);
 
-               const std::string digest = format_digest(format, hash_fn->final());
+               const std::string digest = format_blob(format, hash_fn->final());
 
                if(no_fsname)
                   output() << digest << "\n";
diff --git a/src/scripts/test_cli.py b/src/scripts/test_cli.py
index 15776f6f7..c891d5ffb 100755
--- a/src/scripts/test_cli.py
+++ b/src/scripts/test_cli.py
@@ -390,6 +390,10 @@ def cli_rng_tests():
     test_cli("rng", "16", "D80F88F6ADBE65ACB10C3602E67D985B")
     test_cli("rng", "10 6", "D80F88F6ADBE65ACB10C\n1B119CC068AF")
 
+    test_cli("rng", ['--format=base64', '10'], "2A+I9q2+ZayxDA==")
+    test_cli("rng", ['--format=base58', '10'], "D93XRyVfxqs7oR")
+    test_cli("rng", ['--format=base58check', '10'], "2NS1jYUq92TyGFVnhVLa")
+
 def cli_pk_workfactor_tests():
     test_cli("pk_workfactor", "1024", "80")
     test_cli("pk_workfactor", "2048", "111")