From c2f86d0a0f43aceaede51eab5f4eea63718a1095 Mon Sep 17 00:00:00 2001
From: Trevor Bentley <trevor@yubico.com>
Date: Tue, 24 Oct 2017 15:59:44 +0200
Subject: [PATCH] Move YK4 insecure on-chip key generation prevention from
 yubico-piv-tool to libykpiv

---
 lib/util.c             | 18 ++++++++++++++++++
 tool/yubico-piv-tool.c | 16 ----------------
 2 files changed, 18 insertions(+), 16 deletions(-)

diff --git a/lib/util.c b/lib/util.c
index 4ece743..c0298df 100644
--- a/lib/util.c
+++ b/lib/util.c
@@ -699,6 +699,7 @@ ykpiv_rc ykpiv_util_generate_key(ykpiv_state *state, uint8_t slot, uint8_t algor
   ykpiv_rc res = YKPIV_OK;
   unsigned char in_data[11];
   unsigned char *in_ptr = in_data;
+  char version[7];
   unsigned char data[1024];
   unsigned char templ[] = { 0, YKPIV_INS_GENERATE_ASYMMETRIC, 0, 0 };
   unsigned long recv_len = sizeof(data);
@@ -710,6 +711,23 @@ ykpiv_rc ykpiv_util_generate_key(ykpiv_state *state, uint8_t slot, uint8_t algor
   uint8_t *ptr_point = NULL;
   size_t  cb_point = 0;
 
+  if (ykpiv_util_devicemodel(state) == DEVTYPE_YK4 && (algorithm == YKPIV_ALGO_RSA1024 || algorithm == YKPIV_ALGO_RSA2048)) {
+    if ((res = ykpiv_get_version(state, version, sizeof(version))) == YKPIV_OK) {
+      int major, minor, build;
+      fprintf(stderr, "version: %s\n", version);
+      int match = sscanf(version, "%d.%d.%d", &major, &minor, &build);
+      if (match == 3 && major == 4 && (minor < 3 || (minor == 3 && build < 5))) {
+        fprintf(stderr, "On-chip RSA key generation on this YubiKey has been blocked.\n");
+        fprintf(stderr, "Please see https://yubi.co/ysa201701/ for details.\n");
+        res = YKPIV_NOT_SUPPORTED;
+        goto Cleanup;
+      }
+    } else {
+      fprintf(stderr, "Failed to get device version.\n");
+      goto Cleanup;
+    }
+  }
+
   switch (algorithm) {
   case YKPIV_ALGO_RSA1024:
   case YKPIV_ALGO_RSA2048:
diff --git a/tool/yubico-piv-tool.c b/tool/yubico-piv-tool.c
index d9801f9..4567b16 100644
--- a/tool/yubico-piv-tool.c
+++ b/tool/yubico-piv-tool.c
@@ -141,22 +141,6 @@ static bool generate_key(ykpiv_state *state, enum enum_slot slot,
   size_t mod_len = 0;
   size_t exp_len = 0;
   size_t point_len = 0;
-  char version[7];
-
-  if(algorithm == algorithm_arg_RSA1024 || algorithm == algorithm_arg_RSA2048) {
-    if(ykpiv_get_version(state, version, sizeof(version)) == YKPIV_OK) {
-      int major, minor, build;
-      int match = sscanf(version, "%d.%d.%d", &major, &minor, &build);
-      if(match == 3 && major == 4 && (minor < 3 || (minor == 3 && build < 5))) {
-        fprintf(stderr, "On-chip RSA key generation on this YubiKey has been blocked.\n");
-        fprintf(stderr, "Please see https://yubi.co/ysa201701/ for details.\n");
-        return false;
-      }
-    } else {
-      fprintf(stderr, "Failed to communicate.\n");
-      return false;
-    }
-  }
 
   key = get_slot_hex(slot);