#include "../inc/integretychecker.h"

#include "../cryptopp/base64.h"
#include "../cryptopp/files.h"
#include "../cryptopp/osrng.h"
#include "../cryptopp/rsa.h"

#include <filesystem>
#include <fstream>
#include <iostream>
namespace fs = std::filesystem;

const int IntegretyCheck::KEY_SIZE = 2048;
const std::string IntegretyCheck::HASH_FILE = "hashlist.sha256";
const std::string IntegretyCheck::KEY_FILE = "appKey.pub";
const std::string IntegretyCheck::SIGNATURE_MARKER = "### signature ###\n";
const std::string IntegretyCheck::HASH_FILE_DIVIDER = "  *";

std::string string_to_hex(const std::string &input) {
  static const char *const lut = "0123456789ABCDEF";
  size_t len = input.length();

  std::string output;
  output.reserve(2 * len);
  for (size_t i = 0; i < len; ++i) {
    const unsigned char c = input[i];
    output.push_back(lut[c >> 4]);
    output.push_back(lut[c & 15]);
  }
  return output;
}

IntegretyCheck::IntegretyCheck() {}

IntegretyCheck::IntegretyCheck(const std::string &appPath, bool genKeyPair) {
  if (genKeyPair) {
    generateKeyPair();
  } else {
    loadKeyFile(appPath);
  }
}

bool IntegretyCheck::loadKeyFile(const std::string &app) {
  fs::path appPath(app);
  if (fs::exists(appPath / KEY_FILE)) {
    CryptoPP::FileSource input((appPath / KEY_FILE).c_str(), true);
    m_publicKey.BERDecode(input);
    return true;
  }
  return false;
}

bool IntegretyCheck::saveKeyFile(const std::string &app) {
  ///@todo https://github.com/noloader/cryptopp-pem
  fs::path appPath(app);
  std::cout << (appPath / KEY_FILE).string() << std::endl;
  CryptoPP::FileSink output((appPath / KEY_FILE).c_str());
  m_publicKey.DEREncode(output);
  return true;
}

bool IntegretyCheck::generateKeyPair() {
  CryptoPP::AutoSeededRandomPool rng;

  CryptoPP::InvertibleRSAFunction params;
  params.GenerateRandomWithKeySize(rng, KEY_SIZE);

  m_privateKey = CryptoPP::RSA::PrivateKey(params);
  m_publicKey = CryptoPP::RSA::PublicKey(params);
  return true;
}

std::string IntegretyCheck::generateHashList(const std::string &app) {
  std::stringstream hashlist;
  fs::path appPath(app);
  std::string p = appPath.string();

  for (auto f : fs::recursive_directory_iterator(appPath)) {
    if (f.is_directory()) {
      // skip directory, iterator is recursive
    } else {
      // strip appPath
      std::string filepath = f.path().string();
      filepath.erase(0, p.size() + 1);
      // write hash line to file
      hashlist << generateFileHash(f.path().string()) << HASH_FILE_DIVIDER
               << filepath << std::endl;
    }
  }
  return hashlist.str();
}

std::string IntegretyCheck::generateFileHash(const std::string &filepath) {
  std::ifstream in(filepath);
  std::string contents((std::istreambuf_iterator<char>(in)),
                       std::istreambuf_iterator<char>());
  CryptoPP::byte const *pbData = (CryptoPP::byte *)contents.c_str();
  unsigned long nDataLen = contents.size();
  CryptoPP::byte abDigest[CryptoPP::SHA256::DIGESTSIZE];

  CryptoPP::SHA256().CalculateDigest(abDigest, pbData, nDataLen);

  return string_to_hex(
      std::string((char *)abDigest, CryptoPP::SHA256::DIGESTSIZE));
}

bool IntegretyCheck::extractHashAndFile(const std::string &line,
                                        std::string &hash,
                                        std::string &filename) {
  size_t divider = line.find(HASH_FILE_DIVIDER);
  bool result = false;
  if (divider != std::string::npos) {
    hash = line.substr(0, divider);
    filename = line.substr(divider + HASH_FILE_DIVIDER.size());
    result = true;
  } else {
    std::cerr << "ERROR: malformated hash line '" << line << "'" << std::endl;
  }
  return result;
}

bool IntegretyCheck::checkHashList(const std::string &hashList,
                                   const std::string &appPath) {
  std::stringstream ss(hashList);
  std::string line, hash, filename;
  bool result = true;
  fs::path p(appPath);
  while (std::getline(ss, line, '\n')) {
    if (extractHashAndFile(line, hash, filename)) {
      bool check = verifyFileHash((p / filename).string(), hash);
      result &= check;
      if (!check) {
        std::cerr << "ERROR: Checksum does not match on '" << filename << "'!"
                  << std::endl;
      }
    }
  }
  return result;
}

bool IntegretyCheck::verifyFileHash(const std::string &filePath,
                                    const std::string &hash) {
  return (generateFileHash(filePath) == hash);
}

void IntegretyCheck::signHashList(std::string &hashList) {
  CryptoPP::AutoSeededRandomPool rng;

  CryptoPP::RSASSA_PKCS1v15_SHA_Signer signer(m_privateKey);
  size_t signature_length = signer.MaxSignatureLength();
  CryptoPP::SecByteBlock signature(signature_length);
  signature_length =
      signer.SignMessage(rng, (const CryptoPP::byte *)hashList.c_str(),
                         hashList.length(), signature);
  std::cerr << signature_length << std::endl;
  signature.resize(signature_length);

  // transfrom binary signature to base64 encoded string
  std::string strHexSignature;
  CryptoPP::Base64Encoder base64enc(new CryptoPP::StringSink(strHexSignature));
  base64enc.Put(signature, signature_length);
  base64enc.MessageEnd();

  // add divider and base64 encoded signature to hash list
  hashList.append(SIGNATURE_MARKER);
  hashList.append(strHexSignature);
}

bool IntegretyCheck::verifyHashList(const std::string &signedHashList,
                                    std::string &hashList) {
  size_t startSig = signedHashList.find(SIGNATURE_MARKER);
  if (startSig == std::string::npos) {
    std::cerr << "ERROR: no signature marker found!" << std::endl;
    return false;
  }

  // extract hash list
  hashList = signedHashList.substr(0, startSig);

  // extract base64 encoded signature
  std::string base64Signature =
      signedHashList.substr(startSig + SIGNATURE_MARKER.length());

  // decode base64 encoded signature
  std::string decodedSignature;
  CryptoPP::StringSource decoder(
      base64Signature, true,
      new CryptoPP::Base64Decoder(new CryptoPP::StringSink(decodedSignature)));

  // transform std::string into SecByteBlock
  CryptoPP::SecByteBlock signature(
      (const CryptoPP::byte *)decodedSignature.c_str(),
      decodedSignature.length());

  CryptoPP::RSASSA_PKCS1v15_SHA_Verifier verifier(m_publicKey);
  bool result =
      verifier.VerifyMessage((const CryptoPP::byte *)hashList.c_str(),
                             hashList.length(), signature, signature.size());

  return result;
}

std::string IntegretyCheck::loadHashList(const std::string &app) {
  std::stringstream content;
  fs::path appPath(app);
  std::ifstream in(appPath / HASH_FILE, std::ios::in | std::ios::binary);
  if (in) {
    content << in.rdbuf();
    in.close();
  }
  return content.str();
}

void IntegretyCheck::saveHashList(const std::string &app,
                                  const std::string &hashList) {
  fs::path appPath(app);
  std::ofstream out(appPath / HASH_FILE);
  out << hashList;
  out.close();
}