/** @file
  Main file for attrib shell level 2 function.

  (C) Copyright 2015 Hewlett-Packard Development Company, L.P.<BR>
  Copyright (c) 2009 - 2019, Intel Corporation. All rights reserved.<BR>
  SPDX-License-Identifier: BSD-2-Clause-Patent

**/

#include "UefiShellLevel2CommandsLib.h"

// This function was from from the BdsLib implementation in
// IntelFrameworkModulePkg\Library\GenericBdsLib\BdsConnect.c
// function name: BdsLibConnectAllEfi

/**
  This function will connect all current system handles recursively. The
  connection will finish until every handle's child handle created if it have.

  @retval EFI_SUCCESS           All handles and it's child handle have been
                                connected
  @retval EFI_STATUS            Return the status of gBS->LocateHandleBuffer().

**/
EFI_STATUS
ConnectAllEfi (
  VOID
  )
{
  EFI_STATUS  Status;
  UINTN       HandleCount;
  EFI_HANDLE  *HandleBuffer;
  UINTN       Index;

  Status = gBS->LocateHandleBuffer (
                  AllHandles,
                  NULL,
                  NULL,
                  &HandleCount,
                  &HandleBuffer
                  );
  if (EFI_ERROR (Status)) {
    return Status;
  }

  for (Index = 0; Index < HandleCount; Index++) {
    Status = gBS->ConnectController (HandleBuffer[Index], NULL, NULL, TRUE);
  }

  if (HandleBuffer != NULL) {
    FreePool (HandleBuffer);
  }

  return EFI_SUCCESS;
}

/**
  function to load a .EFI driver into memory and possible connect the driver.

  if FileName is NULL then ASSERT.

  @param[in] FileName           FileName of the driver to load
  @param[in] Connect            Whether to connect or not

  @retval EFI_SUCCESS           the driver was loaded and if Connect was
                                true then connect was attempted. Connection may
                                have failed.
  @retval EFI_OUT_OF_RESOURCES  there was insufficient memory
**/
EFI_STATUS
LoadDriver (
  IN CONST CHAR16   *FileName,
  IN CONST BOOLEAN  Connect
  )
{
  EFI_HANDLE                 LoadedDriverHandle;
  EFI_STATUS                 Status;
  EFI_DEVICE_PATH_PROTOCOL   *FilePath;
  EFI_LOADED_IMAGE_PROTOCOL  *LoadedDriverImage;

  LoadedDriverImage  = NULL;
  FilePath           = NULL;
  LoadedDriverHandle = NULL;
  Status             = EFI_SUCCESS;

  ASSERT (FileName != NULL);

  //
  // Fix local copies of the protocol pointers
  //
  Status = CommandInit ();
  ASSERT_EFI_ERROR (Status);

  //
  // Convert to DEVICE_PATH
  //
  FilePath = gEfiShellProtocol->GetDevicePathFromFilePath (FileName);

  if (FilePath == NULL) {
    ASSERT (FALSE);
    return (EFI_INVALID_PARAMETER);
  }

  //
  // Use LoadImage to get it into memory
  //
  Status = gBS->LoadImage (
                  FALSE,
                  gImageHandle,
                  FilePath,
                  NULL,
                  0,
                  &LoadedDriverHandle
                  );

  if (EFI_ERROR (Status)) {
    //
    // With EFI_SECURITY_VIOLATION retval, the Image was loaded and an ImageHandle was created
    // with a valid EFI_LOADED_IMAGE_PROTOCOL, but the image can not be started right now.
    // If the caller doesn't have the option to defer the execution of an image, we should
    // unload image for the EFI_SECURITY_VIOLATION to avoid resource leak.
    //
    if (Status == EFI_SECURITY_VIOLATION) {
      gBS->UnloadImage (LoadedDriverHandle);
    }

    ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_LOAD_NOT_IMAGE), gShellLevel2HiiHandle, FileName, Status);
  } else {
    //
    // Make sure it is a driver image
    //
    Status = gBS->HandleProtocol (LoadedDriverHandle, &gEfiLoadedImageProtocolGuid, (VOID *)&LoadedDriverImage);

    ASSERT (LoadedDriverImage != NULL);

    if (  EFI_ERROR (Status)
       || (  (LoadedDriverImage->ImageCodeType != EfiBootServicesCode)
          && (LoadedDriverImage->ImageCodeType != EfiRuntimeServicesCode))
          )
    {
      ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_LOAD_NOT_DRIVER), gShellLevel2HiiHandle, FileName);

      //
      // Exit and unload the non-driver image
      //
      gBS->Exit (LoadedDriverHandle, EFI_INVALID_PARAMETER, 0, NULL);
      Status = EFI_INVALID_PARAMETER;
    }
  }

  if (!EFI_ERROR (Status)) {
    //
    // Start the image
    //
    Status = gBS->StartImage (LoadedDriverHandle, NULL, NULL);
    if (EFI_ERROR (Status)) {
      ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_LOAD_ERROR), gShellLevel2HiiHandle, FileName, Status);
    } else {
      ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_LOAD_LOADED), gShellLevel2HiiHandle, FileName, LoadedDriverImage->ImageBase, Status);
    }
  }

  if (!EFI_ERROR (Status) && Connect) {
    //
    // Connect it...
    //
    Status = ConnectAllEfi ();
  }

  //
  // clean up memory...
  //
  if (FilePath != NULL) {
    FreePool (FilePath);
  }

  return (Status);
}

STATIC CONST SHELL_PARAM_ITEM  LoadParamList[] = {
  { L"-nc", TypeFlag },
  { NULL,   TypeMax  }
};

/**
  Function for 'load' command.

  @param[in] ImageHandle  Handle to the Image (NULL if Internal).
  @param[in] SystemTable  Pointer to the System Table (NULL if Internal).
**/
SHELL_STATUS
EFIAPI
ShellCommandRunLoad (
  IN EFI_HANDLE        ImageHandle,
  IN EFI_SYSTEM_TABLE  *SystemTable
  )
{
  EFI_STATUS           Status;
  LIST_ENTRY           *Package;
  CHAR16               *ProblemParam;
  SHELL_STATUS         ShellStatus;
  UINTN                ParamCount;
  EFI_SHELL_FILE_INFO  *ListHead;
  EFI_SHELL_FILE_INFO  *Node;

  ListHead     = NULL;
  ProblemParam = NULL;
  ShellStatus  = SHELL_SUCCESS;

  //
  // initialize the shell lib (we must be in non-auto-init...)
  //
  Status = ShellInitialize ();
  ASSERT_EFI_ERROR (Status);

  //
  // parse the command line
  //
  Status = ShellCommandLineParse (LoadParamList, &Package, &ProblemParam, TRUE);
  if (EFI_ERROR (Status)) {
    if ((Status == EFI_VOLUME_CORRUPTED) && (ProblemParam != NULL)) {
      ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_GEN_PROBLEM), gShellLevel2HiiHandle, L"load", ProblemParam);
      FreePool (ProblemParam);
      ShellStatus = SHELL_INVALID_PARAMETER;
    } else {
      ASSERT (FALSE);
    }
  } else {
    //
    // check for "-?"
    //
    if (ShellCommandLineGetFlag (Package, L"-?")) {
      ASSERT (FALSE);
    } else if (ShellCommandLineGetRawValue (Package, 1) == NULL) {
      //
      // we didnt get a single file to load parameter
      //
      ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_GEN_TOO_FEW), gShellLevel2HiiHandle, L"load");
      ShellStatus = SHELL_INVALID_PARAMETER;
    } else {
      for ( ParamCount = 1
            ; ShellCommandLineGetRawValue (Package, ParamCount) != NULL
            ; ParamCount++
            )
      {
        Status = ShellOpenFileMetaArg ((CHAR16 *)ShellCommandLineGetRawValue (Package, ParamCount), EFI_FILE_MODE_READ, &ListHead);
        if (!EFI_ERROR (Status)) {
          for ( Node = (EFI_SHELL_FILE_INFO *)GetFirstNode (&ListHead->Link)
                ; !IsNull (&ListHead->Link, &Node->Link)
                ; Node = (EFI_SHELL_FILE_INFO *)GetNextNode (&ListHead->Link, &Node->Link)
                )
          {
            //
            // once we have an error preserve that value, but finish the loop.
            //
            if (EFI_ERROR (Status)) {
              LoadDriver (Node->FullName, (BOOLEAN)(ShellCommandLineGetFlag (Package, L"-nc") == FALSE));
            } else {
              Status = LoadDriver (Node->FullName, (BOOLEAN)(ShellCommandLineGetFlag (Package, L"-nc") == FALSE));
            }
          } // for loop for multi-open

          if (EFI_ERROR (Status)) {
            ShellCloseFileMetaArg (&ListHead);
          } else {
            Status = ShellCloseFileMetaArg (&ListHead);
          }
        } else {
          //
          // no files found.
          //
          ShellPrintHiiEx (-1, -1, NULL, STRING_TOKEN (STR_GEN_FILE_NF), gShellLevel2HiiHandle, L"load", (CHAR16 *)ShellCommandLineGetRawValue (Package, ParamCount));
          ShellStatus = SHELL_NOT_FOUND;
        }
      } // for loop for params
    }

    //
    // free the command line package
    //
    ShellCommandLineFreeVarList (Package);
  }

  if (EFI_ERROR (Status) && (ShellStatus == SHELL_SUCCESS)) {
    ShellStatus = SHELL_DEVICE_ERROR;
  }

  return (ShellStatus);
}