// Copyright 2026 Osmium Sorcerer
// SPDX-License-Identifier: MIT

#include <QByteArray>
#include <QCborMap>
#include <QCborStreamReader>
#include <QCborValue>
#include <QDataStream>
#include <QDir>
#include <QIODevice>
#include <QLockFile>
#include <QSaveFile>

#include <sodium.h>

#include "file_functions.h"
#include "keyring.h"

// Encrypted secret key serialized format version.
//
// Version 1: Secret key encrypted with ChaCha20-Poly1305 using password-based
// Argon2-derived wrapping key.
//
// Serialization: Argon2 parameters for key derivation, encrypted key, Poly1305
// MAC over the encrypted payload with the whole header as additional
// authenticated data.
//
// Fields: version (u8), salt[16], opslimit (u32be), memlimit (u32be),
// alg (u8), ciphertext[32], mac[16].
static constexpr quint8 key_format_version = 1;

// Magic that is prepended to a byte sequence to identify it as a 32-byte
// X25519 public key. Don't expect new tags to ever be introduced unless you
// want certified post-quantum-secure lawyer simulator authentication with
// ML-KEM. Simplicity over flexibility.
static constexpr quint8 cert_x25519_tag = 0x26;

// All-zero nonce for use in key wrapping. Nonces are provided to ensure
// uniqueness of multiple keystreams under the same key. In our case, we only
// ever encrypt one message with one derived key, the uniqueness of which is
// provided by the salt that is randomly generated for each new key wrapping.
static constexpr quint8 all_zero_nonce[crypto_aead_chacha20poly1305_NPUBBYTES] = {};

// Domain separator string to isolate generic BLAKE2b specifically to this
// authentication mechanism.
static const auto domain_sep = QByteArrayLiteral("Einsof-Auth-DHCR");

static QString get_keyring_path()
{
  return QDir(get_app_path()).filePath("keyring.cbor");
}

int keyring_initialize()
{
  if (sodium_init() < 0)
  {
    return 1;
  }

  QFile index(get_keyring_path());
  if (!index.exists())
  {
    if (!index.open(QIODevice::NewOnly, QFileDevice::ReadOwner | QFileDevice::WriteOwner))
    {
      return 3;
    }
    index.close();
  }

  return 0;
}

int generate_key(QStringView name, const QByteArray &password)
{
  // If we let users specify arbitrary parameters, additional constraints must
  // be enforced: min/max opslimit, type of memlimit (uint32_t vs size_t).
  // This flexibility is counterproductive. Instead, a simple high/moderate
  // toggle should be used if Argon2 parameter tuning is required.
  if (password.size() < crypto_pwhash_PASSWD_MIN || password.size() > crypto_pwhash_PASSWD_MAX)
  {
    return 1;
  }

  quint8 public_key[crypto_kx_PUBLICKEYBYTES];
  quint8 secret_key[crypto_kx_SECRETKEYBYTES];
  sodium_mlock(secret_key, sizeof(secret_key));
  if (crypto_kx_keypair(public_key, secret_key))
  {
    sodium_munlock(secret_key, sizeof(secret_key));
    return 2;
  }

  // Key wrapping. The crypto_pwhash derives the key from password with the
  // specified parameters (3 iterations over 1 gigabyte of memory). This
  // operation is slow and should ideally be handled asynchronously without
  // blocking the GUI thread.
  quint8 wrap_key[crypto_aead_chacha20poly1305_KEYBYTES];
  sodium_mlock(wrap_key, sizeof(wrap_key));
  quint8 salt[crypto_pwhash_SALTBYTES];
  randombytes_buf(salt, sizeof(salt));
  quint32 pwhash_opslimit = crypto_pwhash_OPSLIMIT_MODERATE;
  quint32 pwhash_memlimit = crypto_pwhash_MEMLIMIT_SENSITIVE;
  quint8 pwhash_alg = crypto_pwhash_ALG_DEFAULT;
  if (crypto_pwhash(wrap_key, sizeof(wrap_key), password.constData(), password.size(), salt, pwhash_opslimit, pwhash_memlimit, pwhash_alg))
  {
    sodium_munlock(wrap_key, sizeof(wrap_key));
    sodium_munlock(secret_key, sizeof(secret_key));
    return 3;
  }

  // Payload: the wrapped secret key. Authenticated together with its header
  // (version tag and Argon2 parameters).
  quint8 payload[sizeof(secret_key) + crypto_aead_chacha20poly1305_ABYTES];
  QByteArray packed_header;
  packed_header.reserve(sizeof(quint8) + sizeof(salt) + sizeof(quint32) + sizeof(quint32) + sizeof(quint8));
  QDataStream hdr_out(&packed_header, QIODevice::WriteOnly);
  hdr_out << key_format_version;
  hdr_out.writeRawData((const char *)salt, sizeof(salt));
  hdr_out << pwhash_opslimit;
  hdr_out << pwhash_memlimit;
  hdr_out << pwhash_alg;
  if (crypto_aead_chacha20poly1305_encrypt(payload, nullptr, secret_key, sizeof(secret_key), (const quint8 *)packed_header.constData(), packed_header.size(), nullptr, all_zero_nonce, wrap_key))
  {
    sodium_munlock(wrap_key, sizeof(wrap_key));
    sodium_munlock(secret_key, sizeof(secret_key));
    return 4;
  }
  sodium_munlock(wrap_key, sizeof(wrap_key));
  sodium_munlock(secret_key, sizeof(secret_key));
  QByteArray public_key_array((const char *)public_key, (qsizetype)sizeof(public_key));
  quint8 tag = cert_x25519_tag;

  // Key fingerprint is used as its unique identifier.
  QByteArray fingerprint;
  fingerprint.append(tag);
  fingerprint.append(public_key_array);
  uchar fpr_hash[crypto_generichash_BYTES];
  crypto_generichash(fpr_hash, sizeof(fpr_hash), (const uchar *)fingerprint.constData(), fingerprint.size(), nullptr, 0);

  QByteArray encrypted_secret(packed_header);
  encrypted_secret.append((const char *)payload, sizeof(payload));

  // Prepare the key CBOR entry and atomically insert it into the keyring.
  // To ensure durability, file locking and atomic overwrite are used as
  // provided by Qt. Frankly, I don't see the keyring being prone to corruption
  // as there's no concurrent access, save for a contrived edge case of a user
  // generating/deleting keys simultaneously from two client instances. It's
  // done regardless for a good measure.
  QCborMap entry;
  entry.insert(0, QCborValue(tag));
  entry.insert(1, QCborValue(name));
  entry.insert(2, QCborValue(public_key_array));
  entry.insert(3, QCborValue(encrypted_secret));
  QFile index_file(get_keyring_path());
  QLockFile index_lock(QDir(get_app_path()).filePath("keyring.lock"));
  if (!index_lock.lock())
  {
    return 9;
  }
  if (!index_file.open(QIODevice::ReadOnly))
  {
    return 6;
  }
  QCborStreamReader idx_in(&index_file);
  QCborValue index = QCborValue::fromCbor(idx_in);
  index_file.close();
  if (!index.isMap())
  {
    index = QCborMap();
  }
  QCborMap index_map = index.toMap();
  index_map.insert(QCborValue(QByteArray((const char *)fpr_hash, (qsizetype)sizeof(fpr_hash))), QCborValue(entry));
  QSaveFile index_file_w(get_keyring_path());
  if (!index_file_w.open(QIODevice::WriteOnly))
  {
    return 8;
  }
  index_file_w.write(QCborValue(index_map).toCbor());
  index_file_w.commit();
  index_lock.unlock();

  return 0;
}

// Look up the key_id in the keyring, unlock it with supplied password using
// parameters from its header, perform a Diffie-Hellman key exchange with the
// server's ephemeral key to derive a shared secret, and hash it along with
// additional data to prove your authenticity to the server.
ResponseResult unlock_and_auth(QByteArrayView key_id, QByteArrayView password, QByteArrayView ephemeral_key, QByteArrayView username, AuthResponse &out)
{
  if (password.size() < crypto_pwhash_PASSWD_MIN || password.size() > crypto_pwhash_PASSWD_MAX)
  {
    return ResponseResult::invalid_password;
  }

  QFile index_file(get_keyring_path());
  if (!index_file.open(QIODevice::ReadOnly))
  {
    return ResponseResult::inaccessible_keyring;
  }
  QCborStreamReader idx_in(&index_file);
  QCborValue index = QCborValue::fromCbor(idx_in);
  index_file.close();
  QCborMap tmp_map = index.toMap().value(QCborValue(QByteArray(key_id))).toMap();
  QCborValue val = tmp_map.value(3);
  if (!val.isByteArray())
  {
    return ResponseResult::corrupted_entry;
  }
  QByteArray encrypted_secret = val.toByteArray();

  val = tmp_map.value(2);
  if (!val.isByteArray())
  {
    return ResponseResult::corrupted_entry;
  }
  QByteArray public_key = val.toByteArray();

  // After retrieving the key from the keyring, the process is the same as
  // in key generation, but in the opposite direction.
  quint8 secret_key[crypto_kx_SECRETKEYBYTES];
  sodium_mlock(secret_key, sizeof(secret_key));

  quint8 sk_version_tag;
  quint8 salt[crypto_pwhash_SALTBYTES];
  quint32 pwhash_opslimit;
  quint32 pwhash_memlimit;
  quint8 pwhash_alg;
  quint8 payload[sizeof(secret_key) + crypto_aead_chacha20poly1305_ABYTES];

  QDataStream key_in(encrypted_secret);
  key_in >> sk_version_tag;
  quint8 supported_version = key_format_version;
  if (sk_version_tag != supported_version)
  {
    sodium_munlock(secret_key, sizeof(secret_key));
    return ResponseResult::unsupported_version;
  }
  key_in.readRawData((char *)salt, sizeof(salt));
  key_in >> pwhash_opslimit;
  key_in >> pwhash_memlimit;
  key_in >> pwhash_alg;
  key_in.readRawData((char *)payload, sizeof(payload));

  quint8 wrap_key[crypto_aead_chacha20poly1305_KEYBYTES];
  sodium_mlock(wrap_key, sizeof(wrap_key));
  if (crypto_pwhash(wrap_key, sizeof(wrap_key), password.constData(), password.size(), salt, pwhash_opslimit, pwhash_memlimit, pwhash_alg))
  {
    sodium_munlock(wrap_key, sizeof(wrap_key));
    sodium_munlock(secret_key, sizeof(secret_key));
    return ResponseResult::derivation_failed;
  }

  const quint8 *packed_header = (quint8 *)encrypted_secret.constData();
  const size_t packed_header_len = sizeof(quint8) + sizeof(salt) + sizeof(quint32) + sizeof(quint32) + sizeof(quint8);

  if (crypto_aead_chacha20poly1305_decrypt(secret_key, nullptr, nullptr, payload, sizeof(payload), packed_header, packed_header_len, all_zero_nonce, wrap_key))
  {
    sodium_munlock(wrap_key, sizeof(wrap_key));
    sodium_munlock(secret_key, sizeof(secret_key));
    return ResponseResult::decryption_failed;
  }
  sodium_munlock(wrap_key, sizeof(wrap_key));

  // Now we've unlocked the key, and ready to perform DH with the server's
  // ephemeral key to prove ourselves.
  quint8 shared_secret[crypto_scalarmult_BYTES];
  if (crypto_scalarmult(shared_secret, secret_key, (const uchar *)ephemeral_key.constData()))
  {
    // We ended up with the point at infinity, something stupid must've
    // happened. Reject.
    sodium_munlock(secret_key, sizeof(secret_key));
    return ResponseResult::bad_curve_point;
  }
  sodium_munlock(secret_key, sizeof(secret_key));

  // The proof is a BLAKE2b-256 hash over the following:
  //   1. Domain-separating constant.
  //   2. Shared DH secret, serving as a proof of posession of the secret key
  //      that the server can verify using our public key.
  //   3. Server's random ephemeral key used specifically in this attempt.
  //   4. Our public key, binding our identity.
  //   5. Our username, binding the transcript to the secret and the rest of the
  //      session context.
  //
  // Since the server possesses our certificate (public key) and its private
  // ephemeral key, it'll end up getting the same DH secret as we did. Assuming
  // everything else matches up in the transcript, it'll independently
  // calculate the same hash and successfully authenticate us.
  //
  // Note that BLAKE2 is resistant to length-extension attacks, and all fields
  // but last are fixed-length, eliminating encoding ambiguity and
  // canonicalization issues.
  quint8 proof[crypto_generichash_BYTES];
  crypto_generichash_state state;
  crypto_generichash_init(&state, nullptr, 0, sizeof(proof));
  crypto_generichash_update(&state, (const uchar *)domain_sep.constData(), domain_sep.size());
  crypto_generichash_update(&state, shared_secret, sizeof(shared_secret));
  crypto_generichash_update(&state, (const uchar *)ephemeral_key.constData(), ephemeral_key.size());
  crypto_generichash_update(&state, (const uchar *)public_key.constData(), public_key.size());
  crypto_generichash_update(&state, (const uchar *)username.constData(), username.size());
  crypto_generichash_final(&state, proof, sizeof(proof));

  out.response = QByteArray((const char *)proof, (qsizetype)sizeof(proof));
  return ResponseResult::success;
}

// The errors aren't handled carefully, there's a hazard of stale locks being
// created. Fortunately, Qt seems to automatically detect these.
void delete_key(const QByteArray &key_id)
{
  QFile index_file(get_keyring_path());
  QLockFile index_lock(QDir(get_app_path()).filePath("keyring.lock"));
  if (!index_lock.lock())
  {
    return;
  }
  if (!index_file.open(QIODevice::ReadOnly))
  {
    return;
  }
  QCborStreamReader idx_in(&index_file);
  QCborValue index = QCborValue::fromCbor(idx_in);
  index_file.close();
  if (!index.isMap())
  {
    return;
  }
  QCborMap index_map = index.toMap();
  index_map.remove(QCborValue(key_id));
  QSaveFile index_file_w(get_keyring_path());
  if (!index_file_w.open(QIODevice::WriteOnly))
  {
    return;
  }
  index_file_w.write(QCborValue(index_map).toCbor());
  index_file_w.commit();
  index_lock.unlock();
}

// Keyring table model methods (for UI).

KeyringModel::KeyringModel(QObject *parent) : QAbstractTableModel(parent)
{}

int KeyringModel::rowCount(const QModelIndex &) const
{
  return m_keys.size();
}

int KeyringModel::columnCount(const QModelIndex &) const
{
  return 2;
}

QVariant KeyringModel::data(const QModelIndex &index, int role) const
{
  if (role == Qt::DisplayRole)
  {
    const KeyInfo entry = m_keys.at(index.row());
    switch (index.column())
    {
    case 0:
      return entry.name;
    case 1:
      return entry.bytes;
    default:
      return {};
    }
  }
  else if (role == Qt::FontRole && index.column() == 1)
  {
    // There's QFontDatabase::systemFont(QFontDatabase::FixedFont)
    QFont font;
    font.setFamily("Monospace");
    return font;
  }
  // "Get key ID" user role. Magic. There ought to be a more ergonomic way to
  // do this.
  else if (role == KeyIDRole)
  {
    return m_keys.at(index.row()).id;
  }
  return {};
}

// I have removed key ID from the table view, I believe it's better to have
// users create unique names, and if not, they still have certificates to
// differentiate. Key ID is purely internal. Make sure it can still be
// accessed by the key unlock dialog.
QVariant KeyringModel::headerData(int section, Qt::Orientation orientation, int role) const
{
  if (role != Qt::DisplayRole || orientation != Qt::Horizontal)
  {
    return {};
  }
  switch (section)
  {
  case 0:
    return "Note";
  case 1:
    return "Certificate";
  default:
    return {};
  }
}

// Load the keyring from the file into the table to be displayed in the Keyring
// tab. Called whenever the keyring is modified.
void KeyringModel::load_keys()
{
  QFile index_file(get_keyring_path());
  if (!index_file.open(QIODevice::ReadOnly))
  {
    return;
  }
  QCborStreamReader idx_in(&index_file);
  const QCborValue index = QCborValue::fromCbor(idx_in);
  index_file.close();
  if (!index.isMap())
  {
    return;
  }
  const QCborMap index_map = index.toMap();
  QVector<KeyInfo> infos;
  infos.reserve(index_map.size());
  for (QCborMap::const_iterator i = index_map.cbegin(); i != index_map.cend(); ++i)
  {
    KeyInfo entry;

    const QCborValue key = i.key();
    if (key.isByteArray())
    {
      entry.id = key.toByteArray();
    }
    else
    {
      entry.id.clear();
    }

    const QCborValue val = i.value();
    if (val.isMap())
    {
      const QCborMap map = val.toMap();
      entry.name = map.value(1).toString();
      quint8 key_tag = (quint8)map.value(0).toInteger();
      if (map.value(2).isByteArray() && key_tag != 0)
      {
        key_tag = key_tag == 0x1 ? 0x56 : key_tag; // TODO: legacy
        const QByteArray pk = map.value(2).toByteArray();
        QByteArray tagged;
        tagged.reserve(sizeof(key_tag) + pk.size());
        tagged.append(key_tag);
        tagged.append(pk);
        entry.bytes = QString::fromLatin1(tagged.toBase64(QByteArray::OmitTrailingEquals));
      }
      else
      {
        entry.bytes.clear();
      }
    }
    infos.append(entry);
  }
  beginResetModel();
  m_keys = infos;
  endResetModel();
}