// 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"

extern "C" {
#include "tpm/tpm.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");

// Internal state, set by tpm2_init if TPM check and self-test succeed, and we
// can use hardware keys.
static int tpm_available = 0;

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();
  }

  if (tpm2_init() == STKB_SUCCESS)
  {
    tpm_available = 1;
  }

  return 0;
}

// Expose a way to query status without making the internal state variable
// global.
// Yes, this is stupid.
int keyring_hardware_available()
{
    return tpm_available;
}

// 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.
static int insert_into_keyring(quint8 tag, QStringView name, const QByteArray &pub, const QByteArray &sec)
{
  QCborMap entry;
  entry.insert(0, QCborValue(tag));
  entry.insert(1, QCborValue(name));
  entry.insert(2, QCborValue(pub));
  entry.insert(3, QCborValue(sec));
  QFile index_file(get_keyring_path());
  QLockFile index_lock(QDir(get_app_path()).filePath("keyring.lock"));
  if (!index_lock.lock())
  {
    return 1;
  }
  if (!index_file.open(QIODevice::ReadOnly))
  {
    return 1;
  }
  QCborStreamReader idx_in(&index_file);
  QCborValue index = QCborValue::fromCbor(idx_in);
  index_file.close();
  if (!index.isMap())
  {
    index = QCborMap();
  }
  QCborMap index_map = index.toMap();
  // Key fingerprint is used as its unique identifier.
  QByteArray fingerprint;
  fingerprint.append(tag);
  fingerprint.append(pub);
  uchar fpr_hash[crypto_generichash_BYTES];
  crypto_generichash(fpr_hash, sizeof(fpr_hash), (const uchar *)fingerprint.constData(), fingerprint.size(), nullptr, 0);
  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 1;
  }
  index_file_w.write(QCborValue(index_map).toCbor());
  index_file_w.commit();
  index_lock.unlock();
  return 0;
}

// Software X25519 keys protected by a passphrase.
int generate_software_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 (2 iterations over 2 gigabytes 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 = 2u;
  quint32 pwhash_memlimit = 2 * 1024 * 1024 * 1024u;
  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));
  QByteArray encrypted_secret(packed_header);
  encrypted_secret.append((const char *)payload, sizeof(payload));

  if (insert_into_keyring(cert_x25519_tag, name, public_key_array, encrypted_secret))
  {
    return 5;
  }

  return 0;
}

static constexpr quint8 cert_hwkey_p256_tag = 0x33;

static std::wstring optional_hardware_key_name(QStringView name)
{
#ifdef SOF_AO_HARDWARE_KEY_REQUIRES_PIN_UI
  (void)name;
  return std::wstring();
#else
  return name.toString().toStdWString();
#endif
}

// Hardware P-256 keys (because most TPMs can only do P-256, unfortunately);
// the TPM verifies the PIN.
int generate_hardware_key(QStringView name, const QByteArray &pin)
{
  // Compressed P-256 public key.
  uint8_t pub[33];
  // Protected secret blobs differ by size depending on the platform (and in
  // case of TSS2 implementation, size is unspecified). Prepare a sufficient
  // buffer to write to, the actual output size will be in `sec_written`.
  uint8_t sec[512];
  size_t sec_written = 0;
  std::wstring optional_name_wstr = optional_hardware_key_name(name);
  stkb_user_input input = {
    .pin = (const uint8_t *)pin.constData(),
    .pin_len = (size_t)pin.size(),
    .key_name = optional_name_wstr.c_str(),
  };
  stkb_rc rc = tpm2_keygen(pub, sec, sizeof(sec), &sec_written, input);
  switch (rc) {
  case STKB_SUCCESS:
    break;
  default:
    return rc;
  }
  QByteArray public_array((const char *)pub, (qsizetype)sizeof(pub));
  QByteArray protected_array((const char *)sec, (qsizetype)sec_written);
  if (insert_into_keyring(cert_hwkey_p256_tag, name, public_array, protected_array))
    return 1;
  return 0;
}

static ResponseResult compute_with_software_key(quint8 *shared_secret, const QByteArray &encrypted_secret, QByteArrayView password, QByteArrayView ephemeral_key)
{
  if (password.size() < crypto_pwhash_PASSWD_MIN || password.size() > crypto_pwhash_PASSWD_MAX)
  {
    return ResponseResult::invalid_password;
  }

  if (ephemeral_key.size() != crypto_kx_PUBLICKEYBYTES)
  {
    return ResponseResult::incompatible_arguments;
  }

  // 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.
  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));

  return ResponseResult::success;
}

static ResponseResult compute_with_hardware_key(quint8 *shared_secret, const QByteArray &blob, QByteArrayView pin, QByteArrayView ephemeral_key)
{
  if (ephemeral_key.size() != 64)
  {
    return ResponseResult::incompatible_arguments;
  }
  stkb_user_input input = {
    .pin = (const uint8_t *)pin.constData(),
    .pin_len = (size_t)pin.size(),
    .key_name = nullptr,
  };
  stkb_rc rc = tpm2_compute_ss(shared_secret, (const uint8_t *)blob.constData(), (size_t)blob.size(), (const uint8_t *)ephemeral_key.constData(), input);
  switch (rc) {
    case STKB_SUCCESS:
      return ResponseResult::success;
    case STKB_ERROR:
      return ResponseResult::hardware_fault;
    case STKB_BAD_PIN:
    case STKB_AUTH_FAILURE:
      return ResponseResult::invalid_pin;
    case STKB_LOCKOUT:
      return ResponseResult::hardware_lockout;
    case STKB_TPM_UNAVAILABLE:
    default:
      return ResponseResult::inaccessible_keyring;
  }
}

KeyringKey acquire_keyring_key(QByteArrayView key_id)
{
  struct KeyringKey key;
  QFile index_file(get_keyring_path());
  if (!index_file.open(QIODevice::ReadOnly))
  {
    return key;
  }
  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(2);
  if (!val.isByteArray())
  {
    return key;
  }
  key.public_key = val.toByteArray();

  val = tmp_map.value(3);
  if (!val.isByteArray())
  {
    return key;
  }
  key.protected_secret = val.toByteArray();

  val = tmp_map.value(0);
  if (!val.isInteger())
  {
    return key;
  }
  qint64 raw_tag = val.toInteger();
  if (raw_tag < 0 || raw_tag > 255)
  {
    return key;
  }
  quint8 tag = (quint8)raw_tag;
  switch (tag) {
  case cert_x25519_tag:
    key.type = KeyringStorageType::software;
    break;
  case cert_hwkey_p256_tag:
    key.type = KeyringStorageType::hardware;
    break;
  default:
    break;
  }

  val = tmp_map.value(1);
  if (!val.isString())
  {
    return key;
  }
  key.name = val.toString();


  key.id = QByteArray(key_id);

  return key;
}

// 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(const KeyringKey &key, QByteArrayView input, QByteArrayView ephemeral_key, QByteArrayView username, AuthResponse &out)
{
  uint8_t shared_secret[32];
  sodium_mlock(shared_secret, sizeof(shared_secret));

  ResponseResult result;
  switch (key.type) {
  case KeyringStorageType::software:
    result = compute_with_software_key(shared_secret, key.protected_secret, input, ephemeral_key);
    break;
  case KeyringStorageType::hardware:
    result = compute_with_hardware_key(shared_secret, key.protected_secret, input, ephemeral_key);
    break;
  default:
    result = ResponseResult::corrupted_entry;
    break;
  }
  if (result != ResponseResult::success)
  {
    sodium_munlock(shared_secret, sizeof(shared_secret));
    return result;
  }

  // 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 *)key.public_key.constData(), key.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));

  sodium_munlock(shared_secret, sizeof(shared_secret));

  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::Base64UrlEncoding | QByteArray::OmitTrailingEquals));
      }
      else
      {
        entry.bytes.clear();
      }
    }
    infos.append(entry);
  }
  beginResetModel();
  m_keys = infos;
  endResetModel();
}