/**
*  @Package     : BeniPkg
*  @FileName    : Exec.c
*  @Date        : 20221210
*  @Author      : Jiangwei
*  @Version     : 1.0
*  @Description :
*    This command is used to execute Shell application.
*
*  @History:
*    20221210: Initialize.
*    20240302: Add security check.
*
*  This program and the accompanying materials
*  are licensed and made available under the terms and conditions of the BSD License
*  which accompanies this distribution. The full text of the license may be found at
*  http://opensource.org/licenses/bsd-license.php
*
*  THE PROGRAM IS DISTRIBUTED UNDER THE BSD LICENSE ON AN "AS IS" BASIS,
*  WITHOUT WARRANTIES OR REPRESENTATIONS OF ANY KIND, EITHER EXPRESS OR IMPLIED.
**/

#include "Exec.h"

#define RSA_LEN 256

STATIC CONST SHELL_PARAM_ITEM ParamList[] = {
  {NULL , TypeMax }
  };

//
// Public modulus of RSA Key. Related private key in BeniPkg\Tools\private.pem.
//
CONST UINT8 mPublicKey[] = {
  0x91, 0xad, 0xec, 0x3f, 0x4d, 0x85, 0x5f, 0xc0, 0xa7, 0x95, 0x14, 0x92, 0x6c, 0x2f, 0x0d, 0x37,
  0x6e, 0x58, 0x2d, 0x9a, 0x06, 0x0b, 0x07, 0xc0, 0x15, 0x90, 0x1e, 0xd9, 0x70, 0x25, 0xa5, 0xfe,
  0x87, 0x68, 0xc3, 0xcd, 0xa2, 0xe5, 0xd4, 0xd7, 0x3c, 0x06, 0x1f, 0x30, 0xa3, 0x81, 0xa7, 0x6a,
  0xf0, 0x27, 0xaa, 0x26, 0x0c, 0xcb, 0x7d, 0xcb, 0xc2, 0x2c, 0xc6, 0x67, 0xb5, 0x76, 0xef, 0x30,
  0x4d, 0x8d, 0x12, 0x6b, 0x4d, 0x20, 0x11, 0x2c, 0xc4, 0x69, 0xa6, 0x9b, 0xdb, 0x0e, 0xc8, 0xae,
  0x3e, 0xcc, 0xa8, 0xe3, 0x83, 0xb9, 0x80, 0x5b, 0xd2, 0x97, 0x3c, 0xe2, 0xe7, 0x85, 0x5a, 0xdb,
  0x53, 0x23, 0x8a, 0xb4, 0xa0, 0xf8, 0x02, 0xf3, 0x03, 0xec, 0x41, 0x37, 0x97, 0xd0, 0xb5, 0x35,
  0xf5, 0x01, 0xd9, 0x3b, 0xe8, 0x24, 0x24, 0xef, 0x39, 0x80, 0x40, 0x5e, 0xc0, 0xc6, 0xb5, 0x3d,
  0x32, 0x3b, 0xf1, 0x4b, 0x80, 0xa9, 0x2d, 0x93, 0x06, 0xd4, 0x8e, 0x06, 0xb6, 0xb0, 0x3e, 0xce,
  0x6a, 0x17, 0x75, 0x28, 0x32, 0x50, 0xa4, 0xc1, 0x86, 0x4c, 0xc0, 0x46, 0xbb, 0x8d, 0x83, 0x6c,
  0x8e, 0x53, 0x96, 0x72, 0x7d, 0x99, 0x85, 0x6f, 0x19, 0xb5, 0x0c, 0x33, 0x1e, 0x00, 0x57, 0x19,
  0x15, 0x59, 0x6b, 0x58, 0x30, 0xdc, 0xc5, 0x00, 0x0d, 0x7c, 0xcc, 0x37, 0x05, 0x00, 0x4f, 0x17,
  0xa7, 0x41, 0x05, 0xe0, 0xd2, 0xf7, 0x67, 0x67, 0xf8, 0xce, 0x77, 0xa3, 0x1b, 0x9a, 0x45, 0xcf,
  0x04, 0x14, 0x04, 0x9a, 0xdf, 0x58, 0x9d, 0x2a, 0x99, 0x00, 0xf7, 0x16, 0x94, 0xad, 0x90, 0x77,
  0x86, 0xff, 0x6e, 0x6b, 0x03, 0xd3, 0x80, 0xf3, 0xf6, 0xde, 0xd9, 0xcc, 0x89, 0xcc, 0xbc, 0x3b,
  0xf9, 0x42, 0x06, 0x5d, 0xba, 0x9b, 0x93, 0x96, 0xb6, 0xf3, 0xe0, 0xfd, 0x98, 0xa1, 0xff, 0x9b,
};

//
// Public exponent of RSA Key.
//
CONST UINT8 mRsaE[] = {
  0x01, 0x00, 0x01
  };

/**
  Check if the application is secure.

  @param[IN]  AppName               The application name.

  @return  TRUE                     It is secure.
  @return  FALSE                    It is not secure.

**/
BOOLEAN
SecureCheck (
  IN  CONST CHAR16                  *AppName
  )
{
  EFI_STATUS        Status = EFI_ABORTED;
  SHELL_FILE_HANDLE Handle = NULL;
  UINT64            FileSize = 0;
  UINT8             *FileBuffer = NULL;
  UINT8             Digest[SHA256_DIGEST_SIZE];
  VOID              *HashContext = NULL;
  BOOLEAN           Result = FALSE;
  BOOLEAN           CryptoStatus = FALSE;
  VOID              *Rsa = NULL;

  //
  // Get the data of application, it consists of two part:
  // 1. The real data for application.
  // 2. The authentic data.
  //
  Status = gEfiShellProtocol->OpenFileByName (AppName, &Handle, EFI_FILE_MODE_READ);
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "OpenFileByName failed - %r\n", Status));
    goto DONE;
  }

  Status = gEfiShellProtocol->GetFileSize (Handle, &FileSize);
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "GetFileSize failed - %r\n", Status));
    goto DONE;
  } else {
    DEBUG ((EFI_D_ERROR, "File size: %d\n", FileSize));
    FileBuffer = AllocateZeroPool (FileSize);
    if (NULL == FileBuffer) {
      goto DONE;
    }
  }

  Status = gEfiShellProtocol->ReadFile (Handle, &FileSize, FileBuffer);
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "ReadFile failed - %r\n", Status));
    goto DONE;
  } else {
    DEBUG ((EFI_D_ERROR, "Read file size: %d\n", FileSize));
  }

  Status = gEfiShellProtocol->CloseFile (Handle);
  if (EFI_ERROR (Status)) {
    goto DONE;
  }

  //
  // Now we have the data, do the check.
  //
  // First we should digest the application real data.
  //
  ZeroMem (Digest, SHA256_DIGEST_SIZE);

  //
  // Allocate hash context buffer required for SHA 256
  //
  HashContext = AllocateZeroPool (Sha256GetContextSize ());
  if (NULL == HashContext) {
    DEBUG ((EFI_D_ERROR, "[%a][%d] Out of memory\n", __FUNCTION__, __LINE__));
    goto DONE;
  }

  CryptoStatus = Sha256Init (HashContext);
  if (!CryptoStatus) {
    DEBUG ((DEBUG_ERROR, "Sha256Init failed\n"));
    goto DONE;
  }

  CryptoStatus = Sha256Update (HashContext, FileBuffer, (FileSize - RSA_LEN));
  if (!CryptoStatus) {
    DEBUG ((DEBUG_ERROR, "Sha256Update failed\n"));
    goto DONE;
  }

  CryptoStatus = Sha256Final (HashContext, Digest);
  if (!CryptoStatus) {
    DEBUG ((DEBUG_ERROR, "Sha256Final failed\n"));
    goto DONE;
  }

  // BeniDumpHex (2, 0, SHA256_DIGEST_SIZE, Digest);

  //
  // Now we have the SHA256 digest, do the RAS check.
  //

  //
  // Generate & Initialize RSA Context.
  //
  Rsa = RsaNew ();
  if (NULL == Rsa) {
    DEBUG ((DEBUG_ERROR, "RsaNew failed\n"));
    goto DONE;
  }

  CryptoStatus = RsaSetKey (Rsa, RsaKeyN, mPublicKey, sizeof (mPublicKey));
  if (!CryptoStatus) {
    DEBUG ((DEBUG_ERROR, "RsaSetKey 1 failed\n"));
    goto DONE;
  }

  CryptoStatus = RsaSetKey (Rsa, RsaKeyE, mRsaE, sizeof (mRsaE));
  if (!CryptoStatus) {
    DEBUG ((DEBUG_ERROR, "RsaSetKey 2 failed\n"));
    goto DONE;
  }

  // BeniDumpHex (2, 0, RSA_LEN, FileBuffer + (FileSize - RSA_LEN));

  CryptoStatus = RsaPkcs1Verify (
                  Rsa,
                  Digest,
                  SHA256_DIGEST_SIZE,
                  FileBuffer + (FileSize - RSA_LEN),
                  RSA_LEN
                  );
  if (!CryptoStatus) {
    DEBUG ((DEBUG_ERROR, "RsaPkcs1Verify failed\n"));
    goto DONE;
  } else {
    //
    // Checked!
    //
    Result = TRUE;
  }

DONE:

  if (NULL != Rsa) {
    RsaFree (Rsa);
  }

  SHELL_FREE_NON_NULL (FileBuffer);
  SHELL_FREE_NON_NULL (HashContext);

  return Result;
}

/**
  Execute Shell application.

  @param[IN]  AppName               The application name.

  @return  NA

**/
VOID
Exec (
  IN  CONST CHAR16                  *AppName
  )
{
  EFI_STATUS                Status = EFI_ABORTED;
  EFI_DEVICE_PATH_PROTOCOL  *DevPath = NULL;
  CHAR16                    *Str = NULL;

  if (!(SecureCheck (AppName))) {
    ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_SECURE_ERROR), mExecHiiHandle);
    return;
  } else {
    ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_SECURE_SUCCESS), mExecHiiHandle);
  }

  DevPath = gEfiShellProtocol->GetDevicePathFromFilePath (AppName);
  if (NULL == DevPath) {
    DEBUG ((EFI_D_ERROR, "Device path not found!\n"));
    return;
  } else {
    Str = ConvertDevicePathToText (DevPath, TRUE, FALSE);
    if (Str) {
      DEBUG ((EFI_D_ERROR, "DevPath: %s\n", Str));
      FreePool (Str);
    }
  }

  Status = gEfiShellProtocol->Execute (&gImageHandle, (CHAR16 *)AppName, NULL, NULL);
  if (EFI_ERROR (Status)) {
    DEBUG ((EFI_D_ERROR, "Execute failed. - %r\n", Status));
  }

  return;
}

/**
  Function for 'exec' command.

  @param[in]  ImageHandle           The image handle.
  @param[in]  SystemTable           The system table.

  @retval  SHELL_SUCCESS            Command completed successfully.
  @retval  SHELL_INVALID_PARAMETER  Command usage error.
  @retval  SHELL_ABORTED            The user aborts the operation.
  @retval  Value                    Unknown error.

**/
SHELL_STATUS
RunExec (
  IN  EFI_HANDLE                    ImageHandle,
  IN  EFI_SYSTEM_TABLE              *SystemTable
  )
{
  SHELL_STATUS  ShellStatus = SHELL_INVALID_PARAMETER;
  EFI_STATUS    Status = EFI_ABORTED;
  LIST_ENTRY    *CheckPackage = NULL;
  CHAR16        *ProblemParam = NULL;
  UINTN         ParamCount = 0;
  CONST CHAR16  *AppName = NULL;

  //
  // Initialize the Shell library (we must be in non-auto-init...).
  //
  Status = ShellInitialize ();
  if (EFI_ERROR (Status)) {
    return SHELL_ABORTED;
  }

  //
  // Parse the command line.
  //
  Status = ShellCommandLineParse (ParamList, &CheckPackage, &ProblemParam, TRUE);
  if (EFI_ERROR (Status)) {
    if ((Status == EFI_VOLUME_CORRUPTED) && (ProblemParam != NULL) ) {
      ShellPrintHiiEx (
        -1, -1, NULL, STRING_TOKEN (STR_GEN_PROBLEM),
        mExecHiiHandle, L"exec", ProblemParam
        );
      FreePool (ProblemParam);
    }
    goto DONE;
  }

  //
  // Check the number of parameters
  //
  ParamCount = ShellCommandLineGetCount (CheckPackage);
  if (ParamCount < 2) {
    ShellPrintHiiEx (
      -1, -1, NULL, STRING_TOKEN (STR_GEN_TOO_FEW),
      mExecHiiHandle, L"exec"
      );
    goto DONE;
  }
  if (ParamCount > 2) {
    ShellPrintHiiEx (
      -1, -1, NULL, STRING_TOKEN (STR_GEN_TOO_MANY),
      mExecHiiHandle, L"exec"
      );
    goto DONE;
  }

  AppName = ShellCommandLineGetRawValue (CheckPackage, 1);
  if (!AppName) {
    ShellPrintHiiEx (
      -1, -1, NULL, STRING_TOKEN (STR_GEN_PROBLEM),
      mExecHiiHandle, L"exec"
      );
    goto DONE;
  }

  ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_EXEC_START), mExecHiiHandle);
  Exec (AppName);
  ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_EXEC_END), mExecHiiHandle);

DONE:

  if ((ShellStatus != SHELL_SUCCESS) && (EFI_ERROR (Status))) {
    ShellStatus = Status & ~MAX_BIT;
  }

  return ShellStatus;
}