// SPDX-License-Identifier: MPL-2.0
#include "../include/vftm.h"
#include "../include/crypto.h"
#include <Uefi.h>
#include <Library/BaseMemoryLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/UefiBootServicesTableLib.h>
#include <Library/UefiLib.h>
#include <Library/BaseLib.h>
#include <Protocol/LoadedImage.h>
#include <Protocol/BlockIo.h>
#include <Protocol/DevicePath.h>
#include <Library/DevicePathLib.h>
#include <Uefi/UefiGpt.h>
#include <Guid/Gpt.h>
bp_vftm_VFTM_ASSETS assets;
static VOID generate_vftm_sectors(const UINT8* part_guid,UINTN* s1,UINTN* s2) {
  UINT32 h=0x811c9dc5;
  for (UINTN i=0;i<16;i++) h=(h^part_guid[i])*0x01000193;
  *s1=8+(h%24);
  *s2=8+(((h>>16)%23+(*s1-7))%24);
}
EFI_STATUS bp_vftm_get_vftm_assets(EFI_HANDLE image_handle) {
  EFI_LOADED_IMAGE_PROTOCOL *loaded_image;
  EFI_STATUS status=gBS->HandleProtocol(image_handle,&gEfiLoadedImageProtocolGuid,(VOID**)&loaded_image);
  if (EFI_ERROR(status)) {
    Print(L"[BlastVFTM] Error: can't get LoadedImage.\n");
    while (TRUE) {};
  } else {
    Print(L"[BlastVFTM] LoadedImage acquired.\n");
  }
  EFI_BLOCK_IO_PROTOCOL *block_io_esp=NULL;
  status=gBS->HandleProtocol(loaded_image->DeviceHandle,&gEfiBlockIoProtocolGuid,(VOID **)&block_io_esp);
  if (EFI_ERROR(status)) {
    Print(L"[BlastVFTM] Error: can't get BlockIo.\n");
    while (TRUE) {};
  } else {
    Print(L"[BlastVFTM] BlockIo on LoadedImage acquired.\n");
  }
  if (!block_io_esp->Media->LogicalPartition) {
    Print(L"[BlastVFTM] Error: BlockIo on LoadedImage isn't a partition.\n");
    while (TRUE) {};
  }
  EFI_DEVICE_PATH_PROTOCOL *full_path;
  status=gBS->OpenProtocol(loaded_image->DeviceHandle,&gEfiDevicePathProtocolGuid,(VOID**)&full_path,image_handle,NULL,EFI_OPEN_PROTOCOL_GET_PROTOCOL);
  if (EFI_ERROR(status)) {
    Print(L"[BlastVFTM] Error: couldn't get full device path for LoadedImage. Status: %d\n",status);
    while (TRUE) {};
  }
  Print(L"[BlastVFTM] Full device path for LoadedImage acquired.\n");
  EFI_DEVICE_PATH_PROTOCOL *parent_path=DuplicateDevicePath(full_path);
  if (parent_path==NULL) {
    Print(L"[BlastVFTM] Error: couldn't duplicate full path to parent path.\n");
    while (TRUE) {};
  }
  EFI_DEVICE_PATH_PROTOCOL *node;
  node=parent_path;
  while (!IsDevicePathEnd(NextDevicePathNode(node))) {
    node=NextDevicePathNode(node);
  }
  if (DevicePathType(node)!=MEDIA_DEVICE_PATH || DevicePathSubType(node)!=MEDIA_HARDDRIVE_DP) {
    FreePool(parent_path);
    Print(L"[BlastVFTM] Error: couldn't find a valid disk device path.\n");
    while (TRUE) {};
  }
  Print(L"[BlastVFTM] Found a valid disk device path.\n");
  SetDevicePathEndNode(node);
  EFI_DEVICE_PATH_PROTOCOL *remaining_path=parent_path;
  EFI_HANDLE disk_handle;
  status=gBS->LocateDevicePath(&gEfiBlockIoProtocolGuid,&remaining_path,&disk_handle);
  FreePool(parent_path);
  if (EFI_ERROR(status)) {
    Print(L"[BlastVFTM] Error: couldn't obtain disk handle.\n");
    while (TRUE) {};
  }
  Print(L"[BlastVFTM] Obtained disk handle.\n");
  EFI_BLOCK_IO_PROTOCOL *disk_blockio=NULL;
  status=gBS->OpenProtocol(disk_handle,&gEfiBlockIoProtocolGuid,(VOID**)&disk_blockio,image_handle,NULL,EFI_OPEN_PROTOCOL_GET_PROTOCOL);
  if (EFI_ERROR(status)) {
    Print(L"[BlastVFTM] Error: couldn't obtain disk BlockIO. Status: %d\n",status);
    while (TRUE) {};
  }
  Print(L"[BlastVFTM] Obtained disk BlockIO.\n");
  if (disk_blockio->Media->LogicalPartition) {
    Print(L"[BlastVFTM] Error: media is a partition.\n");
    while (TRUE) {};
  }
  if (disk_blockio->Media->BlockSize!=512) {
    Print(L"[BlastVFTM] Error: disk doesn't have a block size of 512.\n");
    while (TRUE) {};
  }
  assets.block_io_disk=disk_blockio;
  UINT8 *gpt_buffer=AllocatePool(512);
  if (!gpt_buffer) {
    Print(L"[BlastVFTM] Error: out of memory for GPT Header.\n");
    while (TRUE) {};
  }
  status=disk_blockio->ReadBlocks(disk_blockio,disk_blockio->Media->MediaId,1,512,gpt_buffer);
  if (EFI_ERROR(status)) {
    Print(L"[BlastVFTM] Error: couldn't read first LBA of the disk.\n");
    while (TRUE) {};
  }
  EFI_PARTITION_TABLE_HEADER *gpt_header=(EFI_PARTITION_TABLE_HEADER*)gpt_buffer;
  if (gpt_header->Header.Signature!=EFI_PTAB_HEADER_ID) {
    Print(L"[BlastVFTM] Error: first LBA of the disk isn't a GPT table.\n");
    while (TRUE) {};
  }
  CopyMem(assets.disk_guid,&gpt_header->DiskGUID,sizeof(gpt_header->DiskGUID));
  Print(L"[BlastVFTM] Successfully obtained disk GUID.\n");
  UINTN parts_size=gpt_header->NumberOfPartitionEntries*gpt_header->SizeOfPartitionEntry;
  UINT8 *partition_table_buffer=AllocatePool(parts_size);
  if (!partition_table_buffer) {
    Print(L"[BlastVFTM] Error: out of memory for partitions entries.\n");
    while (TRUE) {};
  }
  status=disk_blockio->ReadBlocks(disk_blockio,disk_blockio->Media->MediaId,gpt_header->PartitionEntryLBA,parts_size,partition_table_buffer);
  if (EFI_ERROR(status)) {
    Print(L"[BlastVFTM] Error: couldn't read partition entries.\n");
    while (TRUE) {};
  }
  assets.partition_table_buffer=partition_table_buffer;
  EFI_PARTITION_ENTRY *entry=(EFI_PARTITION_ENTRY*)partition_table_buffer;
  INTN esp_part_index=-1;
  for (UINTN i=0;i<gpt_header->NumberOfPartitionEntries;i++) {
    EFI_PARTITION_ENTRY *current=&entry[i];
    if (CompareGuid(&current->PartitionTypeGUID,&gEfiPartTypeSystemPartGuid)) {
      esp_part_index=i;
      break;
    }
  }
  if (esp_part_index<0) {
    Print(L"[BlastVFTM] Error: couldn't found any ESP partition.\n");
    while (TRUE) {};
  }
  assets.part_count=gpt_header->NumberOfPartitionEntries;
  EFI_PARTITION_ENTRY esp_entry=entry[esp_part_index];
  CopyMem(assets.efi_part_unique_guid,&esp_entry.UniquePartitionGUID,sizeof(esp_entry.UniquePartitionGUID));
  Print(L"[BlastVFTM] Successfully obtained ESP partition unique GUID.\n");
  UINTN s_manifest,s_pk_manifest;
  generate_vftm_sectors(assets.efi_part_unique_guid,&s_manifest,&s_pk_manifest);
  UINT8 *vftm_buffer=AllocatePool(512);
  if (!vftm_buffer) {
    Print(L"[BlastVFTM] Error: out of memory for manifest.\n");
    while (TRUE) {};
  }
  status=block_io_esp->ReadBlocks(block_io_esp,block_io_esp->Media->MediaId,s_manifest,512,vftm_buffer);
  if (EFI_ERROR(status)) {
    Print(L"[BlastVFTM] Error: couldn't read manifest.\n");
    while (TRUE) {};
  }
  CopyMem(&assets.manifest,vftm_buffer,sizeof(assets.manifest));
  Print(L"[BlastVFTM] Successfully obtained manifest.\n");
  UINT8 *pk_manifest_buffer=AllocatePool(512);
  if (!pk_manifest_buffer) {
    Print(L"[BlastVFTM] Error: out of memory for manifest public key.\n");
    while (TRUE) {};
  }
  status=block_io_esp->ReadBlocks(block_io_esp,block_io_esp->Media->MediaId,s_pk_manifest,512,pk_manifest_buffer);
  if (EFI_ERROR(status)) {
    Print(L"[BlastVFTM] Error: couldn't read manifest public key.\n");
    while (TRUE) {};
  }
  CopyMem(&assets.pk_manifest,pk_manifest_buffer,sizeof(assets.pk_manifest));
  Print(L"[BlastVFTM] Successfully obtained manifest public key.\n");
  UINTN signature_size_padded=(assets.manifest.signature_size+511) & ~511;
  if (assets.manifest.signature_offset_bytes%512!=0) {
    Print(L"[BlastVFTM] Error: provided signature offset isn't a multiple of 512.\n");
    while (TRUE) {};
  }
  UINT8 *signature_buffer=AllocatePool(signature_size_padded);
  if (!signature_buffer) {
    Print(L"[BlastVFTM] Error: out of memory for manifest signature.\n");
    while (TRUE) {};
  }
  status=block_io_esp->ReadBlocks(block_io_esp,block_io_esp->Media->MediaId,assets.manifest.signature_offset_bytes/512,signature_size_padded,signature_buffer);
  if (EFI_ERROR(status)) {
    Print(L"[BlastVFTM] Error: couldn't read manifest signature. Status: %d\n",status);
    while (TRUE) {};
  }
  assets.manifest_signature=signature_buffer;
  Print(L"[BlastVFTM] Successfully obtained manifest signature.\n");
  assets.block_io_esp=block_io_esp;
  return EFI_SUCCESS;
}
EFI_STATUS bp_vftm_check_manifest_authenticity(UINT8 *pk_root,UINT8* sig_root) {
  if (pk_root==NULL || sig_root==NULL) return EFI_INVALID_PARAMETER;
  EFI_STATUS status=bp_crypto_sphincsplus_verify(sig_root,CRYPTO_BYTES,assets.pk_manifest,sizeof(assets.pk_manifest),pk_root);
  if (EFI_ERROR(status)) {
    Print(L"[BlastVFTM] Error: critical security violation in boot process:\n",status);
    Print(L"[BlastVFTM] Provided manifest public key couldn't be authentified.\n",status);
    while (TRUE) {};
  }
  Print(L"[BlastVFTM] Successfully authentified manifest public key.\n",status);
  status=bp_crypto_sphincsplus_verify(assets.manifest_signature,CRYPTO_BYTES,(UINT8*)&assets.manifest,sizeof(assets.manifest),assets.pk_manifest);
  if (EFI_ERROR(status)) {
    Print(L"[BlastVFTM] Error: critical security violation in boot process:\n",status);
    Print(L"[BlastVFTM] Provided Vystem FAT Trusted Manifest couldn't be authentified.\n",status);
    while (TRUE) {};
  }
  Print(L"[BlastVFTM] Successfully authentified provided Vystem FAT Trusted Manifest.\n",status);
  return EFI_SUCCESS;
}
EFI_STATUS bp_vftm_check_manifest_content() {
  UINT8 sig[8]=BP_VFTM_SIGNATURE;
  for (int i=0;i<8;i++) {
    if (sig[i]!=assets.manifest.sig[i]) {
      Print(L"[BlastVFTM] Error: invalid manifest signature.\n");
      while (TRUE) {};
    }
  }
  if (assets.manifest.manifest_version!=BP_VFTM_SUPPORTED_MANIFEST_VERSION) {
    Print(L"[BlastVFTM] Error: non supported manifest version.\n");
    while (TRUE) {};
  }
  if (assets.manifest.os_version!=BP_VFTM_SUPPORTED_OS_VERSION) {
    Print(L"[BlastVFTM] Error: wrong OS version.\n");
    while (TRUE) {};
  }
  if (assets.manifest.bootloader_version!=BP_VFTM_SUPPORTED_BOOTLOADER_VERSION) {
    Print(L"[BlastVFTM] Error: wrong bootloader version.\n");
    while (TRUE) {};
  }
  if (CompareMem(assets.manifest.gpt_disk_guid,assets.disk_guid,16)!=0) {
    Print(L"[BlastVFTM] Error: manifest disk GUID doesn't correspond with GPT table disk GUID.\n");
    while (TRUE) {};
  }
  if (CompareMem(assets.manifest.gpt_efi_part_unique_guid,assets.efi_part_unique_guid,16)!=0) {
    Print(L"[BlastVFTM] Error: manifest ESP unique GUID doesn't correspond with GPT table ESP unique GUID.\n");
    while (TRUE) {};
  }
  Print(L"[BlastVFTM] Successfully verified versions, disk and partitions informations from manifest.\n");
  UINT8* fat32_header_buffer=AllocatePool(512);
  if (!fat32_header_buffer) {
    Print(L"[BlastVFTM] Error: out of memory for FAT32 header buffer.\n");
    while (TRUE) {};
  }
  EFI_STATUS status=assets.block_io_esp->ReadBlocks(assets.block_io_esp,assets.block_io_esp->Media->MediaId,0,512,fat32_header_buffer);
  if (EFI_ERROR(status)) {
    Print(L"[BlastVFTM] Error: couldn't ESP FAT32 header. Status: %d\n",status);
    while (TRUE) {};
  }
  bp_vftm_FAT32_HEADER *fat32_header=(bp_vftm_FAT32_HEADER*)fat32_header_buffer;
  bp_crypto_HASH_DIGEST *header_hash;
  status=bp_crypto_sha3((UINT8*)fat32_header,sizeof(bp_vftm_FAT32_HEADER),64,&header_hash);
  if (EFI_ERROR(status)) {
    Print(L"[BlastVFTM] Error: couldn't hash FAT32 header.\n");
    while (TRUE) {};
  }
  if (CompareMem(header_hash->digest,assets.manifest.fat_header_hash,header_hash->size_bytes)!=0) {
    Print(L"[BlastVFTM] Error: FAT32 header hash didn't correspond with manifest FAT32 header hash.\n");
    while (TRUE) {};
  }
  UINTN fat_total_sectors=fat32_header->BPB_FATSz32*fat32_header->BPB_NumFATs;
  UINTN data_start_sector=fat32_header->BPB_RsvdSecCnt+fat_total_sectors;
  UINTN data_total_sectors=fat32_header->BPB_TotSec32-data_start_sector;
  UINTN data_clusters=data_total_sectors/fat32_header->BPB_SecPerClus;
  if (data_clusters!=assets.manifest.clusters_count) {
    Print(L"[BlastVFTM] Error: FAT32 header indicate a wrong amount of data clusters.\n");
    while (TRUE) {};
  }
  if (fat32_header->BPB_BytsPerSec!=assets.manifest.bytes_per_sector) {
    Print(L"[BlastVFTM] Error: FAT32 header indicate a wrong amount of bytes per sector.\n");
    while (TRUE) {};
  }
  if (fat32_header->BPB_SecPerClus!=assets.manifest.sectors_per_cluster) {
    Print(L"[BlastVFTM] Error: FAT32 header indicate a wrong amount of sector per cluster.\n");
    while (TRUE) {};
  }
  if (fat32_header->BPB_FATSz32*fat32_header->BPB_BytsPerSec!=assets.manifest.fat_size_bytes) {
    Print(L"[BlastVFTM] Error: FAT32 header indicate a wrong amount of bytes per FAT.\n");
    while (TRUE) {};
  }
  if (fat32_header->BPB_RsvdSecCnt*fat32_header->BPB_BytsPerSec!=assets.manifest.fat_offset_bytes) {
    Print(L"[BlastVFTM] Error: FAT32 header indicate a wrong amount of bytes per FAT.\n");
    while (TRUE) {};
  }
  if (CompareMem(fat32_header->BS_OEMName,"VYBUILD1",8)!=0) {
    Print(L"[BlastVFTM] Warning: OEM name indicated in FAT32 header isn't the expected OEM name: %s\n",fat32_header->BS_OEMName);
  }
  if (CompareMem(fat32_header->BS_VolLab,"VYAUTOFAT32",11)!=0) {
    Print(L"[BlastVFTM] Warning: Volume name indicated in FAT32 header isn't the expected volume name: %s\n",fat32_header->BS_VolLab);
  }
  Print(L"[BlastVFTM] Successfully verified ESP FAT32 header.\n");
  UINT8* fat32_fat_buffer=AllocatePool(assets.manifest.fat_size_bytes*2);
  if (!fat32_fat_buffer) {
    Print(L"[BlastVFTM] Error: out of memory for FAT32 double FAT buffer.\n");
    while (TRUE) {};
  }
  status=assets.block_io_esp->ReadBlocks(assets.block_io_esp,assets.block_io_esp->Media->MediaId,assets.manifest.fat_offset_bytes/512,assets.manifest.fat_size_bytes*2,fat32_fat_buffer);
  if (EFI_ERROR(status)) {
    Print(L"[BlastVFTM] Error: couldn't read ESP FAT32 double FAT. Status: %d\n",status);
    while (TRUE) {};
  }
  UINT32 *f1=(UINT32*)fat32_fat_buffer;
  UINT32 *f2=(UINT32*)(fat32_fat_buffer+assets.manifest.fat_size_bytes);
  f1[0]=f1[1]=f2[0]=f2[1]=0;
  bp_crypto_HASH_DIGEST *fat1_hash;
  bp_crypto_HASH_DIGEST *fat2_hash;
  status=bp_crypto_sha3((UINT8*)fat32_fat_buffer,assets.manifest.fat_size_bytes,64,&fat1_hash);
  if (EFI_ERROR(status)) {
    Print(L"[BlastVFTM] Error: couldn't hash FAT32 first FAT.\n");
    while (TRUE) {};
  }
  status=bp_crypto_sha3((UINT8*)fat32_fat_buffer+assets.manifest.fat_size_bytes,assets.manifest.fat_size_bytes,64,&fat2_hash);
  if (EFI_ERROR(status)) {
    Print(L"[BlastVFTM] Error: couldn't hash FAT32 second FAT.\n");
    while (TRUE) {};
  }
  if (CompareMem(fat1_hash->digest,assets.manifest.fat_hash,fat1_hash->size_bytes)!=0) {
    Print(L"[BlastVFTM] Error: FAT32 first FAT hash didn't correspond with manifest FAT32 fat hash.\n");
    while (TRUE) {};
  }
  if (CompareMem(fat2_hash->digest,assets.manifest.fat_hash,fat2_hash->size_bytes)!=0) {
    Print(L"[BlastVFTM] Error: FAT32 second FAT hash didn't correspond with manifest FAT32 fat hash\nA firmware that dump things on the ESP partition can cause the VFTM check to fail.\nConsider adding a VAR file or equivalent to your VM if that is not the case.\n");
    while (TRUE) {};
  }
  Print(L"[BlastVFTM] Successfully verified all ESP FATs.\n");
  UINTN bytes_per_cluster=fat32_header->BPB_BytsPerSec*fat32_header->BPB_SecPerClus;
  UINT8* bootloader_buffer=AllocatePool(assets.manifest.bootloader_cluster_count*bytes_per_cluster);
  if (!bootloader_buffer) {
    Print(L"[BlastVFTM] Error: out of memory for bootloader buffer.\n");
    while (TRUE) {};
  }
  UINTN lba_data_start=fat32_header->BPB_RsvdSecCnt+(fat32_header->BPB_NumFATs*fat32_header->BPB_FATSz32);
  UINTN bootloader_lba=lba_data_start+(assets.manifest.bootloader_cluster_offset-2);
  status=assets.block_io_esp->ReadBlocks(assets.block_io_esp,assets.block_io_esp->Media->MediaId,bootloader_lba,assets.manifest.bootloader_cluster_count*bytes_per_cluster,bootloader_buffer);
  if (EFI_ERROR(status)) {
    Print(L"[BlastVFTM] Error: couldn't read bootloader. Status: %d\n",status);
    while (TRUE) {};
  }
  bp_crypto_HASH_DIGEST *bootloader_hash;
  status=bp_crypto_sha3((UINT8*)bootloader_buffer,assets.manifest.bootloader_cluster_count*bytes_per_cluster,64,&bootloader_hash);
  if (EFI_ERROR(status)) {
    Print(L"[BlastVFTM] Error: couldn't hash bootloader.\n");
    while (TRUE) {};
  }
  if (CompareMem(bootloader_hash->digest,assets.manifest.bootloader_binary_hash,bootloader_hash->size_bytes)!=0) {
    Print(L"[BlastVFTM] Error: bootloader hash doesn't correspond with manifest bootloader hash.\n");
    while (TRUE) {};
  }
  FreePool(bootloader_buffer);
  Print(L"[BlastVFTM] Successfully verified bootloader.\n");
  bp_crypto_HASH_DIGEST *manifest_hash;
  status=bp_crypto_sha3((UINT8*)&assets.manifest,sizeof(assets.manifest)-64,64,&manifest_hash);
  if (EFI_ERROR(status)) {
    Print(L"[BlastVFTM] Error: couldn't hash manifest.\n");
    while (TRUE) {};
  }
  if (CompareMem(manifest_hash->digest,assets.manifest.manifest_hash,manifest_hash->size_bytes)!=0) {
    Print(L"[BlastVFTM] Error: manifest hash doesn't correspond with the indicated hash inside the manifest.\n");
    while (TRUE) {};
  }
  Print(L"[BlastVFTM] Successfully verified manifest.\n");
  Print(L"[BlastVFTM] All VFTM check passed. Launching SBFIE UI...\n");
  return EFI_SUCCESS;
}
bp_vftm_VFTM_ASSETS *bp_vftm_get_assets() {
  return &assets;
}
