// SPDX-License-Identifier: MPL-2.0
#include <exception>
#include <sstream>
extern "C" {
#include "api.h"
#include "argon2.h"
#include "sha3.h"
}
#undef str
#include <sys/socket.h>
#include <sys/un.h>
#include <fstream>
#include <algorithm>
#include <filesystem>
#include <string>
#include <iostream>
#include <vector>
#include <string.h>
#include <pwd.h>
#include <sys/mman.h>
#include <sys/resource.h>
#include <sys/prctl.h>
#include <unistd.h>
using namespace std;
namespace fs=filesystem;
void secure_erase(void *address,size_t size) {
  explicit_bzero(address,size);
}
vector<uint8_t> get_urandom(size_t n) {
  vector<uint8_t> out(n);
  ifstream urandom("/dev/urandom",ios::binary);
  if (!urandom.is_open()) {
    throw runtime_error("open");
  }
  urandom.read(reinterpret_cast<char*>(out.data()),n);
  if (urandom.gcount()!=static_cast<streamsize>(n)) {
    throw runtime_error("error");
  }
  return out;
}
int main(int argc,char **argv) {
  ofstream core("/proc/self/coredump_filter");
  if (!core) {
    cout<<"[Keygen] Error: can't setup security."<<endl;
    return -1;
  }
  core<<"0";
  core.close();
  if (prctl(PR_SET_DUMPABLE,0,0,0,0)!=0) {
    cout<<"[Keygen] Error: can't setup security."<<endl;
    return -1;
  }
  struct rlimit rl_mem;
  rl_mem.rlim_cur=(1<<31);
  rl_mem.rlim_max=(1<<31);
  if (setrlimit(RLIMIT_MEMLOCK,&rl_mem)!=0) {
    cout<<"[Keygen] Error: can't setup security."<<endl;
    return -1;
  }
  if (mlockall(MCL_CURRENT | MCL_FUTURE)!=0) {
    cout<<"[Keygen] Error: can't setup security."<<endl;
    return -1;
  }
  // VFTM stuff
  unsigned char pk_root[CRYPTO_PUBLICKEYBYTES];
  unsigned char sk_root[CRYPTO_SECRETKEYBYTES];
  mlock(pk_root,sizeof(pk_root));
  mlock(sk_root,sizeof(sk_root));
  madvise(pk_root,sizeof(pk_root),MADV_DONTDUMP);
  madvise(sk_root,sizeof(sk_root),MADV_DONTDUMP);
  if (crypto_sign_keypair(pk_root,sk_root)) {
    cout<<"[Keygen] Error: can't generate root keypair."<<endl;
    return -1;
  }
  unsigned char pk_manifest[CRYPTO_PUBLICKEYBYTES];
  unsigned char sk_manifest[CRYPTO_SECRETKEYBYTES];
  mlock(pk_manifest,sizeof(pk_manifest));
  mlock(sk_manifest,sizeof(sk_manifest));
  madvise(pk_manifest,sizeof(pk_manifest),MADV_DONTDUMP);
  madvise(sk_manifest,sizeof(sk_manifest),MADV_DONTDUMP);
  if (crypto_sign_keypair(pk_manifest,sk_manifest)) {
    cout<<"[Keygen] Error: can't generate manifest keypair."<<endl;
    return -1;
  }
  unsigned char sig_root[CRYPTO_BYTES];
  mlock(sig_root,sizeof(sig_root));
  madvise(sig_root,sizeof(sig_root),MADV_DONTDUMP);
  size_t siglen;
  if (crypto_sign_signature(sig_root,&siglen,pk_manifest,sizeof(pk_manifest),sk_root)) {
    cout<<"[Keygen] Error: can't generate signature for manifest public key."<<endl;
    secure_erase(pk_root,sizeof(pk_root));
    secure_erase(sk_root,sizeof(sk_root));
    secure_erase(pk_manifest,sizeof(pk_manifest));
    secure_erase(sk_manifest,sizeof(sk_manifest));
    secure_erase(sig_root,siglen);
    return -1;
  }
  if (siglen!=CRYPTO_BYTES) {
    cout<<"[Keygen] Error: signature isn't the expected size for manifest public key."<<endl;
    secure_erase(pk_root,sizeof(pk_root));
    secure_erase(sk_root,sizeof(sk_root));
    secure_erase(pk_manifest,sizeof(pk_manifest));
    secure_erase(sk_manifest,sizeof(sk_manifest));
    secure_erase(sig_root,siglen);
    return -1;
  }
  secure_erase(sk_root,sizeof(sk_root));
  uint8_t payload[CRYPTO_PUBLICKEYBYTES+CRYPTO_SECRETKEYBYTES];
  mlock(payload,sizeof(payload));
  madvise(payload,sizeof(payload),MADV_DONTDUMP);
  memcpy(payload,pk_manifest,CRYPTO_PUBLICKEYBYTES);
  memcpy(payload+CRYPTO_PUBLICKEYBYTES,sk_manifest,CRYPTO_SECRETKEYBYTES);
  int fd=socket(AF_UNIX,SOCK_SEQPACKET,0);
  if (fd<0) {
    cout<<"[Keygen] Error: couldn't create client socket."<<endl;
    secure_erase(pk_root,sizeof(pk_root));
    secure_erase(pk_manifest,sizeof(pk_manifest));
    secure_erase(sk_manifest,sizeof(sk_manifest));
    secure_erase(sig_root,siglen);
    return -1;
  }
  sockaddr_un addr{};
  addr.sun_family=AF_UNIX;
  strcpy(addr.sun_path,"./vftm.sock");
  int r=connect(fd,(sockaddr*)&addr,sizeof(sockaddr_un));
  if (r<0) {
    cout<<"[Keygen] Error: couldn't connect to server socket."<<endl;
    secure_erase(pk_root,sizeof(pk_root));
    secure_erase(pk_manifest,sizeof(pk_manifest));
    secure_erase(sk_manifest,sizeof(sk_manifest));
    secure_erase(sig_root,siglen);
    return -1;
  }
  size_t s=send(fd,payload,sizeof(payload),0);
  if (s<0 || s!=sizeof(payload)) {
    perror("send");
    cout<<"[Keygen] Error: couldn't send keys to caller."<<endl;
    secure_erase(pk_root,sizeof(pk_root));
    secure_erase(pk_manifest,sizeof(pk_manifest));
    secure_erase(sk_manifest,sizeof(sk_manifest));
    secure_erase(sig_root,siglen);
    return -1;
  }
  secure_erase(pk_manifest,sizeof(pk_manifest));
  secure_erase(sk_manifest,sizeof(sk_manifest));
  secure_erase(payload,sizeof(payload));
  // SBFIE stuff
  bool set_files_to_user=true;
  const char* original_user=getenv("SUDO_USER");
  const char* original_uid=getenv("SUDO_UID");
  const char* original_gid=getenv("SUDO_GID");
  if (!(original_gid && original_uid && original_user)) {
    cout<<"[Keygen] Warning: it look like keygen is running without sudo or as root. It won't be able to make the file it produces accessible to original user."<<endl;
    set_files_to_user=false;
  }
  struct passwd* user_info=nullptr;
  if (set_files_to_user) {
    user_info=getpwnam(original_user);
    if (user_info==nullptr) {
      cout<<"[Keygen] Warning: couldn't found user named \""<<string(original_user)<<"\"."<<endl;
      set_files_to_user=false;
    }
  }
  const char* password_env=getenv("VYSTEM_KEYGEN_PASSWORD");
  string password="";
  if (password_env) {
    password=string(password_env);
    if (password=="") {
      cout<<"[Keygen] Warning: VYSTEM_KEYGEN_PASSWORD defined but empty, defaulting to manual password setup."<<endl;
      password_env=nullptr;
    }
  }
  vector<string> files;
  if (argc>1) {
    for (int i=1;i<argc;++i) {
      files.push_back(string(argv[i]));
    }
  } else {
    cout<<"[Keygen] Error: no file provided."<<endl;
    return -1;
  }
  for (int i=0;i<files.size();i++) {
    if (!fs::exists(files[i])) {
      cout<<"[Keygen] Error: file "<<files[i]<<" don't exist."<<endl;
      return -1;
    }
  }
  vector<unsigned char> fbuf;
  vector<unsigned char> sigsbuf(files.size()*CRYPTO_BYTES);
  vector<unsigned char> pksbuf(files.size()*CRYPTO_PUBLICKEYBYTES);
  if (mlock(sigsbuf.data(),sigsbuf.size())!=0) {
    cout<<"[Keygen] Error: can't setup security."<<endl;
    return -1;
  }
  if (mlock(pksbuf.data(),pksbuf.size())!=0) {
    cout<<"[Keygen] Error: can't setup security."<<endl;
    return -1;
  }
  unsigned char pk[CRYPTO_PUBLICKEYBYTES];
  unsigned char sk[CRYPTO_SECRETKEYBYTES];
  mlock(pk,sizeof(pk));
  mlock(sk,sizeof(sk));
  madvise(pk,sizeof(pk),MADV_DONTDUMP);
  madvise(sk,sizeof(sk),MADV_DONTDUMP);
  unsigned char sig[CRYPTO_BYTES];
  mlock(sig,sizeof(sig));
  madvise(sig,sizeof(sig),MADV_DONTDUMP);
  for (int i=0;i<files.size();++i) {
    ifstream file_content(files[i],ios::binary);
    if (!file_content) {
      cout<<"[Keygen] Error: can't read file: "<<files[i]<<endl;
      return -1;
    }
    file_content.seekg(0,ios::end);
    streamsize size=file_content.tellg();
    file_content.seekg(0,ios::beg);
    fbuf.resize(size);
    if (!file_content.read(reinterpret_cast<char*>(fbuf.data()),size)) {
      cout<<"[Keygen] Error: can't read file: "<<files[i]<<endl;
      return -1;
    }
    if (crypto_sign_keypair(pk,sk)) {
      cout<<"[Keygen] Error: can't generate keypair for file: "<<files[i]<<endl;
      return -1;
    }
    size_t siglen;
    if (size<1024*1024) {
      if (crypto_sign_signature(sig,&siglen,fbuf.data(),fbuf.size(),sk)) {
        cout<<"[Keygen] Error: can't generate signature for file: "<<files[i]<<endl;
        secure_erase(pk,sizeof(pk));
        secure_erase(sk,sizeof(sk));
        secure_erase(sig,siglen);
        return -1;
      }
    } else {
      vector<unsigned char> hash(64);
      sha3(fbuf.data(),size,hash.data(),hash.size());
      if (crypto_sign_signature(sig,&siglen,hash.data(),hash.size(),sk)) {
        cout<<"[Keygen] Error: can't generate signature for file: "<<files[i]<<endl;
        secure_erase(pk,sizeof(pk));
        secure_erase(sk,sizeof(sk));
        secure_erase(sig,siglen);
        return -1;
      }
    }
    if (siglen!=CRYPTO_BYTES) {
      cout<<"[Keygen] Error: signature isn't the expected size for file: "<<files[i]<<endl;
      secure_erase(pk,sizeof(pk));
      secure_erase(sk,sizeof(sk));
      secure_erase(sig,siglen);
      return -1;
    }
    memcpy(&pksbuf[i*CRYPTO_PUBLICKEYBYTES],pk,CRYPTO_PUBLICKEYBYTES);
    memcpy(&sigsbuf[i*CRYPTO_BYTES],sig,CRYPTO_BYTES);
    secure_erase(sk,sizeof(sk));
    secure_erase(pk,sizeof(pk));
    secure_erase(sig,siglen);
    if (!fbuf.empty()) {
      secure_erase(fbuf.data(),fbuf.size());
      fbuf.clear();
    }
  }
  if (!filesystem::exists("./sign")) {
    filesystem::create_directory("./sign");
  }
  for (int i=0;i<files.size();i++) {
    fs::path sigpath="./sign/"+fs::path(files[i]).filename().string();
    sigpath.replace_extension(".sig");
    ofstream sigfile(sigpath,ios::binary);
    if (!sigfile) {
      cout<<"[Keygen] Error: can't open signature file for file: "<<files[i]<<endl;
      return -1;
    }
    sigfile.write(reinterpret_cast<const char*>(sigsbuf.data()+i*CRYPTO_BYTES),CRYPTO_BYTES);
    sigfile.close();
    if (set_files_to_user && user_info) {
      if (chown(sigpath.c_str(),user_info->pw_uid,user_info->pw_gid)!=0) {
        cout<<"[Keygen] Warning: couldn't change file "<<sigpath<<" to owner "<<string(original_user)<<"."<<endl;
      }
    }
  }
  secure_erase(sigsbuf.data(),sigsbuf.size());
  if (!password_env) {
    cout<<"[Keygen] You are about to define a boot password. This password will be asked everytime you boot this instance of Vystem B."<<endl;
    cout<<"[Keygen] It will be used to sign every public key on this instance of Vystem B. Make sure it's secure. Here are the requirements:\n  - It should be at least 12 characters.\n  - Password can only include ASCII characters.\n  - Password shouldn't be longer than 512 characters.\n  - DO NOT use the numerical pad at all. It's not supported in the EFI environment."<<endl;
    cout<<"[Keygen] Leave empty for us to generate entropy as a password."<<endl;
    char* pwd=getpass("[Keygen] Enter boot password: ");
    if (strlen(pwd)==0) {
      cout<<"[Keygen] You asked for entropy. It will be provided as hexadecimal characters. Please enter the amount of hexadecimal characters you want."<<endl;
      cout<<"[Keygen] The amount of hex characters shouldn't be bigger than 512."<<endl;
      cout<<"[Keygen] Amount of hex characters (should be a multiple of 2 and at least 16 characters to be a minimum secure): ";
      string amount;
      getline(cin,amount);
      uint64_t num;
      try {
        num=stoi(amount);
      } catch (const exception e) {
        cout<<"[Keygen] Value entered isn't a number. Using default amount of 16 characters."<<endl;
        num=8;
      }
      if (num%2!=0) {
        cout<<"[Keygen] Number entered isn't a multiple of 2. Using default amount of 16 characters."<<endl;
        num=8;
      } else if (num<513) {
        num=num/2;
      } else {
        cout<<"[Keygen] Error: password is too long."<<endl;
        return -1;
      }
      vector<uint8_t> passhex;
      passhex.resize(num);
      try {
        passhex=get_urandom(num);
      } catch (const exception e) {
        if (string(e.what())=="open") {
          cout<<"[Keygen] Error: can't open secure source of randomness."<<endl;
          return -1;
        } else if (string(e.what())=="error") {
          cout<<"[Keygen] Error: can't read enought secure entropy."<<endl;
          return -1;
        }
      }
      ostringstream oss;
      oss<<hex<<setfill('0');
      for (uint8_t byte:passhex) {
        oss<<setw(2)<<std::uppercase<<static_cast<int>(byte);
      }
      password=oss.str();
      secure_erase(passhex.data(),passhex.size());
    } else if (strlen(pwd)<=512){
      password=string(pwd);
      char* confirm_pwd=getpass("[Keygen] Confirm password: ");
      string confirm_password(confirm_pwd);
      if (confirm_password!=password) {
        cout<<"[Keygen] Error: password don't correspond. Please try again."<<endl;
        return -1;
      }
      secure_erase(confirm_pwd,strlen(confirm_pwd));
    } else {
      cout<<"[Keygen] Error: password is too long."<<endl;
      return -1;
    }
    string pass=password;
    cout<<"[Keygen] Your password is set to be: "<<pass<<endl;
    cout<<endl;
    secure_erase(pwd,strlen(pwd));
    for (unsigned char c:password) {
      if (c>=0x80) {
        cout<<"[Keygen] Error: one character isn't ASCII."<<endl;
        return -1;
      }
    }
  } else {
    cout<<"[Keygen] Password provided from VYSTEM_KEYGEN_PASSWORD: "+password<<endl;
    for (unsigned char c:password) {
      if (c>=0x80) {
        cout<<"[Keygen] Error: one character isn't ASCII."<<endl;
        return -1;
      }
    }
  }
  vector<unsigned char> passutf16(password.size()*2);
  for (int i=0;i<password.size();++i) {
    passutf16[i*2]=password[i];
    passutf16[i*2+1]=0x00;
  }
  vector<uint8_t> salt(32);
  salt=get_urandom(32);
  vector<unsigned char> seed;
  seed.resize(CRYPTO_SEEDBYTES);
  if (argon2id_hash_raw(3,262144,1,passutf16.data(),passutf16.size(),salt.data(),salt.size(),seed.data(),CRYPTO_SEEDBYTES)!=0) {
    cout<<"[Keygen] Error: can't generate seed from password."<<endl;
    return -1;
  }
  if (crypto_sign_seed_keypair(pk,sk,seed.data())!=0) {
    cout<<"[Keygen] Error: can't generate keys from seed."<<endl;
    return -1;
  }
  secure_erase(pk,sizeof(pk));
  secure_erase(seed.data(),seed.size());
  if (crypto_sign_signature(sig,&siglen,pksbuf.data(),pksbuf.size(),sk)!=0) {
    cout<<"[Keygen] Error: can't generate signature from all publics keys."<<endl;
    return -1;
  }
  secure_erase(sk,sizeof(sk));
  ofstream keyfile("key.h");
  keyfile<<"#ifndef BP_LIB_KEY_H\n#define BP_LIB_KEY_H\n#include <Uefi.h>\nUINT8 bp_key_mainsig["<<to_string(CRYPTO_BYTES)<<"]={";
  for (int i=0;i<CRYPTO_BYTES;++i) {
    keyfile<<"0x"<<hex<<setw(2)<<setfill('0')<<(int)sig[i];
    if (i+1!=CRYPTO_BYTES) keyfile<<",";
  }
  keyfile<<"};\nUINT8 bp_key_pkblob["<<to_string(pksbuf.size())<<"]={";
  for (int i=0;i<pksbuf.size();++i) {
    keyfile<<"0x"<<hex<<setw(2)<<setfill('0')<<(int)pksbuf[i];
    if (i+1!=pksbuf.size()) keyfile<<",";
  }
  keyfile<<"};\nUINT8 bp_key_pwdsalt["<<to_string(salt.size())<<"]={";
  for (int i=0;i<salt.size();++i) {
    keyfile<<"0x"<<hex<<setw(2)<<setfill('0')<<(int)salt[i];
    if (i+1!=salt.size()) keyfile<<",";
  }
  keyfile<<"};\nUINT8 bp_key_pk_root["<<to_string(sizeof(pk_root))<<"]={";
  for (int i=0;i<sizeof(pk_root);++i) {
    keyfile<<"0x"<<hex<<setw(2)<<setfill('0')<<(int)pk_root[i];
    if (i+1!=sizeof(pk_root)) keyfile<<",";
  }
  keyfile<<"};\nUINT8 bp_key_sig_root["<<to_string(sizeof(sig_root))<<"]={";
  for (int i=0;i<sizeof(sig_root);++i) {
    keyfile<<"0x"<<hex<<setw(2)<<setfill('0')<<(int)sig_root[i];
    if (i+1!=sizeof(sig_root)) keyfile<<",";
  }
  keyfile<<"};\nSTATIC const CHAR16 * const bp_key_files["<<files.size()<<"]={";
  for (int i=0;i<files.size();++i) {
    string filename=fs::path(files[i]).filename().string();
    std::transform(filename.begin(),filename.end(),filename.begin(),
      [](unsigned char c) {
        return toupper(c);
      }
    );
    keyfile<<"L\""<<filename+"\"";
    if (i+1!=files.size()) keyfile<<",";
  }
  keyfile<<"};\nSTATIC const CHAR16 * const bp_key_files_sign["<<files.size()<<"]={";
  for (int i=0;i<files.size();++i) {
    string filename=fs::path(files[i]).filename().replace_extension(".sig").string();
    std::transform(filename.begin(),filename.end(),filename.begin(),
      [](unsigned char c) {
        return toupper(c);
      }
    );
    keyfile<<"L\""<<filename+"\"";
    if (i+1!=files.size()) keyfile<<",";
  }
  keyfile<<"};\n#endif";
  keyfile.close();
  if (set_files_to_user && user_info) {
    if (chown(fs::absolute("key.h").c_str(),user_info->pw_uid,user_info->pw_gid)!=0) {
      cout<<"[Keygen] Warning: couldn't change file \""<<fs::absolute("key.h")<<"\" to owner "<<string(original_user)<<"."<<endl;
    }
    if (chown(fs::absolute("sign").c_str(),user_info->pw_uid,user_info->pw_gid)!=0) {
      cout<<"[Keygen] Warning: couldn't change file \""<<fs::absolute("sign")<<"\" to owner "<<string(original_user)<<"."<<endl;
    }
  }
  return 0;
}