/** @file

;******************************************************************************
;* Copyright (c) 2018, Insyde Software Corp. All Rights Reserved.
;*
;* You may not reproduce, distribute, publish, display, perform, modify, adapt,
;* transmit, broadcast, present, recite, release, license or otherwise exploit
;* any part of this publication in any form, by any means, without the prior
;* written permission of Insyde Software Corporation.
;*
;******************************************************************************
*/
#include "WmiSetupUnderOsSmm.h"

BOOLEAN                                 mAdminPasswordExist = FALSE;
BOOLEAN                                 mUserPasswordExist  = FALSE;
BOOLEAN                                 mHddPasswordExist   = FALSE;
BOOLEAN                                 mPasswordChange     = FALSE;
UINTN                                   mPasswordErrorCount = 0;
WMI_PASSWORD_DATA_MAP                   mPasswordCheckData;

//
// Hexadecimal to Unicode char
//
STATIC UINT16                           HexToChar16[] = {
  L'0',
  L'1',
  L'2',
  L'3',
  L'4',
  L'5',
  L'6',
  L'7',
  L'8',
  L'9',
  L'A',
  L'B',
  L'C',
  L'D',
  L'E',
  L'F'
};

//
// Refer from EFI_KEYBOARD_SCANCODE_LIST at SetupMenuService.c
//
// WMI keyboard scancode list
//
WMI_KEYBOARD_SCANCODE_LIST              mScanCodeMap[] = {
  {'a', 0x1E},
  {'b', 0x30},
  {'c', 0x2E},
  {'d', 0x20},
  {'e', 0x12},
  {'f', 0x21},
  {'g', 0x22},
  {'h', 0x23},
  {'i', 0x17},
  {'j', 0x24},
  {'k', 0x25},
  {'l', 0x26},
  {'m', 0x32},
  {'n', 0x31},
  {'o', 0x18},
  {'p', 0x19},
  {'q', 0x10},
  {'r', 0x13},
  {'s', 0x1F},
  {'t', 0x14},
  {'u', 0x16},
  {'v', 0x2F},
  {'w', 0x11},
  {'x', 0x2D},
  {'y', 0x15},
  {'z', 0x2C},
  {'1', 0x02},
  {'2', 0x03},
  {'3', 0x04},
  {'4', 0x05},
  {'5', 0x06},
  {'6', 0x07},
  {'7', 0x08},
  {'8', 0x09},
  {'9', 0x0A},
  {'0', 0x0B},
#ifdef L05_SMB_BIOS_ENABLE
  {' ', 0x39},
#endif
};

//
// WMI password type map
//
WMI_PASSWORD_TYPE_MAP                   mWmiPasswordTypeMap[] = {
// PasswordType,       PasswordTypeStr
  {WmiSystemAdminType, WMI_BIOS_PASSWORD_ADMIN_TYPE     },
  {WmiSystemUserType,  WMI_BIOS_PASSWORD_USER_TYPE      },
  {WmiHddUserType,     WMI_BIOS_PASSWORD_HDD_USER_TYPE  },
  {WmiHddMasterType,   WMI_BIOS_PASSWORD_HDD_MASTER_TYPE},
};

//
// WMI password encoding map
//
WMI_PASSWORD_ENCODING_MAP               mWmiPasswordEncodingMap[] = {
// PasswordEncoding,    PasswordEncodingStr
  {WmiAsciiEncoding,    WMI_BIOS_PASSWORD_ASCII_ENCODING   },
  {WmiScancodeEncoding, WMI_BIOS_PASSWORD_SCANCODE_ENCODING},
};

//
// WMI password language map
//
WMI_PASSWORD_LANGUAGE_MAP               mWmiPasswordLanguageMap[] = {
// PasswordLanguage, PasswordLanguageStr
  {WmiUsLanguage,    WMI_BIOS_PASSWORD_US_LANGUAGE},
  {WmiFrLanguage,    WMI_BIOS_PASSWORD_FR_LANGUAGE},
  {WmiGrLanguage,    WMI_BIOS_PASSWORD_GR_LANGUAGE},
};

//
// Map List Count
//
UINT16                                  mWmiPasswordTypeMapListCount     = sizeof (mWmiPasswordTypeMap)     / sizeof (WMI_PASSWORD_TYPE_MAP);
UINT16                                  mWmiPasswordEncodingMapListCount = sizeof (mWmiPasswordEncodingMap) / sizeof (WMI_PASSWORD_ENCODING_MAP);
UINT16                                  mWmiPasswordLanguageMapListCount = sizeof (mWmiPasswordLanguageMap) / sizeof (WMI_PASSWORD_LANGUAGE_MAP);

/**
  Refer from ConvertUnicodeCharToKeyboardScancode() at SetupMenuService.c

  Convert Ascii Char to Keyboard Scancode(Normal).

  @param  KeyData                       Point to Key Data(Ascii).

  @retval EFI_SUCCESS                   The function completed successfully.
  @retval EFI_NOT_FOUND                 The Ascii Char can't map to Keyboard Scancode.
**/
EFI_STATUS
ConvertAsciiCharToKeyboardScancode (
  OUT CHAR8                             *KeyData
  )
{
  UINTN                                 Index;

  for (Index = 0; Index < (sizeof (mScanCodeMap) / sizeof (WMI_KEYBOARD_SCANCODE_LIST)); Index++) {
#ifdef L05_SMB_BIOS_ENABLE
    //
    // [Lenovo SMB BIOS Special Requirements V1.1]
    //   2.1 BIOS Password Characters
    //     Supported characters
    //     - Alphabet (case insensitive)
    //
    if (*KeyData >= L'A' && *KeyData <= L'Z') {
      *KeyData += 0x20;
    }
#endif
    if (*KeyData == mScanCodeMap[Index].AsciiChar) {
      *KeyData = mScanCodeMap[Index].KeyboardScanCode;
      return EFI_SUCCESS;
    }
  }

  return EFI_NOT_FOUND;
}

/**
  This function uses to check password type.

  @param  PasswordTypeStr               A pointer to the password type string.
  @param  PasswordType                  A pointer to the system admin password or user password.

  @retval TRUE                          Password type is valid.
  @retval FALSE                         Password type is not valid.
**/
BOOLEAN
CheckPasswordTypeIsValid (
  IN  CHAR8                             *PasswordTypeStr,
  OUT WMI_PASSWORD_TYPE                 *PasswordType
  )
{
  BOOLEAN                               IsValid;
  UINTN                                 Index;

  //
  // Initialization
  //
  IsValid = FALSE;

  //
  // Compare password type string to get password type
  //
  for (Index = 0; Index < mWmiPasswordTypeMapListCount; Index++) {
    if (AsciiStrnCmp (PasswordTypeStr, mWmiPasswordTypeMap[Index].PasswordTypeStr, AsciiStrLen (mWmiPasswordTypeMap[Index].PasswordTypeStr)) == 0) {
      *PasswordType = mWmiPasswordTypeMap[Index].PasswordType;
      IsValid = TRUE;
      break;
    }
  }

  return IsValid;
}

/**
  Refer from IsValidKeyData() at SetupMenuService.c

  Check password is valid.

  @param  Password                      A pointer to the Password buffer.
  @param  PasswordType                  Type of password.

  @retval TRUE                          Password is valid.
  @retval FALSE                         Password is not valid.
**/
BOOLEAN
CheckPasswordIsValid (
  IN CHAR8                              **Password,
  IN WMI_PASSWORD_TYPE                  PasswordType
  )
{
  BOOLEAN                               IsValid;
  EFI_STATUS                            Status;
  UINTN                                 MaxPasswordLen;
  UINTN                                 Index;

  //
  // Initialization
  //
  IsValid        = TRUE;
  MaxPasswordLen = 0;

  //
  // Step 1. Check password length
  //
  switch (PasswordType) {

  case WmiSystemAdminType:
  case WmiSystemUserType:
    MaxPasswordLen = L05_WMI_PASSWORD_MAX_LENGTH;
    break;

  case WmiHddUserType:
  case WmiHddMasterType:
    MaxPasswordLen = HDD_PASSWORD_MAX_NUMBER;
    break;
  }

  if (AsciiStrLen (*Password) > MaxPasswordLen) {
    IsValid = FALSE;
  }

  //
  // Step 2. Only 'a'~'z', '0'~'9' are valid key data for password Character
  //
  for (Index = 0; Index < AsciiStrLen (*Password); Index++) {
    if (((*Password)[Index] < L'a' || (*Password)[Index] > L'z')
      && ((*Password)[Index] < L'0' || (*Password)[Index] > L'9')
#ifdef L05_SMB_BIOS_ENABLE
      //
      // [Lenovo SMB BIOS Special Requirements V1.1]
      //   2.1 BIOS Password Characters
      //     Supported characters
      //     - Alphabet (case insensitive)
      //     - Space ' '
      //
      && ((*Password)[Index] < L'A' || (*Password)[Index] > L'Z')
      && ((*Password)[Index] != L' ')
#endif
    ) {

      IsValid = FALSE;
      break;
    }

    //
    // Step 3. Convert ASCII char to keyboard scancode
    //
    Status = ConvertAsciiCharToKeyboardScancode (&((*Password)[Index]));

    if (EFI_ERROR (Status)) {
      IsValid = FALSE;
      break;
    }
  }

  return IsValid;
}

/**
  This function uses to check password encoding.

  @param  PasswordEncodingStr           A pointer to the password encoding string.
  @param  PasswordEncoding              A pointer to the password encoding.

  @retval TRUE                          Password encoding is valid.
  @retval FALSE                         Password encoding is not valid.
**/
BOOLEAN
CheckPasswordEncodingIsValid (
  IN  CHAR8                             *PasswordEncodingStr,
  OUT WMI_PASSWORD_TYPE                 *PasswordEncoding
  )
{
  BOOLEAN                               IsValid;
  UINTN                                 Index;

  //
  // Initialization
  //
  IsValid = FALSE;

  //
  // Compare password encoding string to get password encoding
  //
  for (Index = 0; Index < mWmiPasswordEncodingMapListCount; Index++) {
    if (AsciiStrCmp (PasswordEncodingStr, mWmiPasswordEncodingMap[Index].PasswordEncodingStr) == 0) {
      *PasswordEncoding = mWmiPasswordEncodingMap[Index].PasswordEncoding;
      IsValid = TRUE;
      break;
    }
  }

  return IsValid;
}

/**
  This function uses to check password language.

  @param  PasswordLanguageStr           A pointer to the password language string.
  @param  PasswordLanguage              A pointer to the password language.

  @retval TRUE                          Password language is valid.
  @retval FALSE                         Password language is not valid.
**/
BOOLEAN
CheckPasswordLanguageIsValid (
  IN  CHAR8                             *PasswordLanguageStr,
  OUT WMI_PASSWORD_LANGUAGE             *PasswordLanguage
  )
{
  BOOLEAN                               IsValid;
  UINTN                                 Index;

  //
  // Initialization
  //
  IsValid = FALSE;

  //
  // Compare password language string to get password language
  //
  for (Index = 0; Index < mWmiPasswordLanguageMapListCount; Index++) {
    if (AsciiStrCmp (PasswordLanguageStr, mWmiPasswordLanguageMap[Index].PasswordLanguageStr) == 0) {
      *PasswordLanguage = mWmiPasswordLanguageMap[Index].PasswordLanguage;
      IsValid = TRUE;
      break;
    }
  }

  return IsValid;
}

/**
  Analyze WMI set password input buffer.

  @param  PasswordType                  Assign a pointer to the password type.
  @param  CurrentPassword               Assign a pointer to the current password.
  @param  NewPassword                   Assign a pointer to the new password.
  @param  PasswordEncodingStr           Assign a pointer to the password encoding string.
  @param  PasswordLanguageStr           Assign a pointer to the password keyboard language string.

  @retval EFI_SUCCESS                   This function execute successfully.
  @retval Others                        An unexpected error occurred.
**/
EFI_STATUS
AnalyzeWmiSetPasswordInputBuffer (
  OUT CHAR8                             **PasswordType,
  OUT CHAR8                             **CurrentPassword,
  OUT CHAR8                             **NewPassword,
  OUT CHAR8                             **PasswordEncodingStr,
  OUT CHAR8                             **PasswordLanguageStr
  )
{
  EFI_STATUS                            Status;
  UINTN                                 BufferOffset;
  UINTN                                 TempBufferOffset;
  UINTN                                 SectionCount;

  //
  // Initialization
  //
  Status       = EFI_INVALID_PARAMETER;
  BufferOffset = 0;

  //
  // Part section
  //
  SectionCount = 0;
  TempBufferOffset = 0;

  for (BufferOffset = 0; BufferOffset < mL05GlobalNVSArea->L05WmiBufferLen; BufferOffset++) {

    if ((mL05GlobalNVSArea->L05WmiBuffer[BufferOffset] != ',')
      &&(mL05GlobalNVSArea->L05WmiBuffer[BufferOffset] != ';')) {
      continue;
    }

    mL05GlobalNVSArea->L05WmiBuffer[BufferOffset] = 0;

    switch (SectionCount) {

    case SetBiosPasswordType:
      *PasswordType = &mL05GlobalNVSArea->L05WmiBuffer[TempBufferOffset];
      break;

    case SetBiosPasswordCurrentPassword:
      *CurrentPassword = &mL05GlobalNVSArea->L05WmiBuffer[TempBufferOffset];
      break;

    case SetBiosPasswordNewPassword:
      *NewPassword = &mL05GlobalNVSArea->L05WmiBuffer[TempBufferOffset];
      break;

    case SetBiosPasswordEncoding:
      *PasswordEncodingStr = &mL05GlobalNVSArea->L05WmiBuffer[TempBufferOffset];
      break;

    case SetBiosPasswordLanguage:
      *PasswordLanguageStr = &mL05GlobalNVSArea->L05WmiBuffer[TempBufferOffset];
      break;
    }

    TempBufferOffset = BufferOffset;
    TempBufferOffset++;
    SectionCount++;
  }

  //
  // Check section number is valid
  //
  if (SectionCount == SetBiosPasswordMaxSection) {
    Status = EFI_SUCCESS;
  }

  return Status;
}

/**
  Check WMI password format.

  @param  PasswordTypeStr               A pointer to the password type string.
  @param  PasswordType                  A pointer to the password type.
  @param  PasswordEncodingStr           A pointer to the password encoding string.
  @param  PasswordEncoding              A pointer to the password encoding.
  @param  PasswordLanguageStr           A pointer to the password keyboard language string.
  @param  PasswordLanguage              A pointer to the password keyboard language.

  @retval EFI_SUCCESS                   This function execute successfully.
  @retval EFI_INVALID_PARAMETER         Input invalid password format.
  @retval EFI_ACCESS_DENIED             Input invalid password rule.
**/
EFI_STATUS
CheckWmiPasswordFormat (
  IN     CHAR8                          *PasswordTypeStr,
  OUT    WMI_PASSWORD_TYPE              *PasswordType,
  IN     CHAR8                          *PasswordEncodingStr,
  IN OUT WMI_PASSWORD_ENCODING          *PasswordEncoding,
  IN     CHAR8                          *PasswordLanguageStr,
  IN OUT WMI_PASSWORD_LANGUAGE          *PasswordLanguage
  )
{
  EFI_STATUS                            Status;

  //
  // Initialization
  //
  Status = EFI_SUCCESS;

  //
  // Check password type is valid
  //
  if (!CheckPasswordTypeIsValid (PasswordTypeStr, PasswordType)) {
    Status = EFI_INVALID_PARAMETER;
    return Status;
  }

  //
  // Check password encoding is valid
  //
  if (!CheckPasswordEncodingIsValid (PasswordEncodingStr, PasswordEncoding)) {
    Status = EFI_INVALID_PARAMETER;
    return Status;
  }

  //
  // L05 only support ascii encoding
  //
  if (*PasswordEncoding != WmiAsciiEncoding) {
    Status = EFI_UNSUPPORTED;
    return Status;
  }

  //
  // Check password keyboard languages is valid
  //
  if (!CheckPasswordLanguageIsValid (PasswordLanguageStr, PasswordLanguage)) {
    Status = EFI_INVALID_PARAMETER;
    return Status;
  }

  //
  // L05 only support us language
  //
  if (*PasswordLanguage != WmiUsLanguage) {
    Status = EFI_UNSUPPORTED;
    return Status;
  }

  return Status;
}

/**
  Check WMI password input buffer.

  @param  PasswordTypeStr               A pointer to the password type string.
  @param  PasswordType                  A pointer to the password type.
  @param  CurrentPassword               A pointer to the current password.
  @param  NewPassword                   A pointer to the new password.
  @param  PasswordEncodingStr           A pointer to the password encoding string.
  @param  PasswordEncoding              A pointer to the password encoding.
  @param  PasswordLanguageStr           A pointer to the password keyboard language string.
  @param  PasswordLanguage              A pointer to the password keyboard language.

  @retval EFI_SUCCESS                   This function execute successfully.
  @retval EFI_INVALID_PARAMETER         Input invalid password format.
  @retval EFI_ACCESS_DENIED             Input invalid password rule.
**/
EFI_STATUS
CheckWmiPasswordInputBuffer (
  IN     CHAR8                          *PasswordTypeStr,
  OUT    WMI_PASSWORD_TYPE              *PasswordType,
  IN OUT CHAR8                          *CurrentPassword,
  IN OUT CHAR8                          *NewPassword,
  IN     CHAR8                          *PasswordEncodingStr,
  IN OUT WMI_PASSWORD_ENCODING          *PasswordEncoding,
  IN     CHAR8                          *PasswordLanguageStr,
  IN OUT WMI_PASSWORD_LANGUAGE          *PasswordLanguage
  )
{
  EFI_STATUS                            Status;

  //
  // Initialization
  //
  Status = EFI_SUCCESS;

  //
  // Check WMI password format
  //
  Status = CheckWmiPasswordFormat (
             PasswordTypeStr,
             PasswordType,
             PasswordEncodingStr,
             PasswordEncoding,
             PasswordLanguageStr,
             PasswordLanguage
             );

  if (EFI_ERROR (Status)) {
    return Status;
  }

  //
  // Check current password is valid
  //
  if (!CheckPasswordIsValid (&CurrentPassword, *PasswordType)) {
    Status = EFI_INVALID_PARAMETER;
    return Status;
  }

  //
  // Check NewPassword password is valid
  //
  if (!CheckPasswordIsValid (&NewPassword, *PasswordType)) {
    Status = EFI_INVALID_PARAMETER;
    return Status;
  }

  //
  // [Lenovo BIOS Setup using WMI Deployment Guide - Third Edition (February 2016)]
  //   2. A password cannot be set using this method when one does not already exist.
  //      Passwords can only be updated or cleared.
  //
  if (AsciiStrLen (CurrentPassword) == 0) {
    Status = EFI_ACCESS_DENIED;
    return Status;
  }

  //
  // [Lenovo China Minimum BIOS Spec Version 1.37]
  //   [3.4.1 Rules for BIOS CMOS password]
  //     When user did not set any password,the security menu have one submenu:
  //     Set Administrator Password.
  //
  if (!mAdminPasswordExist && (*PasswordType == WmiSystemUserType)) {
    Status = EFI_ACCESS_DENIED;
    return Status;
  }

  return Status;
}

/**
  Check WMI input password is repeat.

  @param  None

  @retval TRUE                          WMI input password is repeat.
  @retval FALSE                         WMI input password is not repeat.
**/
BOOLEAN
WmiInputPasswordIsRepeat (
  VOID
  )
{
  EFI_STATUS                            Status;
  CHAR8                                 *PasswordTypeStr;
  WMI_PASSWORD_TYPE                     PasswordType;
  CHAR8                                 *CurrentPassword;
  CHAR8                                 *NewPassword;
  CHAR8                                 *PasswordEncodingStr;
  WMI_PASSWORD_ENCODING                 PasswordEncoding;
  CHAR8                                 *PasswordLanguageStr;
  WMI_PASSWORD_LANGUAGE                 PasswordLanguage;

  //
  // Initialization
  //
  Status              = EFI_SUCCESS;
  PasswordTypeStr     = NULL;
  PasswordType        = WmiPasswoedMaxType;
  CurrentPassword     = NULL;
  NewPassword         = NULL;
  PasswordEncodingStr = NULL;
  PasswordEncoding    = WmiPasswoedMaxEncoding;
  PasswordLanguageStr = NULL;
  PasswordLanguage    = WmiPasswoedMaxLanguage;

  //
  // Analyze WMI set password input buffer
  //
  Status = AnalyzeWmiSetPasswordInputBuffer (
             &PasswordTypeStr,
             &CurrentPassword,
             &NewPassword,
             &PasswordEncodingStr,
             &PasswordLanguageStr
             );


  if (EFI_ERROR (Status)) {
    return FALSE;
  }

  //
  // Check WMI password input buffer is valid
  //
  Status = CheckWmiPasswordInputBuffer (
             PasswordTypeStr,
             &PasswordType,
             CurrentPassword,
             NewPassword,
             PasswordEncodingStr,
             &PasswordEncoding,
             PasswordLanguageStr,
             &PasswordLanguage
             );

  if (EFI_ERROR (Status)) {
    return FALSE;
  }

  //
  // Check password type is the same
  //
  if (PasswordType != mPasswordCheckData.PasswordType) {
    return FALSE;
  }

  //
  // Check current password is the same
  //
  if (AsciiStrCmp (CurrentPassword, mPasswordCheckData.CurrentPassword) != 0) {
    return FALSE;
  }

  //
  // Check new password is the same
  //
  if (AsciiStrCmp (NewPassword, mPasswordCheckData.NewPassword) != 0) {
    return FALSE;
  }

  //
  // Clean password check data
  //
  ZeroMem (&mPasswordCheckData, sizeof (WMI_PASSWORD_DATA_MAP));

  return TRUE;
}

/**
  Refer from ReadPassword() at SysPasswordDxe.c

  WMI get admin or user password.

  @param  PasswordType                  System admin or user password.
  @param  Password                      Assign a pointer to the password buffer.

  @retval EFI_SUCCESS                   The function completed successfully.
  @retval EFI_UNSUPPORTED               Invalid system password type.
**/
EFI_STATUS
EFIAPI
WmiReadPassword (
  IN  WMI_PASSWORD_TYPE                 PasswordType,
  OUT UINT8                             **Password
  )
{
  EFI_STATUS                            Status;
  UINTN                                 DataSize;

  //
  // Initialization
  //
#ifndef L05_NOTEBOOK_PASSWORD_ENABLE
  DataSize = L05_WMI_PASSWORD_MAX_LENGTH + 2;

#else
  DataSize = SHA256_DIGEST_SIZE;
#endif

  //
  // Allocate resources
  //
  *Password = AllocateRuntimeZeroPool (DataSize);

  //
  // Get system passwords by password type
  //
  Status = OemSvcGetSystemPasswords (PasswordType, DataSize, *Password);

  if (Status != EFI_MEDIA_CHANGED) {
    return EFI_NOT_FOUND;
  }

  return EFI_SUCCESS;
}

/**
  Refer from GetPasswordStatus() at SysPasswordDxe.c

  Check admin or user password is set or not.

  @param  PasswordType                  System admin or user password.

  @retval EFI_SUCCESS                   Password is set.
  @retval EFI_NOT_FOUND                 Password isn't set.
**/
EFI_STATUS
EFIAPI
WmiGetPasswordStatus (
  IN  WMI_PASSWORD_TYPE                 PasswordType
  )
{
  EFI_STATUS                            Status;
  UINT8                                 *Password;

  //
  // Initialization
  //
  Password = NULL;

  //
  // Read password by password type
  //
  Status = WmiReadPassword (PasswordType, &Password);

  if (EFI_ERROR (Status) || (Password == NULL)) {
    if (Password != NULL) {
      FreePool (Password);
    }

    return Status;
  }

#ifndef L05_NOTEBOOK_PASSWORD_ENABLE
  //
  // Check password Checksum is valid
  //
  if (!IsSystemPasswordChecksumValid (Password)) {
#else
  //
  //  Check system has system password or not.
  //
  if (!L05HaveSystemPassword (Password)) {
#endif
    FreePool (Password);
    return EFI_NOT_FOUND;
  }

  //
  // Free resources
  //
  FreePool (Password);

  return EFI_SUCCESS;
}

/**
  Refer from IsPasswordExist() at Password.c

  This function uses to check admin password is exist.

  @param  None

  @retval TRUE                          Admin password is in locked state.
  @retval FALSE                         Admin password is in unloced state.
**/
BOOLEAN
WmiIsAdminPasswordExist (
  VOID
  )
{
  EFI_STATUS                            Status;

  Status = WmiGetPasswordStatus (WmiSystemAdminType);

  return Status == EFI_SUCCESS ? TRUE : FALSE;
}

/**
  Refer from IsPasswordExist() at Password.c

  This function uses to check user password is exist.

  @param  None

  @retval TRUE                          User password is in locked state.
  @retval FALSE                         User password is in unloced state.
**/
BOOLEAN
WmiIsUserPasswordExist (
  VOID
  )
{
  EFI_STATUS                            Status;

  Status = WmiGetPasswordStatus (WmiSystemUserType);

  return Status == EFI_SUCCESS ? TRUE : FALSE;
}

/**
  This function uses to check HDD password is exist.

  @param  None

  @retval TRUE                          Hdd password is in locked state.
  @retval FALSE                         Hdd password is in unloced state.
**/
BOOLEAN
WmiIsHddPasswordExist (
  VOID
  )
{
  EFI_STATUS                            Status;
  HDD_PASSWORD_TABLE                    *HddPasswordTable;
  UINTN                                 HddPasswordTableSize;

  //
  // Initialization
  //
  Status = EFI_SUCCESS;
  HddPasswordTableSize = 0;
  HddPasswordTable = NULL;

  //
  // Get save HDD password
  //
  HddPasswordTable = SmmGetVariableAndSize (
                       SAVE_HDD_PASSWORD_VARIABLE_NAME,
                       &gSaveHddPasswordGuid,
                       &HddPasswordTableSize
                       );

  if (HddPasswordTable == NULL) {
    Status = EFI_NOT_FOUND;
  }

  //
  // Free resources
  //
  if (HddPasswordTable != NULL) {
    FreePool (HddPasswordTable);
  }

  return Status == EFI_SUCCESS ? TRUE : FALSE;
}

/**
  Refer from CheckPassword() at SysPasswordDxe.c

  Check admin or user password is correct or not.

  @param  None

  @param  PasswordType                  System admin or user password.
  @param  Password                      A pointer to the password.

  @retval
**/
EFI_STATUS
WmiCheckPassword (
  WMI_PASSWORD_TYPE                     PasswordType,
  CHAR8                                 *Password
  )
{
  EFI_STATUS                            Status;
  UINT8                                 *SavePassword;

  //
  // Initialization
  //
  SavePassword = NULL;

  //
  // Read password by password type
  //
  Status = WmiReadPassword (PasswordType, &SavePassword);

  if (EFI_ERROR (Status) || (SavePassword == NULL)) {
    if (SavePassword != NULL) {
      FreePool (SavePassword);
    }

    return Status;
  }

#ifndef L05_NOTEBOOK_PASSWORD_ENABLE
  //
  // Check password Checksum is valid
  //
  if (!IsSystemPasswordChecksumValid (SavePassword)) {
#else
  //
  //  Check system has system password or not.
  //
  if (!L05HaveSystemPassword (SavePassword)) {
#endif
    FreePool (SavePassword);
    return EFI_NOT_FOUND;
  }

  //
  // Compare input password with save password
  //
  if (!L05PasswordCmp (Password, SavePassword)) {
    Status = EFI_CRC_ERROR;
  }

  //
  // Free resources
  //
  FreePool (SavePassword);

  return Status;
}

/**
  Clear System password.

  @param  PasswordType                  System admin or user password.

  @retval EFI_SUCCESS                   This function execute successfully.
  @retval Others                        An unexpected error occurred.
**/
EFI_STATUS
ClearSystemPassword (
  WMI_PASSWORD_TYPE                     PasswordType
  )
{
  EFI_STATUS                            Status;
#ifndef L05_NOTEBOOK_PASSWORD_ENABLE
  UINT8                                 EmptyPassword[2];
#else
  UINT8                                 EmptyPassword[SHA256_DIGEST_SIZE];
#endif

  //
  // Initialization
  //
  Status = EFI_SUCCESS;

  //
  //  Clear Password
  //
  ZeroMem (EmptyPassword, sizeof (EmptyPassword));
  Status = OemSvcSetSystemPasswords (PasswordType, sizeof (EmptyPassword), EmptyPassword);

  if (Status == EFI_MEDIA_CHANGED) {
    Status = EFI_SUCCESS;
  }

  //
  // Re-init security item
  //
  Status = SecurityDataInit (FALSE);

  return Status;
}

/**
  Check admin password.

  @param  Password                      A pointer to the password buffer.

  @retval EFI_SUCCESS                   This function execute successfully.
  @retval EFI_ACCESS_DENIED             WMI input admin password is invalid.
  @retval Others                        An unexpected error occurred.
**/
EFI_STATUS
WmiCheckAdminPassword (
  CHAR8                                 *Password,
  CHAR8                                 *PasswordEncodingStr,
  CHAR8                                 *PasswordLanguageStr
  )
{
  EFI_STATUS                            Status;
  CHAR8                                 *NewPassword;
  UINT8                                 *EncodeAdminPassword;
  UINT8                                 *EncodeAdminPasswordTemp;
  CHAR8                                 *PasswordTypeStr = WMI_BIOS_PASSWORD_ADMIN_TYPE;
  WMI_PASSWORD_TYPE                     PasswordType;
  WMI_PASSWORD_ENCODING                 PasswordEncoding;
  WMI_PASSWORD_LANGUAGE                 PasswordLanguage;

  //
  // Initialization
  //
  Status                  = EFI_SUCCESS;
  NewPassword             = "";
  EncodeAdminPassword     = NULL;
  EncodeAdminPasswordTemp = NULL;
  PasswordType            = WmiPasswoedMaxType;
  PasswordEncoding        = WmiPasswoedMaxEncoding;
  PasswordLanguage        = WmiPasswoedMaxLanguage;

  //
  // Check WMI password input buffer is valid
  //
  Status = CheckWmiPasswordInputBuffer (
             PasswordTypeStr,
             &PasswordType,
             Password,
             NewPassword,
             PasswordEncodingStr,
             &PasswordEncoding,
             PasswordLanguageStr,
             &PasswordLanguage
             );

  if (EFI_ERROR (Status)) {
    return Status;
  }

  //
  // Encode password
  //
  EncodeAdminPasswordTemp = L05SysPasswordPacket (Password, AsciiStrLen (Password));
  EncodeAdminPassword = L05SysPasswordEncode (EncodeAdminPasswordTemp, AsciiStrLen (Password));

  if (EncodeAdminPassword == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }

  //
  // Compare current password with save password
  //
  Status = WmiCheckPassword (WmiSystemAdminType, EncodeAdminPassword);

  if (EFI_ERROR (Status)) {
    mPasswordErrorCount++;
    Status = EFI_ACCESS_DENIED;

  } else {
    //
    // Clean password error count
    //
    mPasswordErrorCount = 0;
  }

  //
  // Free resources
  //
  FreePool (EncodeAdminPassword);

  return Status;
}

UINTN
WmiGetNumOfHdd (
  VOID
  )
{
  EFI_STATUS                            Status;
  EFI_HDD_PASSWORD_SERVICE_PROTOCOL     *HddPasswordService;
  HDD_PASSWORD_HDD_INFO                 *HddInfoArray;
  UINTN                                 NumOfHdd;

  HddPasswordService = NULL;
  HddInfoArray       = NULL;

  Status = gSmst->SmmLocateProtocol (
                    &gEfiHddPasswordServiceProtocolGuid,
                    NULL,
                    &HddPasswordService
                    );

  if (EFI_ERROR (Status)) {
    return 0;
  }

  //
  // Get HDD info
  //
  Status = HddPasswordService->GetHddInfo (
                                 HddPasswordService,
                                 &HddInfoArray,
                                 &NumOfHdd
                                 );

  if (EFI_ERROR (Status)) {
    return 0;
  }

    return NumOfHdd;
}


/**
  This function uses to get HDD password type number.

  @param  PasswordType                  HDD master or user password.
  @param  PasswordTypeStr               A pointer to the password type string.
  @param  HddPasswordTypeNumber         A pointer to the HDD password type number.

  @retval EFI_SUCCESS                   This function execute successfully.
  @retval Others                        An unexpected error occurred.
**/
EFI_STATUS
WmiGetHddPasswordTypeNumber (
  IN WMI_PASSWORD_TYPE                  PasswordType,
  IN  CHAR8                             *PasswordTypeStr,
  OUT UINT8                             *HddPasswordTypeNumber
  )
{
  EFI_STATUS                            Status;
  UINTN                                 Index;
  CHAR8                                 *HddPasswordTypeNumberStr;
  UINTN                                 NumOfHdd;

  Status                   = EFI_INVALID_PARAMETER;
  HddPasswordTypeNumberStr = NULL;

  if (PasswordTypeStr == NULL || HddPasswordTypeNumber == NULL) {
    return EFI_INVALID_PARAMETER;
  }

  NumOfHdd = WmiGetNumOfHdd ();
  *HddPasswordTypeNumber = 0;

  for (Index = 0; Index < mWmiPasswordTypeMapListCount; Index++) {
    if (PasswordType == mWmiPasswordTypeMap[Index].PasswordType) {
      HddPasswordTypeNumberStr = PasswordTypeStr + AsciiStrLen (mWmiPasswordTypeMap[Index].PasswordTypeStr);

      while ((*HddPasswordTypeNumberStr >= L'0' && *HddPasswordTypeNumberStr <= L'9')) {
        *HddPasswordTypeNumber *= 10;
        *HddPasswordTypeNumber += *HddPasswordTypeNumberStr - L'0';
        HddPasswordTypeNumberStr++;
      }

      if (*HddPasswordTypeNumber == 0 || *HddPasswordTypeNumberStr != 0) {
        *HddPasswordTypeNumber = 0;
        break;
      }

      Status = EFI_SUCCESS;
      break;
    }
  }

  if (*HddPasswordTypeNumber > NumOfHdd) {
    *HddPasswordTypeNumber = 0;
    Status = EFI_INVALID_PARAMETER;
  }

  return Status;
}

/**
  WMI Set HDD Passwords.

  @param  PasswordType                  HDD master or user password.
  @param  HddPasswordTypeNumber         HDD password type number.
  @param  HddPasswordBuffer             A pointer to the HDD current password.
  @param  HddNewPassword                A pointer to the HDD new password.

  @retval EFI_SUCCESS                   This function execute successfully.
  @retval Others                        An unexpected error occurred.
**/
EFI_STATUS
WmiSetHddPasswords (
  IN WMI_PASSWORD_TYPE                  PasswordType,
  IN UINT8                              HddPasswordTypeNumber,
  IN UINT8                              *HddCurrentPassword,
  IN UINT8                              *HddNewPassword
  )
{
  EFI_STATUS                            Status;
  HDD_PASSWORD_TABLE                    HddPasswordTable[WmiHddPasswordMaxSection];
  UINTN                                 NumberLength;
  CHAR16                                *VariableName;
  UINTN                                 BufferSize;

  //
  // Initialization
  //
  VariableName = NULL;
  BufferSize   = 0;
  ZeroMem (HddPasswordTable, sizeof (HddPasswordTable));

  //
  // Put HDD current password
  //
  AsciiStrToUnicodeStrS (HddCurrentPassword, HddPasswordTable[WmiHddPasswordCurrentSection].PasswordStr, HDD_PASSWORD_MAX_NUMBER + 1);
  HddPasswordTable[WmiHddPasswordCurrentSection].PasswordType = (UINT8) (PasswordType - WmiHddUserType);

  //
  // Put HDD new password
  //
  AsciiStrToUnicodeStrS (HddNewPassword, HddPasswordTable[WmiHddPasswordNewSection].PasswordStr, HDD_PASSWORD_MAX_NUMBER + 1);
  HddPasswordTable[WmiHddPasswordNewSection].PasswordType = (UINT8) (PasswordType - WmiHddUserType);

  //
  // Allocate resources
  //
  L05_GET_NUMBER_LENGTH (HddPasswordTypeNumber, NumberLength);
  BufferSize = StrSize (L05_WMI_HDD_PASWORD_DATA_VARIABLE_NAME) + (NumberLength * sizeof (CHAR16));
  VariableName = AllocateZeroPool (BufferSize);

  if (VariableName == NULL) {
    return EFI_OUT_OF_RESOURCES;
  }

  //
  // Set WMI HDD password variable name
  //
  UnicodeSPrint (VariableName, BufferSize, L"%s%d", L05_WMI_HDD_PASWORD_DATA_VARIABLE_NAME, HddPasswordTypeNumber);

  //
  // Set WMI HDD password variable
  //
  Status = mSmmVariable->SmmSetVariable (
                           VariableName,
                           &gL05WmiSetupUnderOsSetupVariableGuid,
                           EFI_VARIABLE_NON_VOLATILE | EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS,
                           (sizeof (HDD_PASSWORD_TABLE) * WmiHddPasswordMaxSection),
                           (VOID *) HddPasswordTable
                           );

  //
  // Free resources
  //
  FreePool (VariableName);

  return Status;
}

/**
  WMI Method for Set BIOS password -

    Mapping to WMI class  : Lenovo_SetBiosPassword
    Mapping to ASL method : WMA6

  @param  None

  @retval EFI_SUCCESS                   This function execute successfully.
  @retval EFI_INVALID_PARAMETER         Input invalid password format.
  @retval EFI_ACCESS_DENIED             Input invalid password rule or password is invalid.
  @retval Others                        An unexpected error occurred.
**/
EFI_STATUS
EFIAPI
WmiSetBiosPasswordA6 (
  VOID
  )
{
  EFI_STATUS                            Status;
  CHAR8                                 *PasswordTypeStr;
  WMI_PASSWORD_TYPE                     PasswordType;
  CHAR8                                 *CurrentPassword;
  CHAR8                                 *NewPassword;
  CHAR8                                 *PasswordEncodingStr;
  WMI_PASSWORD_ENCODING                 PasswordEncoding;
  CHAR8                                 *PasswordLanguageStr;
  WMI_PASSWORD_LANGUAGE                 PasswordLanguage;
  UINT8                                 *EncodeCurrentPassword;
  UINT8                                 *EncodeCurrentPasswordTemp;
  UINT8                                 *EncodeNewPassword;
  UINT8                                 *EncodeNewPasswordTemp;
  UINT8                                 HddPasswordTypeNumber;

  //
  // Initialization
  //
  PasswordTypeStr           = NULL;
  PasswordType              = WmiPasswoedMaxType;
  CurrentPassword           = NULL;
  NewPassword               = NULL;
  PasswordEncodingStr       = NULL;
  PasswordEncoding          = WmiPasswoedMaxEncoding;
  PasswordLanguageStr       = NULL;
  PasswordLanguage          = WmiPasswoedMaxLanguage;
  EncodeCurrentPassword     = NULL;
  EncodeCurrentPasswordTemp = NULL;
  EncodeNewPassword         = NULL;
  EncodeNewPasswordTemp     = NULL;

  mAdminPasswordExist = WmiIsAdminPasswordExist ();
  mUserPasswordExist  = WmiIsUserPasswordExist ();
  mHddPasswordExist   = WmiIsHddPasswordExist ();

  //
  // Analyze WMI set password input buffer
  //
  Status = AnalyzeWmiSetPasswordInputBuffer (
             &PasswordTypeStr,
             &CurrentPassword,
             &NewPassword,
             &PasswordEncodingStr,
             &PasswordLanguageStr
             );

  if (EFI_ERROR (Status)) {
    return Status;
  }

  //
  // Check WMI password input buffer is valid
  //
  Status = CheckWmiPasswordInputBuffer (
             PasswordTypeStr,
             &PasswordType,
             CurrentPassword,
             NewPassword,
             PasswordEncodingStr,
             &PasswordEncoding,
             PasswordLanguageStr,
             &PasswordLanguage
             );

  if (EFI_ERROR (Status)) {
    return Status;
  }

  switch (PasswordType) {

  case WmiSystemAdminType:
  case WmiSystemUserType:

    //
    // Encode password
    //
    EncodeCurrentPasswordTemp = L05SysPasswordPacket (CurrentPassword, AsciiStrLen (CurrentPassword));
    EncodeCurrentPassword = L05SysPasswordEncode (EncodeCurrentPasswordTemp, AsciiStrLen (CurrentPassword));

    EncodeNewPasswordTemp = L05SysPasswordPacket (NewPassword, AsciiStrLen (NewPassword));
    EncodeNewPassword = L05SysPasswordEncode (EncodeNewPasswordTemp, AsciiStrLen (NewPassword));

    if (EncodeCurrentPassword == NULL) {
      return EFI_OUT_OF_RESOURCES;
    }

    if (EncodeNewPassword == NULL) {
      return EFI_OUT_OF_RESOURCES;
    }

    //
    // Compare current password with old password
    //
    Status = WmiCheckPassword (PasswordType, EncodeCurrentPassword);

    if (EFI_ERROR (Status) && (Status != EFI_NOT_FOUND)) {
      mPasswordErrorCount++;
      Status = EFI_ACCESS_DENIED;
      break;
    }

    //
    // Avoid password don't exist but when current password exist can set new password
    //
    if ((!(PasswordType ? mUserPasswordExist : mAdminPasswordExist)) && (AsciiStrLen (CurrentPassword) != 0)) {
      mPasswordErrorCount++;
      Status = EFI_ACCESS_DENIED;
      break;
    }

    //
    // [Lenovo China Minimum BIOS Spec Version 1.37]
    //    When the user password set, there is an additional submenu added.
    //    Clear user password menu just appeared in the security menu when
    //    user entered CMOS setup menu with administrator password.
    //
    if ((PasswordType == WmiSystemUserType) && (AsciiStrLen (NewPassword) == 0)) {
      Status = EFI_ACCESS_DENIED;
      break;
    }

    //
    // Set new password.
    //
    if (AsciiStrLen (NewPassword) != 0) {
#ifndef L05_NOTEBOOK_PASSWORD_ENABLE
      Status = OemSvcSetSystemPasswords (PasswordType, (EncodeNewPassword[0] + 2), EncodeNewPassword);

#else
      Status = OemSvcSetSystemPasswords (PasswordType, SHA256_DIGEST_SIZE, EncodeNewPassword);
#endif

      if (Status == EFI_MEDIA_CHANGED) {
        Status = EFI_SUCCESS;
      }
    }

    //
    // When clear admin password also need clear user password
    //
    if ((PasswordType == WmiSystemAdminType) && (AsciiStrLen (NewPassword) == 0)) {
      Status = ClearSystemPassword (WmiSystemAdminType);
      Status = ClearSystemPassword (WmiSystemUserType);
    }

    //
    // Re-init security item
    //
    Status = SecurityDataInit (FALSE);
    break;

  case WmiHddUserType:
  case WmiHddMasterType:

#ifdef L05_HDD_PASSWORD_ENABLE
    //
    // Avoid HDD password don't exist but when current password exist can set new password
    //
    if (!mHddPasswordExist) {
      Status = EFI_ACCESS_DENIED;
      break;
    }

    //
    // By L05 Spec., User hdd password can not disable hdd password function.
    // return EFI_ACCESS_DENIED for L05 only.
    //
    if ((PasswordType == WmiHddUserType) && (AsciiStrLen (NewPassword) == 0)) {
      Status = EFI_ACCESS_DENIED;
      break;
    }

    //
    // Get HDD Number
    //
    Status = WmiGetHddPasswordTypeNumber (PasswordType, PasswordTypeStr, &HddPasswordTypeNumber);
    if (EFI_ERROR (Status) || HddPasswordTypeNumber == 0) {
      Status = EFI_ACCESS_DENIED;
      break;
    }

    //
    // Set HDD password.
    //
    Status = WmiSetHddPasswords (PasswordType, HddPasswordTypeNumber, CurrentPassword, NewPassword);
    break;

#else
    Status = EFI_UNSUPPORTED;
    break;
    
#endif
  }

  if (!EFI_ERROR (Status)) {
    //
    // [Lenovo BIOS Setup using WMI Deployment Guide - Third Edition (February 2016)]
    //   Limitations and Notes:
    //     1.  BIOS settings cannot be changed at the same boot as power-on passwords (POP) and hard disk passwords
    //         (HDP). If you want to change BIOS settings and POP or HDP, you must reboot the system after changing
    //         one of them.
    //
    mPasswordChange = TRUE;

    //
    // Init password check data
    //
    mPasswordCheckData.PasswordType = PasswordType;
    AsciiStrCpyS (mPasswordCheckData.CurrentPassword, sizeof (mPasswordCheckData.CurrentPassword), CurrentPassword);
    AsciiStrCpyS (mPasswordCheckData.NewPassword, sizeof (mPasswordCheckData.NewPassword), NewPassword);
  }

  //
  // Free resources
  //
  if (EncodeCurrentPassword != NULL) {
    FreePool (EncodeCurrentPassword);
  }

  if (EncodeNewPassword != NULL) {
    FreePool (EncodeNewPassword);
  }

  return Status;
}