// SPDX-License-Identifier: MPL-2.0
extern "C" {
#include "api.h"
#include "sha3.h"
}
#undef str
#include <cstdint>
#include <iostream>
#include <filesystem>
#include <fstream>
#include <string>
#include <vector>
#include <cstring>
#include <algorithm>
using namespace std;
namespace fs=filesystem;
#pragma pack(push,1)
struct initfs_header {
  uint8_t sign[8]={'I','n','i','t','F','i','S','y'};
  uint16_t bootloader_version=0x0001;
  uint16_t initfs_version=0x0001;
  uint32_t os_version=0x00000001;
  uint8_t installation_id[48]={0};
  uint64_t initfs_size=0;
  uint64_t table_size=0;
  uint64_t files_area_size=0;
  uint64_t entries_width=256;
  uint64_t entries_count=0;
  uint64_t files_area_offset=0;
  uint64_t entropy_check1=0;
  uint64_t check1=0;
  uint8_t entry_table_hash[64]={0};
  uint8_t files_area_hash[64]={0};
  uint8_t installation_id_hash_hash[64]={0};
  uint8_t padding[128]={0};
  uint8_t header_hash[64]={0};
};
#pragma pack(pop)
#pragma pack(push,1)
struct file_entry {
  uint64_t file_offset=0;
  uint64_t file_size=0;
  uint8_t pk[64]={0};
  uint8_t hash[64]={0};
  char file_name[112]={0};
};
#pragma pack(pop)
#pragma pack(push,1)
struct signsyst_header {
  uint8_t sign[8]={'S','i','g','n','S','y','s','t'};
  uint16_t bootloader_version=0x0001;
  uint16_t initfs_version=0x0001;
  uint32_t os_version=0x00000001;
  uint8_t installation_id[48]={0};
  uint64_t signature_size=0;
  uint64_t signature_count=0;
  uint64_t signsyst_size=0;
  uint64_t signature_block_size=0;
  uint8_t signature_block_hash[64];
  uint8_t padding[288]={0};
  uint8_t header_hash[64]={0};
};
#pragma pack(pop)
void secure_erase(void *address,size_t size) {
  explicit_bzero(address,size);
}
vector<unsigned char> generate_padding(size_t x,unsigned char p) {
  vector<unsigned char> result(x,0);
  size_t f0=0;
  size_t f1=1;
  while (f0<x) {
    result[f0]=p;
    size_t next=f0+f1;
    f0=f1;
    f1=next;
  }
  return result;
}
int main(int argc,char **argv) {
  if (argc!=2) {
    cout<<"[InitFSGen] Error: invalid number of argument."<<endl;
    return -1;
  }
  string folder_path=string(argv[1]);
  if (!fs::exists(folder_path)) {
    cout<<"[InitFSGen] Error: provided path doesn't exists."<<endl;
    return -1;
  }
  if (!fs::is_directory(folder_path)) {
    cout<<"[InitFSGen] Error: provided path isn't a directory"<<endl;
    return -1;
  }
  vector<string> initfs_files_name;
  vector<size_t> initfs_files_size;
  size_t total_size=0;
  for (auto &entry:fs::directory_iterator(folder_path)) {
    if (!entry.is_regular_file()) {
      cout<<"[InitFSGen] Error: InitFS only support files. Following entry isn't a regular file : "<<entry.path()<<endl;
      return -1;
    }
    initfs_files_name.push_back(string(fs::absolute(entry.path())));
    initfs_files_size.push_back(entry.file_size());
    total_size+=entry.file_size();
  }
  vector<unsigned char> entropy_id;
  entropy_id.resize(48);
  ifstream random("/dev/urandom",ios::binary);
  if (!random) {
    cout<<"[InitFSGen] Error: Can't open secure entropy source."<<endl;
    return -1;
  }
  random.read(reinterpret_cast<char*>(entropy_id.data()),entropy_id.size());
  if (random.gcount()!=entropy_id.size()) {
    cout<<"[InitFSGen] Error: Can't read enougth entropy for installation id."<<endl;
    return -1;
  }
  vector<unsigned char> entropy;
  entropy.resize(8);
  random.read(reinterpret_cast<char*>(entropy.data()),entropy.size());
  if (random.gcount()!=entropy.size()) {
    cout<<"[InitFSGen] Error: Can't read enougth entropy."<<endl;
    return -1;
  }
  initfs_header header{};
  signsyst_header sign_header={};
  memcpy(header.installation_id,entropy_id.data(),entropy_id.size());
  memcpy(sign_header.installation_id,entropy_id.data(),entropy_id.size());
  sign_header.signature_size=CRYPTO_BYTES;
  sign_header.signature_count=initfs_files_name.size();
  sign_header.signature_block_size=sign_header.signature_count*sign_header.signature_size;
  sign_header.signsyst_size=512+sign_header.signature_block_size;
  memcpy(&header.entropy_check1,entropy.data(),entropy.size());
  header.entries_count=initfs_files_name.size();
  header.table_size=header.entries_count*header.entries_width;
  header.files_area_size=total_size;
  header.initfs_size=512+header.table_size+header.files_area_size;
  header.files_area_offset=512+header.table_size;
  uint64_t value=(header.initfs_size+header.table_size+header.files_area_size+header.entries_width+header.entries_count+header.files_area_offset)%UINT64_MAX;
  header.check1=(value*0x9E3779B185EBCA87)^header.entropy_check1;
  sha3(header.installation_id,sizeof(header.installation_id),header.installation_id_hash_hash,sizeof(header.installation_id_hash_hash));
  ofstream initfs_footprint("initfs-footprint.bin",ios::binary);
  if (!initfs_footprint) {
    cout<<"[InitFSGen] Error: Can't open initfs-footprint.bin."<<endl;
    return -1;
  }
  initfs_footprint.write(reinterpret_cast<char*>(header.installation_id_hash_hash),sizeof(header.installation_id_hash_hash));
  sha3(header.installation_id_hash_hash,sizeof(header.installation_id_hash_hash),header.installation_id_hash_hash,sizeof(header.installation_id_hash_hash));
  vector<unsigned char> files_area;
  vector<file_entry> files_entries_table;
  vector<unsigned char> sig_list;
  sig_list.resize(CRYPTO_BYTES*header.entries_count);
  size_t cursor=0;
  files_area.resize(header.files_area_size);
  uint8_t sk[CRYPTO_SECRETKEYBYTES]={0};
  for (size_t i=0;i<initfs_files_name.size();++i) {
    ifstream file(initfs_files_name[i],ios::binary);
    if (!file) {
      cout<<"[InitFSGen] Error: Can't read file: "<<initfs_files_name[i]<<endl;
      return -1;
    }
    file.read(reinterpret_cast<char*>(files_area.data()+cursor),initfs_files_size[i]);
    if (file.gcount()!=initfs_files_size[i]) {
      cout<<"[InitFSGen] Error: Couldn't read full file: "<<initfs_files_name[i]<<endl;
      return -1;
    }
    file_entry entry;
    string filename=string(fs::path(initfs_files_name[i]).filename());
    copy(filename.data(),filename.data()+min(filename.size(),static_cast<size_t>(112)),entry.file_name);
    entry.file_size=initfs_files_size[i];
    entry.file_offset=cursor;
    sha3(files_area.data()+cursor,initfs_files_size[i],entry.hash,64);
    if (crypto_sign_keypair(entry.pk,sk)) {
      cout<<"[InitFSGen] Error: can't generate keypair for file: "<<initfs_files_name[i]<<endl;
      secure_erase(sk,sizeof(sk));
      return -1;
    }
    size_t siglen=0;
    if (crypto_sign_signature(sig_list.data()+i*CRYPTO_BYTES,&siglen,files_area.data()+cursor,initfs_files_size[i],sk)) {
      cout<<"[InitFSGen] Error: can't generate signature for file: "<<initfs_files_name[i]<<endl;
      secure_erase(sk,sizeof(sk));
      return -1;
    }
    secure_erase(sk,sizeof(sk));
    if (siglen!=CRYPTO_BYTES) {
      cout<<"[InitFSGen] Error: Signature isn't the expected size for file: "<<initfs_files_name[i]<<endl;
      return -1;
    }
    files_entries_table.push_back(entry);
    cursor+=initfs_files_size[i];
  }
  vector<unsigned char> entries_table;
  entries_table.resize(header.table_size);
  for (size_t i=0;i<files_entries_table.size();++i) {
    memcpy(entries_table.data()+i*header.entries_width,&files_entries_table[i],sizeof(file_entry));
  }
  sha3(entries_table.data(),entries_table.size(),header.entry_table_hash,sizeof(header.entry_table_hash));
  for (size_t i=0;i<48;++i) {
    header.entry_table_hash[i+16]^=header.installation_id[i];
  }
  sha3(files_area.data(),files_area.size(),header.files_area_hash,sizeof(header.files_area_hash));
  for (size_t i=0;i<48;++i) {
    header.files_area_hash[i+16]^=header.installation_id[i];
  }
  sha3(sig_list.data(),sig_list.size(),sign_header.signature_block_hash,sizeof(sign_header.signature_block_hash));
  for (size_t i=0;i<48;++i) {
    sign_header.signature_block_hash[i+16]^=sign_header.installation_id[i];
  }
  auto padding=generate_padding(sizeof(header.padding),header.installation_id[0]);
  memcpy(header.padding,padding.data(),padding.size());
  auto padding_sign=generate_padding(sizeof(sign_header.padding),sign_header.installation_id[sizeof(sign_header.installation_id)-1]);
  memcpy(sign_header.padding,padding_sign.data(),padding_sign.size());
  sha3(&header,sizeof(header)-sizeof(header.header_hash),header.header_hash,sizeof(header.header_hash));
  sha3(&sign_header,sizeof(sign_header)-sizeof(sign_header.header_hash),sign_header.header_hash,sizeof(sign_header.header_hash));
  ofstream initfs_bin("initfs.bin",ios::binary);
  if (!initfs_bin) {
    cout<<"[InitFSGen] Error: Can't open initfs.bin."<<endl;
    return -1;
  }
  initfs_bin.write(reinterpret_cast<char*>(&header),sizeof(header));
  initfs_bin.write(reinterpret_cast<char*>(entries_table.data()),entries_table.size());
  initfs_bin.write(reinterpret_cast<char*>(files_area.data()),files_area.size());
  initfs_bin.close();
  ofstream signsyst_hash("signsyst-hash.bin",ios::binary);
  if (!signsyst_hash) {
    cout<<"[InitFSGen] Error: Can't open signsyst-hash.bin."<<endl;
    return -1;
  }
  vector<unsigned char> signsysthash(64,0);
  sha3(&sign_header,sizeof(sign_header),signsysthash.data(),signsysthash.size());
  signsyst_hash.write(reinterpret_cast<char*>(signsysthash.data()),signsysthash.size());
  signsyst_hash.close();
  ofstream signsyst_bin("signsyst.bin",ios::binary);
  if (!signsyst_bin) {
    cout<<"[InitFSGen] Error: Can't open signsyst.bin."<<endl;
    return -1;
  }
  signsyst_bin.write(reinterpret_cast<char*>(&sign_header),sizeof(sign_header));
  signsyst_bin.write(reinterpret_cast<char*>(sig_list.data()),sig_list.size());
  signsyst_bin.close();
  return 0;
}