]> git.proxmox.com Git - efi-boot-shim.git/blobdiff - MokManager.c
Import upstream version 0.9+1474479173.6c180c6
[efi-boot-shim.git] / MokManager.c
index d764bf6e7509b32a69d173a701221a57a55d60f2..ebc85db98beab50b2994e1fbd261b4b76bce8744 100644 (file)
@@ -1,20 +1,35 @@
 #include <efi.h>
 #include <efilib.h>
+#include <stdarg.h>
 #include <Library/BaseCryptLib.h>
 #include <openssl/x509.h>
+#include <openssl/x509v3.h>
+#include <openssl/asn1.h>
+#include <openssl/bn.h>
 #include "shim.h"
-#include "signature.h"
 #include "PeImage.h"
+#include "PasswordCrypt.h"
 
-#define PASSWORD_MAX 16
-#define PASSWORD_MIN 8
-#define SB_PASSWORD_LEN 8
+#include "guid.h"
+#include "console.h"
+#include "variables.h"
+#include "simple_file.h"
+#include "efiauthenticated.h"
+
+#define PASSWORD_MAX 256
+#define PASSWORD_MIN 1
+#define SB_PASSWORD_LEN 16
+
+#define NAME_LINE_MAX 70
 
 #ifndef SHIM_VENDOR
 #define SHIM_VENDOR L"Shim"
 #endif
 
-#define EFI_VARIABLE_APPEND_WRITE 0x00000040
+EFI_GUID SHIM_LOCK_GUID = { 0x605dab50, 0xe046, 0x4300, {0xab, 0xb6, 0x3d, 0xd8, 0x10, 0xdd, 0x8b, 0x23} };
+EFI_GUID EFI_CERT_SHA224_GUID = { 0xb6e5233, 0xa65c, 0x44c9, {0x94, 0x7, 0xd9, 0xab, 0x83, 0xbf, 0xc8, 0xbd} };
+EFI_GUID EFI_CERT_SHA384_GUID = { 0xff3e5307, 0x9fd0, 0x48c9, {0x85, 0xf1, 0x8a, 0xd5, 0x6c, 0x70, 0x1e, 0x1} };
+EFI_GUID EFI_CERT_SHA512_GUID = { 0x93e0fae, 0xa6c4, 0x4f50, {0x9f, 0x1b, 0xd4, 0x1e, 0x2b, 0x89, 0xc1, 0x9a} };
 
 #define CERT_STRING L"Select an X509 certificate to enroll:\n\n"
 #define HASH_STRING L"Select a file to trust:\n\n"
@@ -31,25 +46,20 @@ struct menu_item {
 typedef struct {
        UINT32 MokSize;
        UINT8 *Mok;
+       EFI_GUID Type;
 } __attribute__ ((packed)) MokListNode;
 
 typedef struct {
        UINT32 MokSBState;
        UINT32 PWLen;
-       CHAR16 Password[PASSWORD_MAX];
+       CHAR16 Password[SB_PASSWORD_LEN];
 } __attribute__ ((packed)) MokSBvar;
 
-static EFI_INPUT_KEY get_keystroke (void)
-{
-       EFI_INPUT_KEY key;
-       UINTN EventIndex;
-
-       uefi_call_wrapper(BS->WaitForEvent, 3, 1, &ST->ConIn->WaitForKey,
-                         &EventIndex);
-       uefi_call_wrapper(ST->ConIn->ReadKeyStroke, 2, ST->ConIn, &key);
-
-       return key;
-}
+typedef struct {
+       UINT32 MokDBState;
+       UINT32 PWLen;
+       CHAR16 Password[SB_PASSWORD_LEN];
+} __attribute__ ((packed)) MokDBvar;
 
 static EFI_STATUS get_sha1sum (void *Data, int DataSize, UINT8 *hash)
 {
@@ -61,24 +71,24 @@ static EFI_STATUS get_sha1sum (void *Data, int DataSize, UINT8 *hash)
        ctx = AllocatePool(ctxsize);
 
        if (!ctx) {
-               Print(L"Unable to allocate memory for hash context\n");
+               console_notify(L"Unable to allocate memory for hash context");
                return EFI_OUT_OF_RESOURCES;
        }
 
        if (!Sha1Init(ctx)) {
-               Print(L"Unable to initialise hash\n");
+               console_notify(L"Unable to initialise hash");
                status = EFI_OUT_OF_RESOURCES;
                goto done;
        }
 
        if (!(Sha1Update(ctx, Data, DataSize))) {
-               Print(L"Unable to generate hash\n");
+               console_notify(L"Unable to generate hash");
                status = EFI_OUT_OF_RESOURCES;
                goto done;
        }
 
        if (!(Sha1Final(ctx, hash))) {
-               Print(L"Unable to finalise hash\n");
+               console_notify(L"Unable to finalise hash");
                status = EFI_OUT_OF_RESOURCES;
                goto done;
        }
@@ -88,44 +98,154 @@ done:
        return status;
 }
 
+static BOOLEAN is_sha2_hash (EFI_GUID Type)
+{
+       EFI_GUID Sha224 = EFI_CERT_SHA224_GUID;
+       EFI_GUID Sha256 = EFI_CERT_SHA256_GUID;
+       EFI_GUID Sha384 = EFI_CERT_SHA384_GUID;
+       EFI_GUID Sha512 = EFI_CERT_SHA512_GUID;
+
+       if (CompareGuid(&Type, &Sha224) == 0)
+               return TRUE;
+       else if (CompareGuid(&Type, &Sha256) == 0)
+               return TRUE;
+       else if (CompareGuid(&Type, &Sha384) == 0)
+               return TRUE;
+       else if (CompareGuid(&Type, &Sha512) == 0)
+               return TRUE;
+
+       return FALSE;
+}
+
+static UINT32 sha_size (EFI_GUID Type)
+{
+       EFI_GUID Sha1 = EFI_CERT_SHA1_GUID;
+       EFI_GUID Sha224 = EFI_CERT_SHA224_GUID;
+       EFI_GUID Sha256 = EFI_CERT_SHA256_GUID;
+       EFI_GUID Sha384 = EFI_CERT_SHA384_GUID;
+       EFI_GUID Sha512 = EFI_CERT_SHA512_GUID;
+
+       if (CompareGuid(&Type, &Sha1) == 0)
+               return SHA1_DIGEST_SIZE;
+       else if (CompareGuid(&Type, &Sha224) == 0)
+               return SHA224_DIGEST_LENGTH;
+       else if (CompareGuid(&Type, &Sha256) == 0)
+               return SHA256_DIGEST_SIZE;
+       else if (CompareGuid(&Type, &Sha384) == 0)
+               return SHA384_DIGEST_LENGTH;
+       else if (CompareGuid(&Type, &Sha512) == 0)
+               return SHA512_DIGEST_LENGTH;
+
+       return 0;
+}
+
+static BOOLEAN is_valid_siglist (EFI_GUID Type, UINT32 SigSize)
+{
+       EFI_GUID CertType = X509_GUID;
+       UINT32 hash_sig_size;
+
+       if (CompareGuid (&Type, &CertType) == 0 && SigSize != 0)
+               return TRUE;
+
+       if (!is_sha2_hash (Type))
+               return FALSE;
+
+       hash_sig_size = sha_size (Type) + sizeof(EFI_GUID);
+       if (SigSize != hash_sig_size)
+               return FALSE;
+
+       return TRUE;
+}
+
+static UINT32 count_keys(void *Data, UINTN DataSize)
+{
+       EFI_SIGNATURE_LIST *CertList = Data;
+       UINTN dbsize = DataSize;
+       UINT32 MokNum = 0;
+       void *end = Data + DataSize;
+
+       while ((dbsize > 0) && (dbsize >= CertList->SignatureListSize)) {
+
+               /* Use ptr arithmetics to ensure bounded access. Do not allow 0
+                * SignatureListSize that will cause endless loop.
+                */
+               if ((void *)(CertList + 1) > end || CertList->SignatureListSize == 0) {
+                       console_notify(L"Invalid MOK detected! Ignoring MOK List.");
+                       return 0;
+               }
+
+               if (CertList->SignatureListSize == 0 ||
+                   CertList->SignatureListSize <= CertList->SignatureSize) {
+                       console_errorbox(L"Corrupted signature list");
+                       return 0;
+               }
+
+               if (!is_valid_siglist(CertList->SignatureType, CertList->SignatureSize)) {
+                       console_errorbox(L"Invalid signature list found");
+                       return 0;
+               }
+
+               MokNum++;
+               dbsize -= CertList->SignatureListSize;
+               CertList = (EFI_SIGNATURE_LIST *) ((UINT8 *) CertList +
+                                                 CertList->SignatureListSize);
+       }
+
+       return MokNum;
+}
+
 static MokListNode *build_mok_list(UINT32 num, void *Data, UINTN DataSize) {
        MokListNode *list;
        EFI_SIGNATURE_LIST *CertList = Data;
        EFI_SIGNATURE_DATA *Cert;
-       EFI_GUID CertType = EfiCertX509Guid;
-       EFI_GUID HashType = EfiHashSha256Guid;
+       EFI_GUID CertType = X509_GUID;
        UINTN dbsize = DataSize;
        UINTN count = 0;
+       void *end = Data + DataSize;
 
        list = AllocatePool(sizeof(MokListNode) * num);
 
        if (!list) {
-               Print(L"Unable to allocate MOK list\n");
+               console_notify(L"Unable to allocate MOK list");
                return NULL;
        }
 
        while ((dbsize > 0) && (dbsize >= CertList->SignatureListSize)) {
-               if ((CompareGuid (&CertList->SignatureType, &CertType) != 0) &&
-                   (CompareGuid (&CertList->SignatureType, &HashType) != 0)) {
-                       dbsize -= CertList->SignatureListSize;
-                       CertList = (EFI_SIGNATURE_LIST *)((UINT8 *) CertList +
-                                                 CertList->SignatureListSize);
-                       continue;
+               /* CertList out of bounds? */
+               if ((void *)(CertList + 1) > end || CertList->SignatureListSize == 0) {
+                       FreePool(list);
+                       return NULL;
                }
 
-               if ((CompareGuid (&CertList->SignatureType, &HashType) == 0) &&
-                   (CertList->SignatureSize != 48)) {
-                       dbsize -= CertList->SignatureListSize;
-                       CertList = (EFI_SIGNATURE_LIST *)((UINT8 *) CertList +
-                                                 CertList->SignatureListSize);
-                       continue;
-               }
+               /* Omit the signature check here since we already did it
+                  in count_keys() */
 
                Cert = (EFI_SIGNATURE_DATA *) (((UINT8 *) CertList) +
                  sizeof (EFI_SIGNATURE_LIST) + CertList->SignatureHeaderSize);
 
-               list[count].MokSize = CertList->SignatureSize;
-               list[count].Mok = (void *)Cert->SignatureData;
+               /* Cert out of bounds? */
+               if ((void *)(Cert + 1) > end || CertList->SignatureSize <= sizeof(EFI_GUID)) {
+                       FreePool(list);
+                       return NULL;
+               }
+
+               list[count].Type = CertList->SignatureType;
+               if (CompareGuid (&CertList->SignatureType, &CertType) == 0) {
+                       list[count].MokSize = CertList->SignatureSize -
+                                             sizeof(EFI_GUID);
+                       list[count].Mok = (void *)Cert->SignatureData;
+               } else {
+                       list[count].MokSize = CertList->SignatureListSize -
+                                             sizeof(EFI_SIGNATURE_LIST);
+                       list[count].Mok = (void *)Cert;
+               }
+
+               /* MOK out of bounds? */
+               if (list[count].MokSize > (unsigned long)end -
+                                         (unsigned long)list[count].Mok) {
+                       FreePool(list);
+                       return NULL;
+               }
 
                count++;
                dbsize -= CertList->SignatureListSize;
@@ -136,349 +256,434 @@ static MokListNode *build_mok_list(UINT32 num, void *Data, UINTN DataSize) {
        return list;
 }
 
-static void print_x509_name (X509_NAME *X509Name, CHAR16 *name)
-{
-       char *str;
-
-       str = X509_NAME_oneline(X509Name, NULL, 0);
-       if (str) {
-               Print(L"  %s:\n    %a\n", name, str);
-               OPENSSL_free(str);
-       }
-}
-
-static const char *mon[12]= {
-"Jan","Feb","Mar","Apr","May","Jun",
-"Jul","Aug","Sep","Oct","Nov","Dec"
+typedef struct {
+       int nid;
+       CHAR16 *name;
+} NidName;
+
+static NidName nidname[] = {
+       {NID_commonName, L"CN"},
+       {NID_organizationName, L"O"},
+       {NID_countryName, L"C"},
+       {NID_stateOrProvinceName, L"ST"},
+       {NID_localityName, L"L"},
+       {-1, NULL}
 };
 
-static void print_x509_GENERALIZEDTIME_time (ASN1_TIME *time, CHAR16 *time_string)
+static CHAR16* get_x509_name (X509_NAME *X509Name)
 {
-       char *v;
-       int gmt = 0;
-       int i;
-       int y = 0,M = 0,d = 0,h = 0,m = 0,s = 0;
-       char *f = NULL;
-       int f_len = 0;
-
-       i=time->length;
-       v=(char *)time->data;
-
-       if (i < 12)
-               goto error;
-
-       if (v[i-1] == 'Z')
-               gmt=1;
-
-       for (i=0; i<12; i++) {
-               if ((v[i] > '9') || (v[i] < '0'))
-                       goto error;
-       }
+       CHAR16 name[NAME_LINE_MAX+1];
+       CHAR16 part[NAME_LINE_MAX+1];
+       char str[NAME_LINE_MAX];
+       int i, len, rest, first;
+
+       name[0] = '\0';
+       rest = NAME_LINE_MAX;
+       first = 1;
+       for (i = 0; nidname[i].name != NULL; i++) {
+               int add;
+               len = X509_NAME_get_text_by_NID (X509Name, nidname[i].nid,
+                                                str, NAME_LINE_MAX);
+               if (len <= 0)
+                       continue;
 
-       y = (v[0]-'0')*1000+(v[1]-'0')*100 + (v[2]-'0')*10+(v[3]-'0');
-       M = (v[4]-'0')*10+(v[5]-'0');
+               if (first)
+                       add = len + (int)StrLen(nidname[i].name) + 1;
+               else
+                       add = len + (int)StrLen(nidname[i].name) + 3;
 
-       if ((M > 12) || (M < 1))
-               goto error;
+               if (add > rest)
+                       continue;
 
-       d = (v[6]-'0')*10+(v[7]-'0');
-       h = (v[8]-'0')*10+(v[9]-'0');
-       m = (v[10]-'0')*10+(v[11]-'0');
-
-       if (time->length >= 14 &&
-           (v[12] >= '0') && (v[12] <= '9') &&
-           (v[13] >= '0') && (v[13] <= '9')) {
-               s = (v[12]-'0')*10+(v[13]-'0');
-               /* Check for fractions of seconds. */
-               if (time->length >= 15 && v[14] == '.') {
-                       int l = time->length;
-                       f = &v[14];     /* The decimal point. */
-                       f_len = 1;
-                       while (14 + f_len < l && f[f_len] >= '0' &&
-                              f[f_len] <= '9')
-                               ++f_len;
+               if (first) {
+                       SPrint(part, NAME_LINE_MAX * sizeof(CHAR16), L"%s=%a",
+                              nidname[i].name, str);
+               } else {
+                       SPrint(part, NAME_LINE_MAX * sizeof(CHAR16), L", %s=%a",
+                              nidname[i].name, str);
                }
+               StrCat(name, part);
+               rest -= add;
+               first = 0;
        }
 
-       SPrint(time_string, 0, L"%a %2d %02d:%02d:%02d%.*a %d%a",
-              mon[M-1], d, h, m, s, f_len, f, y, (gmt)?" GMT":"");
-error:
-       return;
-}
-
-static void print_x509_UTCTIME_time (ASN1_TIME *time, CHAR16 *time_string)
-{
-       char *v;
-       int gmt=0;
-       int i;
-       int y = 0,M = 0,d = 0,h = 0,m = 0,s = 0;
-
-       i=time->length;
-       v=(char *)time->data;
-
-       if (i < 10)
-               goto error;
-
-       if (v[i-1] == 'Z')
-               gmt=1;
-
-       for (i=0; i<10; i++)
-               if ((v[i] > '9') || (v[i] < '0'))
-                       goto error;
-
-       y = (v[0]-'0')*10+(v[1]-'0');
-
-       if (y < 50)
-               y+=100;
-
-       M = (v[2]-'0')*10+(v[3]-'0');
-
-       if ((M > 12) || (M < 1))
-               goto error;
-
-       d = (v[4]-'0')*10+(v[5]-'0');
-       h = (v[6]-'0')*10+(v[7]-'0');
-       m = (v[8]-'0')*10+(v[9]-'0');
-
-       if (time->length >=12 &&
-           (v[10] >= '0') && (v[10] <= '9') &&
-           (v[11] >= '0') && (v[11] <= '9'))
-               s = (v[10]-'0')*10+(v[11]-'0');
+       if (rest >= 0 && rest < NAME_LINE_MAX)
+               return PoolPrint(L"%s", name);
 
-       SPrint(time_string, 0, L"%a %2d %02d:%02d:%02d %d%a",
-              mon[M-1], d, h, m, s, y+1900, (gmt)?" GMT":"");
-error:
-       return;
+       return NULL;
 }
 
-static void print_x509_time (ASN1_TIME *time, CHAR16 *name)
+static CHAR16* get_x509_time (ASN1_TIME *time)
 {
-       CHAR16 time_string[30];
-
-       if (time->type == V_ASN1_UTCTIME) {
-               print_x509_UTCTIME_time(time, time_string);
-       } else if (time->type == V_ASN1_GENERALIZEDTIME) {
-               print_x509_GENERALIZEDTIME_time(time, time_string);
-       } else {
-               time_string[0] = '\0';
-       }
-
-       Print(L"  %s:\n    %s\n", name, time_string);
+       BIO *bio = BIO_new (BIO_s_mem());
+       char str[30];
+       int len;
+
+       ASN1_TIME_print (bio, time);
+       len = BIO_read(bio, str, 29);
+       if (len < 0)
+               len = 0;
+       str[len] = '\0';
+       BIO_free (bio);
+
+       return PoolPrint(L"%a", str);
 }
 
-static void show_x509_info (X509 *X509Cert)
+static void show_x509_info (X509 *X509Cert, UINT8 *hash)
 {
        ASN1_INTEGER *serial;
        BIGNUM *bnser;
        unsigned char hexbuf[30];
        X509_NAME *X509Name;
        ASN1_TIME *time;
+       CHAR16 *issuer = NULL;
+       CHAR16 *subject = NULL;
+       CHAR16 *from = NULL;
+       CHAR16 *until = NULL;
+       EXTENDED_KEY_USAGE *extusage;
+       POOL_PRINT hash_string1;
+       POOL_PRINT hash_string2;
+       POOL_PRINT serial_string;
+       int fields = 0;
+       CHAR16 **text;
+       int i = 0;
+
+       ZeroMem(&hash_string1, sizeof(hash_string1));
+       ZeroMem(&hash_string2, sizeof(hash_string2));
+       ZeroMem(&serial_string, sizeof(serial_string));
 
        serial = X509_get_serialNumber(X509Cert);
        if (serial) {
                int i, n;
                bnser = ASN1_INTEGER_to_BN(serial, NULL);
                n = BN_bn2bin(bnser, hexbuf);
-               Print(L"  Serial Number:\n    ");
-               for (i = 0; i < n-1; i++) {
-                       Print(L"%02x:", hexbuf[i]);
+               for (i = 0; i < n; i++) {
+                       CatPrint(&serial_string, L"%02x:", hexbuf[i]);
                }
-               Print(L"%02x\n", hexbuf[n-1]);
        }
 
+       if (serial_string.str)
+               fields++;
+
        X509Name = X509_get_issuer_name(X509Cert);
        if (X509Name) {
-               print_x509_name(X509Name, L"Issuer");
+               issuer = get_x509_name(X509Name);
+               if (issuer)
+                       fields++;
        }
 
        X509Name = X509_get_subject_name(X509Cert);
        if (X509Name) {
-               print_x509_name(X509Name, L"Subject");
+               subject = get_x509_name(X509Name);
+               if (subject)
+                       fields++;
        }
 
        time = X509_get_notBefore(X509Cert);
        if (time) {
-               print_x509_time(time, L"Validity from");
+               from = get_x509_time(time);
+               if (from)
+                       fields++;
        }
 
        time = X509_get_notAfter(X509Cert);
        if (time) {
-               print_x509_time(time, L"Validity till");
+               until = get_x509_time(time);
+               if (until)
+                       fields++;
        }
-}
 
-static void show_mok_info (void *Mok, UINTN MokSize)
-{
-       EFI_STATUS efi_status;
-       UINT8 hash[SHA1_DIGEST_SIZE];
-       unsigned int i;
-       X509 *X509Cert;
+       for (i=0; i<10; i++)
+               CatPrint(&hash_string1, L"%02x ", hash[i]);
+       for (i=10; i<20; i++)
+               CatPrint(&hash_string2, L"%02x ", hash[i]);
 
-       if (!Mok || MokSize == 0)
+       if (hash_string1.str)
+               fields++;
+
+       if (hash_string2.str)
+               fields++;
+
+       if (!fields)
                return;
 
-       if (MokSize != 48) {
-               if (X509ConstructCertificate(Mok, MokSize,
-                                (UINT8 **) &X509Cert) && X509Cert != NULL) {
-                       show_x509_info(X509Cert);
-                       X509_free(X509Cert);
-               } else {
-                       Print(L"  Not a valid X509 certificate: %x\n\n",
-                             ((UINT32 *)Mok)[0]);
-                       return;
-               }
+       i = 0;
 
-               efi_status = get_sha1sum(Mok, MokSize, hash);
+       extusage = X509_get_ext_d2i(X509Cert, NID_ext_key_usage, NULL, NULL);
+       text = AllocateZeroPool(sizeof(CHAR16 *) * (fields*3 + sk_ASN1_OBJECT_num(extusage) + 3));
 
-               if (efi_status != EFI_SUCCESS) {
-                       Print(L"Failed to compute MOK fingerprint\n");
-                       return;
-               }
+       if (extusage) {
+               int j = 0;
 
-               Print(L"  Fingerprint (SHA1):\n    ");
-               for (i = 0; i < SHA1_DIGEST_SIZE; i++) {
-                       Print(L" %02x", hash[i]);
-                       if (i % 10 == 9)
-                               Print(L"\n    ");
-               }
-       } else {
-               Print(L"SHA256 hash:\n   ");
-               for (i = 0; i < SHA256_DIGEST_SIZE; i++) {
-                       Print(L" %02x", ((UINT8 *)Mok)[i]);
-                       if (i % 10 == 9)
-                               Print(L"\n   ");
+               text[i++] = StrDuplicate(L"[Extended Key Usage]");
+
+               for (j = 0; j < sk_ASN1_OBJECT_num(extusage); j++) {
+                       POOL_PRINT extkeyusage;
+                       ASN1_OBJECT *obj = sk_ASN1_OBJECT_value(extusage, j);
+                       int buflen = 80;
+                       char buf[buflen];
+
+                       ZeroMem(&extkeyusage, sizeof(extkeyusage));
+
+                       OBJ_obj2txt(buf, buflen, obj, 0);
+                       CatPrint(&extkeyusage, L"OID: %a", buf);
+                       text[i++] = StrDuplicate(extkeyusage.str);
+                       FreePool(extkeyusage.str);
                }
-               Print(L"\n");
+               text[i++] = StrDuplicate(L"");
+               EXTENDED_KEY_USAGE_free(extusage);
        }
 
-       Print(L"\n");
+       if (serial_string.str) {
+               text[i++] = StrDuplicate(L"[Serial Number]");
+               text[i++] = serial_string.str;
+               text[i++] = StrDuplicate(L"");
+       }
+       if (issuer) {
+               text[i++] = StrDuplicate(L"[Issuer]");
+               text[i++] = issuer;
+               text[i++] = StrDuplicate(L"");
+       }
+       if (subject) {
+               text[i++] = StrDuplicate(L"[Subject]");
+               text[i++] = subject;
+               text[i++] = StrDuplicate(L"");
+       }
+       if (from) {
+               text[i++] = StrDuplicate(L"[Valid Not Before]");
+               text[i++] = from;
+               text[i++] = StrDuplicate(L"");
+       }
+       if (until) {
+               text[i++] = StrDuplicate(L"[Valid Not After]");
+               text[i++] = until;
+               text[i++] = StrDuplicate(L"");
+       }
+       if (hash_string1.str) {
+               text[i++] = StrDuplicate(L"[Fingerprint]");
+               text[i++] = hash_string1.str;
+       }
+       if (hash_string2.str) {
+               text[i++] = hash_string2.str;
+               text[i++] = StrDuplicate(L"");
+       }
+       text[i] = NULL;
+
+       console_print_box(text, -1);
+
+       for (i=0; text[i] != NULL; i++)
+               FreePool(text[i]);
+
+       FreePool(text);
 }
 
-static INTN get_number ()
+static void show_sha_digest (EFI_GUID Type, UINT8 *hash)
 {
-       EFI_INPUT_KEY input_key;
-       CHAR16 input[10];
-       int count = 0;
+       EFI_GUID Sha1 = EFI_CERT_SHA1_GUID;
+       EFI_GUID Sha224 = EFI_CERT_SHA224_GUID;
+       EFI_GUID Sha256 = EFI_CERT_SHA256_GUID;
+       EFI_GUID Sha384 = EFI_CERT_SHA384_GUID;
+       EFI_GUID Sha512 = EFI_CERT_SHA512_GUID;
+       CHAR16 *text[5];
+       POOL_PRINT hash_string1;
+       POOL_PRINT hash_string2;
+       int i;
+       int length;
+
+       if (CompareGuid(&Type, &Sha1) == 0) {
+               length = SHA1_DIGEST_SIZE;
+               text[0] = L"SHA1 hash";
+       } else if (CompareGuid(&Type, &Sha224) == 0) {
+               length = SHA224_DIGEST_LENGTH;
+               text[0] = L"SHA224 hash";
+       } else if (CompareGuid(&Type, &Sha256) == 0) {
+               length = SHA256_DIGEST_SIZE;
+               text[0] = L"SHA256 hash";
+       } else if (CompareGuid(&Type, &Sha384) == 0) {
+               length = SHA384_DIGEST_LENGTH;
+               text[0] = L"SHA384 hash";
+       } else if (CompareGuid(&Type, &Sha512) == 0) {
+               length = SHA512_DIGEST_LENGTH;
+               text[0] = L"SHA512 hash";
+       } else {
+               return;
+       }
 
-       do {
-               input_key = get_keystroke();
+       ZeroMem(&hash_string1, sizeof(hash_string1));
+       ZeroMem(&hash_string2, sizeof(hash_string2));
 
-               if ((input_key.UnicodeChar < '0' ||
-                    input_key.UnicodeChar > '9' ||
-                    count >= 10) &&
-                   input_key.UnicodeChar != CHAR_BACKSPACE) {
-                       continue;
-               }
+       text[1] = L"";
 
-               if (count == 0 && input_key.UnicodeChar == CHAR_BACKSPACE)
-                       continue;
+       for (i=0; i<length/2; i++)
+               CatPrint(&hash_string1, L"%02x ", hash[i]);
+       for (i=length/2; i<length; i++)
+               CatPrint(&hash_string2, L"%02x ", hash[i]);
 
-               Print(L"%c", input_key.UnicodeChar);
+       text[2] = hash_string1.str;
+       text[3] = hash_string2.str;
+       text[4] = NULL;
 
-               if (input_key.UnicodeChar == CHAR_BACKSPACE) {
-                       input[--count] = '\0';
-                       continue;
-               }
+       console_print_box(text, -1);
 
-               input[count++] = input_key.UnicodeChar;
-       } while (input_key.UnicodeChar != CHAR_CARRIAGE_RETURN);
+       if (hash_string1.str)
+               FreePool(hash_string1.str);
 
-       if (count == 0)
-               return -1;
+       if (hash_string2.str)
+               FreePool(hash_string2.str);
+}
+
+static void show_efi_hash (EFI_GUID Type, void *Mok, UINTN MokSize)
+{
+       UINTN sig_size;
+       UINTN hash_num;
+       UINT8 *hash;
+       CHAR16 **menu_strings;
+       UINTN key_num = 0;
+       UINTN i;
+
+       sig_size = sha_size(Type) + sizeof(EFI_GUID);
+       if ((MokSize % sig_size) != 0) {
+               console_errorbox(L"Corrupted Hash List");
+               return;
+       }
+       hash_num = MokSize / sig_size;
+
+       if (hash_num == 1) {
+               hash = (UINT8 *)Mok + sizeof(EFI_GUID);
+               show_sha_digest(Type, hash);
+               return;
+       }
+
+       menu_strings = AllocateZeroPool(sizeof(CHAR16 *) * (hash_num + 2));
+       if (!menu_strings) {
+               console_errorbox(L"Out of Resources");
+               return;
+       }
+       for (i=0; i<hash_num; i++) {
+               menu_strings[i] = PoolPrint(L"View hash %d", i);
+       }
+       menu_strings[i] = StrDuplicate(L"Back");
+       menu_strings[i+1] = NULL;
+
+       while (key_num < hash_num) {
+               key_num = console_select((CHAR16 *[]){ L"[Hash List]", NULL },
+                                        menu_strings, key_num);
+
+               if (key_num < 0 || key_num >= hash_num)
+                       break;
+
+               hash = (UINT8 *)Mok + sig_size*key_num + sizeof(EFI_GUID);
+               show_sha_digest(Type, hash);
+       }
 
-       input[count] = '\0';
+       for (i=0; menu_strings[i] != NULL; i++)
+               FreePool(menu_strings[i]);
 
-       return (INTN)Atoi(input);
+       FreePool(menu_strings);
 }
 
-static UINT8 list_keys (void *MokNew, UINTN MokNewSize)
+static void show_mok_info (EFI_GUID Type, void *Mok, UINTN MokSize)
 {
-       UINT32 MokNum = 0;
-       MokListNode *keys = NULL;
-       INTN key_num = 0;
-       UINT8 initial = 1;
-       EFI_SIGNATURE_LIST *CertList = MokNew;
-       EFI_GUID CertType = EfiCertX509Guid;
-       EFI_GUID HashType = EfiHashSha256Guid;
-       UINTN dbsize = MokNewSize;
-
-       if (MokNewSize < (sizeof(EFI_SIGNATURE_LIST) +
-                         sizeof(EFI_SIGNATURE_DATA))) {
-               Print(L"No keys\n");
-               Pause();
-               return 0;
-       }
+       EFI_STATUS efi_status;
+       EFI_GUID CertType = X509_GUID;
 
-       while ((dbsize > 0) && (dbsize >= CertList->SignatureListSize)) {
-               if ((CompareGuid (&CertList->SignatureType, &CertType) != 0) &&
-                   (CompareGuid (&CertList->SignatureType, &HashType) != 0)) {
-                       Print(L"Doesn't look like a key or hash\n");
-                       dbsize -= CertList->SignatureListSize;
-                       CertList = (EFI_SIGNATURE_LIST *) ((UINT8 *) CertList +
-                                                 CertList->SignatureListSize);
-                       continue;
-               }
+       if (!Mok || MokSize == 0)
+               return;
 
-               if ((CompareGuid (&CertList->SignatureType, &CertType) != 0) &&
-                   (CertList->SignatureSize != 48)) {
-                       Print(L"Doesn't look like a valid hash\n");
-                       dbsize -= CertList->SignatureListSize;
-                       CertList = (EFI_SIGNATURE_LIST *) ((UINT8 *) CertList +
-                                                 CertList->SignatureListSize);
-                       continue;
+       if (CompareGuid (&Type, &CertType) == 0) {
+               UINT8 hash[SHA1_DIGEST_SIZE];
+               X509 *X509Cert;
+               efi_status = get_sha1sum(Mok, MokSize, hash);
+
+               if (efi_status != EFI_SUCCESS) {
+                       console_notify(L"Failed to compute MOK fingerprint");
+                       return;
                }
 
-               MokNum++;
-               dbsize -= CertList->SignatureListSize;
-               CertList = (EFI_SIGNATURE_LIST *) ((UINT8 *) CertList +
-                                                 CertList->SignatureListSize);
+               if (X509ConstructCertificate(Mok, MokSize,
+                                (UINT8 **) &X509Cert) && X509Cert != NULL) {
+                       show_x509_info(X509Cert, hash);
+                       X509_free(X509Cert);
+               } else {
+                       console_notify(L"Not a valid X509 certificate");
+                       return;
+               }
+       } else if (is_sha2_hash(Type)) {
+               show_efi_hash(Type, Mok, MokSize);
        }
+}
+
+static EFI_STATUS list_keys (void *KeyList, UINTN KeyListSize, CHAR16 *title)
+{
+       UINTN MokNum = 0;
+       MokListNode *keys = NULL;
+       UINT32 key_num = 0;
+       CHAR16 **menu_strings;
+       unsigned int i;
 
-       keys = build_mok_list(MokNum, MokNew, MokNewSize);
+       if (KeyListSize < (sizeof(EFI_SIGNATURE_LIST) +
+                          sizeof(EFI_SIGNATURE_DATA))) {
+               console_notify(L"No MOK keys found");
+               return EFI_NOT_FOUND;
+       }
 
+       MokNum = count_keys(KeyList, KeyListSize);
+       if (MokNum == 0) {
+               console_errorbox(L"Invalid key list");
+               return EFI_ABORTED;
+       }
+       keys = build_mok_list(MokNum, KeyList, KeyListSize);
        if (!keys) {
-               Print(L"Failed to construct key list in MokNew\n");
-               return 0;
+               console_errorbox(L"Failed to construct key list");
+               return EFI_ABORTED;
        }
 
-       do {
-               uefi_call_wrapper(ST->ConOut->ClearScreen, 1, ST->ConOut);
-               Print(L"Input the key number to show the details of the key or\n"
-                     L"type \'0\' to continue\n\n");
-               Print(L"%d key(s) in the new key list\n\n", MokNum);
-
-               if (key_num > MokNum) {
-                       Print(L"[Key %d]\n", key_num);
-                       Print(L"No such key\n\n");
-               } else if (initial != 1 && key_num > 0){
-                       Print(L"[Key %d]\n", key_num);
-                       show_mok_info(keys[key_num-1].Mok, keys[key_num-1].MokSize);
-               }
+       menu_strings = AllocateZeroPool(sizeof(CHAR16 *) * (MokNum + 2));
 
-               Print(L"Key Number: ");
+       if (!menu_strings)
+               return EFI_OUT_OF_RESOURCES;
+
+       for (i=0; i<MokNum; i++) {
+               menu_strings[i] = PoolPrint(L"View key %d", i);
+       }
+       menu_strings[i] = StrDuplicate(L"Continue");
 
-               key_num = get_number();
+       menu_strings[i+1] = NULL;
 
-               Print(L"\n\n");
+       while (key_num < MokNum) {
+               key_num = console_select((CHAR16 *[]){ title, NULL },
+                                        menu_strings, key_num);
 
-               if (key_num == -1)
-                       continue;
+               if (key_num < 0 || key_num >= MokNum)
+                       break;
+
+               show_mok_info(keys[key_num].Type, keys[key_num].Mok,
+                             keys[key_num].MokSize);
+       }
+
+       for (i=0; menu_strings[i] != NULL; i++)
+               FreePool(menu_strings[i]);
 
-               initial = 0;
-       } while (key_num != 0);
+       FreePool(menu_strings);
 
        FreePool(keys);
 
-       return 1;
+       return EFI_SUCCESS;
 }
 
-static UINT8 get_line (UINT32 *length, CHAR16 *line, UINT32 line_max, UINT8 show)
+static EFI_STATUS get_line (UINT32 *length, CHAR16 *line, UINT32 line_max, UINT8 show)
 {
        EFI_INPUT_KEY key;
-       int count = 0;
+       EFI_STATUS status;
+       unsigned int count = 0;
 
        do {
-               key = get_keystroke();
+               status = console_get_keystroke(&key);
+               if (EFI_ERROR (status)) {
+                       console_error(L"Failed to read the keystroke", status);
+                       *length = 0;
+                       return status;
+               }
 
                if ((count >= line_max &&
                     key.UnicodeChar != CHAR_BACKSPACE) ||
@@ -509,11 +714,11 @@ static UINT8 get_line (UINT32 *length, CHAR16 *line, UINT32 line_max, UINT8 show
 
        *length = count;
 
-       return 1;
+       return EFI_SUCCESS;
 }
 
-static EFI_STATUS compute_pw_hash (void *MokNew, UINTN MokNewSize, CHAR16 *password,
-                            UINT32 pw_length, UINT8 *hash)
+static EFI_STATUS compute_pw_hash (void *Data, UINTN DataSize, UINT8 *password,
+                                  UINT32 pw_length, UINT8 *hash)
 {
        EFI_STATUS status;
        unsigned int ctxsize;
@@ -523,32 +728,32 @@ static EFI_STATUS compute_pw_hash (void *MokNew, UINTN MokNewSize, CHAR16 *passw
        ctx = AllocatePool(ctxsize);
 
        if (!ctx) {
-               Print(L"Unable to allocate memory for hash context\n");
+               console_notify(L"Unable to allocate memory for hash context");
                return EFI_OUT_OF_RESOURCES;
        }
 
        if (!Sha256Init(ctx)) {
-               Print(L"Unable to initialise hash\n");
+               console_notify(L"Unable to initialise hash");
                status = EFI_OUT_OF_RESOURCES;
                goto done;
        }
 
-       if (MokNew && MokNewSize) {
-               if (!(Sha256Update(ctx, MokNew, MokNewSize))) {
-                       Print(L"Unable to generate hash\n");
+       if (Data && DataSize) {
+               if (!(Sha256Update(ctx, Data, DataSize))) {
+                       console_notify(L"Unable to generate hash");
                        status = EFI_OUT_OF_RESOURCES;
                        goto done;
                }
        }
 
-       if (!(Sha256Update(ctx, password, pw_length * sizeof(CHAR16)))) {
-               Print(L"Unable to generate hash\n");
+       if (!(Sha256Update(ctx, password, pw_length))) {
+               console_notify(L"Unable to generate hash");
                status = EFI_OUT_OF_RESOURCES;
                goto done;
        }
 
        if (!(Sha256Final(ctx, hash))) {
-               Print(L"Unable to finalise hash\n");
+               console_notify(L"Unable to finalise hash");
                status = EFI_OUT_OF_RESOURCES;
                goto done;
        }
@@ -558,445 +763,1083 @@ done:
        return status;
 }
 
-static EFI_STATUS store_keys (void *MokNew, UINTN MokNewSize, int authenticate)
+static void console_save_and_set_mode (SIMPLE_TEXT_OUTPUT_MODE *SavedMode)
 {
-       EFI_GUID shim_lock_guid = SHIM_LOCK_GUID;
-       EFI_STATUS efi_status;
-       UINT8 hash[SHA256_DIGEST_SIZE];
-       UINT8 auth[SHA256_DIGEST_SIZE];
-       UINTN auth_size;
-       UINT32 attributes;
-       CHAR16 password[PASSWORD_MAX];
-       UINT32 pw_length;
-       UINT8 fail_count = 0;
+       if (!SavedMode) {
+               Print(L"Invalid parameter: SavedMode\n");
+               return;
+       }
 
-       if (authenticate) {
-               auth_size = SHA256_DIGEST_SIZE;
-               efi_status = uefi_call_wrapper(RT->GetVariable, 5, L"MokAuth",
-                                              &shim_lock_guid,
-                                              &attributes, &auth_size, auth);
+       CopyMem(SavedMode, ST->ConOut->Mode, sizeof(SIMPLE_TEXT_OUTPUT_MODE));
+       uefi_call_wrapper(ST->ConOut->EnableCursor, 2, ST->ConOut, FALSE);
+       uefi_call_wrapper(ST->ConOut->SetAttribute, 2, ST->ConOut,
+                         EFI_LIGHTGRAY | EFI_BACKGROUND_BLUE);
+}
 
+static void console_restore_mode (SIMPLE_TEXT_OUTPUT_MODE *SavedMode)
+{
+       uefi_call_wrapper(ST->ConOut->EnableCursor, 2, ST->ConOut,
+                         SavedMode->CursorVisible);
+       uefi_call_wrapper(ST->ConOut->SetCursorPosition, 3, ST->ConOut,
+                         SavedMode->CursorColumn, SavedMode->CursorRow);
+       uefi_call_wrapper(ST->ConOut->SetAttribute, 2, ST->ConOut,
+                         SavedMode->Attribute);
+}
 
-               if (efi_status != EFI_SUCCESS || auth_size != SHA256_DIGEST_SIZE) {
-                       Print(L"Failed to get MokAuth %d\n", efi_status);
-                       return efi_status;
-               }
+static UINT32 get_password (CHAR16 *prompt, CHAR16 *password, UINT32 max)
+{
+       SIMPLE_TEXT_OUTPUT_MODE SavedMode;
+       CHAR16 *str;
+       CHAR16 *message[2];
+       UINTN length;
+       UINT32 pw_length;
 
-               while (fail_count < 3) {
-                       Print(L"Password(%d-%d characters): ",
-                             PASSWORD_MIN, PASSWORD_MAX);
-                       get_line(&pw_length, password, PASSWORD_MAX, 0);
+       if (!prompt)
+               prompt = L"Password:";
 
-                       if (pw_length < 8) {
-                               Print(L"At least %d characters for the password\n",
-                                     PASSWORD_MIN);
-                       }
+       console_save_and_set_mode(&SavedMode);
 
-                       efi_status = compute_pw_hash(MokNew, MokNewSize, password,
-                                                    pw_length, hash);
+       str = PoolPrint(L"%s            ", prompt);
+       if (!str) {
+               console_errorbox(L"Failed to allocate prompt");
+               return 0;
+       }
 
-                       if (efi_status != EFI_SUCCESS) {
-                               return efi_status;
-                       }
+       message[0] = str;
+       message[1] = NULL;
+       length = StrLen(message[0]);
+       console_print_box_at(message, -1, -length-4, -5, length+4, 3, 0, 1);
+       get_line(&pw_length, password, max, 0);
 
-                       if (CompareMem(auth, hash, SHA256_DIGEST_SIZE) != 0) {
-                               Print(L"Password doesn't match\n");
-                               fail_count++;
-                       } else {
-                               break;
-                       }
-               }
+       console_restore_mode(&SavedMode);
 
-               if (fail_count >= 3)
+       FreePool(str);
+
+       return pw_length;
+}
+
+static EFI_STATUS match_password (PASSWORD_CRYPT *pw_crypt,
+                                 void *Data, UINTN DataSize,
+                                 UINT8 *auth, CHAR16 *prompt)
+{
+       EFI_STATUS status;
+       UINT8 hash[128];
+       UINT8 *auth_hash;
+       UINT32 auth_size;
+       CHAR16 password[PASSWORD_MAX];
+       UINT32 pw_length;
+       UINT8 fail_count = 0;
+       unsigned int i;
+
+       if (pw_crypt) {
+               auth_hash = pw_crypt->hash;
+               auth_size = get_hash_size (pw_crypt->method);
+               if (auth_size == 0)
+                       return EFI_INVALID_PARAMETER;
+       } else if (auth) {
+               auth_hash = auth;
+               auth_size = SHA256_DIGEST_SIZE;
+       } else {
+               return EFI_INVALID_PARAMETER;
+       }
+
+       while (fail_count < 3) {
+               pw_length = get_password(prompt, password, PASSWORD_MAX);
+
+               if (pw_length < PASSWORD_MIN || pw_length > PASSWORD_MAX) {
+                       console_errorbox(L"Invalid password length");
+                       fail_count++;
+                       continue;
+               }
+
+               /*
+                * Compute password hash
+                */
+               if (pw_crypt) {
+                       char pw_ascii[PASSWORD_MAX + 1];
+                       for (i = 0; i < pw_length; i++)
+                               pw_ascii[i] = (char)password[i];
+                       pw_ascii[pw_length] = '\0';
+
+                       status = password_crypt(pw_ascii, pw_length, pw_crypt, hash);
+               } else {
+                       /*
+                        * For backward compatibility
+                        */
+                       status = compute_pw_hash(Data, DataSize, (UINT8 *)password,
+                                                pw_length * sizeof(CHAR16), hash);
+               }
+               if (status != EFI_SUCCESS) {
+                       console_errorbox(L"Unable to generate password hash");
+                       fail_count++;
+                       continue;
+               }
+
+               if (CompareMem(auth_hash, hash, auth_size) != 0) {
+                       console_errorbox(L"Password doesn't match");
+                       fail_count++;
+                       continue;
+               }
+
+               break;
+       }
+
+       if (fail_count >= 3)
+               return EFI_ACCESS_DENIED;
+
+       return EFI_SUCCESS;
+}
+
+static EFI_STATUS write_db (CHAR16 *db_name, void *MokNew, UINTN MokNewSize)
+{
+       EFI_GUID shim_lock_guid = SHIM_LOCK_GUID;
+       EFI_STATUS status;
+       UINT32 attributes;
+       void *old_data = NULL;
+       void *new_data = NULL;
+       UINTN old_size;
+       UINTN new_size;
+
+       status = uefi_call_wrapper(RT->SetVariable, 5, db_name,
+                                  &shim_lock_guid,
+                                  EFI_VARIABLE_NON_VOLATILE
+                                  | EFI_VARIABLE_BOOTSERVICE_ACCESS
+                                  | EFI_VARIABLE_APPEND_WRITE,
+                                  MokNewSize, MokNew);
+       if (status == EFI_SUCCESS || status != EFI_INVALID_PARAMETER) {
+               return status;
+       }
+
+       status = get_variable_attr(db_name, (UINT8 **)&old_data, &old_size,
+                                  shim_lock_guid, &attributes);
+       if (EFI_ERROR(status) && status != EFI_NOT_FOUND) {
+               return status;
+       }
+
+       /* Check if the old db is compromised or not */
+       if (attributes & EFI_VARIABLE_RUNTIME_ACCESS) {
+               FreePool(old_data);
+               old_data = NULL;
+               old_size = 0;
+       }
+
+       new_size = old_size + MokNewSize;
+       new_data = AllocatePool(new_size);
+       if (new_data == NULL) {
+               status = EFI_OUT_OF_RESOURCES;
+               goto out;
+       }
+
+       CopyMem(new_data, old_data, old_size);
+       CopyMem(new_data + old_size, MokNew, MokNewSize);
+
+       status = uefi_call_wrapper(RT->SetVariable, 5, db_name,
+                                  &shim_lock_guid,
+                                  EFI_VARIABLE_NON_VOLATILE
+                                  | EFI_VARIABLE_BOOTSERVICE_ACCESS,
+                                  new_size, new_data);
+
+out:
+       if (old_size > 0) {
+               FreePool(old_data);
+       }
+
+       if (new_data != NULL) {
+               FreePool(new_data);
+       }
+
+       return status;
+}
+
+static EFI_STATUS store_keys (void *MokNew, UINTN MokNewSize, int authenticate,
+                             BOOLEAN MokX)
+{
+       EFI_GUID shim_lock_guid = SHIM_LOCK_GUID;
+       EFI_STATUS efi_status;
+       CHAR16 *db_name;
+       CHAR16 *auth_name;
+       UINT8 auth[PASSWORD_CRYPT_SIZE];
+       UINTN auth_size = PASSWORD_CRYPT_SIZE;
+       UINT32 attributes;
+
+       if (MokX) {
+               db_name = L"MokListX";
+               auth_name = L"MokXAuth";
+       } else {
+               db_name = L"MokList";
+               auth_name = L"MokAuth";
+       }
+
+       if (authenticate) {
+               efi_status = uefi_call_wrapper(RT->GetVariable, 5, auth_name,
+                                              &shim_lock_guid,
+                                              &attributes, &auth_size, auth);
+
+               if (efi_status != EFI_SUCCESS ||
+                   (auth_size != SHA256_DIGEST_SIZE &&
+                    auth_size != PASSWORD_CRYPT_SIZE)) {
+                       if (MokX)
+                               console_error(L"Failed to get MokXAuth", efi_status);
+                       else
+                               console_error(L"Failed to get MokAuth", efi_status);
+                       return efi_status;
+               }
+
+               if (auth_size == PASSWORD_CRYPT_SIZE) {
+                       efi_status = match_password((PASSWORD_CRYPT *)auth,
+                                                   NULL, 0, NULL, NULL);
+               } else {
+                       efi_status = match_password(NULL, MokNew, MokNewSize,
+                                                   auth, NULL);
+               }
+               if (efi_status != EFI_SUCCESS)
                        return EFI_ACCESS_DENIED;
        }
 
        if (!MokNewSize) {
                /* Delete MOK */
-               efi_status = uefi_call_wrapper(RT->SetVariable, 5, L"MokList",
+               efi_status = uefi_call_wrapper(RT->SetVariable, 5, db_name,
                                               &shim_lock_guid,
                                               EFI_VARIABLE_NON_VOLATILE
                                               | EFI_VARIABLE_BOOTSERVICE_ACCESS,
                                               0, NULL);
        } else {
                /* Write new MOK */
-               efi_status = uefi_call_wrapper(RT->SetVariable, 5, L"MokList",
-                                              &shim_lock_guid,
-                                              EFI_VARIABLE_NON_VOLATILE
-                                              | EFI_VARIABLE_BOOTSERVICE_ACCESS
-                                              | EFI_VARIABLE_APPEND_WRITE,
-                                              MokNewSize, MokNew);
+               efi_status = write_db(db_name, MokNew, MokNewSize);
        }
 
        if (efi_status != EFI_SUCCESS) {
-               Print(L"Failed to set variable %d\n", efi_status);
+               console_error(L"Failed to set variable", efi_status);
                return efi_status;
        }
 
        return EFI_SUCCESS;
 }
 
-static UINTN mok_enrollment_prompt (void *MokNew, UINTN MokNewSize, int auth) {
-       CHAR16 line[1];
-       UINT32 length;
+static INTN mok_enrollment_prompt (void *MokNew, UINTN MokNewSize, int auth,
+                                   BOOLEAN MokX)
+{
+       EFI_GUID shim_lock_guid = SHIM_LOCK_GUID;
        EFI_STATUS efi_status;
+       CHAR16 *title;
 
-       do {
-               if (!list_keys(MokNew, MokNewSize)) {
-                       return 0;
-               }
+       if (MokX)
+               title = L"[Enroll MOKX]";
+       else
+               title = L"[Enroll MOK]";
 
-               Print(L"Enroll the key(s)? (y/n): ");
+       if (list_keys(MokNew, MokNewSize, title) != EFI_SUCCESS)
+               return 0;
 
-               get_line (&length, line, 1, 1);
+       if (console_yes_no((CHAR16 *[]){L"Enroll the key(s)?", NULL}) == 0)
+               return 0;
 
-               if (line[0] == 'Y' || line[0] == 'y') {
-                       efi_status = store_keys(MokNew, MokNewSize, auth);
+       efi_status = store_keys(MokNew, MokNewSize, auth, MokX);
 
-                       if (efi_status != EFI_SUCCESS) {
-                               Print(L"Failed to enroll keys\n");
-                               return -1;
-                       }
-                       return 0;
+       if (efi_status != EFI_SUCCESS) {
+               console_notify(L"Failed to enroll keys\n");
+               return -1;
+       }
+
+       if (auth) {
+               if (MokX) {
+                       LibDeleteVariable(L"MokXNew", &shim_lock_guid);
+                       LibDeleteVariable(L"MokXAuth", &shim_lock_guid);
+               } else {
+                       LibDeleteVariable(L"MokNew", &shim_lock_guid);
+                       LibDeleteVariable(L"MokAuth", &shim_lock_guid);
                }
-       } while (line[0] != 'N' && line[0] != 'n');
-       return -1;
-}
 
-static INTN mok_enrollment_prompt_callback (void *MokNew, void *data2,
-                                           void *data3) {
-       return mok_enrollment_prompt(MokNew, (UINTN)data2, TRUE);
+               console_notify(L"The system must now be rebooted");
+               uefi_call_wrapper(RT->ResetSystem, 4, EfiResetWarm,
+                                 EFI_SUCCESS, 0, NULL);
+               console_notify(L"Failed to reboot");
+               return -1;
+       }
+
+       return 0;
 }
 
-static INTN mok_deletion_prompt (void *MokNew, void *data2, void *data3) {
-       CHAR16 line[1];
-       UINT32 length;
+static INTN mok_reset_prompt (BOOLEAN MokX)
+{
+       EFI_GUID shim_lock_guid = SHIM_LOCK_GUID;
        EFI_STATUS efi_status;
+       CHAR16 *prompt;
+
+       uefi_call_wrapper(ST->ConOut->ClearScreen, 1, ST->ConOut);
 
-       Print(L"Erase all stored keys? (y/N): ");
+       if (MokX)
+               prompt = L"Erase all stored keys in MokListX?";
+       else
+               prompt = L"Erase all stored keys in MokList?";
+       if (console_yes_no((CHAR16 *[]){prompt, NULL }) == 0)
+               return 0;
 
-       get_line (&length, line, 1, 1);
+       efi_status = store_keys(NULL, 0, TRUE, MokX);
 
-       if (line[0] == 'Y' || line[0] == 'y') {
-               efi_status = store_keys(NULL, 0, TRUE);
+       if (efi_status != EFI_SUCCESS) {
+               console_notify(L"Failed to erase keys\n");
+               return -1;
+       }
 
-               if (efi_status != EFI_SUCCESS) {
-                       Print(L"Failed to erase keys\n");
-                       return -1;
-               }
+       if (MokX) {
+               LibDeleteVariable(L"MokXNew", &shim_lock_guid);
+               LibDeleteVariable(L"MokXAuth", &shim_lock_guid);
+       } else {
+               LibDeleteVariable(L"MokNew", &shim_lock_guid);
+               LibDeleteVariable(L"MokAuth", &shim_lock_guid);
        }
 
-       return 0;
+       console_notify(L"The system must now be rebooted");
+       uefi_call_wrapper(RT->ResetSystem, 4, EfiResetWarm,
+                         EFI_SUCCESS, 0, NULL);
+       console_notify(L"Failed to reboot\n");
+       return -1;
 }
 
-static INTN mok_sb_prompt (void *MokSB, void *data2, void *data3) {
+static EFI_STATUS write_back_mok_list (MokListNode *list, INTN key_num,
+                                      BOOLEAN MokX)
+{
        EFI_GUID shim_lock_guid = SHIM_LOCK_GUID;
+       EFI_GUID CertType = X509_GUID;
        EFI_STATUS efi_status;
-       UINTN MokSBSize = (UINTN)data2;
-       MokSBvar *var = MokSB;
-       CHAR16 password[1];
-       UINT8 correct = 0, fail_count = 0;
-       UINT8 hash[SHA256_DIGEST_SIZE];
-       UINT32 length;
-       CHAR16 line[1];
-       UINT8 sbval = 1;
-       UINT8 pos;
+       EFI_SIGNATURE_LIST *CertList;
+       EFI_SIGNATURE_DATA *CertData;
+       void *Data = NULL, *ptr;
+       INTN DataSize = 0;
+       int i;
+       CHAR16 *db_name;
 
-       LibDeleteVariable(L"MokSB", &shim_lock_guid);
+       if (MokX)
+               db_name = L"MokListX";
+       else
+               db_name = L"MokList";
 
-       if (MokSBSize != sizeof(MokSBvar)) {
-               Print(L"Invalid MokSB variable contents\n");
-               return -1;
+       for (i = 0; i < key_num; i++) {
+               if (list[i].Mok == NULL)
+                       continue;
+
+               DataSize += sizeof(EFI_SIGNATURE_LIST);
+               if (CompareGuid(&(list[i].Type), &CertType) == 0)
+                       DataSize += sizeof(EFI_GUID);
+               DataSize += list[i].MokSize;
        }
 
-       uefi_call_wrapper(ST->ConOut->ClearScreen, 1, ST->ConOut);
+       Data = AllocatePool(DataSize);
+       if (Data == NULL && DataSize != 0)
+               return EFI_OUT_OF_RESOURCES;
 
-       while (correct < 3) {
-               RandomBytes (&pos, sizeof(pos));
+       ptr = Data;
 
-               pos = pos % var->PWLen;
+       for (i = 0; i < key_num; i++) {
+               if (list[i].Mok == NULL)
+                       continue;
 
-               Print(L"Enter password character %d: ", pos + 1);
-               get_line(&length, password, 1, 0);
+               CertList = (EFI_SIGNATURE_LIST *)ptr;
+               CertData = (EFI_SIGNATURE_DATA *)(((uint8_t *)ptr) +
+                          sizeof(EFI_SIGNATURE_LIST));
 
-               if (password[0] != var->Password[pos]) {
-                       Print(L"Invalid character\n");
-                       fail_count++;
+               CertList->SignatureType = list[i].Type;
+               CertList->SignatureHeaderSize = 0;
+
+               if (CompareGuid(&(list[i].Type), &CertType) == 0) {
+                       CertList->SignatureListSize = list[i].MokSize +
+                                                     sizeof(EFI_SIGNATURE_LIST) +
+                                                     sizeof(EFI_GUID);
+                       CertList->SignatureSize = list[i].MokSize + sizeof(EFI_GUID);
+
+                       CertData->SignatureOwner = shim_lock_guid;
+                       CopyMem(CertData->SignatureData, list[i].Mok, list[i].MokSize);
                } else {
-                       correct++;
+                       CertList->SignatureListSize = list[i].MokSize +
+                                                     sizeof(EFI_SIGNATURE_LIST);
+                       CertList->SignatureSize = sha_size(list[i].Type) + sizeof(EFI_GUID);
+
+                       CopyMem(CertData, list[i].Mok, list[i].MokSize);
                }
+               ptr = (uint8_t *)ptr + CertList->SignatureListSize;
+       }
 
-               if (fail_count >= 3)
-                       break;
+       efi_status = uefi_call_wrapper(RT->SetVariable, 5, db_name,
+                                      &shim_lock_guid,
+                                      EFI_VARIABLE_NON_VOLATILE
+                                      | EFI_VARIABLE_BOOTSERVICE_ACCESS,
+                                      DataSize, Data);
+       if (Data)
+               FreePool(Data);
+
+       if (efi_status != EFI_SUCCESS) {
+               console_error(L"Failed to set variable", efi_status);
+               return efi_status;
        }
 
-       if (fail_count >= 3) {
-               Print(L"Password limit reached\n");
-               return -1;
+       return EFI_SUCCESS;
+}
+
+static void delete_cert (void *key, UINT32 key_size,
+                        MokListNode *mok, INTN mok_num)
+{
+       EFI_GUID CertType = X509_GUID;
+       int i;
+
+       for (i = 0; i < mok_num; i++) {
+               if (CompareGuid(&(mok[i].Type), &CertType) != 0)
+                       continue;
+
+               if (mok[i].MokSize == key_size &&
+                   CompareMem(key, mok[i].Mok, key_size) == 0) {
+                       /* Remove the key */
+                       mok[i].Mok = NULL;
+                       mok[i].MokSize = 0;
+               }
        }
+}
 
-       if (var->MokSBState == 0) {
-               Print(L"Disable Secure Boot? (y/n): ");
-       } else {
-               Print(L"Enable Secure Boot? (y/n): ");
+static int match_hash (UINT8 *hash, UINT32 hash_size, int start,
+                      void *hash_list, UINT32 list_num)
+{
+       UINT8 *ptr;
+       UINTN i;
+
+       ptr = hash_list + sizeof(EFI_GUID);
+       for (i = start; i < list_num; i++) {
+               if (CompareMem(hash, ptr, hash_size) == 0)
+                       return i;
+               ptr += hash_size + sizeof(EFI_GUID);
        }
 
-       do {
-               get_line (&length, line, 1, 1);
+       return -1;
+}
 
-               if (line[0] == 'Y' || line[0] == 'y') {
-                       if (var->MokSBState == 0) {
-                               efi_status = uefi_call_wrapper(RT->SetVariable,
-                                                              5, L"MokSBState",
-                                                              &shim_lock_guid,
-                                              EFI_VARIABLE_NON_VOLATILE |
-                                              EFI_VARIABLE_BOOTSERVICE_ACCESS,
-                                                              1, &sbval);
-                               if (efi_status != EFI_SUCCESS) {
-                                       Print(L"Failed to set Secure Boot state\n");
-                                       return -1;
-                               }
-                       } else {
-                               LibDeleteVariable(L"MokSBState",
-                                                 &shim_lock_guid);
+static void mem_move (void *dest, void *src, UINTN size)
+{
+       UINT8 *d, *s;
+       UINTN i;
+
+       d = (UINT8 *)dest;
+       s = (UINT8 *)src;
+       for (i = 0; i < size; i++)
+               d[i] = s[i];
+}
+
+static void delete_hash_in_list (EFI_GUID Type, UINT8 *hash, UINT32 hash_size,
+                                MokListNode *mok, INTN mok_num)
+{
+       UINT32 sig_size;
+       UINT32 list_num;
+       int i, del_ind;
+       void *start, *end;
+       UINT32 remain;
+
+       sig_size = hash_size + sizeof(EFI_GUID);
+
+       for (i = 0; i < mok_num; i++) {
+               if ((CompareGuid(&(mok[i].Type), &Type) != 0) ||
+                   (mok[i].MokSize < sig_size))
+                       continue;
+
+               list_num = mok[i].MokSize / sig_size;
+
+               del_ind = match_hash(hash, hash_size, 0, mok[i].Mok,
+                                    list_num);
+               while (del_ind >= 0) {
+                       /* Remove the hash */
+                       if (sig_size == mok[i].MokSize) {
+                               mok[i].Mok = NULL;
+                               mok[i].MokSize = 0;
+                               break;
                        }
 
-                       Print(L"Press a key to reboot system\n");
-                       Pause();
-                       uefi_call_wrapper(RT->ResetSystem, 4, EfiResetWarm,
-                                         EFI_SUCCESS, 0, NULL);
-                       Print(L"Failed to reboot\n");
-                       return -1;
+                       start = mok[i].Mok + del_ind * sig_size;
+                       end = start + sig_size;
+                       remain = mok[i].MokSize - (del_ind + 1)*sig_size;
+
+                       mem_move(start, end, remain);
+                       mok[i].MokSize -= sig_size;
+                       list_num--;
+
+                       del_ind = match_hash(hash, hash_size, del_ind,
+                                            mok[i].Mok, list_num);
+               }
+       }
+}
+
+static void delete_hash_list (EFI_GUID Type, void *hash_list, UINT32 list_size,
+                             MokListNode *mok, INTN mok_num)
+{
+       UINT32 hash_size;
+       UINT32 hash_num;
+       UINT32 sig_size;
+       UINT8 *hash;
+       UINT32 i;
+
+       hash_size = sha_size (Type);
+       sig_size = hash_size + sizeof(EFI_GUID);
+       if (list_size < sig_size)
+               return;
+
+       hash_num = list_size / sig_size;
+
+       hash = hash_list + sizeof(EFI_GUID);
+
+       for (i = 0; i < hash_num; i++) {
+               delete_hash_in_list (Type, hash, hash_size, mok, mok_num);
+               hash += sig_size;
+       }
+}
+
+static EFI_STATUS delete_keys (void *MokDel, UINTN MokDelSize, BOOLEAN MokX)
+{
+       EFI_GUID shim_lock_guid = SHIM_LOCK_GUID;
+       EFI_GUID CertType = X509_GUID;
+       EFI_STATUS efi_status;
+       CHAR16 *db_name;
+       CHAR16 *auth_name;
+       CHAR16 *err_str1;
+       CHAR16 *err_str2;
+       UINT8 auth[PASSWORD_CRYPT_SIZE];
+       UINTN auth_size = PASSWORD_CRYPT_SIZE;
+       UINT32 attributes;
+       UINT8 *MokListData = NULL;
+       UINTN MokListDataSize = 0;
+       MokListNode *mok, *del_key;
+       INTN mok_num, del_num;
+       int i;
+
+       if (MokX) {
+               db_name = L"MokListX";
+               auth_name = L"MokXDelAuth";
+       } else {
+               db_name = L"MokList";
+               auth_name = L"MokDelAuth";
+       }
+
+       efi_status = uefi_call_wrapper(RT->GetVariable, 5, auth_name,
+                                      &shim_lock_guid,
+                                      &attributes, &auth_size, auth);
+
+       if (efi_status != EFI_SUCCESS ||
+           (auth_size != SHA256_DIGEST_SIZE && auth_size != PASSWORD_CRYPT_SIZE)) {
+               if (MokX)
+                       console_error(L"Failed to get MokXDelAuth", efi_status);
+               else
+                       console_error(L"Failed to get MokDelAuth", efi_status);
+               return efi_status;
+       }
+
+       if (auth_size == PASSWORD_CRYPT_SIZE) {
+               efi_status = match_password((PASSWORD_CRYPT *)auth, NULL, 0,
+                                           NULL, NULL);
+       } else {
+               efi_status = match_password(NULL, MokDel, MokDelSize, auth, NULL);
+       }
+       if (efi_status != EFI_SUCCESS)
+               return EFI_ACCESS_DENIED;
+
+       efi_status = get_variable_attr (db_name, &MokListData, &MokListDataSize,
+                                       shim_lock_guid, &attributes);
+       if (efi_status != EFI_SUCCESS) {
+               if (MokX)
+                       console_errorbox(L"Failed to retrieve MokListX");
+               else
+                       console_errorbox(L"Failed to retrieve MokList");
+               return EFI_ABORTED;
+       } else if (attributes & EFI_VARIABLE_RUNTIME_ACCESS) {
+               if (MokX) {
+                       err_str1 = L"MokListX is compromised!";
+                       err_str2 = L"Erase all keys in MokListX!";
+               } else {
+                       err_str1 = L"MokList is compromised!";
+                       err_str2 = L"Erase all keys in MokList!";
                }
-       } while (line[0] != 'N' && line[0] != 'n');
+               console_alertbox((CHAR16 *[]){err_str1, err_str2, NULL});
+               uefi_call_wrapper(RT->SetVariable, 5, db_name,
+                                 &shim_lock_guid,
+                                 EFI_VARIABLE_NON_VOLATILE |
+                                 EFI_VARIABLE_BOOTSERVICE_ACCESS,
+                                 0, NULL);
+               return EFI_ACCESS_DENIED;
+       }
+
+       /* Nothing to do */
+       if (!MokListData || MokListDataSize == 0)
+               return EFI_SUCCESS;
 
+       /* Construct lists */
+       mok_num = count_keys(MokListData, MokListDataSize);
+       if (mok_num == 0) {
+               if (MokX) {
+                       err_str1 = L"Failed to construct the key list of MokListX";
+                       err_str2 = L"Reset MokListX!";
+               } else {
+                       err_str1 = L"Failed to construct the key list of MokList";
+                       err_str2 = L"Reset MokList!";
+               }
+               console_alertbox((CHAR16 *[]){err_str1, err_str2, NULL});
+               uefi_call_wrapper(RT->SetVariable, 5, db_name,
+                                 &shim_lock_guid,
+                                 EFI_VARIABLE_NON_VOLATILE |
+                                 EFI_VARIABLE_BOOTSERVICE_ACCESS,
+                                 0, NULL);
+               efi_status = EFI_ABORTED;
+               goto error;
+       }
+       mok = build_mok_list(mok_num, MokListData, MokListDataSize);
+       if (!mok) {
+               console_errorbox(L"Failed to construct key list");
+               efi_status = EFI_ABORTED;
+               goto error;
+       }
+       del_num = count_keys(MokDel, MokDelSize);
+       if (del_num == 0) {
+               console_errorbox(L"Invalid key delete list");
+               efi_status = EFI_ABORTED;
+               goto error;
+       }
+       del_key = build_mok_list(del_num, MokDel, MokDelSize);
+       if (!del_key) {
+               console_errorbox(L"Failed to construct key list");
+               efi_status = EFI_ABORTED;
+               goto error;
+       }
+
+       /* Search and destroy */
+       for (i = 0; i < del_num; i++) {
+               if (CompareGuid(&(del_key[i].Type), &CertType) == 0) {
+                       delete_cert(del_key[i].Mok, del_key[i].MokSize,
+                                   mok, mok_num);
+               } else if (is_sha2_hash(del_key[i].Type)) {
+                       delete_hash_list(del_key[i].Type, del_key[i].Mok,
+                                        del_key[i].MokSize, mok, mok_num);
+               }
+       }
+
+       efi_status = write_back_mok_list(mok, mok_num, MokX);
+
+error:
+       if (MokListData)
+               FreePool(MokListData);
+       if (mok)
+               FreePool(mok);
+       if (del_key)
+               FreePool(del_key);
+
+       return efi_status;
+}
+
+static INTN mok_deletion_prompt (void *MokDel, UINTN MokDelSize, BOOLEAN MokX)
+{
+       EFI_GUID shim_lock_guid = SHIM_LOCK_GUID;
+       EFI_STATUS efi_status;
+       CHAR16 *title;
+
+       if (MokX)
+               title = L"[Delete MOKX]";
+       else
+               title = L"[Delete MOK]";
+
+       if (list_keys(MokDel, MokDelSize, title) != EFI_SUCCESS) {
+               return 0;
+       }
+
+        if (console_yes_no((CHAR16 *[]){L"Delete the key(s)?", NULL}) == 0)
+                return 0;
+
+       efi_status = delete_keys(MokDel, MokDelSize, MokX);
+
+       if (efi_status != EFI_SUCCESS) {
+               console_notify(L"Failed to delete keys");
+               return -1;
+       }
+
+       if (MokX) {
+               LibDeleteVariable(L"MokXDel", &shim_lock_guid);
+               LibDeleteVariable(L"MokXDelAuth", &shim_lock_guid);
+       } else {
+               LibDeleteVariable(L"MokDel", &shim_lock_guid);
+               LibDeleteVariable(L"MokDelAuth", &shim_lock_guid);
+       }
+
+       console_notify(L"The system must now be rebooted");
+       uefi_call_wrapper(RT->ResetSystem, 4, EfiResetWarm,
+                         EFI_SUCCESS, 0, NULL);
+       console_notify(L"Failed to reboot");
        return -1;
 }
 
+static CHAR16 get_password_charater (CHAR16 *prompt)
+{
+       SIMPLE_TEXT_OUTPUT_MODE SavedMode;
+       EFI_STATUS status;
+       CHAR16 *message[2];
+       CHAR16 character;
+       UINTN length;
+       UINT32 pw_length;
+
+       if (!prompt)
+               prompt = L"Password charater: ";
+
+       console_save_and_set_mode(&SavedMode);
+
+       message[0] = prompt;
+       message[1] = NULL;
+       length = StrLen(message[0]);
+       console_print_box_at(message, -1, -length-4, -5, length+4, 3, 0, 1);
+       status = get_line(&pw_length, &character, 1, 0);
+       if (EFI_ERROR(status))
+               character = 0;
+
+       console_restore_mode(&SavedMode);
+
+       return character;
+}
 
-static INTN mok_pw_prompt (void *MokPW, void *data2, void *data3) {
+static INTN mok_sb_prompt (void *MokSB, UINTN MokSBSize) {
        EFI_GUID shim_lock_guid = SHIM_LOCK_GUID;
        EFI_STATUS efi_status;
-       UINTN MokPWSize = (UINTN)data2;
+       SIMPLE_TEXT_OUTPUT_MODE SavedMode;
+       MokSBvar *var = MokSB;
+       CHAR16 *message[4];
+       CHAR16 pass1, pass2, pass3;
+       CHAR16 *str;
        UINT8 fail_count = 0;
-       UINT8 hash[SHA256_DIGEST_SIZE];
-       CHAR16 password[PASSWORD_MAX];
-       UINT32 length;
-       CHAR16 line[1];
+       UINT8 sbval = 1;
+       UINT8 pos1, pos2, pos3;
+       int ret;
 
-       if (MokPWSize != SHA256_DIGEST_SIZE) {
-               Print(L"Invalid MokPW variable contents\n");
+       if (MokSBSize != sizeof(MokSBvar)) {
+               console_notify(L"Invalid MokSB variable contents");
                return -1;
        }
 
-       LibDeleteVariable(L"MokPW", &shim_lock_guid);
-
        uefi_call_wrapper(ST->ConOut->ClearScreen, 1, ST->ConOut);
 
-       while (fail_count < 3) {
-               Print(L"Confirm MOK passphrase: ");
-               get_line(&length, password, PASSWORD_MAX, 0);
+       message[0] = L"Change Secure Boot state";
+       message[1] = NULL;
 
-               if ((length < PASSWORD_MIN) || (length > PASSWORD_MAX)) {
-                       Print(L"Invalid password length\n");
-                       fail_count++;
-                       continue;
+       console_save_and_set_mode(&SavedMode);
+       console_print_box_at(message, -1, 0, 0, -1, -1, 1, 1);
+       console_restore_mode(&SavedMode);
+
+       while (fail_count < 3) {
+               RandomBytes (&pos1, sizeof(pos1));
+               pos1 = (pos1 % var->PWLen);
+
+               do {
+                       RandomBytes (&pos2, sizeof(pos2));
+                       pos2 = (pos2 % var->PWLen);
+               } while (pos2 == pos1);
+
+               do {
+                       RandomBytes (&pos3, sizeof(pos3));
+                       pos3 = (pos3 % var->PWLen) ;
+               } while (pos3 == pos2 || pos3 == pos1);
+
+               str = PoolPrint(L"Enter password character %d: ", pos1 + 1);
+               if (!str) {
+                       console_errorbox(L"Failed to allocate buffer");
+                       return -1;
                }
+               pass1 = get_password_charater(str);
+               FreePool(str);
 
-               efi_status = compute_pw_hash(NULL, 0, password,
-                                            SB_PASSWORD_LEN, hash);
+               str = PoolPrint(L"Enter password character %d: ", pos2 + 1);
+               if (!str) {
+                       console_errorbox(L"Failed to allocate buffer");
+                       return -1;
+               }
+               pass2 = get_password_charater(str);
+               FreePool(str);
 
-               if (efi_status != EFI_SUCCESS) {
-                       Print(L"Unable to generate password hash\n");
-                       fail_count++;
-                       continue;
+               str = PoolPrint(L"Enter password character %d: ", pos3 + 1);
+               if (!str) {
+                       console_errorbox(L"Failed to allocate buffer");
+                       return -1;
                }
+               pass3 = get_password_charater(str);
+               FreePool(str);
 
-               if (CompareMem(MokPW, hash, SHA256_DIGEST_SIZE) != 0) {
-                       Print(L"Password doesn't match\n");
+               if (pass1 != var->Password[pos1] ||
+                   pass2 != var->Password[pos2] ||
+                   pass3 != var->Password[pos3]) {
+                       Print(L"Invalid character\n");
                        fail_count++;
-                       continue;
+               } else {
+                       break;
                }
-
-               break;
        }
 
        if (fail_count >= 3) {
-               Print(L"Password limit reached\n");
+               console_notify(L"Password limit reached");
                return -1;
        }
 
-       Print(L"Set MOK password? (y/n): ");
+       if (var->MokSBState == 0)
+               ret = console_yes_no((CHAR16 *[]){L"Disable Secure Boot", NULL});
+       else
+               ret = console_yes_no((CHAR16 *[]){L"Enable Secure Boot", NULL});
 
-       do {
-               get_line (&length, line, 1, 1);
+       if (ret == 0) {
+               LibDeleteVariable(L"MokSB", &shim_lock_guid);
+               return -1;
+       }
 
-               if (line[0] == 'Y' || line[0] == 'y') {
-                       efi_status = uefi_call_wrapper(RT->SetVariable, 5,
-                                                      L"MokPWStore",
-                                                      &shim_lock_guid,
+       if (var->MokSBState == 0) {
+               efi_status = uefi_call_wrapper(RT->SetVariable,
+                                              5, L"MokSBState",
+                                              &shim_lock_guid,
                                               EFI_VARIABLE_NON_VOLATILE |
                                               EFI_VARIABLE_BOOTSERVICE_ACCESS,
-                                                      MokPWSize, MokPW);
-                       if (efi_status != EFI_SUCCESS) {
-                               Print(L"Failed to set MOK password\n");
-                               return -1;
-                       }
-
-                       Print(L"Press a key to reboot system\n");
-                       Pause();
-                       uefi_call_wrapper(RT->ResetSystem, 4, EfiResetWarm,
-                                         EFI_SUCCESS, 0, NULL);
-                       Print(L"Failed to reboot\n");
+                                              1, &sbval);
+               if (efi_status != EFI_SUCCESS) {
+                       console_notify(L"Failed to set Secure Boot state");
+                       return -1;
+               }
+       } else {
+               efi_status = uefi_call_wrapper(RT->SetVariable,
+                                              5, L"MokSBState",
+                                              &shim_lock_guid,
+                                              EFI_VARIABLE_NON_VOLATILE |
+                                              EFI_VARIABLE_BOOTSERVICE_ACCESS,
+                                              0, NULL);
+               if (efi_status != EFI_SUCCESS) {
+                       console_notify(L"Failed to delete Secure Boot state");
                        return -1;
                }
-       } while (line[0] != 'N' && line[0] != 'n');
+       }
 
-       return 0;
+       console_notify(L"The system must now be rebooted");
+       uefi_call_wrapper(RT->ResetSystem, 4, EfiResetWarm,
+                         EFI_SUCCESS, 0, NULL);
+       console_notify(L"Failed to reboot");
+       return -1;
 }
 
-static UINTN draw_menu (CHAR16 *header, UINTN lines, struct menu_item *items,
-                       UINTN count) {
-       UINTN i;
-
-       uefi_call_wrapper(ST->ConOut->ClearScreen, 1, ST->ConOut);
+static INTN mok_db_prompt (void *MokDB, UINTN MokDBSize) {
+       EFI_GUID shim_lock_guid = SHIM_LOCK_GUID;
+       EFI_STATUS efi_status;
+       SIMPLE_TEXT_OUTPUT_MODE SavedMode;
+       MokDBvar *var = MokDB;
+       CHAR16 *message[4];
+       CHAR16 pass1, pass2, pass3;
+       CHAR16 *str;
+       UINT8 fail_count = 0;
+       UINT8 dbval = 1;
+       UINT8 pos1, pos2, pos3;
+       int ret;
 
-       uefi_call_wrapper(ST->ConOut->SetAttribute, 2, ST->ConOut,
-                         EFI_WHITE | EFI_BACKGROUND_BLACK);
+       if (MokDBSize != sizeof(MokDBvar)) {
+               console_notify(L"Invalid MokDB variable contents");
+               return -1;
+       }
 
-       Print(L"%s UEFI key management\n\n", SHIM_VENDOR);
+       uefi_call_wrapper(ST->ConOut->ClearScreen, 1, ST->ConOut);
 
-       if (header)
-               Print(L"%s", header);
+       message[0] = L"Change DB state";
+       message[1] = NULL;
 
-       for (i = 0; i < count; i++) {
-               uefi_call_wrapper(ST->ConOut->SetAttribute, 2, ST->ConOut,
-                                 items[i].colour | EFI_BACKGROUND_BLACK);
-               Print(L"  %s\n", items[i].text);
-       }
+       console_save_and_set_mode(&SavedMode);
+       console_print_box_at(message, -1, 0, 0, -1, -1, 1, 1);
+       console_restore_mode(&SavedMode);
 
-       uefi_call_wrapper(ST->ConOut->SetCursorPosition, 3, ST->ConOut, 0, 0);
-       uefi_call_wrapper(ST->ConOut->EnableCursor, 2, ST->ConOut, TRUE);
+       while (fail_count < 3) {
+               RandomBytes (&pos1, sizeof(pos1));
+               pos1 = (pos1 % var->PWLen);
+
+               do {
+                       RandomBytes (&pos2, sizeof(pos2));
+                       pos2 = (pos2 % var->PWLen);
+               } while (pos2 == pos1);
+
+               do {
+                       RandomBytes (&pos3, sizeof(pos3));
+                       pos3 = (pos3 % var->PWLen) ;
+               } while (pos3 == pos2 || pos3 == pos1);
+
+               str = PoolPrint(L"Enter password character %d: ", pos1 + 1);
+               if (!str) {
+                       console_errorbox(L"Failed to allocate buffer");
+                       return -1;
+               }
+               pass1 = get_password_charater(str);
+               FreePool(str);
 
-       return 2 + lines;
-}
+               str = PoolPrint(L"Enter password character %d: ", pos2 + 1);
+               if (!str) {
+                       console_errorbox(L"Failed to allocate buffer");
+                       return -1;
+               }
+               pass2 = get_password_charater(str);
+               FreePool(str);
 
-static void free_menu (struct menu_item *items, UINTN count) {
-       UINTN i;
+               str = PoolPrint(L"Enter password character %d: ", pos3 + 1);
+               if (!str) {
+                       console_errorbox(L"Failed to allocate buffer");
+                       return -1;
+               }
+               pass3 = get_password_charater(str);
+               FreePool(str);
 
-       for (i=0; i<count; i++) {
-               if (items[i].text)
-                       FreePool(items[i].text);
+               if (pass1 != var->Password[pos1] ||
+                   pass2 != var->Password[pos2] ||
+                   pass3 != var->Password[pos3]) {
+                       Print(L"Invalid character\n");
+                       fail_count++;
+               } else {
+                       break;
+               }
        }
 
-       FreePool(items);
-}
-
-static void update_time (UINTN position, UINTN timeout)
-{
-       uefi_call_wrapper(ST->ConOut->SetCursorPosition, 3, ST->ConOut, 0,
-                         position);
-
-       uefi_call_wrapper(ST->ConOut->SetAttribute, 2, ST->ConOut,
-                         EFI_BLACK | EFI_BACKGROUND_BLACK);
+       if (fail_count >= 3) {
+               console_notify(L"Password limit reached");
+               return -1;
+       }
 
-       Print(L"                       ", timeout);
+       if (var->MokDBState == 0)
+               ret = console_yes_no((CHAR16 *[]){L"Ignore DB certs/hashes", NULL});
+       else
+               ret = console_yes_no((CHAR16 *[]){L"Use DB certs/hashes", NULL});
 
-       uefi_call_wrapper(ST->ConOut->SetCursorPosition, 3, ST->ConOut, 0,
-                         position);
+       if (ret == 0) {
+               LibDeleteVariable(L"MokDB", &shim_lock_guid);
+               return -1;
+       }
 
-       uefi_call_wrapper(ST->ConOut->SetAttribute, 2, ST->ConOut,
-                         EFI_WHITE | EFI_BACKGROUND_BLACK);
+       if (var->MokDBState == 0) {
+               efi_status = uefi_call_wrapper(RT->SetVariable,
+                                              5, L"MokDBState",
+                                              &shim_lock_guid,
+                                              EFI_VARIABLE_NON_VOLATILE |
+                                              EFI_VARIABLE_BOOTSERVICE_ACCESS,
+                                              1, &dbval);
+               if (efi_status != EFI_SUCCESS) {
+                       console_notify(L"Failed to set DB state");
+                       return -1;
+               }
+       } else {
+               efi_status = uefi_call_wrapper(RT->SetVariable, 5,
+                                              L"MokDBState",
+                                              &shim_lock_guid,
+                                              EFI_VARIABLE_NON_VOLATILE |
+                                              EFI_VARIABLE_BOOTSERVICE_ACCESS,
+                                              0, NULL);
+               if (efi_status != EFI_SUCCESS) {
+                       console_notify(L"Failed to delete DB state");
+                       return -1;
+               }
+       }
 
-       if (timeout > 1)
-               Print(L"Booting in %d seconds\n", timeout);
-       else if (timeout)
-               Print(L"Booting in %d second\n", timeout);
+       console_notify(L"The system must now be rebooted");
+       uefi_call_wrapper(RT->ResetSystem, 4, EfiResetWarm,
+                         EFI_SUCCESS, 0, NULL);
+       console_notify(L"Failed to reboot");
+       return -1;
 }
 
-static void run_menu (CHAR16 *header, UINTN lines, struct menu_item *items,
-                     UINTN count, UINTN timeout) {
-       UINTN index, pos = 0, wait = 0, offset;
-       EFI_INPUT_KEY key;
-       EFI_STATUS status;
-       INTN ret;
+static INTN mok_pw_prompt (void *MokPW, UINTN MokPWSize) {
+       EFI_GUID shim_lock_guid = SHIM_LOCK_GUID;
+       EFI_STATUS efi_status;
+       UINT8 hash[PASSWORD_CRYPT_SIZE];
+       UINT8 clear = 0;
 
-       if (timeout)
-               wait = 10000000;
+       if (MokPWSize != SHA256_DIGEST_SIZE && MokPWSize != PASSWORD_CRYPT_SIZE) {
+               console_notify(L"Invalid MokPW variable contents");
+               return -1;
+       }
 
-       offset = draw_menu (header, lines, items, count);
+       uefi_call_wrapper(ST->ConOut->ClearScreen, 1, ST->ConOut);
 
-       while (1) {
-               update_time(count + offset + 1, timeout);
+       SetMem(hash, PASSWORD_CRYPT_SIZE, 0);
 
-               uefi_call_wrapper(ST->ConOut->SetCursorPosition, 3, ST->ConOut,
-                                 0, pos + offset);
-               status = WaitForSingleEvent(ST->ConIn->WaitForKey, wait);
+       if (MokPWSize == PASSWORD_CRYPT_SIZE) {
+               if (CompareMem(MokPW, hash, PASSWORD_CRYPT_SIZE) == 0)
+                       clear = 1;
+       } else {
+               if (CompareMem(MokPW, hash, SHA256_DIGEST_SIZE) == 0)
+                       clear = 1;
+       }
 
-               if (status == EFI_TIMEOUT) {
-                       timeout--;
-                       if (!timeout) {
-                               free_menu(items, count);
-                               return;
-                       }
-                       continue;
-               }
+       if (clear) {
+               if (console_yes_no((CHAR16 *[]){L"Clear MOK password?", NULL}) == 0)
+                       return 0;
 
-               wait = 0;
-               timeout = 0;
+               uefi_call_wrapper(RT->SetVariable, 5, L"MokPWStore",
+                                 &shim_lock_guid,
+                                 EFI_VARIABLE_NON_VOLATILE
+                                 | EFI_VARIABLE_BOOTSERVICE_ACCESS,
+                                 0, NULL);
+               LibDeleteVariable(L"MokPW", &shim_lock_guid);
+               console_notify(L"The system must now be rebooted");
+               uefi_call_wrapper(RT->ResetSystem, 4, EfiResetWarm, EFI_SUCCESS, 0,
+                                 NULL);
+               console_notify(L"Failed to reboot");
+               return -1;
+       }
 
-               uefi_call_wrapper(BS->WaitForEvent, 3, 1,
-                                 &ST->ConIn->WaitForKey, &index);
-               uefi_call_wrapper(ST->ConIn->ReadKeyStroke, 2, ST->ConIn,
-                                 &key);
+       if (MokPWSize == PASSWORD_CRYPT_SIZE) {
+               efi_status = match_password((PASSWORD_CRYPT *)MokPW, NULL, 0,
+                                           NULL, L"Confirm MOK passphrase: ");
+       } else {
+               efi_status = match_password(NULL, NULL, 0, MokPW,
+                                           L"Confirm MOK passphrase: ");
+       }
 
-               switch(key.ScanCode) {
-               case SCAN_UP:
-                       if (pos == 0)
-                               continue;
-                       pos--;
-                       continue;
-                       break;
-               case SCAN_DOWN:
-                       if (pos == (count - 1))
-                               continue;
-                       pos++;
-                       continue;
-                       break;
-               }
+       if (efi_status != EFI_SUCCESS) {
+               console_notify(L"Password limit reached");
+               return -1;
+       }
 
-               switch(key.UnicodeChar) {
-               case CHAR_LINEFEED:
-               case CHAR_CARRIAGE_RETURN:
-                       if (items[pos].callback == NULL) {
-                               free_menu(items, count);
-                               return;
-                       }
+       if (console_yes_no((CHAR16 *[]){L"Set MOK password?", NULL}) == 0)
+               return 0;
 
-                       ret = items[pos].callback(items[pos].data,
-                                                 items[pos].data2,
-                                                 items[pos].data3);
-                       if (ret < 0) {
-                               Print(L"Press a key to continue\n");
-                               Pause();
-                       }
-                       draw_menu (header, lines, items, count);
-                       pos = 0;
-                       break;
-               }
+       efi_status = uefi_call_wrapper(RT->SetVariable, 5,
+                                      L"MokPWStore",
+                                      &shim_lock_guid,
+                                      EFI_VARIABLE_NON_VOLATILE |
+                                      EFI_VARIABLE_BOOTSERVICE_ACCESS,
+                                      MokPWSize, MokPW);
+       if (efi_status != EFI_SUCCESS) {
+               console_notify(L"Failed to set MOK password");
+               return -1;
        }
+
+       LibDeleteVariable(L"MokPW", &shim_lock_guid);
+
+       console_notify(L"The system must now be rebooted");
+       uefi_call_wrapper(RT->ResetSystem, 4, EfiResetWarm, EFI_SUCCESS, 0,
+                         NULL);
+       console_notify(L"Failed to reboot");
+       return -1;
 }
 
-static UINTN verify_certificate(void *cert, UINTN size)
+static BOOLEAN verify_certificate(UINT8 *cert, UINTN size)
 {
        X509 *X509Cert;
-       if (!cert || size == 0)
+       UINTN length;
+       if (!cert || size < 0)
                return FALSE;
 
+       /*
+        * A DER encoding x509 certificate starts with SEQUENCE(0x30),
+        * the number of length bytes, and the number of value bytes.
+        * The size of a x509 certificate is usually between 127 bytes
+        * and 64KB. For convenience, assume the number of value bytes
+        * is 2, i.e. the second byte is 0x82.
+        */
+       if (cert[0] != 0x30 || cert[1] != 0x82) {
+               console_notify(L"Not a DER encoding X509 certificate");
+               return FALSE;
+       }
+
+       length = (cert[2]<<8 | cert[3]);
+       if (length != (size - 4)) {
+               console_notify(L"Invalid X509 certificate: Inconsistent size");
+               return FALSE;
+       }
+
        if (!(X509ConstructCertificate(cert, size, (UINT8 **) &X509Cert)) ||
            X509Cert == NULL) {
-               Print(L"Invalid X509 certificate\n");
-               Pause();
+               console_notify(L"Invalid X509 certificate");
                return FALSE;
        }
 
@@ -1004,49 +1847,22 @@ static UINTN verify_certificate(void *cert, UINTN size)
        return TRUE;
 }
 
-static INTN file_callback (void *data, void *data2, void *data3) {
-       EFI_FILE_INFO *buffer = NULL;
-       UINTN buffersize = 0, mokbuffersize;
-       EFI_STATUS status;
-       EFI_FILE *file;
-       CHAR16 *filename = data;
-       EFI_FILE *parent = data2;
-       BOOLEAN hash = !!data3;
-       EFI_GUID file_info_guid = EFI_FILE_INFO_ID;
-       EFI_GUID shim_lock_guid = SHIM_LOCK_GUID;
-       EFI_SIGNATURE_LIST *CertList;
-       EFI_SIGNATURE_DATA *CertData;
-       void *mokbuffer = NULL;
-
-       status = uefi_call_wrapper(parent->Open, 5, parent, &file, filename,
-                                  EFI_FILE_MODE_READ, 0);
-
-       if (status != EFI_SUCCESS)
-               return 1;
-
-       status = uefi_call_wrapper(file->GetInfo, 4, file, &file_info_guid,
-                                  &buffersize, buffer);
-
-       if (status == EFI_BUFFER_TOO_SMALL) {
-               buffer = AllocatePool(buffersize);
-               status = uefi_call_wrapper(file->GetInfo, 4, file,
-                                          &file_info_guid, &buffersize,
-                                          buffer);
-       }
-
-       if (!buffer)
-               return 0;
-
-       buffersize = buffer->FileSize;
+static EFI_STATUS enroll_file (void *data, UINTN datasize, BOOLEAN hash)
+{
+       EFI_STATUS status = EFI_SUCCESS;
+       EFI_GUID shim_lock_guid = SHIM_LOCK_GUID;
+       EFI_SIGNATURE_LIST *CertList;
+       EFI_SIGNATURE_DATA *CertData;
+       UINTN mokbuffersize;
+       void *mokbuffer = NULL;
 
        if (hash) {
-               void *binary;
                UINT8 sha256[SHA256_DIGEST_SIZE];
                UINT8 sha1[SHA1_DIGEST_SIZE];
                SHIM_LOCK *shim_lock;
                EFI_GUID shim_guid = SHIM_LOCK_GUID;
                PE_COFF_LOADER_IMAGE_CONTEXT context;
-
+       
                status = LibLocateProtocol(&shim_guid, (VOID **)&shim_lock);
 
                if (status != EFI_SUCCESS)
@@ -1060,33 +1876,25 @@ static INTN file_callback (void *data, void *data2, void *data3) {
                if (!mokbuffer)
                        goto out;
 
-               binary = AllocatePool(buffersize);
-
-               status = uefi_call_wrapper(file->Read, 3, file, &buffersize,
-                                          binary);
-
-               if (status != EFI_SUCCESS)
-                       goto out;
-
-               status = shim_lock->Context(binary, buffersize, &context);
+               status = shim_lock->Context(data, datasize, &context);
 
                if (status != EFI_SUCCESS)
                        goto out;
-
-               status = shim_lock->Hash(binary, buffersize, &context, sha256,
+       
+               status = shim_lock->Hash(data, datasize, &context, sha256,
                                         sha1);
 
                if (status != EFI_SUCCESS)
                        goto out;
 
                CertList = mokbuffer;
-               CertList->SignatureType = EfiHashSha256Guid;
+               CertList->SignatureType = EFI_CERT_SHA256_GUID;
                CertList->SignatureSize = 16 + SHA256_DIGEST_SIZE;
                CertData = (EFI_SIGNATURE_DATA *)(((UINT8 *)mokbuffer) +
                                                  sizeof(EFI_SIGNATURE_LIST));
                CopyMem(CertData->SignatureData, sha256, SHA256_DIGEST_SIZE);
        } else {
-               mokbuffersize = buffersize + sizeof(EFI_SIGNATURE_LIST) +
+               mokbuffersize = datasize + sizeof(EFI_SIGNATURE_LIST) +
                        sizeof(EFI_GUID);
                mokbuffer = AllocatePool(mokbuffersize);
 
@@ -1094,13 +1902,12 @@ static INTN file_callback (void *data, void *data2, void *data3) {
                        goto out;
 
                CertList = mokbuffer;
-               CertList->SignatureType = EfiCertX509Guid;
-               CertList->SignatureSize = 16 + buffersize;
-               status = uefi_call_wrapper(file->Read, 3, file, &buffersize,
-                                mokbuffer + sizeof(EFI_SIGNATURE_LIST) + 16);
+               CertList->SignatureType = X509_GUID;
+               CertList->SignatureSize = 16 + datasize;
+
+               memcpy(mokbuffer + sizeof(EFI_SIGNATURE_LIST) + 16, data,
+                      datasize);
 
-               if (status != EFI_SUCCESS)
-                       goto out;
                CertData = (EFI_SIGNATURE_DATA *)(((UINT8 *)mokbuffer) +
                                                  sizeof(EFI_SIGNATURE_LIST));
        }
@@ -1110,329 +1917,150 @@ static INTN file_callback (void *data, void *data2, void *data3) {
        CertData->SignatureOwner = shim_lock_guid;
 
        if (!hash) {
-               if (!verify_certificate(CertData->SignatureData, buffersize))
+               if (!verify_certificate(CertData->SignatureData, datasize))
                        goto out;
        }
 
-       mok_enrollment_prompt(mokbuffer, mokbuffersize, FALSE);
+       mok_enrollment_prompt(mokbuffer, mokbuffersize, FALSE, FALSE);
 out:
-       if (buffer)
-               FreePool(buffer);
-
        if (mokbuffer)
                FreePool(mokbuffer);
 
-       return 0;
+       return status;
 }
 
-static INTN directory_callback (void *data, void *data2, void *data3) {
-       EFI_FILE_INFO *buffer = NULL;
-       UINTN buffersize = 0;
-       EFI_STATUS status;
-       UINTN dircount = 0, i = 0;
-       struct menu_item *dircontent;
-       EFI_FILE *dir;
-       CHAR16 *filename = data;
-       EFI_FILE *root = data2;
-       BOOLEAN hash = !!data3;
-
-       status = uefi_call_wrapper(root->Open, 5, root, &dir, filename,
-                                  EFI_FILE_MODE_READ, 0);
-
-       if (status != EFI_SUCCESS)
-               return 1;
-
-       while (1) {
-               status = uefi_call_wrapper(dir->Read, 3, dir, &buffersize,
-                                          buffer);
-
-               if (status == EFI_BUFFER_TOO_SMALL) {
-                       buffer = AllocatePool(buffersize);
-                       status = uefi_call_wrapper(dir->Read, 3, dir,
-                                                  &buffersize, buffer);
-               }
-
-               if (status != EFI_SUCCESS)
-                       return 1;
+static void mok_hash_enroll(void)
+{
+       EFI_STATUS efi_status;
+        CHAR16 *file_name = NULL;
+       EFI_HANDLE im = NULL;
+       EFI_FILE *file = NULL;
+       UINTN filesize;
+       void *data;
 
-               if (!buffersize)
-                       break;
+       simple_file_selector(&im, (CHAR16 *[]){
+      L"Select Binary",
+      L"",
+      L"The Selected Binary will have its hash Enrolled",
+      L"This means it will Subsequently Boot with no prompting",
+      L"Remember to make sure it is a genuine binary before Enroling its hash",
+      NULL
+             }, L"\\", L"", &file_name);
 
-               if ((StrCmp(buffer->FileName, L".") == 0) ||
-                   (StrCmp(buffer->FileName, L"..") == 0))
-                       continue;
+       if (!file_name)
+               return;
 
-               dircount++;
+       efi_status = simple_file_open(im, file_name, &file, EFI_FILE_MODE_READ);
 
-               FreePool(buffer);
-               buffersize = 0;
+       if (efi_status != EFI_SUCCESS) {
+               console_error(L"Unable to open file", efi_status);
+               return;
        }
 
-       dircount++;
+       simple_file_read_all(file, &filesize, &data);
+       simple_file_close(file);
 
-       dircontent = AllocatePool(sizeof(struct menu_item) * dircount);
+       if (!filesize) {
+               console_error(L"Unable to read file", efi_status);
+               return;
+       }
 
-       dircontent[0].text = StrDuplicate(L"..");
-       dircontent[0].callback = NULL;
-       dircontent[0].colour = EFI_YELLOW;
-       i++;
+       efi_status = enroll_file(data, filesize, TRUE);
 
-       uefi_call_wrapper(dir->SetPosition, 2, dir, 0);
+       if (efi_status != EFI_SUCCESS)
+               console_error(L"Hash failed (did you select a valid EFI binary?)", efi_status);
 
-       while (1) {
-               status = uefi_call_wrapper(dir->Read, 3, dir, &buffersize,
-                                          buffer);
+       FreePool(data);
+}
 
-               if (status == EFI_BUFFER_TOO_SMALL) {
-                       buffer = AllocatePool(buffersize);
-                       status = uefi_call_wrapper(dir->Read, 3, dir,
-                                                  &buffersize, buffer);
-               }
+static CHAR16 *der_suffix[] = {
+       L".cer",
+       L".der",
+       L".crt",
+       NULL
+};
 
-               if (status != EFI_SUCCESS)
-                       return 1;
+static BOOLEAN check_der_suffix (CHAR16 *file_name)
+{
+       CHAR16 suffix[5];
+       int i;
 
-               if (!buffersize)
-                       break;
+       if (!file_name || StrLen(file_name) <= 4)
+               return FALSE;
 
-               if ((StrCmp(buffer->FileName, L".") == 0) ||
-                   (StrCmp(buffer->FileName, L"..") == 0))
-                       continue;
+       suffix[0] = '\0';
+       StrCat(suffix, file_name + StrLen(file_name) - 4);
 
-               if (buffer->Attribute & EFI_FILE_DIRECTORY) {
-                       dircontent[i].text = StrDuplicate(buffer->FileName);
-                       dircontent[i].callback = directory_callback;
-                       dircontent[i].data = dircontent[i].text;
-                       dircontent[i].data2 = dir;
-                       dircontent[i].data3 = data3;
-                       dircontent[i].colour = EFI_YELLOW;
-               } else {
-                       dircontent[i].text = StrDuplicate(buffer->FileName);
-                       dircontent[i].callback = file_callback;
-                       dircontent[i].data = dircontent[i].text;
-                       dircontent[i].data2 = dir;
-                       dircontent[i].data3 = data3;
-                       dircontent[i].colour = EFI_WHITE;
+       StrLwr (suffix);
+       for (i = 0; der_suffix[i] != NULL; i++) {
+               if (StrCmp(suffix, der_suffix[i]) == 0) {
+                       return TRUE;
                }
-
-               i++;
-               FreePool(buffer);
-               buffersize = 0;
-               buffer = NULL;
        }
 
-       if (hash)
-               run_menu(HASH_STRING, 2, dircontent, dircount, 0);
-       else
-               run_menu(CERT_STRING, 2, dircontent, dircount, 0);
-
-       return 0;
+       return FALSE;
 }
 
-static INTN filesystem_callback (void *data, void *data2, void *data3) {
-       EFI_FILE_INFO *buffer = NULL;
-       UINTN buffersize = 0;
-       EFI_STATUS status;
-       UINTN dircount = 0, i = 0;
-       struct menu_item *dircontent;
-       EFI_FILE *root = data;
-       BOOLEAN hash = !!data3;
-
-       uefi_call_wrapper(root->SetPosition, 2, root, 0);
-
-       while (1) {
-               status = uefi_call_wrapper(root->Read, 3, root, &buffersize,
-                                          buffer);
-
-               if (status == EFI_BUFFER_TOO_SMALL) {
-                       buffer = AllocatePool(buffersize);
-                       status = uefi_call_wrapper(root->Read, 3, root,
-                                                  &buffersize, buffer);
-               }
-
-               if (status != EFI_SUCCESS)
-                       return 1;
-
-               if (!buffersize)
-                       break;
-
-               if ((StrCmp(buffer->FileName, L".") == 0) ||
-                   (StrCmp(buffer->FileName, L"..") == 0))
-                       continue;
-
-               dircount++;
-
-               FreePool(buffer);
-               buffersize = 0;
-       }
-
-       dircount++;
-
-       dircontent = AllocatePool(sizeof(struct menu_item) * dircount);
-
-       dircontent[0].text = StrDuplicate(L"Return to filesystem list");
-       dircontent[0].callback = NULL;
-       dircontent[0].colour = EFI_YELLOW;
-       i++;
-
-       uefi_call_wrapper(root->SetPosition, 2, root, 0);
-
-       while (1) {
-               status = uefi_call_wrapper(root->Read, 3, root, &buffersize,
-                                          buffer);
-
-               if (status == EFI_BUFFER_TOO_SMALL) {
-                       buffer = AllocatePool(buffersize);
-                       status = uefi_call_wrapper(root->Read, 3, root,
-                                                  &buffersize, buffer);
-               }
-
-               if (status != EFI_SUCCESS)
-                       return 1;
-
-               if (!buffersize)
-                       break;
+static void mok_key_enroll(void)
+{
+       EFI_STATUS efi_status;
+        CHAR16 *file_name = NULL;
+       EFI_HANDLE im = NULL;
+       EFI_FILE *file = NULL;
+       UINTN filesize;
+       void *data;
 
-               if ((StrCmp(buffer->FileName, L".") == 0) ||
-                   (StrCmp(buffer->FileName, L"..") == 0))
-                       continue;
+       simple_file_selector(&im, (CHAR16 *[]){
+      L"Select Key",
+      L"",
+      L"The selected key will be enrolled into the MOK database",
+      L"This means any binaries signed with it will be run without prompting",
+      L"Remember to make sure it is a genuine key before Enroling it",
+      NULL
+             }, L"\\", L"", &file_name);
 
-               if (buffer->Attribute & EFI_FILE_DIRECTORY) {
-                       dircontent[i].text = StrDuplicate(buffer->FileName);
-                       dircontent[i].callback = directory_callback;
-                       dircontent[i].data = dircontent[i].text;
-                       dircontent[i].data2 = root;
-                       dircontent[i].data3 = data3;
-                       dircontent[i].colour = EFI_YELLOW;
-               } else {
-                       dircontent[i].text = StrDuplicate(buffer->FileName);
-                       dircontent[i].callback = file_callback;
-                       dircontent[i].data = dircontent[i].text;
-                       dircontent[i].data2 = root;
-                       dircontent[i].data3 = data3;
-                       dircontent[i].colour = EFI_WHITE;
-               }
+       if (!file_name)
+               return;
 
-               i++;
-               FreePool(buffer);
-               buffer = NULL;
-               buffersize = 0;
+       if (!check_der_suffix(file_name)) {
+               console_alertbox((CHAR16 *[]){
+                       L"Unsupported Format",
+                       L"",
+                       L"Only DER encoded certificate (*.cer/der/crt) is supported",
+                       NULL});
+               return;
        }
 
-       if (hash)
-               run_menu(HASH_STRING, 2, dircontent, dircount, 0);
-       else
-               run_menu(CERT_STRING, 2, dircontent, dircount, 0);
-
-       return 0;
-}
-
-static INTN find_fs (void *data, void *data2, void *data3) {
-       EFI_GUID fs_guid = SIMPLE_FILE_SYSTEM_PROTOCOL;
-       UINTN count, i;
-       UINTN OldSize, NewSize;
-       EFI_HANDLE **filesystem_handles;
-       struct menu_item *filesystems;
-       BOOLEAN hash = !!data3;
-
-       uefi_call_wrapper(BS->LocateHandleBuffer, 5, ByProtocol, &fs_guid,
-                         NULL, &count, &filesystem_handles);
+       efi_status = simple_file_open(im, file_name, &file, EFI_FILE_MODE_READ);
 
-       if (!count || !filesystem_handles) {
-               Print(L"No filesystems?\n");
-               return 1;
+       if (efi_status != EFI_SUCCESS) {
+               console_error(L"Unable to open file", efi_status);
+               return;
        }
 
-       count++;
-
-       filesystems = AllocatePool(sizeof(struct menu_item) * count);
-
-       filesystems[0].text = StrDuplicate(L"Exit");
-       filesystems[0].callback = NULL;
-       filesystems[0].colour = EFI_YELLOW;
-
-       for (i=1; i<count; i++) {
-               EFI_HANDLE *fs = filesystem_handles[i-1];
-               EFI_FILE_IO_INTERFACE *fs_interface;
-               EFI_DEVICE_PATH *path;
-               EFI_FILE *root;
-               EFI_STATUS status;
-               CHAR16 *VolumeLabel = NULL;
-               EFI_FILE_SYSTEM_INFO *buffer = NULL;
-               UINTN buffersize = 0;
-               EFI_GUID file_info_guid = EFI_FILE_INFO_ID;
-
-               status = uefi_call_wrapper(BS->HandleProtocol, 3, fs, &fs_guid,
-                                          &fs_interface);
-
-               if (status != EFI_SUCCESS || !fs_interface)
-                       continue;
-
-               path = DevicePathFromHandle(fs);
-
-               status = uefi_call_wrapper(fs_interface->OpenVolume, 2,
-                                          fs_interface, &root);
-
-               if (status != EFI_SUCCESS || !root)
-                       continue;
-
-               status = uefi_call_wrapper(root->GetInfo, 4, root,
-                                          &file_info_guid, &buffersize,
-                                          buffer);
-
-               if (status == EFI_BUFFER_TOO_SMALL) {
-                       buffer = AllocatePool(buffersize);
-                       status = uefi_call_wrapper(root->GetInfo, 4, root,
-                                                  &file_info_guid,
-                                                  &buffersize, buffer);
-               }
-
-               if (status == EFI_SUCCESS)
-                       VolumeLabel = buffer->VolumeLabel;
-
-               if (path)
-                       filesystems[i].text = DevicePathToStr(path);
-               else
-                       filesystems[i].text = StrDuplicate(L"Unknown device\n");
-               if (VolumeLabel) {
-                       OldSize = (StrLen(filesystems[i].text) + 1) * sizeof(CHAR16);
-                       NewSize = OldSize + StrLen(VolumeLabel) * sizeof(CHAR16);
-                       filesystems[i].text = ReallocatePool(filesystems[i].text,
-                                                            OldSize, NewSize);
-                       StrCat(filesystems[i].text, VolumeLabel);
-               }
-
-               if (buffersize)
-                       FreePool(buffer);
+       simple_file_read_all(file, &filesize, &data);
+       simple_file_close(file);
 
-               filesystems[i].data = root;
-               filesystems[i].data2 = NULL;
-               filesystems[i].data3 = data3;
-               filesystems[i].callback = filesystem_callback;
-               filesystems[i].colour = EFI_YELLOW;
+       if (!filesize) {
+               console_error(L"Unable to read file", efi_status);
+               return;
        }
 
-       uefi_call_wrapper(BS->FreePool, 1, filesystem_handles);
-
-       if (hash)
-               run_menu(HASH_STRING, 2, filesystems, count, 0);
-       else
-               run_menu(CERT_STRING, 2, filesystems, count, 0);
-
-       return 0;
+       enroll_file(data, filesize, FALSE);
+       FreePool(data);
 }
 
-static BOOLEAN verify_pw(void)
+static BOOLEAN verify_pw(BOOLEAN *protected)
 {
        EFI_GUID shim_lock_guid = SHIM_LOCK_GUID;
        EFI_STATUS efi_status;
-       CHAR16 password[PASSWORD_MAX];
-       UINT8 fail_count = 0;
-       UINT8 hash[SHA256_DIGEST_SIZE];
-       UINT8 pwhash[SHA256_DIGEST_SIZE];
-       UINTN size = SHA256_DIGEST_SIZE;
-       UINT32 length;
+       SIMPLE_TEXT_OUTPUT_MODE SavedMode;
+       UINT8 pwhash[PASSWORD_CRYPT_SIZE];
+       UINTN size = PASSWORD_CRYPT_SIZE;
        UINT32 attributes;
+       CHAR16 *message[2];
+
+       *protected = FALSE;
 
        efi_status = uefi_call_wrapper(RT->GetVariable, 5, L"MokPWStore",
                                       &shim_lock_guid, &attributes, &size,
@@ -1443,7 +2071,8 @@ static BOOLEAN verify_pw(void)
         * known value, so there's no safety advantage in failing to validate
         * purely because of a failure to read the variable
         */
-       if (efi_status != EFI_SUCCESS)
+       if (efi_status != EFI_SUCCESS ||
+           (size != SHA256_DIGEST_SIZE && size != PASSWORD_CRYPT_SIZE))
                return TRUE;
 
        if (attributes & EFI_VARIABLE_RUNTIME_ACCESS)
@@ -1451,173 +2080,427 @@ static BOOLEAN verify_pw(void)
 
        uefi_call_wrapper(ST->ConOut->ClearScreen, 1, ST->ConOut);
 
-       while (fail_count < 3) {
-               Print(L"Enter MOK password: ");
-               get_line(&length, password, PASSWORD_MAX, 0);
+       /* Draw the background */
+       console_save_and_set_mode(&SavedMode);
+       message[0] = PoolPrint (L"%s UEFI key management", SHIM_VENDOR);
+       message[1] = NULL;
+       console_print_box_at(message, -1, 0, 0, -1, -1, 1, 1);
+       FreePool(message[0]);
+       console_restore_mode(&SavedMode);
+
+       if (size == PASSWORD_CRYPT_SIZE) {
+               efi_status = match_password((PASSWORD_CRYPT *)pwhash, NULL, 0,
+                                           NULL, L"Enter MOK password:");
+       } else {
+               efi_status = match_password(NULL, NULL, 0, pwhash,
+                                           L"Enter MOK password:");
+       }
+       if (efi_status != EFI_SUCCESS) {
+               console_notify(L"Password limit reached");
+               return FALSE;
+       }
 
-               if (length < PASSWORD_MIN || length > PASSWORD_MAX) {
-                       Print(L"Invalid password length\n");
-                       fail_count++;
-                       continue;
-               }
+       *protected = TRUE;
 
-               efi_status = compute_pw_hash(NULL, 0, password, length, hash);
+       return TRUE;
+}
 
-               if (efi_status != EFI_SUCCESS) {
-                       Print(L"Unable to generate password hash\n");
-                       fail_count++;
-                       continue;
-               }
+static int draw_countdown()
+{
+       SIMPLE_TEXT_OUTPUT_MODE SavedMode;
+       EFI_INPUT_KEY key;
+       EFI_STATUS status;
+       UINTN cols, rows;
+       CHAR16 *title[2];
+       CHAR16 *message =  L"Press any key to perform MOK management";
+       int timeout = 10, wait = 10000000;
 
-               if (CompareMem(pwhash, hash, SHA256_DIGEST_SIZE) != 0) {
-                       Print(L"Password doesn't match\n");
-                       fail_count++;
-                       continue;
+       console_save_and_set_mode (&SavedMode);
+
+       title[0] = PoolPrint (L"%s UEFI key management", SHIM_VENDOR);
+       title[1] = NULL;
+
+       console_print_box_at(title, -1, 0, 0, -1, -1, 1, 1);
+
+       uefi_call_wrapper(ST->ConOut->QueryMode, 4, ST->ConOut,
+                         ST->ConOut->Mode->Mode, &cols, &rows);
+
+       PrintAt((cols - StrLen(message))/2, rows/2, message);
+       while (1) {
+               if (timeout > 1)
+                       PrintAt(2, rows - 3, L"Booting in %d seconds  ", timeout);
+               else if (timeout)
+                       PrintAt(2, rows - 3, L"Booting in %d second   ", timeout);
+
+               status = WaitForSingleEvent(ST->ConIn->WaitForKey, wait);
+
+               if (status != EFI_TIMEOUT) {
+                       /* Clear the key in the queue */
+                       uefi_call_wrapper(ST->ConIn->ReadKeyStroke, 2,
+                                         ST->ConIn, &key);
+                       break;
                }
 
-               return TRUE;
+               timeout--;
+               if (!timeout)
+                       break;
        }
 
-       Print(L"Password limit reached\n");
-       return FALSE;
+       FreePool(title[0]);
+
+       console_restore_mode(&SavedMode);
+
+       return timeout;
 }
 
-static EFI_STATUS enter_mok_menu(EFI_HANDLE image_handle, void *MokNew,
-                                UINTN MokNewSize, void *MokSB,
-                                UINTN MokSBSize, void *MokPW, UINTN MokPWSize)
+typedef enum {
+       MOK_CONTINUE_BOOT,
+       MOK_RESET_MOK,
+       MOK_RESET_MOKX,
+       MOK_ENROLL_MOK,
+       MOK_ENROLL_MOKX,
+       MOK_DELETE_MOK,
+       MOK_DELETE_MOKX,
+       MOK_CHANGE_SB,
+       MOK_SET_PW,
+       MOK_CHANGE_DB,
+       MOK_KEY_ENROLL,
+       MOK_HASH_ENROLL
+} mok_menu_item;
+
+static EFI_STATUS enter_mok_menu(EFI_HANDLE image_handle,
+                                void *MokNew, UINTN MokNewSize,
+                                void *MokDel, UINTN MokDelSize,
+                                void *MokSB, UINTN MokSBSize,
+                                void *MokPW, UINTN MokPWSize,
+                                void *MokDB, UINTN MokDBSize,
+                                void *MokXNew, UINTN MokXNewSize,
+                                void *MokXDel, UINTN MokXDelSize)
 {
-       struct menu_item *menu_item;
+       CHAR16 **menu_strings;
+       mok_menu_item *menu_item;
+       int choice = 0;
        UINT32 MokAuth = 0;
+       UINT32 MokDelAuth = 0;
+       UINT32 MokXAuth = 0;
+       UINT32 MokXDelAuth = 0;
        UINTN menucount = 3, i = 0;
        EFI_STATUS efi_status;
        EFI_GUID shim_lock_guid = SHIM_LOCK_GUID;
-       UINT8 auth[SHA256_DIGEST_SIZE];
-       UINTN auth_size = SHA256_DIGEST_SIZE;
+       UINT8 auth[PASSWORD_CRYPT_SIZE];
+       UINTN auth_size = PASSWORD_CRYPT_SIZE;
        UINT32 attributes;
+       BOOLEAN protected;
+       EFI_STATUS ret = EFI_SUCCESS;
 
-       if (verify_pw() == FALSE)
+       if (verify_pw(&protected) == FALSE)
                return EFI_ACCESS_DENIED;
-
+       
        efi_status = uefi_call_wrapper(RT->GetVariable, 5, L"MokAuth",
                                               &shim_lock_guid,
                                               &attributes, &auth_size, auth);
 
-       if ((efi_status == EFI_SUCCESS) && (auth_size == SHA256_DIGEST_SIZE))
+       if ((efi_status == EFI_SUCCESS) &&
+           (auth_size == SHA256_DIGEST_SIZE || auth_size == PASSWORD_CRYPT_SIZE))
                MokAuth = 1;
 
+       efi_status = uefi_call_wrapper(RT->GetVariable, 5, L"MokDelAuth",
+                                              &shim_lock_guid,
+                                              &attributes, &auth_size, auth);
+
+       if ((efi_status == EFI_SUCCESS) &&
+          (auth_size == SHA256_DIGEST_SIZE || auth_size == PASSWORD_CRYPT_SIZE))
+               MokDelAuth = 1;
+
+       efi_status = uefi_call_wrapper(RT->GetVariable, 5, L"MokXAuth",
+                                              &shim_lock_guid,
+                                              &attributes, &auth_size, auth);
+
+       if ((efi_status == EFI_SUCCESS) &&
+           (auth_size == SHA256_DIGEST_SIZE || auth_size == PASSWORD_CRYPT_SIZE))
+               MokXAuth = 1;
+
+       efi_status = uefi_call_wrapper(RT->GetVariable, 5, L"MokXDelAuth",
+                                              &shim_lock_guid,
+                                              &attributes, &auth_size, auth);
+
+       if ((efi_status == EFI_SUCCESS) &&
+          (auth_size == SHA256_DIGEST_SIZE || auth_size == PASSWORD_CRYPT_SIZE))
+               MokXDelAuth = 1;
+
        if (MokNew || MokAuth)
                menucount++;
 
+       if (MokDel || MokDelAuth)
+               menucount++;
+
+       if (MokXNew || MokXAuth)
+               menucount++;
+
+       if (MokXDel || MokXDelAuth)
+               menucount++;
+
        if (MokSB)
                menucount++;
 
        if (MokPW)
                menucount++;
 
-       menu_item = AllocateZeroPool(sizeof(struct menu_item) * menucount);
+       if (MokDB)
+               menucount++;
+
+       menu_strings = AllocateZeroPool(sizeof(CHAR16 *) * (menucount + 1));
+
+       if (!menu_strings)
+               return EFI_OUT_OF_RESOURCES;
+
+       menu_item = AllocateZeroPool(sizeof(mok_menu_item) * menucount);
 
-       if (!menu_item)
+       if (!menu_item) {
+               FreePool(menu_strings);
                return EFI_OUT_OF_RESOURCES;
+       }
 
-       menu_item[i].text = StrDuplicate(L"Continue boot");
-       menu_item[i].colour = EFI_WHITE;
-       menu_item[i].callback = NULL;
+       menu_strings[i] = L"Continue boot";
+       menu_item[i] = MOK_CONTINUE_BOOT;
 
        i++;
 
        if (MokNew || MokAuth) {
                if (!MokNew) {
-                       menu_item[i].text = StrDuplicate(L"Delete MOK");
-                       menu_item[i].colour = EFI_WHITE;
-                       menu_item[i].callback = mok_deletion_prompt;
+                       menu_strings[i] = L"Reset MOK";
+                       menu_item[i] = MOK_RESET_MOK;
+               } else {
+                       menu_strings[i] = L"Enroll MOK";
+                       menu_item[i] = MOK_ENROLL_MOK;
+               }
+               i++;
+       }
+
+       if (MokDel || MokDelAuth) {
+               menu_strings[i] = L"Delete MOK";
+               menu_item[i] = MOK_DELETE_MOK;
+               i++;
+       }
+
+       if (MokXNew || MokXAuth) {
+               if (!MokXNew) {
+                       menu_strings[i] = L"Reset MOKX";
+                       menu_item[i] = MOK_RESET_MOKX;
                } else {
-                       menu_item[i].text = StrDuplicate(L"Enroll MOK");
-                       menu_item[i].colour = EFI_WHITE;
-                       menu_item[i].data = MokNew;
-                       menu_item[i].data2 = (void *)MokNewSize;
-                       menu_item[i].callback = mok_enrollment_prompt_callback;
+                       menu_strings[i] = L"Enroll MOKX";
+                       menu_item[i] = MOK_ENROLL_MOKX;
                }
                i++;
        }
 
+       if (MokXDel || MokXDelAuth) {
+               menu_strings[i] = L"Delete MOKX";
+               menu_item[i] = MOK_DELETE_MOKX;
+               i++;
+       }
+
        if (MokSB) {
-               menu_item[i].text = StrDuplicate(L"Change Secure Boot state");
-               menu_item[i].colour = EFI_WHITE;
-               menu_item[i].callback = mok_sb_prompt;
-               menu_item[i].data = MokSB;
-               menu_item[i].data2 = (void *)MokSBSize;
+               menu_strings[i] = L"Change Secure Boot state";
+               menu_item[i] = MOK_CHANGE_SB;
                i++;
        }
 
        if (MokPW) {
-               menu_item[i].text = StrDuplicate(L"Set MOK password");
-               menu_item[i].colour = EFI_WHITE;
-               menu_item[i].callback = mok_pw_prompt;
-               menu_item[i].data = MokPW;
-               menu_item[i].data2 = (void *)MokPWSize;
+               menu_strings[i] = L"Set MOK password";
+               menu_item[i] = MOK_SET_PW;
                i++;
        }
 
-       menu_item[i].text = StrDuplicate(L"Enroll key from disk");
-       menu_item[i].colour = EFI_WHITE;
-       menu_item[i].callback = find_fs;
-       menu_item[i].data3 = (void *)FALSE;
+       if (MokDB) {
+               menu_strings[i] = L"Change DB state";
+               menu_item[i] = MOK_CHANGE_DB;
+               i++;
+       }
 
+       menu_strings[i] = L"Enroll key from disk";
+       menu_item[i] = MOK_KEY_ENROLL;
        i++;
 
-       menu_item[i].text = StrDuplicate(L"Enroll hash from disk");
-       menu_item[i].colour = EFI_WHITE;
-       menu_item[i].callback = find_fs;
-       menu_item[i].data3 = (void *)TRUE;
-
+       menu_strings[i] = L"Enroll hash from disk";
+       menu_item[i] = MOK_HASH_ENROLL;
        i++;
 
-       run_menu(NULL, 0, menu_item, menucount, 10);
+       menu_strings[i] = NULL;
 
-       uefi_call_wrapper(ST->ConOut->ClearScreen, 1, ST->ConOut);
+       if (protected == FALSE && draw_countdown() == 0)
+               goto out;
 
-       return 0;
+       while (choice >= 0) {
+               choice = console_select((CHAR16 *[]){ L"Perform MOK management", NULL },
+                                       menu_strings, 0);
+
+               if (choice < 0)
+                       goto out;
+
+               switch (menu_item[choice]) {
+               case MOK_CONTINUE_BOOT:
+                       goto out;
+               case MOK_RESET_MOK:
+                       mok_reset_prompt(FALSE);
+                       break;
+               case MOK_ENROLL_MOK:
+                       mok_enrollment_prompt(MokNew, MokNewSize, TRUE, FALSE);
+                       break;
+               case MOK_DELETE_MOK:
+                       mok_deletion_prompt(MokDel, MokDelSize, FALSE);
+                       break;
+               case MOK_RESET_MOKX:
+                       mok_reset_prompt(TRUE);
+                       break;
+               case MOK_ENROLL_MOKX:
+                       mok_enrollment_prompt(MokXNew, MokXNewSize, TRUE, TRUE);
+                       break;
+               case MOK_DELETE_MOKX:
+                       mok_deletion_prompt(MokXDel, MokXDelSize, TRUE);
+                       break;
+               case MOK_CHANGE_SB:
+                       mok_sb_prompt(MokSB, MokSBSize);
+                       break;
+               case MOK_SET_PW:
+                       mok_pw_prompt(MokPW, MokPWSize);
+                       break;
+               case MOK_CHANGE_DB:
+                       mok_db_prompt(MokDB, MokDBSize);
+                       break;
+               case MOK_KEY_ENROLL:
+                       mok_key_enroll();
+                       break;
+               case MOK_HASH_ENROLL:
+                       mok_hash_enroll();
+                       break;
+               }
+       }
+
+out:
+       console_reset();
+
+       FreePool(menu_strings);
+
+       if (menu_item)
+               FreePool(menu_item);
+
+       return ret;
 }
 
 static EFI_STATUS check_mok_request(EFI_HANDLE image_handle)
 {
        EFI_GUID shim_lock_guid = SHIM_LOCK_GUID;
-       UINTN MokNewSize = 0, MokSBSize = 0, MokPWSize = 0;
+       UINTN MokNewSize = 0, MokDelSize = 0, MokSBSize = 0, MokPWSize = 0;
+       UINTN MokDBSize = 0, MokXNewSize = 0, MokXDelSize = 0;
        void *MokNew = NULL;
+       void *MokDel = NULL;
        void *MokSB = NULL;
        void *MokPW = NULL;
+       void *MokDB = NULL;
+       void *MokXNew = NULL;
+       void *MokXDel = NULL;
+       EFI_STATUS status;
 
-       MokNew = LibGetVariableAndSize(L"MokNew", &shim_lock_guid, &MokNewSize);
+       status = get_variable(L"MokNew", (UINT8 **)&MokNew, &MokNewSize,
+                               shim_lock_guid);
+       if (status == EFI_SUCCESS) {
+               if (LibDeleteVariable(L"MokNew", &shim_lock_guid) != EFI_SUCCESS) {
+                       console_notify(L"Failed to delete MokNew");
+               }
+       } else if (EFI_ERROR(status) && status != EFI_NOT_FOUND) {
+               console_error(L"Could not retrieve MokNew", status);
+       }
 
-       MokSB = LibGetVariableAndSize(L"MokSB", &shim_lock_guid, &MokSBSize);
+       status = get_variable(L"MokDel", (UINT8 **)&MokDel, &MokDelSize,
+                               shim_lock_guid);
+       if (status == EFI_SUCCESS) {
+               if (LibDeleteVariable(L"MokDel", &shim_lock_guid) != EFI_SUCCESS) {
+                       console_notify(L"Failed to delete MokDel");
+               }
+       } else if (EFI_ERROR(status) && status != EFI_NOT_FOUND) {
+               console_error(L"Could not retrieve MokDel", status);
+       }
 
-       MokPW = LibGetVariableAndSize(L"MokPW", &shim_lock_guid, &MokPWSize);
+       status = get_variable(L"MokSB", (UINT8 **)&MokSB, &MokSBSize,
+                               shim_lock_guid);
+       if (status == EFI_SUCCESS) {
+               if (LibDeleteVariable(L"MokSB", &shim_lock_guid) != EFI_SUCCESS) {
+                       console_notify(L"Failed to delete MokSB");
+               }
+       } else if (EFI_ERROR(status) && status != EFI_NOT_FOUND) {
+               console_error(L"Could not retrieve MokSB", status);
+       }
 
-       enter_mok_menu(image_handle, MokNew, MokNewSize, MokSB, MokSBSize,
-                      MokPW, MokPWSize);
+       status = get_variable(L"MokPW", (UINT8 **)&MokPW, &MokPWSize,
+                               shim_lock_guid);
+       if (status == EFI_SUCCESS) {
+               if (LibDeleteVariable(L"MokPW", &shim_lock_guid) != EFI_SUCCESS) {
+                       console_notify(L"Failed to delete MokPW");
+               }
+       } else if (EFI_ERROR(status) && status != EFI_NOT_FOUND) {
+               console_error(L"Could not retrieve MokPW", status);
+       }
 
-       if (MokNew) {
-               if (LibDeleteVariable(L"MokNew", &shim_lock_guid) != EFI_SUCCESS) {
-                       Print(L"Failed to delete MokNew\n");
+       status = get_variable(L"MokDB", (UINT8 **)&MokDB, &MokDBSize,
+                               shim_lock_guid);
+       if (status == EFI_SUCCESS) {
+               if (LibDeleteVariable(L"MokDB", &shim_lock_guid) != EFI_SUCCESS) {
+                       console_notify(L"Failed to delete MokDB");
                }
-               FreePool (MokNew);
+       } else if (EFI_ERROR(status) && status != EFI_NOT_FOUND) {
+               console_error(L"Could not retrieve MokDB", status);
        }
 
-       if (MokSB) {
-               if (LibDeleteVariable(L"MokSB", &shim_lock_guid) != EFI_SUCCESS) {
-                       Print(L"Failed to delete MokSB\n");
+       status = get_variable(L"MokXNew", (UINT8 **)&MokXNew, &MokXNewSize,
+                               shim_lock_guid);
+       if (status == EFI_SUCCESS) {
+               if (LibDeleteVariable(L"MokXNew", &shim_lock_guid) != EFI_SUCCESS) {
+                       console_notify(L"Failed to delete MokXNew");
                }
-               FreePool (MokNew);
+       } else if (EFI_ERROR(status) && status != EFI_NOT_FOUND) {
+               console_error(L"Could not retrieve MokXNew", status);
        }
 
-       if (MokPW) {
-               if (LibDeleteVariable(L"MokPW", &shim_lock_guid) != EFI_SUCCESS) {
-                       Print(L"Failed to delete MokPW\n");
+       status = get_variable(L"MokXDel", (UINT8 **)&MokXDel, &MokXDelSize,
+                               shim_lock_guid);
+       if (status == EFI_SUCCESS) {
+               if (LibDeleteVariable(L"MokXDel", &shim_lock_guid) != EFI_SUCCESS) {
+                       console_notify(L"Failed to delete MokXDel");
                }
-               FreePool (MokNew);
+       } else if (EFI_ERROR(status) && status != EFI_NOT_FOUND) {
+               console_error(L"Could not retrieve MokXDel", status);
        }
 
+       enter_mok_menu(image_handle, MokNew, MokNewSize, MokDel, MokDelSize,
+                      MokSB, MokSBSize, MokPW, MokPWSize, MokDB, MokDBSize,
+                      MokXNew, MokXNewSize, MokXDel, MokXDelSize);
+
+       if (MokNew)
+               FreePool (MokNew);
+
+       if (MokDel)
+               FreePool (MokDel);
+
+       if (MokSB)
+               FreePool (MokSB);
+
+       if (MokPW)
+               FreePool (MokPW);
+
+       if (MokDB)
+               FreePool (MokDB);
+
+       if (MokXNew)
+               FreePool (MokXNew);
+
+       if (MokXDel)
+               FreePool (MokXDel);
+
        LibDeleteVariable(L"MokAuth", &shim_lock_guid);
+       LibDeleteVariable(L"MokDelAuth", &shim_lock_guid);
+       LibDeleteVariable(L"MokXAuth", &shim_lock_guid);
+       LibDeleteVariable(L"MokXDelAuth", &shim_lock_guid);
 
        return EFI_SUCCESS;
 }
@@ -1653,9 +2536,12 @@ EFI_STATUS efi_main (EFI_HANDLE image_handle, EFI_SYSTEM_TABLE *systab)
 
        InitializeLib(image_handle, systab);
 
+       setup_console(1);
+
        setup_rand();
 
        efi_status = check_mok_request(image_handle);
 
+       setup_console(0);
        return efi_status;
 }