// SPDX-License-Identifier: MPL-2.0
#include "../include/initfs.h"
#include "../include/crypto.h"
#include <Uefi.h>
#include <Base.h>
#include <Library/BaseMemoryLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Protocol/SimpleFileSystem.h>
#include <Protocol/BlockIo.h>
EFI_STATUS bp_initfs_parse_initfs_header(UINT8 *header_data,UINTN buffer_size,bp_initfs_INITFS_HEADER **header) {
  if (header==NULL || header_data==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  if (buffer_size<sizeof(bp_initfs_INITFS_HEADER)) {
    return EFI_BUFFER_TOO_SMALL;
  }
  UINT8 sig[8]={'I','n','i','t','F','i','S','y'};
  if (CompareMem(sig,header_data,8)!=0) {
    return EFI_COMPROMISED_DATA;
  }
  *header=(bp_initfs_INITFS_HEADER *)header_data;
  if ((*header)->initfs_version!=BP_INITFS_INITFS_SUPPORTED_VERSION) {
    return EFI_UNSUPPORTED;
  }
  return EFI_SUCCESS;
}
EFI_STATUS bp_initfs_generate_padding(UINTN size,UINT8 byte,UINT8 **result) {
  if (size==0 || result==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  *result=AllocateZeroPool(size);
  if (!*result) {
    *result=NULL;
    return EFI_OUT_OF_RESOURCES;
  }
  UINTN f0=0;
  UINTN f1=1;
  while (f0<size) {
    (*result)[f0]=byte;
    UINTN next=f0+f1;
    f0=f1;
    f1=next;
  }
  return EFI_SUCCESS;
}
EFI_STATUS bp_initfs_check_initfs_header(bp_initfs_INITFS_HEADER *header,EFI_FILE_PROTOCOL *root) {
  if (header==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  UINT64 value=(header->initfs_size+header->table_size+header->files_area_size+header->entries_width+header->entries_count+header->files_area_offset)%MAX_UINT64;
  UINT64 check1=(value*0x9E3779B185EBCA87)^header->entropy_check1;
  if (check1!=header->check1) {
    return EFI_SECURITY_VIOLATION;
  }
  EFI_FILE_PROTOCOL *opened_file=NULL;
  EFI_STATUS status;
  status=bp_disk_open_file(root,L"\\EFI\\BPFILES\\INITFSFP.BIN",&opened_file);
  if (EFI_ERROR(status)) {
    return status;
  }
  bp_disk_FILE *file=NULL;
  status=bp_disk_read_file(opened_file,&file);
  if (EFI_ERROR(status)) {
    return status;
  }
  if (file->file_size_bytes!=64) {
    return EFI_SECURITY_VIOLATION;
  }
  bp_crypto_HASH_DIGEST *installation_id_hash=NULL;
  status=bp_crypto_sha3(header->installation_id,sizeof(header->installation_id),64,&installation_id_hash);
  if (EFI_ERROR(status)) {
    return status;
  }
  if (CompareMem(installation_id_hash->digest,file->content,64)!=0) {
    return EFI_SECURITY_VIOLATION;
  }
  bp_crypto_HASH_DIGEST *installation_id_hash_hash=NULL;
  status=bp_crypto_sha3(file->content,file->file_size_bytes,64,&installation_id_hash_hash);
  if (EFI_ERROR(status)) {
    return status;
  }
  if (CompareMem(installation_id_hash_hash->digest,header->installation_id_hash_hash,64)!=0) {
    return EFI_SECURITY_VIOLATION;
  }
  UINT8 *padding=NULL;
  status=bp_initfs_generate_padding(sizeof(header->padding),header->installation_id[0],&padding);
  if (EFI_ERROR(status)) {
    return status;
  }
  if (CompareMem(padding,header->padding,sizeof(header->padding))!=0) {
    return EFI_SECURITY_VIOLATION;
  }
  bp_crypto_HASH_DIGEST *header_hash=NULL;
  status=bp_crypto_sha3((UINT8 *)header,sizeof(bp_initfs_INITFS_HEADER)-sizeof(header->header_hash),64,&header_hash);
  if (EFI_ERROR(status)) {
    return status;
  }
  if (CompareMem(header_hash->digest,header->header_hash,sizeof(header->header_hash))!=0) {
    return EFI_VOLUME_CORRUPTED;
  }
  if (header->table_size!=header->entries_count*header->entries_width) {
    return EFI_VOLUME_CORRUPTED;
  }
  if (header->initfs_size!=512+header->table_size+header->files_area_size) {
    return EFI_VOLUME_CORRUPTED;
  }
  if (header->files_area_offset!=512+header->table_size) {
    return EFI_VOLUME_CORRUPTED;
  }
  if (header->files_area_size+header->table_size>header->initfs_size) {
    return EFI_VOLUME_CORRUPTED;
  }
  return EFI_SUCCESS;
}
EFI_STATUS bp_initfs_parse_signsyst_header(UINT8 *header_data,UINTN buffer_size,bp_initfs_SIGNSYST_HEADER **header) {
  if (header==NULL || header_data==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  if (buffer_size<sizeof(bp_initfs_SIGNSYST_HEADER)) {
    return EFI_BUFFER_TOO_SMALL;
  }
  UINT8 sig[8]={'S','i','g','n','S','y','s','t'};
  if (CompareMem(sig,header_data,8)!=0) {
    return EFI_COMPROMISED_DATA;
  }
  *header=(bp_initfs_SIGNSYST_HEADER *)header_data;
  if ((*header)->initfs_version!=BP_INITFS_INITFS_SUPPORTED_VERSION) {
    return EFI_UNSUPPORTED;
  }
  return EFI_SUCCESS;
}
EFI_STATUS bp_initfs_check_signsyst_header(bp_initfs_SIGNSYST_HEADER *header,EFI_FILE_PROTOCOL *root) {
  if (header==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  EFI_STATUS status;
  UINT8 *padding=NULL;
  status=bp_initfs_generate_padding(sizeof(header->padding),header->installation_id[sizeof(header->installation_id)-1],&padding);
  if (EFI_ERROR(status)) {
    return status;
  }
  if (CompareMem(padding,header->padding,sizeof(header->padding))!=0) {
    return EFI_SECURITY_VIOLATION;
  }
  bp_crypto_HASH_DIGEST *header_hash=NULL;
  status=bp_crypto_sha3((UINT8 *)header,sizeof(bp_initfs_SIGNSYST_HEADER)-sizeof(header->header_hash),64,&header_hash);
  if (EFI_ERROR(status)) {
    return status;
  }
  if (CompareMem(header_hash->digest,header->header_hash,sizeof(header->header_hash))!=0) {
    return EFI_SECURITY_VIOLATION;
  }
  EFI_FILE_PROTOCOL *opened_file=NULL;
  status=bp_disk_open_file(root,L"\\EFI\\BPFILES\\SSHASH.BIN",&opened_file);
  if (EFI_ERROR(status)) {
    return status;
  }
  bp_disk_FILE *file=NULL;
  status=bp_disk_read_file(opened_file,&file);
  if (EFI_ERROR(status)) {
    return status;
  }
  if (file->file_size_bytes!=64) {
    return EFI_SECURITY_VIOLATION;
  }
  bp_crypto_HASH_DIGEST *signsyst_header_hash=NULL;
  status=bp_crypto_sha3((UINT8 *)header,sizeof(bp_initfs_SIGNSYST_HEADER),64,&signsyst_header_hash);
  if (EFI_ERROR(status)) {
    return status;
  }
  if (CompareMem(signsyst_header_hash->digest,file->content,64)!=0) {
    return EFI_SECURITY_VIOLATION;
  }
  if (header->signature_block_size!=header->signature_count*header->signature_size) {
    return EFI_VOLUME_CORRUPTED;
  }
  if (header->signsyst_size!=512+header->signature_block_size) {
    return EFI_VOLUME_CORRUPTED;
  }
  if (header->signature_block_size>header->signsyst_size) {
    return EFI_VOLUME_CORRUPTED;
  }
  return EFI_SUCCESS;
}
EFI_STATUS bp_initfs_cross_check_headers(bp_initfs_INITFS_HEADER *initfs,bp_initfs_SIGNSYST_HEADER *signsyst) {
  if (initfs==NULL || signsyst==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  if (initfs->bootloader_version!=signsyst->bootloader_version || initfs->initfs_version!=signsyst->initfs_version || initfs->os_version!=signsyst->os_version) {
    return EFI_SECURITY_VIOLATION;
  }
  if (CompareMem(initfs->installation_id,signsyst->installation_id,sizeof(initfs->installation_id))!=0) {
    return EFI_SECURITY_VIOLATION;
  }
  if (initfs->entries_count!=signsyst->signature_count) {
    return EFI_SECURITY_VIOLATION;
  }
  return EFI_SUCCESS;
}
EFI_STATUS bp_initfs_create_initfs_structure(bp_initfs_INITFS_HEADER *supposed_header,UINT8 *loaded_initfs,bp_initfs_INITFS **initfs) {
  if (supposed_header==NULL || loaded_initfs==NULL || initfs==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  bp_initfs_INITFS *fs=AllocatePool(sizeof(bp_initfs_INITFS));
  if (!fs) {
    return EFI_OUT_OF_RESOURCES;
  }
  if (CompareMem(supposed_header,loaded_initfs,512)!=0) {
    return EFI_VOLUME_CORRUPTED;
  }
  fs->header=supposed_header;
  if (fs->header->entries_count==0) {
    return EFI_NOT_FOUND;
  }
  bp_initfs_FILES_TABLE *files_table=AllocatePool(sizeof(bp_initfs_FILES_TABLE));
  if (!files_table) {
    FreePool(fs);
    return EFI_OUT_OF_RESOURCES;
  }
  files_table->files_entries=(bp_initfs_FILE_ENTRY *)(loaded_initfs+512);
  files_table->file_count=fs->header->entries_count;
  fs->files_table=files_table;
  fs->files_area=loaded_initfs+512+(256*fs->header->entries_count);
  *initfs=fs;
  return EFI_SUCCESS;
}
EFI_STATUS bp_initfs_create_signsyst_structure(bp_initfs_SIGNSYST_HEADER *supposed_header,UINT8 *loaded_signsyst,bp_initfs_SIGNSYST **signsyst) {
  if (supposed_header==NULL || loaded_signsyst==NULL || signsyst==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  bp_initfs_SIGNSYST *fs=AllocatePool(sizeof(bp_initfs_SIGNSYST));
  if (!fs) {
    return EFI_OUT_OF_RESOURCES;
  }
  if (CompareMem(supposed_header,loaded_signsyst,512)!=0) {
    return EFI_VOLUME_CORRUPTED;
  }
  fs->header=supposed_header;
  if (fs->header->signature_count==0) {
    return EFI_NOT_FOUND;
  }
  fs->signature_area=loaded_signsyst+512;
  *signsyst=fs;
  return EFI_SUCCESS;
}
EFI_STATUS bp_initfs_check_initfs_integrity(bp_initfs_INITFS *initfs,bp_initfs_SIGNSYST *signsyst) {
  if (initfs==NULL || signsyst==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  EFI_STATUS status;
  bp_crypto_HASH_DIGEST *table_hash=NULL;
  status=bp_crypto_sha3((UINT8 *)initfs->files_table->files_entries,sizeof(bp_initfs_FILE_ENTRY)*initfs->header->entries_count,64,&table_hash);
  if (EFI_ERROR(status)) {
    return status;
  }
  for (UINTN i=0;i<48;i++) {
    table_hash->digest[i+16]^=initfs->header->installation_id[i];
  }
  if (CompareMem(table_hash->digest,initfs->header->entries_table_hash,64)) {
    return EFI_VOLUME_CORRUPTED;
  }
  bp_crypto_HASH_DIGEST *files_area_hash=NULL;
  status=bp_crypto_sha3(initfs->files_area,initfs->header->files_area_size,64,&files_area_hash);
  if (EFI_ERROR(status)) {
    return status;
  }
  for (UINTN i=0;i<48;i++) {
    files_area_hash->digest[i+16]^=initfs->header->installation_id[i];
  }
  if (CompareMem(files_area_hash->digest,initfs->header->files_area_hash,64)!=0) {
    return EFI_VOLUME_CORRUPTED;
  }
  bp_crypto_HASH_DIGEST *signature_area_hash=NULL;
  status=bp_crypto_sha3(signsyst->signature_area,signsyst->header->signature_block_size,64,&signature_area_hash);
  if (EFI_ERROR(status)) {
    return status;
  }
  for (UINTN i=0;i<48;i++) {
    signature_area_hash->digest[i+16]^=initfs->header->installation_id[i];
  }
  if (CompareMem(signature_area_hash->digest,signsyst->header->signature_block_hash,64)!=0) {
    return EFI_VOLUME_CORRUPTED;
  }
  UINTN cursor=0;
  for (UINTN i=0;i<initfs->header->entries_count;++i) {
    bp_crypto_HASH_DIGEST *file_hash=NULL;
    status=bp_crypto_sha3(initfs->files_area+cursor,initfs->files_table->files_entries[i].file_size,64,&file_hash);
    if (EFI_ERROR(status)) {
      return status;
    }
    if (CompareMem(file_hash->digest,initfs->files_table->files_entries[i].hash,64)!=0) {
      return EFI_SECURITY_VIOLATION;
    }
    status=bp_crypto_sphincsplus_verify(signsyst->signature_area+i*signsyst->header->signature_size,signsyst->header->signature_size,initfs->files_area+cursor,initfs->files_table->files_entries[i].file_size,initfs->files_table->files_entries[i].pk);
    if (EFI_ERROR(status)) {
      return EFI_SECURITY_VIOLATION;
    }
    FreePool(file_hash->digest);
    FreePool(file_hash);
    if (cursor+initfs->files_table->files_entries[i].file_size>initfs->header->files_area_size) {
      return EFI_VOLUME_CORRUPTED;
    }
    cursor+=initfs->files_table->files_entries[i].file_size;
  }
  return EFI_SUCCESS;
}
EFI_STATUS bp_initfs_generate_file_name(CHAR8 *file_name_in,UINT8 **file_name_out) {
  if (file_name_in==NULL || file_name_out==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  UINTN file_name_len=AsciiStrnLenS(file_name_in,113);
  if (file_name_len==113) {
    return EFI_INVALID_PARAMETER;
  }
  *file_name_out=AllocateZeroPool(112);
  if (!*file_name_out) {
    return EFI_OUT_OF_RESOURCES;
  }
  CopyMem(*file_name_out,file_name_in,file_name_len);
  return EFI_SUCCESS;
}
EFI_STATUS bp_initfs_read_file(bp_initfs_INITFS *initfs,CHAR8 *file_name,bp_disk_FILE **file) {
  if (initfs==NULL || file_name==NULL || file==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  UINT8 *file_name_raw=NULL;
  EFI_STATUS status=bp_initfs_generate_file_name(file_name,&file_name_raw);
  if (EFI_ERROR(status)) {
    return status;
  }
  UINTN index=0;
  for (UINTN i=0;i<initfs->header->entries_count;++i) {
    if (CompareMem(file_name_raw,initfs->files_table->files_entries[i].file_name,112)==0) {
      index=i+1;
      break;
    }
  }
  if (index==0) {
    return EFI_NOT_FOUND;
  }
  index-=1;
  if (initfs->files_table->files_entries[index].file_offset+initfs->files_table->files_entries[index].file_size>initfs->header->files_area_size) {
    return EFI_BAD_BUFFER_SIZE;
  }
  bp_disk_FILE *file_out=AllocatePool(sizeof(bp_disk_FILE));
  if (!file_out) {
    return EFI_OUT_OF_RESOURCES;
  }
  file_out->file_size_bytes=initfs->files_table->files_entries[index].file_size;
  file_out->content=AllocateZeroPool(file_out->file_size_bytes);
  if (!file_out->content) {
    FreePool(file_out);
    return EFI_OUT_OF_RESOURCES;
  }
  CopyMem(file_out->content,initfs->files_area+initfs->files_table->files_entries[index].file_offset,file_out->file_size_bytes);
  *file=file_out;
  return EFI_SUCCESS;
}
