From 801fcc59dbf04dfce4323d2a3cad99f9d643e525 Mon Sep 17 00:00:00 2001
From: Szczepan Zalega <szczepan@nitrokey.com>
Date: Fri, 14 Oct 2016 15:00:04 +0200
Subject: Support regenerating AES keys on NK Storage

Signed-off-by: Szczepan Zalega <szczepan@nitrokey.com>
---
 NitrokeyManager.cc         | 18 +++++++++++++++---
 include/stick20_commands.h | 26 ++++++++++++++++++--------
 2 files changed, 33 insertions(+), 11 deletions(-)

diff --git a/NitrokeyManager.cc b/NitrokeyManager.cc
index 20657a5..47b68d4 100644
--- a/NitrokeyManager.cc
+++ b/NitrokeyManager.cc
@@ -396,9 +396,21 @@ namespace nitrokey{
     }
 
     void NitrokeyManager::build_aes_key(const char *admin_password) {
-        auto p = get_payload<BuildAESKey>();
-        strcpyT(p.admin_password, admin_password);
-        BuildAESKey::CommandTransaction::run(*device, p);
+        switch (device->get_device_model()) {
+            case DeviceModel::PRO: {
+                auto p = get_payload<BuildAESKey>();
+                strcpyT(p.admin_password, admin_password);
+                BuildAESKey::CommandTransaction::run(*device, p);
+                break;
+            }
+            case DeviceModel::STORAGE : {
+                auto p = get_payload<stick20::CreateNewKeys>();
+                strcpyT(p.admin_password, admin_password);
+                p.setKindPrefixed();
+                stick20::CreateNewKeys::CommandTransaction::run(*device, p);
+                break;
+            }
+        }
     }
 
     void NitrokeyManager::factory_reset(const char *admin_password) {
diff --git a/include/stick20_commands.h b/include/stick20_commands.h
index 03761a5..f4e7500 100644
--- a/include/stick20_commands.h
+++ b/include/stick20_commands.h
@@ -146,15 +146,25 @@ class ExportFirmware : semantics::non_constructible {
                       struct EmptyPayload> CommandTransaction;
 };
 
-class CreateNewKeys : semantics::non_constructible {
- public:
-  struct CommandPayload {
-    uint8_t password[30];
-  };
+    class CreateNewKeys : Command<CommandID::GENERATE_NEW_KEYS> {
+    public:
+        struct CommandPayload {
+            uint8_t kind;
+            uint8_t admin_password[30]; //CS20_MAX_PASSWORD_LEN
+            std::string dissect() const {
+              std::stringstream ss;
+              ss << " admin_password:\t" <<  admin_password<< std::endl;
+              return ss.str();
+            }
+            void setKindPrefixed(){
+              kind = 'P';
+            }
+        } __packed;
+
+        typedef Transaction<command_id(), struct CommandPayload, struct EmptyPayload>
+            CommandTransaction;
+    };
 
-  typedef Transaction<CommandID::GENERATE_NEW_KEYS, struct CommandPayload,
-                      struct EmptyPayload> CommandTransaction;
-};
 
 class FillSDCardWithRandomChars : semantics::non_constructible {
  public:
-- 
cgit v1.2.1