From b38ae6a6542d1c8dbf7687541e691636027a1cc3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Sosth=C3=A8ne=20Gu=C3=A9don?= <sosthene@nitrokey.com>
Date: Tue, 15 Nov 2022 11:20:37 +0100
Subject: [PATCH] Add support for Aes key management in tests

---
 Cargo.toml                | 1 +
 src/state.rs              | 6 ++++--
 tests/command_response.rs | 6 +++++-
 3 files changed, 10 insertions(+), 3 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index b939b9d..2711082 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -41,6 +41,7 @@ hex  = "0.4"
 test-log = "0.2"
 ron  = "0.8"
 des = "0.8"
+aes = "0.8.2"
 
 
 [features]
diff --git a/src/state.rs b/src/state.rs
index c4fccb4..07ae7c4 100644
--- a/src/state.rs
+++ b/src/state.rs
@@ -9,7 +9,7 @@ use trussed::{
     api::reply::Metadata,
     config::MAX_MESSAGE_LENGTH,
     syscall, try_syscall,
-    types::{KeyId, Location, Mechanism, PathBuf},
+    types::{KeyId, KeySerialization, Location, Mechanism, PathBuf},
 };
 
 use crate::constants::*;
@@ -455,9 +455,11 @@ impl Persistent {
     pub fn initialize(client: &mut impl trussed::Client) -> Self {
         info!("initializing PIV state");
         let management_key = ManagementKey {
-            id: syscall!(client.unsafe_inject_shared_key(
+            id: syscall!(client.unsafe_inject_key(
+                YUBICO_DEFAULT_MANAGEMENT_KEY_ALG.mechanism(),
                 YUBICO_DEFAULT_MANAGEMENT_KEY,
                 trussed::types::Location::Internal,
+                KeySerialization::Raw
             ))
             .key,
             alg: YUBICO_DEFAULT_MANAGEMENT_KEY_ALG,
diff --git a/tests/command_response.rs b/tests/command_response.rs
index 4a0ed79..340ae1c 100644
--- a/tests/command_response.rs
+++ b/tests/command_response.rs
@@ -6,6 +6,7 @@ mod setup;
 
 use std::borrow::Cow;
 
+use aes::Aes256Enc;
 use hex_literal::hex;
 use serde::Deserialize;
 use trussed::types::GenericArray;
@@ -348,7 +349,10 @@ impl IoCmd {
                 let cipher = TdesEde3::new(GenericArray::from_slice(&key));
                 cipher.encrypt_block(GenericArray::from_mut_slice(challenge));
             }
-            Algorithm::Aes256 => todo!(),
+            Algorithm::Aes256 => {
+                let cipher = Aes256Enc::new(GenericArray::from_slice(&key));
+                cipher.encrypt_block(GenericArray::from_mut_slice(challenge));
+            }
             _ => panic!(),
         }
         let second_data = tlv(&[0x7C], &tlv(&[0x82], challenge));