#include "connectors/connector.h"

#include <openssl/evp.h>
#include <openssl/rand.h>
#include <stdio.h>
#include <string.h>

static const char BASE64URL[] =
    "ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789-_";

boolean Connector_Base64Url_Encode(
    const uint8 *input, size_t input_length, char *output, size_t output_size)
{
  size_t required = (input_length * 4 + 2) / 3;
  if (!input || !output || output_size <= required)
    return FALSE;
  size_t i = 0, j = 0;
  while (i + 3 <= input_length) {
    uint32 value = ((uint32)input[i] << 16) |
                   ((uint32)input[i + 1] << 8) | input[i + 2];
    output[j++] = BASE64URL[(value >> 18) & 63];
    output[j++] = BASE64URL[(value >> 12) & 63];
    output[j++] = BASE64URL[(value >> 6) & 63];
    output[j++] = BASE64URL[value & 63];
    i += 3;
  }
  if (i < input_length) {
    uint32 value = (uint32)input[i] << 16;
    output[j++] = BASE64URL[(value >> 18) & 63];
    if (i + 1 < input_length) {
      value |= (uint32)input[i + 1] << 8;
      output[j++] = BASE64URL[(value >> 12) & 63];
      output[j++] = BASE64URL[(value >> 6) & 63];
    } else {
      output[j++] = BASE64URL[(value >> 12) & 63];
    }
  }
  output[j] = '\0';
  return TRUE;
}

static int32 base64url_value(char c)
{
  const char *position = strchr(BASE64URL, c);
  return position ? (int32)(position - BASE64URL) : -1;
}

boolean Connector_Base64Url_Decode(
    const char *input, uint8 *output, size_t output_size, size_t *output_length)
{
  if (!input || !output || !output_length)
    return FALSE;
  size_t length = strlen(input);
  if ((length % 4) == 1 || output_size < (length * 3) / 4)
    return FALSE;
  uint32 accumulator = 0;
  int32 bits = 0;
  size_t written = 0;
  for (size_t i = 0; i < length; ++i) {
    int32 value = base64url_value(input[i]);
    if (value < 0)
      return FALSE;
    accumulator = (accumulator << 6) | (uint32)value;
    bits += 6;
    if (bits >= 8) {
      bits -= 8;
      if (written >= output_size)
        return FALSE;
      output[written++] = (uint8)((accumulator >> bits) & 255);
    }
  }
  *output_length = written;
  return TRUE;
}

boolean Connector_Encrypt(
    const Connector_Master_Key *key, const char *plaintext,
    char *encoded, size_t encoded_size)
{
  if (!key || !plaintext || !encoded)
    return FALSE;
  size_t plaintext_length = strlen(plaintext);
  if (plaintext_length > 4096)
    return FALSE;
  uint8 nonce[12], tag[16], ciphertext[4096];
  if (RAND_bytes(nonce, sizeof(nonce)) != 1)
    return FALSE;
  EVP_CIPHER_CTX *context = EVP_CIPHER_CTX_new();
  int32 length = 0, total = 0;
  boolean ok = context &&
      EVP_EncryptInit_ex(context, EVP_aes_256_gcm(), NULL, NULL, NULL) == 1 &&
      EVP_CIPHER_CTX_ctrl(context, EVP_CTRL_GCM_SET_IVLEN, sizeof(nonce), NULL) == 1 &&
      EVP_EncryptInit_ex(context, NULL, NULL, key->key, nonce) == 1 &&
      EVP_EncryptUpdate(context, ciphertext, &length,
                        (const uint8 *)plaintext, (int32)plaintext_length) == 1;
  total = length;
  ok = ok && EVP_EncryptFinal_ex(context, ciphertext + total, &length) == 1;
  total += length;
  ok = ok && EVP_CIPHER_CTX_ctrl(context, EVP_CTRL_GCM_GET_TAG, sizeof(tag), tag) == 1;
  EVP_CIPHER_CTX_free(context);
  if (!ok)
    return FALSE;
  uint8 envelope[4 + 12 + 16 + 4096];
  envelope[0] = (uint8)(key->version >> 24);
  envelope[1] = (uint8)(key->version >> 16);
  envelope[2] = (uint8)(key->version >> 8);
  envelope[3] = (uint8)key->version;
  memcpy(envelope + 4, nonce, sizeof(nonce));
  memcpy(envelope + 16, tag, sizeof(tag));
  memcpy(envelope + 32, ciphertext, (size_t)total);
  return Connector_Base64Url_Encode(
      envelope, 32 + (size_t)total, encoded, encoded_size);
}

boolean Connector_Decrypt(
    const Connector_Master_Key *key, const char *encoded,
    char *plaintext, size_t plaintext_size)
{
  uint8 envelope[4 + 12 + 16 + 4096];
  size_t envelope_length = 0;
  if (!key || !encoded || !plaintext ||
      !Connector_Base64Url_Decode(
          encoded, envelope, sizeof(envelope), &envelope_length) ||
      envelope_length < 32)
    return FALSE;
  uint32 version = ((uint32)envelope[0] << 24) |
      ((uint32)envelope[1] << 16) | ((uint32)envelope[2] << 8) | envelope[3];
  size_t ciphertext_length = envelope_length - 32;
  if (version != key->version || plaintext_size <= ciphertext_length)
    return FALSE;
  EVP_CIPHER_CTX *context = EVP_CIPHER_CTX_new();
  int32 length = 0, total = 0;
  boolean ok = context &&
      EVP_DecryptInit_ex(context, EVP_aes_256_gcm(), NULL, NULL, NULL) == 1 &&
      EVP_CIPHER_CTX_ctrl(context, EVP_CTRL_GCM_SET_IVLEN, 12, NULL) == 1 &&
      EVP_DecryptInit_ex(context, NULL, NULL, key->key, envelope + 4) == 1 &&
      EVP_DecryptUpdate(context, (uint8 *)plaintext, &length,
                        envelope + 32, (int32)ciphertext_length) == 1;
  total = length;
  ok = ok && EVP_CIPHER_CTX_ctrl(
      context, EVP_CTRL_GCM_SET_TAG, 16, envelope + 16) == 1 &&
      EVP_DecryptFinal_ex(context, (uint8 *)plaintext + total, &length) == 1;
  total += length;
  EVP_CIPHER_CTX_free(context);
  if (!ok) {
    memset(plaintext, 0, plaintext_size);
    return FALSE;
  }
  plaintext[total] = '\0';
  return TRUE;
}

boolean Connector_Form_Encode(
    const char *input, char *output, size_t output_size)
{
  static const char hex[] = "0123456789ABCDEF";
  if (!input || !output)
    return FALSE;
  size_t j = 0;
  for (size_t i = 0; input[i]; ++i) {
    uint8 c = (uint8)input[i];
    boolean safe = (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') ||
        (c >= '0' && c <= '9') || c == '-' || c == '_' || c == '.' || c == '~';
    size_t needed = safe ? 1 : 3;
    if (j + needed >= output_size)
      return FALSE;
    if (safe)
      output[j++] = (char)c;
    else {
      output[j++] = '%';
      output[j++] = hex[c >> 4];
      output[j++] = hex[c & 15];
    }
  }
  output[j] = '\0';
  return TRUE;
}

const char *Connector_Operation_Name(Connector_Operation operation)
{
  static const char *names[] = {
    "drive.list", "drive.get", "drive.download", "drive.changes",
    "drive.create", "drive.upload", "drive.update", "gmail.list",
    "gmail.get", "gmail.attachment", "gmail.history", "gmail.draft.create",
    "gmail.send"
  };
  return operation >= CONNECTOR_OP_DRIVE_LIST &&
      operation <= CONNECTOR_OP_GMAIL_SEND ? names[operation] : "unknown";
}

boolean Connector_Operation_Is_Mutation(Connector_Operation operation)
{
  return operation == CONNECTOR_OP_DRIVE_CREATE ||
      operation == CONNECTOR_OP_DRIVE_UPLOAD ||
      operation == CONNECTOR_OP_DRIVE_UPDATE ||
      operation == CONNECTOR_OP_GMAIL_DRAFT_CREATE ||
      operation == CONNECTOR_OP_GMAIL_SEND;
}

boolean Connector_Operation_Is_Allowed(Connector_Operation operation)
{
  return operation >= CONNECTOR_OP_DRIVE_LIST &&
      operation <= CONNECTOR_OP_GMAIL_SEND;
}

boolean Connector_Request_Digest(
    const char *user_id, const char *account_id, Connector_Operation operation,
    const Connector_Provider_Request *request, char output[65])
{
  if (!user_id || !account_id || !request || !output ||
      request->body_length > CONNECTOR_MAX_JSON_BYTES)
    return FALSE;
  EVP_MD_CTX *context = EVP_MD_CTX_new();
  uint8 digest[32];
  uint32 digest_length = 0;
  const char separator = '\0';
#define HASH_FIELD(value, length) do { \
  EVP_DigestUpdate(context, (value) ? (value) : "", (value) ? (length) : 0); \
  EVP_DigestUpdate(context, &separator, 1); \
} while (0)
  boolean ok = context &&
      EVP_DigestInit_ex(context, EVP_sha256(), NULL) == 1;
  if (!ok) {
    EVP_MD_CTX_free(context);
    return FALSE;
  }
  const char *name = Connector_Operation_Name(operation);
  HASH_FIELD(user_id, strlen(user_id));
  HASH_FIELD(account_id, strlen(account_id));
  HASH_FIELD(name, strlen(name));
  HASH_FIELD(request->method, request->method ? strlen(request->method) : 0);
  HASH_FIELD(request->path, request->path ? strlen(request->path) : 0);
  HASH_FIELD(request->query, request->query ? strlen(request->query) : 0);
  HASH_FIELD(request->content_type,
             request->content_type ? strlen(request->content_type) : 0);
  HASH_FIELD(request->body, request->body_length);
  uint8 overwrite = request->overwrite ? 1 : 0;
  EVP_DigestUpdate(context, &overwrite, 1);
  ok = EVP_DigestFinal_ex(context, digest, &digest_length) == 1;
  EVP_MD_CTX_free(context);
#undef HASH_FIELD
  if (!ok || digest_length != 32)
    return FALSE;
  for (size_t i = 0; i < sizeof(digest); ++i)
    snprintf(output + i * 2, 3, "%02x", digest[i]);
  output[64] = '\0';
  return TRUE;
}

boolean Connector_OAuth_PKCE_Start(Connector_OAuth_Start *start)
{
  uint8 state[32], verifier[48], digest[32];
  uint32 digest_length = 0;
  if (!start || RAND_bytes(state, sizeof(state)) != 1 ||
      RAND_bytes(verifier, sizeof(verifier)) != 1 ||
      !Connector_Base64Url_Encode(
          state, sizeof(state), start->state, sizeof(start->state)) ||
      !Connector_Base64Url_Encode(
          verifier, sizeof(verifier), start->code_verifier,
          sizeof(start->code_verifier)))
    return FALSE;
  if (EVP_Digest(
          start->code_verifier, strlen(start->code_verifier), digest,
          &digest_length, EVP_sha256(), NULL) != 1 || digest_length != 32)
    return FALSE;
  return Connector_Base64Url_Encode(
      digest, sizeof(digest), start->code_challenge,
      sizeof(start->code_challenge));
}