1 files changed, 63 insertions, 24 deletions

diff --git a/test/encrypt_block_aes.cpp b/test/encrypt_block_aes.cpp
index 8be1fef..fc56019 100644
--- a/test/encrypt_block_aes.cpp
+++ b/test/encrypt_block_aes.cpp
@@ -6,56 +6,90 @@
  *            See LICENSE.txt for details.
  */
 
+#include "common_aes.hpp"
+
 #include <aesni/all.h>
 
 #include <aesnixx/all.hpp>
 
-#include <cstdlib>
-#include <cstring>
-
 #include <exception>
 #include <iostream>
 
-namespace
-{
-    void exit_with_usage()
-    {
-        std::cout << "Usage: encrypt_block_aes.exe KEY0 IV0 [PLAIN0...] [-- KEY1 IV1 [PLAIN1...]...]\n";
-        std::exit(EXIT_FAILURE);
-    }
-}
-
 int main(int argc, char** argv)
 {
     try
     {
-        for (--argc, ++argv; argc > -1; --argc, ++argv)
-        {
-            if (argc < 2)
-                exit_with_usage();
+        CommandLineParser cmd_parser("encrypt_block_aes.exe");
 
+        if (!cmd_parser.parse_options(argc, argv))
+            return 0;
+
+        auto args = cmd_parser.get_args();
+
+        while (!args.empty())
+        {
             AesNI_BoxAlgorithmParams algorithm_params;
-            aesni::aes::from_string(algorithm_params.aes128_key, argv[0]);
+
+            switch (cmd_parser.get_algorithm())
+            {
+                case AESNI_AES128:
+                    aesni::aes::from_string(algorithm_params.aes128_key, args.front());
+                    break;
+
+                case AESNI_AES192:
+                    aesni::aes::from_string(algorithm_params.aes192_key, args.front());
+                    break;
+
+                case AESNI_AES256:
+                    aesni::aes::from_string(algorithm_params.aes256_key, args.front());
+                    break;
+            }
+
+            args.pop_front();
 
             AesNI_BoxBlock iv;
-            aesni::aes::from_string(iv.aes_block, argv[1]);
+            AesNI_BoxBlock* iv_ptr = nullptr;
+
+            switch (cmd_parser.get_mode())
+            {
+                case AESNI_ECB:
+                    break;
+
+                case AESNI_CBC:
+                case AESNI_CFB:
+                case AESNI_OFB:
+                case AESNI_CTR:
+                    if (args.empty())
+                    {
+                        cmd_parser.print_usage();
+                        return 1;
+                    }
+                    aesni::aes::from_string(iv.aes_block, args.front());
+                    iv_ptr = &iv;
+                    args.pop_front();
+                    break;
+            }
 
             AesNI_Box box;
             aesni_box_init(
                 &box,
-                AESNI_AES128,
+                cmd_parser.get_algorithm(),
                 &algorithm_params,
-                AESNI_OFB,
-                &iv,
+                cmd_parser.get_mode(),
+                iv_ptr,
                 aesni::ErrorDetailsThrowsInDestructor());
 
-            for (argc -= 2, argv += 2; argc > 0; --argc, ++argv)
+            while (!args.empty())
             {
-                if (std::strcmp("--", argv[0]) == 0)
+                if (args.front() == "--")
+                {
+                    args.pop_front();
                     break;
+                }
 
                 AesNI_BoxBlock plaintext;
-                aesni::aes::from_string(plaintext.aes_block, argv[0]);
+                aesni::aes::from_string(plaintext.aes_block, args.front());
+                args.pop_front();
 
                 AesNI_BoxBlock ciphertext;
                 aesni_box_encrypt(
@@ -70,6 +104,11 @@ int main(int argc, char** argv)
 
         return 0;
     }
+    catch (const boost::program_options::error& e)
+    {
+        std::cerr << "Usage error: " << e.what() << "\n";
+        return 1;
+    }
     catch (const std::exception& e)
     {
         std::cerr << e.what() << "\n";