diff --git a/src/crypto_utils.c b/src/crypto_utils.c
index cdc2b44..e394ebc 100644
--- a/src/crypto_utils.c
+++ b/src/crypto_utils.c
@@ -92,7 +92,7 @@ void pin_derive_kenc2(const uint8_t pin_token[32], uint8_t kenc[32]) {
 // Encrypt 32-byte device key using AES-256-GCM
 // Output: [nonce|ciphertext|tag]  =  12 + in_len + 16 = 60 bytes
 // ------------------------------------------------------------------
-int encrypt_with_aad(const uint8_t key[32], const uint8_t *in_buf, size_t in_len, uint8_t version, uint8_t *out_buf) {
+int encrypt_with_aad(const uint8_t key[32], const uint8_t *in_buf, size_t in_len, const pin_kdf_version_t version, uint8_t *out_buf) {
     uint8_t *nonce = out_buf;
     uint8_t *ct    = out_buf + 12;
     uint8_t *tag   = out_buf + 12 + in_len;
@@ -102,11 +102,15 @@ int encrypt_with_aad(const uint8_t key[32], const uint8_t *in_buf, size_t in_len
     mbedtls_gcm_context gcm;
     mbedtls_gcm_init(&gcm);
     uint8_t kenc[32];
-    if (version == 2) {
+    if (version == PIN_KDF_V2) {
         pin_derive_kenc2(key, kenc);
-    } else {
+    }
+    else if (version == PIN_KDF_V1) {
         pin_derive_kenc(key, kenc);
     }
+    else {
+        return PICOKEY_WRONG_DATA;
+    }
     int rc = mbedtls_gcm_setkey(&gcm, MBEDTLS_CIPHER_ID_AES, kenc, 256);
     mbedtls_platform_zeroize(kenc, sizeof(kenc));
     if (rc != 0) {
@@ -123,7 +127,7 @@ int encrypt_with_aad(const uint8_t key[32], const uint8_t *in_buf, size_t in_len
 // Input: [nonce|ciphertext|tag]  =  in_len bytes
 // Output: decrypted = in_len - 12 - 16 bytes
 // ------------------------------------------------------------------
-int decrypt_with_aad(const uint8_t key[32], const uint8_t *in_buf, size_t in_len, uint8_t version, uint8_t *out_buf) {
+int decrypt_with_aad(const uint8_t key[32], const uint8_t *in_buf, size_t in_len, const pin_kdf_version_t version, uint8_t *out_buf) {
     const uint8_t *nonce = in_buf;
     const uint8_t *ct    = in_buf + 12;
     const uint8_t *tag   = in_buf + in_len - 16;
@@ -131,11 +135,14 @@ int decrypt_with_aad(const uint8_t key[32], const uint8_t *in_buf, size_t in_len
     mbedtls_gcm_context gcm;
     mbedtls_gcm_init(&gcm);
     uint8_t kenc[32];
-    if (version == 2) {
+    if (version == PIN_KDF_V2) {
         pin_derive_kenc2(key, kenc);
-    } else {
+    } else if (version == PIN_KDF_V1) {
         pin_derive_kenc(key, kenc);
     }
+    else {
+        return PICOKEY_WRONG_DATA;
+    }
     int rc = mbedtls_gcm_setkey(&gcm, MBEDTLS_CIPHER_ID_AES, kenc, 256);
     mbedtls_platform_zeroize(kenc, sizeof(kenc));
     if (rc != 0) {
@@ -157,7 +164,6 @@ void double_hash_pin(const uint8_t *pin, uint16_t len, uint8_t output[32]) {
     hash_multi(o1, sizeof(o1), output);
 }
 
-
 void hash_multi(const uint8_t *input, uint16_t len, uint8_t output[32]) {
     mbedtls_sha256_context ctx;
     mbedtls_sha256_init(&ctx);
diff --git a/src/crypto_utils.h b/src/crypto_utils.h
index 8975bb0..df1f30d 100644
--- a/src/crypto_utils.h
+++ b/src/crypto_utils.h
@@ -38,6 +38,14 @@
 
 #define IV_SIZE 16
 
+typedef enum {
+    PIN_KDF_V1 = 1,
+    PIN_KDF_V2 = 2,
+    PIN_KDF_UNKNOWN = 0xff
+} pin_kdf_version_t;
+
+#define PIN_KDF_DEFAULT_VERSION PIN_KDF_V2
+
 extern int ct_memcmp(const void *a, const void *b, size_t n);
 // Newer and safe functions
 extern void derive_kbase(uint8_t kbase[32]);
@@ -46,8 +54,8 @@ extern void pin_derive_kenc(const uint8_t pin_token[32], uint8_t kenc[32]);
 extern void pin_derive_kenc2(const uint8_t pin_token[32], uint8_t kenc[32]);
 extern void pin_derive_session(const uint8_t *pin, size_t pin_len, uint8_t pin_token[32]);
 extern void pin_derive_verifier(const uint8_t *pin, size_t pin_len, uint8_t verifier[32]);
-extern int encrypt_with_aad(const uint8_t key[32], const uint8_t *in_buf, size_t in_len, uint8_t version, uint8_t *out_buf);
-extern int decrypt_with_aad(const uint8_t key[32], const uint8_t *in_buf, size_t in_len, uint8_t version, uint8_t *out_buf);
+extern int encrypt_with_aad(const uint8_t key[32], const uint8_t *in_buf, size_t in_len, const pin_kdf_version_t version, uint8_t *out_buf);
+extern int decrypt_with_aad(const uint8_t key[32], const uint8_t *in_buf, size_t in_len, const pin_kdf_version_t version, uint8_t *out_buf);
 extern void double_hash_pin(const uint8_t *pin, uint16_t len, uint8_t output[32]);
 extern void hash_multi(const uint8_t *input, uint16_t len, uint8_t output[32]);
 extern void hash256(const uint8_t *input, size_t len, uint8_t output[32]);
@@ -59,4 +67,6 @@ extern int aes_decrypt_cfb_256(const uint8_t *key, const uint8_t *iv, uint8_t *d
 extern mbedtls_ecp_group_id ec_get_curve_from_prime(const uint8_t *prime, size_t prime_len);
 extern uint32_t crc32c(const uint8_t *buf, size_t len);
 
+#define PIN_KDF_SIZE(x)  (12 + (x) + 16)
+
 #endif