/** @file

;******************************************************************************
;* Copyright (c) 2016 - 2021, 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.
;*
;******************************************************************************
// Abstract:
//
//  Create a notify event, then it will change secure boot mode and lock specific variable
//
*/

#include <Library/BaseLib.h>
#include <Library/BaseMemoryLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/DebugLib.h>
#include <Library/PcdLib.h>
#include <Library/UefiLib.h>

#include <Library/VariableLib.h>
#include <Library/UefiRuntimeServicesTableLib.h>
#include <Library/UefiBootServicesTableLib.h>

#include <Protocol/SetupUtility.h>
#include <Protocol/VariableLock.h>
#include <Protocol/NonVolatileVariable.h>
#include <Protocol/SmmCommunication.h>
#include <Protocol/SmmVariable.h>

#include <Guid/AdmiSecureBoot.h>
#include <Guid/DebugMask.h>
#include <Guid/EventGroup.h>
#include <Guid/AsfSecureBootVariable.h>
#include <Guid/AuthenticatedVariableFormat.h>

#include <ChipsetSmiTable.h>
#include <Protocol/AmtWrapperProtocol.h>

#define H2O_SIGNATURE        SIGNATURE_32 ('O', '2', 'H', '$')

extern EFI_GUID gActiveManagementProtocolGuid;


UINT8    mDefaultAthenData[] = {
  0x01, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x28, 0x02, 0x00, 0x00, 0x00, 0x02, 0x02, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
  0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0xCC, 0xdd, 0x26, 0x02, 0x97, 0x9d, 0x44, 0x1f, 0x64, 0x95, 0xf8, 0x88, 0x46, 0xef, 0x9f, 0xb2,
  0xa1, 0x6c, 0xbc, 0x89, 0x1c, 0x29, 0xad, 0x4e, 0x5f, 0x67, 0xe0, 0xdb, 0xeb, 0x2f, 0xa4, 0xc0, 0x3f, 0x44, 0x72, 0x5c, 0x93, 0x39, 0x6d, 0xe5, 0x1c, 0xe2, 0x2b, 0xcb, 0x76, 0x11, 0x71, 0x4f,
  0xc7, 0xe7, 0xa2, 0xde, 0x1e, 0x07, 0x56, 0x92, 0xfc, 0x63, 0x61, 0xf1, 0x06, 0xbb, 0x77, 0x26, 0xb8, 0x99, 0x7c, 0x0e, 0xb8, 0x81, 0x24, 0x2e, 0x72, 0x78, 0x8c, 0xfa, 0x6a, 0x71, 0x5d, 0x08,
  0x42, 0xda, 0xde, 0x86, 0xfe, 0x99, 0x92, 0xf7, 0xf8, 0xa0, 0x3d, 0x1f, 0x52, 0x0b, 0xb8, 0xc6, 0x1b, 0xd5, 0x39, 0xcb, 0x34, 0xdb, 0x2b, 0x25, 0x57, 0x44, 0x72, 0x02, 0x93, 0x67, 0x2a, 0x82,
  0x90, 0xa2, 0x5a, 0x76, 0xcb, 0x81, 0x7d, 0x42, 0x39, 0x1a, 0xf8, 0x54, 0xe0, 0x42, 0xbf, 0x49, 0x36, 0x1c, 0x4e, 0x88, 0xda, 0x5c, 0x7b, 0xd4, 0x67, 0xf6, 0x82, 0x30, 0xe4, 0x18, 0xf9, 0xb3,
  0xe8, 0xab, 0xa4, 0xca, 0x53, 0x82, 0xcf, 0xf4, 0xac, 0x83, 0x4b, 0xd6, 0xd4, 0xa1, 0xc0, 0x01, 0xfc, 0x16, 0x9d, 0xca, 0xc5, 0x17, 0x2e, 0x70, 0xff, 0xaf, 0x86, 0xc8, 0x43, 0x02, 0xde, 0x52,
  0xb1, 0x23, 0xf0, 0x2f, 0xd8, 0x0a, 0x46, 0xd3, 0xaf, 0x6b, 0x7e, 0x77, 0xf6, 0x72, 0x00, 0x66, 0x69, 0x66, 0x37, 0x71, 0x19, 0xc7, 0x7b, 0xf0, 0xd1, 0x41, 0xfd, 0x2b, 0x9b, 0xa1, 0x14, 0xb5,
  0xb3, 0x55, 0xb9, 0x96, 0x53, 0x9f, 0x29, 0x63, 0xb2, 0x38, 0x2e, 0x29, 0x10, 0x16, 0x8a, 0xbd, 0x2e, 0xed, 0x3b, 0xdc, 0xb5, 0x40, 0x6f, 0xab, 0x26, 0xea, 0xb5, 0x28, 0x36, 0x6b, 0x4a, 0x68,
  0x85, 0xd2, 0xb1, 0xff, 0xbf, 0xe6, 0x0a, 0xd3, 0x8c, 0xe5, 0x01, 0xda, 0xfd, 0x65, 0x75, 0x41, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
  0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
  0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
  0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
  0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
  0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
  0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
  0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
  0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x01
};

/**
  Initial communication buffer header.

  @param[in] CommunicationBuffer         Input communication buffer.
  @param[in] CommunicationBufferSize     Size of communication buffer in bytes.
**/
STATIC
VOID
InitCommunicationBufferHeader (
  IN VOID         *CommunicationBuffer,
  IN UINTN        CommunicationBufferSize
  )
{
  EFI_SMM_COMMUNICATE_HEADER    *SmmCommunicateHeader;

  if (CommunicationBuffer == NULL || CommunicationBufferSize < SMM_COMMUNICATE_HEADER_SIZE + sizeof (SMM_VAR_BUFFER)) {
    return;
  }
  SmmCommunicateHeader = (EFI_SMM_COMMUNICATE_HEADER *) CommunicationBuffer;
  ZeroMem (SmmCommunicateHeader, SMM_COMMUNICATE_HEADER_SIZE + sizeof (SMM_VAR_BUFFER));
  CopyGuid (&SmmCommunicateHeader->HeaderGuid, &gEfiSmmVariableProtocolGuid);
  SmmCommunicateHeader->MessageLength = SMM_COMMUNICATE_BUFFER_SIZE;
}

VOID
EFIAPI
AsfSecureBootChanged (
  IN EFI_EVENT    Event,
  IN VOID         *Context
  )
{
  EFI_STATUS                     Status;
  UINTN                          VarSize;
  UINT8                          AsfData;
  UINTN                          BufferSize;
  EDKII_VARIABLE_LOCK_PROTOCOL   *VariableLock;
  AMT_WRAPPER_PROTOCOL           *AmtWrapper;
  NON_VOLATILE_VARIABLE_PROTOCOL *NonVolatileVariableProtocol;
  VOID                           *VariableBuffer;
  UINTN                           VariableBufferSize;
  EFI_SMM_COMMUNICATION_PROTOCOL *SmmCommunication;
  SMM_VAR_BUFFER                 *SmmVarBuffer;

  AmtWrapper = NULL;
  BufferSize = sizeof (mDefaultAthenData);

  Status = gBS->LocateProtocol (&gAmtWrapperProtocolGuid, NULL, (VOID **) &AmtWrapper);

  if (EFI_ERROR (Status)) {
    return;
  } else {
    gBS->CloseEvent (Event);
  }

  Status = gBS->LocateProtocol (
               &gEfiNonVolatileVariableProtocolGuid,
               NULL,
               (VOID **) &NonVolatileVariableProtocol
               );
  if (EFI_ERROR (Status)) {
    return;
  }

  Status = gBS->LocateProtocol (
               &gEfiSmmCommunicationProtocolGuid,
               NULL,
               (VOID **) &SmmCommunication
               );
  if (EFI_ERROR (Status)) {
    return;
  }

  Status = NonVolatileVariableProtocol->GetRuntimeVariableBuffer (&VariableBuffer, &VariableBufferSize);
  if (EFI_ERROR (Status) || VariableBufferSize < BufferSize) {
    return;
  }

  Status = gBS->LocateProtocol (&gEdkiiVariableLockProtocolGuid, NULL, (VOID **) &VariableLock);
  if (!EFI_ERROR (Status)) {
    Status = VariableLock->RequestToLock (VariableLock, ASF_SECURE_BOOT_VARIABLE_NAME, &gAsfSecureBootVariableGuid);
    ASSERT_EFI_ERROR (Status);
  } else {
    return;
  }

  VarSize = sizeof(UINT8);
  Status = CommonGetVariable (
                  ASF_SECURE_BOOT_VARIABLE_NAME,
                  &gAsfSecureBootVariableGuid,
                  &VarSize,
                  &AsfData
                  );
  if (Status == EFI_BUFFER_TOO_SMALL) {

    Status = CommonGetVariable (
                    ASF_SECURE_BOOT_VARIABLE_NAME,
                    &gAsfSecureBootVariableGuid,
                    &VarSize,
                    &AsfData
                    );
  }

  if(!Status) {

    switch (AsfData) {
    case 0:
      //
      //  Set Enforce Secure Boot to Disable
      //
      BufferSize = sizeof (mDefaultAthenData);
      mDefaultAthenData[BufferSize - 1] = 0;
      InitCommunicationBufferHeader (VariableBuffer, VariableBufferSize);
      SmmVarBuffer = (SMM_VAR_BUFFER *)((UINTN) VariableBuffer + SMM_COMMUNICATE_HEADER_SIZE);
      SmmVarBuffer->DataSize   = BufferSize;
      SmmVarBuffer->AccessType = SECURE_BOOT_ENFORCE_FUN_NUM;
      SmmVarBuffer->Status = EFI_UNSUPPORTED;
      CopyMem (SmmVarBuffer + 1, mDefaultAthenData, BufferSize);
      Status = SmmCommunication->Communicate (SmmCommunication, VariableBuffer, &VariableBufferSize);
    break;
    case 1:
      //
      //  Set Enforce Secure Boot to Enable
      //
      BufferSize = sizeof (mDefaultAthenData);
      mDefaultAthenData[BufferSize - 1] = 1;
      InitCommunicationBufferHeader (VariableBuffer, VariableBufferSize);
      SmmVarBuffer = (SMM_VAR_BUFFER *)((UINTN) VariableBuffer + SMM_COMMUNICATE_HEADER_SIZE);
      SmmVarBuffer->DataSize   = BufferSize;
      SmmVarBuffer->AccessType = SECURE_BOOT_ENFORCE_FUN_NUM;
      SmmVarBuffer->Status = EFI_UNSUPPORTED;
      CopyMem (SmmVarBuffer + 1, mDefaultAthenData, BufferSize);
      Status = SmmCommunication->Communicate (SmmCommunication, VariableBuffer, &VariableBufferSize);
    break;

    default :
    break;
    }

    //
    // Del Variable ASF-SB
    //
    Status = CommonSetVariable (
                ASF_SECURE_BOOT_VARIABLE_NAME,
                &gAsfSecureBootVariableGuid,
                0,
                0,
                NULL
                );
    //
    // Del Variable of ASBN
    //
    Status = CommonSetVariable (
                    EFI_ADMINISTER_SECURE_BOOT_NAME,
                    &gEfiGenericVariableGuid,
                    EFI_VARIABLE_NON_VOLATILE | EFI_VARIABLE_BOOTSERVICE_ACCESS | EFI_VARIABLE_RUNTIME_ACCESS,
                    0,
                    NULL
                    );

    //
    // Reset
    //
    gRT->ResetSystem(EfiResetCold, EFI_SUCCESS, 0, NULL);
    CpuDeadLoop ();
  }
  return;
}

/**
 Crated the callback function

 @param[in]         ImageHandle         Pointer to the loaded image protocol for this driver
 @param[in]         SystemTable         Pointer to the EFI System Table

 @retval            EFI_SUCCESS
 @retval            Assert, otherwise.
*/
EFI_STATUS
EFIAPI
AsfSecureBootDxeEntrypoint (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE  *SystemTable
  )
{
  VOID                 *Registration;

  EfiCreateProtocolNotifyEvent (
        &gAmtWrapperProtocolGuid,
        TPL_CALLBACK,
        AsfSecureBootChanged,
        NULL,
        &Registration
        );

 return EFI_SUCCESS;
}