diff --git a/Makefile b/Makefile
index 33ede26e5..ba3d5535f 100644
--- a/Makefile
+++ b/Makefile
@@ -86,7 +86,7 @@ ifeq ($(TARGET),ti_hercules)
 endif
 
 # Environment variables for sign tool
-SIGN_ENV=IMAGE_HEADER_SIZE=$(IMAGE_HEADER_SIZE) WOLFBOOT_SECTOR_SIZE=$(WOLFBOOT_SECTOR_SIZE)
+SIGN_ENV=IMAGE_HEADER_SIZE=$(IMAGE_HEADER_SIZE) WOLFBOOT_SECTOR_SIZE=$(WOLFBOOT_SECTOR_SIZE) ML_DSA_LEVEL=$(ML_DSA_LEVEL) IMAGE_SIGNATURE_SIZE=$(IMAGE_SIGNATURE_SIZE)
 
 
 MAIN_TARGET=factory.bin
@@ -208,7 +208,7 @@ keytools_check: keytools
 
 $(PRIVATE_KEY):
 	$(Q)$(MAKE) keytools_check
-	$(Q)(test $(SIGN) = NONE) || ("$(KEYGEN_TOOL)" $(KEYGEN_OPTIONS) -g $(PRIVATE_KEY)) || true
+	$(Q)(test $(SIGN) = NONE) || ($(SIGN_ENV) "$(KEYGEN_TOOL)" $(KEYGEN_OPTIONS) -g $(PRIVATE_KEY)) || true
 	$(Q)(test $(SIGN) = NONE) && (echo "// SIGN=NONE" >  src/keystore.c) || true
 	$(Q)(test "$(FLASH_OTP_KEYSTORE)" = "1") && (make -C tools/keytools/otp) || true
 
@@ -216,7 +216,7 @@ $(SECONDARY_PRIVATE_KEY): $(PRIVATE_KEY) keystore.der
 	$(Q)$(MAKE) keytools_check
 	$(Q)rm -f src/keystore.c
 	$(Q)dd if=keystore.der of=pubkey_1.der bs=1 skip=16
-	$(Q)(test $(SIGN_SECONDARY) = NONE) || ("$(KEYGEN_TOOL)" \
+	$(Q)(test $(SIGN_SECONDARY) = NONE) || ($(SIGN_ENV) "$(KEYGEN_TOOL)" \
 		$(KEYGEN_OPTIONS) -i pubkey_1.der $(SECONDARY_KEYGEN_OPTIONS) \
 		-g $(SECONDARY_PRIVATE_KEY)) || true
 	$(Q)(test "$(FLASH_OTP_KEYSTORE)" = "1") && (make -C tools/keytools/otp) || true
diff --git a/include/wolfboot/wolfboot.h b/include/wolfboot/wolfboot.h
index 6d44e9e94..e697c38bf 100644
--- a/include/wolfboot/wolfboot.h
+++ b/include/wolfboot/wolfboot.h
@@ -138,21 +138,17 @@ extern "C" {
 /* ML-DSA pub key size is a function of parameters.
  * This needs to be configurable. Default to security
  * category 2. */
-#ifdef ML_DSA_LEVEL
-    #if ML_DSA_LEVEL == 2
-        #define KEYSTORE_PUBKEY_SIZE_ML_DSA  1312
-    #elif ML_DSA_LEVEL == 3
-        #define KEYSTORE_PUBKEY_SIZE_ML_DSA  1952
-    #elif ML_DSA_LEVEL == 5
-        #define KEYSTORE_PUBKEY_SIZE_ML_DSA  2592
-    #endif
-#else
-    #ifdef SIGN_ML_DSA
-        #error "ML_DSA_LEVEL not defined"
-    #endif
-    /* Default to max size for keystore */
-    #define KEYSTORE_PUBKEY_SIZE_ML_DSA  2592
-#endif /* defined ML_DSA_LEVEL */
+
+#define ML_DSA_L2_PUBKEY_SIZE   1312
+#define ML_DSA_L3_PUBKEY_SIZE   1952
+#define ML_DSA_L5_PUBKEY_SIZE   2592
+
+
+#if defined(SIGN_ML_DSA) && !defined(ML_DSA_LEVEL)
+    #define ML_DSA_LEVEL 5
+#endif
+/* Default to max size for keystore */
+#define KEYSTORE_PUBKEY_SIZE_ML_DSA  2592
 
 /* Mask for key permissions */
 #define KEY_VERIFY_ALL         (0xFFFFFFFFU)
diff --git a/tools/keytools/Makefile b/tools/keytools/Makefile
index d1c387b5a..e1e3be552 100644
--- a/tools/keytools/Makefile
+++ b/tools/keytools/Makefile
@@ -17,7 +17,8 @@ LDFLAGS =
 OBJDIR = ./
 LIBS =
 
-ML_DSA_LEVEL?=2
+ML_DSA_LEVEL?=5
+CFLAGS+=-DML_DSA_LEVEL=$(ML_DSA_LEVEL)
 
 LMS_LEVELS?=1
 LMS_HEIGHT?=10
@@ -33,6 +34,9 @@ CFLAGS +=-DWOLFBOOT_SIGN_LMS -DWOLFSSL_HAVE_LMS \
 # LMS flags
 CFLAGS +=-DWOLFSSL_WC_LMS
 
+# ML_DSA flags
+CFLAGS +=-DWOLFSSL_HAVE_DILITHIUM
+
 # XMSS flags
 CFLAGS +=-DWOLFBOOT_SIGN_XMSS -DWOLFSSL_HAVE_XMSS \
 		 -D"IMAGE_SIGNATURE_SIZE"=$(IMAGE_SIGNATURE_SIZE) \
@@ -111,10 +115,6 @@ OBJS_REAL+=\
 	$(WOLFDIR)/wolfcrypt/src/wc_xmss_impl.o
 OBJS_REAL+=$(WOLFDIR)/wolfcrypt/src/dilithium.o
 
-CFLAGS += -D"WOLFBOOT_SIGN_ML_DSA" \
-            -D"IMAGE_SIGNATURE_SIZE"=$(IMAGE_SIGNATURE_SIZE) \
-            -D"ML_DSA_LEVEL"=$(ML_DSA_LEVEL)
-
 OBJS_VIRT=$(addprefix $(OBJDIR), $(notdir $(OBJS_REAL)))
 vpath %.c $(WOLFDIR)/wolfcrypt/src/
 vpath %.c $(WOLFBOOTDIR)/src/
diff --git a/tools/keytools/keygen.c b/tools/keytools/keygen.c
index b517fdcef..60c253aef 100644
--- a/tools/keytools/keygen.c
+++ b/tools/keytools/keygen.c
@@ -451,15 +451,35 @@ static uint32_t get_pubkey_size(uint32_t keyType)
         case KEYGEN_XMSS:
             size = KEYSTORE_PUBKEY_SIZE_XMSS;
             break;
-#ifdef KEYSTORE_PUBKEY_SIZE_ML_DSA
         case KEYGEN_ML_DSA:
-            size = KEYSTORE_PUBKEY_SIZE_ML_DSA;
+        {
+            char *env_ml_dsa_level = getenv("ML_DSA_LEVEL");
+            if (env_ml_dsa_level == NULL) {
+                fprintf(stderr, "warning: ML_DSA_LEVEL environment variable"
+                       " not set, assuming level 2\n");
+                size = ML_DSA_L2_PUBKEY_SIZE;
+            } else {
+                int level = atoi(env_ml_dsa_level);
+                switch (level) {
+                    case 2:
+                        size = ML_DSA_L2_PUBKEY_SIZE;
+                        break;
+                    case 3:
+                        size = ML_DSA_L3_PUBKEY_SIZE;
+                        break;
+                    case 5:
+                        size = ML_DSA_L5_PUBKEY_SIZE;
+                        break;
+                    default:
+                        fprintf(stderr, "error: invalid ML_DSA_LEVEL: %d\n", level);
+                        exit(1);
+                }
+            }
             break;
-#endif
         default:
             size = 0;
+        }
     }
-
     return size;
 }
 
@@ -520,7 +540,6 @@ void keystore_add(uint32_t ktype, uint8_t *key, uint32_t sz, const char *keyfile
 }
 
 
-#if !defined(NO_RSA) && defined(WOLFSSL_KEY_GEN)
 static void keygen_rsa(const char *keyfile, int kbits, uint32_t id_mask)
 {
     RsaKey k;
@@ -570,9 +589,7 @@ static void keygen_rsa(const char *keyfile, int kbits, uint32_t id_mask)
     else if (kbits == 4096)
         keystore_add(KEYGEN_RSA4096, pub_der, publen, keyfile, id_mask);
 }
-#endif
 
-#ifdef HAVE_ECC
 #define MAX_ECC_KEY_SIZE 66
 
 static void keygen_ecc(const char *priv_fname, uint16_t ecc_key_size,
@@ -676,10 +693,8 @@ static void keygen_ecc(const char *priv_fname, uint16_t ecc_key_size,
     else if (ecc_key_size == 66)
         keystore_add(KEYGEN_ECC521, k_buffer, 2 * ecc_key_size, priv_fname, id_mask);
 }
-#endif
 
 
-#ifdef HAVE_ED25519
 static void keygen_ed25519(const char *privkey, uint32_t id_mask)
 {
     ed25519_key k;
@@ -716,9 +731,7 @@ static void keygen_ed25519(const char *privkey, uint32_t id_mask)
 
     keystore_add(KEYGEN_ED25519, pub, ED25519_PUB_KEY_SIZE, privkey, id_mask);
 }
-#endif
 
-#ifdef HAVE_ED448
 static void keygen_ed448(const char *privkey, uint32_t id_mask)
 {
     ed448_key k;
@@ -755,9 +768,7 @@ static void keygen_ed448(const char *privkey, uint32_t id_mask)
 
     keystore_add(KEYGEN_ED448, pub, ED448_PUB_KEY_SIZE, privkey, id_mask);
 }
-#endif
 
-#if defined(WOLFSSL_HAVE_LMS)
 #include "../lms/lms_common.h"
 
 static void keygen_lms(const char *priv_fname, uint32_t id_mask)
@@ -844,9 +855,7 @@ static void keygen_lms(const char *priv_fname, uint32_t id_mask)
 
     wc_LmsKey_Free(&key);
 }
-#endif /* if defined(WOLFSSL_HAVE_LMS) */
 
-#if defined(WOLFSSL_HAVE_XMSS)
 #include "../xmss/xmss_common.h"
 
 static void keygen_xmss(const char *priv_fname, uint32_t id_mask)
@@ -942,9 +951,7 @@ static void keygen_xmss(const char *priv_fname, uint32_t id_mask)
 
     wc_XmssKey_Free(&key);
 }
-#endif /* if defined(WOLFSSL_HAVE_XMSS) */
 
-#if defined(WOLFSSL_WC_DILITHIUM)
 
 static void keygen_ml_dsa(const char *priv_fname, uint32_t id_mask)
 {
@@ -957,6 +964,13 @@ static void keygen_ml_dsa(const char *priv_fname, uint32_t id_mask)
     word32   pub_len = 0;
     int      ml_dsa_priv_len = 0;
     int      ml_dsa_pub_len = 0;
+    int      ml_dsa_level = ML_DSA_LEVEL;
+    char *   env_ml_dsa_level = getenv("ML_DSA_LEVEL");
+    if (env_ml_dsa_level != NULL) {
+        ml_dsa_level = atoi(env_ml_dsa_level);
+    }
+
+    fprintf(stderr, "info: using DSA level %d\n", ml_dsa_level);
 
     ret = wc_MlDsaKey_Init(&key, NULL, INVALID_DEVID);
     if (ret != 0) {
@@ -964,10 +978,10 @@ static void keygen_ml_dsa(const char *priv_fname, uint32_t id_mask)
         exit(1);
     }
 
-    ret = wc_MlDsaKey_SetParams(&key, ML_DSA_LEVEL);
+    ret = wc_MlDsaKey_SetParams(&key, ml_dsa_level);
     if (ret != 0) {
         fprintf(stderr, "error: wc_MlDsaKey_SetParams(%d) returned %d\n",
-                ML_DSA_LEVEL, ret);
+                ml_dsa_level, ret);
         exit(1);
     }
 
@@ -985,6 +999,7 @@ static void keygen_ml_dsa(const char *priv_fname, uint32_t id_mask)
                 ret);
         exit(1);
     }
+    printf("info: ml-dsa public key length: %d\n", ml_dsa_pub_len);
 
     /* Get the ML-DSA private key length. This API returns
      * the public + private length. */
@@ -994,6 +1009,7 @@ static void keygen_ml_dsa(const char *priv_fname, uint32_t id_mask)
                 ret);
         exit(1);
     }
+    printf("info: ml-dsa private key length: %d\n", ml_dsa_priv_len);
 
     if (ml_dsa_priv_len <= ml_dsa_pub_len) {
         printf("error: ml-dsa: unexpected key lengths: %d, %d",
@@ -1026,9 +1042,9 @@ static void keygen_ml_dsa(const char *priv_fname, uint32_t id_mask)
         exit(1);
     }
 
-    if (pub_len != sizeof(pub)) {
+    if ((int)pub_len != ml_dsa_pub_len) {
         fprintf(stderr, "error: wc_MlDsaKey_ExportPubRaw returned pub_len=%d, " \
-                        "expected %zu\n", pub_len, sizeof(pub));
+                        "expected %d\n", pub_len, ml_dsa_pub_len);
         exit(1);
     }
 
@@ -1050,14 +1066,13 @@ static void keygen_ml_dsa(const char *priv_fname, uint32_t id_mask)
     fwrite(pub, pub_len, 1, fpriv);
     fclose(fpriv);
 
-    keystore_add(KEYGEN_ML_DSA, pub, KEYSTORE_PUBKEY_SIZE_ML_DSA,
+    keystore_add(KEYGEN_ML_DSA, pub, pub_len,
                  priv_fname, id_mask);
 
     wc_MlDsaKey_Free(&key);
     free(priv);
     priv = NULL;
 }
-#endif /* if defined(WOLFSSL_WC_DILITHIUM) */
 
 static void key_gen_check(const char *kfilename)
 {
diff --git a/tools/keytools/user_settings.h b/tools/keytools/user_settings.h
index ca66ac942..0759d1a3f 100644
--- a/tools/keytools/user_settings.h
+++ b/tools/keytools/user_settings.h
@@ -74,19 +74,17 @@
 #undef  NO_SHA256
 
 /* ML-DSA (dilithium) */
-#if defined(WOLFBOOT_SIGN_ML_DSA)
-#   define HAVE_DILITHIUM
-#   define WOLFSSL_WC_DILITHIUM
-#   define WOLFSSL_EXPERIMENTAL_SETTINGS
-    /* Wolfcrypt builds ML-DSA (dilithium) to the FIPS 204 final
-     * standard by default. Uncomment this if you want the draft
-     * version instead. */
-    #if 0
-    #define WOLFSSL_DILITHIUM_FIPS204_DRAFT
-    #endif
-    /* dilithium needs these sha functions. */
-#   define WOLFSSL_SHAKE128
-#endif /* WOLFBOOT_SIGN_ML_DSA */
+#define HAVE_DILITHIUM
+#define WOLFSSL_WC_DILITHIUM
+#define WOLFSSL_EXPERIMENTAL_SETTINGS
+/* Wolfcrypt builds ML-DSA (dilithium) to the FIPS 204 final
+* standard by default. Uncomment this if you want the draft
+* version instead. */
+#if 0
+  #define WOLFSSL_DILITHIUM_FIPS204_DRAFT
+#endif
+/* dilithium needs these sha functions. */
+#define WOLFSSL_SHAKE128
 
 /* ASN */
 #define WOLFSSL_ASN_TEMPLATE
diff --git a/tools/test-renode.mk b/tools/test-renode.mk
index a3b00e294..de40b3a66 100644
--- a/tools/test-renode.mk
+++ b/tools/test-renode.mk
@@ -36,8 +36,6 @@ else
 	SIGN_TOOL?=$(WOLFBOOT_ROOT)/tools/keytools/sign
 endif
 
-SIGN_ENV=IMAGE_HEADER_SIZE=$(IMAGE_HEADER_SIZE) WOLFBOOT_SECTOR_SIZE=$(WOLFBOOT_SECTOR_SIZE)
-
 ifeq ($(TARGET),stm32f7)
   RENODE_CONFIG=tools/renode/stm32f746_wolfboot.resc
   POFF=393211