#include "../include/crypto.h"
#include "../sha3/sha3.h"
#include "../argon2/argon2.h"
#include <Uefi.h>
#include <Library/BaseMemoryLib.h>
#include <Library/MemoryAllocationLib.h>
#include <Library/UefiBootServicesTableLib.h>
EFI_STATUS bp_crypto_secure_erase(UINT8 *ptr,UINTN size) {
  if (ptr==NULL || size==0) {
    return EFI_INVALID_PARAMETER;
  }
  ZeroMem(ptr,size);
  return EFI_SUCCESS;
}
EFI_STATUS bp_crypto_sha3(UINT8 *input,UINTN size_input,UINTN size_output,bp_crypto_HASH_DIGEST **output) {
  EFI_STATUS status=EFI_SUCCESS;
  if (size_input<1 || input==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  if (size_output!=32 && size_output!=48 && size_output!=64) {
    return EFI_INVALID_PARAMETER;
  }
  *output=AllocateZeroPool(sizeof(bp_crypto_HASH_DIGEST));
  if (!*output) {
    return EFI_OUT_OF_RESOURCES;
  }
  (*output)->digest=AllocateZeroPool(64);
  if (!(*output)->digest) {
    FreePool(*output);
    *output=NULL;
    return EFI_OUT_OF_RESOURCES;
  }
  (*output)->size_bytes=size_output;
  sha3(input,size_input,(*output)->digest,size_output);
  return status;
}
EFI_STATUS bp_crypto_argon2id_raw(bp_crypto_ARGON2_CONTEXT *context,bp_crypto_HASH_DIGEST **output) {
  if (context==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  if (context->password.password==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  if (context->password.size==0) {
    return EFI_INVALID_PARAMETER;
  }
  if (context->memory_cost==0 || context->time_cost==0) {
    return EFI_INVALID_PARAMETER;
  }
  if (context->hash_size==0) {
    return EFI_INVALID_PARAMETER;
  }
  if (output==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  *output=NULL;
  *output=AllocateZeroPool(sizeof(bp_crypto_HASH_DIGEST));
  if (*output==NULL) {
    return EFI_OUT_OF_RESOURCES;
  }
  (*output)->digest=AllocateZeroPool(context->hash_size);
  if ((*output)->digest==NULL) {
    FreePool(*output);
    *output=NULL;
    return EFI_OUT_OF_RESOURCES;
  }
  INTN res=argon2id_hash_raw(context->time_cost,context->memory_cost,1,context->password.password,context->password.size*2,context->salt,context->salt_size,(*output)->digest,context->hash_size);
  if (res!=0) {
    FreePool((*output)->digest);
    FreePool(*output);
    *output=NULL;
    return EFI_INVALID_PARAMETER;
  }
  (*output)->size_bytes=context->hash_size;
  return EFI_SUCCESS;
}
EFI_STATUS bp_crypto_argon2d_raw(bp_crypto_ARGON2_CONTEXT *context,bp_crypto_HASH_DIGEST **output) {
  if (context==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  if (context->password.password==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  if (context->password.size==0) {
    return EFI_INVALID_PARAMETER;
  }
  if (context->memory_cost==0 || context->time_cost==0) {
    return EFI_INVALID_PARAMETER;
  }
  if (context->hash_size==0) {
    return EFI_INVALID_PARAMETER;
  }
  if (output==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  *output=NULL;
  *output=AllocateZeroPool(sizeof(bp_crypto_HASH_DIGEST));
  if (*output==NULL) {
    return EFI_OUT_OF_RESOURCES;
  }
  (*output)->digest=AllocateZeroPool(context->hash_size);
  if ((*output)->digest==NULL) {
    FreePool(*output);
    *output=NULL;
    return EFI_OUT_OF_RESOURCES;
  }
  INTN res=argon2d_hash_raw(context->time_cost,context->memory_cost,1,context->password.password,context->password.size*2,context->salt,context->salt_size,(*output)->digest,context->hash_size);
  if (res!=0) {
    FreePool((*output)->digest);
    FreePool(*output);
    *output=NULL;
    return EFI_INVALID_PARAMETER;
  }
  (*output)->size_bytes=context->hash_size;
  return EFI_SUCCESS;
}
EFI_STATUS bp_crypto_argon2i_raw(bp_crypto_ARGON2_CONTEXT *context,bp_crypto_HASH_DIGEST **output) {
  if (context==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  if (context->password.password==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  if (context->password.size==0) {
    return EFI_INVALID_PARAMETER;
  }
  if (context->memory_cost==0 || context->time_cost==0) {
    return EFI_INVALID_PARAMETER;
  }
  if (context->hash_size==0) {
    return EFI_INVALID_PARAMETER;
  }
  if (output==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  *output=NULL;
  *output=AllocateZeroPool(sizeof(bp_crypto_HASH_DIGEST));
  if (*output==NULL) {
    return EFI_OUT_OF_RESOURCES;
  }
  (*output)->digest=AllocateZeroPool(context->hash_size);
  if ((*output)->digest==NULL) {
    FreePool(*output);
    *output=NULL;
    return EFI_OUT_OF_RESOURCES;
  }
  INTN res=argon2i_hash_raw(context->time_cost,context->memory_cost,1,context->password.password,context->password.size*2,context->salt,context->salt_size,(*output)->digest,context->hash_size);
  if (res!=0) {
    FreePool((*output)->digest);
    FreePool(*output);
    *output=NULL;
    return EFI_INVALID_PARAMETER;
  }
  (*output)->size_bytes=context->hash_size;
  return EFI_SUCCESS;
}
EFI_STATUS bp_crypto_sphincsplus_keypair_seed(UINT8 *seed,UINTN seed_size,BOOLEAN erase_sk,BOOLEAN erase_pk,bp_crypto_SPHINCSPLUS_KEYPAIR **output) {
  if (seed==NULL || seed_size!=CRYPTO_SEEDBYTES) {
    return EFI_INVALID_PARAMETER;
  }
  if (output==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  *output=NULL;
  *output=AllocatePool(sizeof(bp_crypto_SPHINCSPLUS_KEYPAIR));
  if (*output==NULL) {
    return EFI_OUT_OF_RESOURCES;
  }
  (*output)->erase_pk=erase_pk;
  (*output)->erase_sk=erase_sk;
  INTN res=crypto_sign_seed_keypair((*output)->pk,(*output)->sk,seed);
  if (res!=0) {
    FreePool(*output);
    *output=NULL;
    return EFI_INVALID_PARAMETER;
  }
  if (erase_pk) {
    bp_crypto_secure_erase((*output)->pk,CRYPTO_PUBLICKEYBYTES);
  }
  if (erase_sk) {
    bp_crypto_secure_erase((*output)->sk,CRYPTO_SECRETKEYBYTES);
  }
  return EFI_SUCCESS;
}
EFI_STATUS bp_crypto_sphincsplus_verify(UINT8 *signature,UINTN signature_size,UINT8 *message,UINTN message_size,UINT8 *pk) {
  if (signature==NULL || signature_size!=CRYPTO_BYTES) {
    return EFI_INVALID_PARAMETER;
  }
  if (message==NULL || message_size==0) {
    return EFI_INVALID_PARAMETER;
  }
  if (pk==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  INTN res=crypto_sign_verify(signature,signature_size,message,message_size,pk);
  if (res!=0) {
    return EFI_INVALID_PARAMETER;
  }
  return EFI_SUCCESS;
}
EFI_STATUS bp_crypto_sphincsplus_verify_file(EFI_FILE_PROTOCOL *root,CHAR16 *file,CHAR16 *sigfile,UINT8 *pk) {
  EFI_STATUS status;
  if (root==NULL || file==NULL) {
    return EFI_INVALID_PARAMETER;
  }
  EFI_FILE_PROTOCOL *opened_file=NULL;
  EFI_FILE_PROTOCOL *opened_sig_file=NULL;
  bp_disk_FILE *file_content=NULL;
  bp_disk_FILE *sig_file_content=NULL;
  if (EFI_ERROR(bp_disk_check_file_existence(root,file))) {
    return EFI_NOT_FOUND;
  }
  if (EFI_ERROR(bp_disk_check_file_existence(root,sigfile))) {
    return EFI_NOT_FOUND;
  }
  status=bp_disk_open_file(root,file,&opened_file);
  if (EFI_ERROR(status)) goto cleanup;
  EFI_FILE_INFO *file_info;
  status=bp_disk_get_file_info(opened_file,&file_info);
  if (EFI_ERROR(status)) goto cleanup;
  status=bp_disk_open_file(root,sigfile,&opened_sig_file);
  if (EFI_ERROR(status)) goto cleanup;
  status=bp_disk_read_file(opened_file,&file_content);
  if (EFI_ERROR(status)) goto cleanup;
  status=bp_disk_read_file(opened_sig_file,&sig_file_content);
  if (EFI_ERROR(status)) goto cleanup;
  if (sig_file_content->file_size_bytes!=CRYPTO_BYTES) {
    status=EFI_SECURITY_VIOLATION;
    goto cleanup;
  }
  bp_crypto_HASH_DIGEST *hash=NULL;
  if (file_info->FileSize>1024*1024) {
    status=bp_crypto_sha3(file_content->content,file_content->file_size_bytes,64,&hash);
    if (EFI_ERROR(status)) goto cleanup;
    INTN res=crypto_sign_verify(sig_file_content->content,sig_file_content->file_size_bytes,hash->digest,hash->size_bytes,pk);
    if (res!=0) {
      status=EFI_SECURITY_VIOLATION;
      goto cleanup;
    }
    status=EFI_SUCCESS;
  } else {
    INTN res=crypto_sign_verify(sig_file_content->content,sig_file_content->file_size_bytes,file_content->content,file_content->file_size_bytes,pk);
    if (res!=0) {
      status=EFI_SECURITY_VIOLATION;
      goto cleanup;
    }
    status=EFI_SUCCESS;
  }
cleanup:
  if (opened_file!=NULL) {
    opened_file->Close(opened_file);
  }
  if (opened_sig_file!=NULL) {
    opened_sig_file->Close(opened_sig_file);
  }
  if (file_content->content) {
    bp_crypto_secure_erase(file_content->content,file_content->file_size_bytes);
    FreePool(file_content->content);
  }
  if (sig_file_content->content) {
    bp_crypto_secure_erase(sig_file_content->content,sig_file_content->file_size_bytes);
    FreePool(sig_file_content->content);
  }
  if (file_info) {
    FreePool(file_info);
  }
  if (hash) {
    bp_crypto_secure_erase(hash->digest,hash->size_bytes);
    FreePool(hash->digest);
    FreePool(hash);
  }
  return status;
}
