#include <platform.h>
#include "common/ilib/json.h"
#include "common/ilib/base64.h"
#include "common/ilib/hash.h"
#include "common/ilib/jwk.h"
#include "common/ilib/ecc.h"

NAMESPACE_BEGIN

const char * jwa::alg_to_str(jwa_alg alg)
{
    if (alg == JWA_NONE) return "none";
    else if (alg == JWA_HS256) return "HS256";
    else if (alg == JWA_HS384) return "HS384";
    else if (alg == JWA_HS512) return "HS512";
    else if (alg == JWA_RS256) return "RS256";
    else if (alg == JWA_RS384) return "RS384";
    else if (alg == JWA_RS512) return "RS512";
    else if (alg == JWA_ES256) return "ES256";
    else if (alg == JWA_ES384) return "ES384";
    else if (alg == JWA_ES512) return "ES512";
    else if (alg == JWA_PS256) return "PS256";
    else if (alg == JWA_PS384) return "PS384";
    else if (alg == JWA_PS512) return "PS512";
    else return "none";
}

jwa_alg jwa::alg_from_str(const char * str)
{
    if (strcmp(str, "none") == 0) return JWA_NONE;
    else if (strcmp(str, "HS256") == 0) return JWA_HS256;
    else if (strcmp(str, "HS384") == 0) return JWA_HS384;
    else if (strcmp(str, "HS512") == 0) return JWA_HS512;
    else if (strcmp(str, "RS256") == 0) return JWA_RS256;
    else if (strcmp(str, "RS384") == 0) return JWA_RS384;
    else if (strcmp(str, "RS512") == 0) return JWA_RS512;
    else if (strcmp(str, "ES256") == 0) return JWA_ES256;
    else if (strcmp(str, "ES384") == 0) return JWA_ES384;
    else if (strcmp(str, "ES512") == 0) return JWA_ES512;
    else if (strcmp(str, "PS256") == 0) return JWA_PS256;
    else if (strcmp(str, "PS384") == 0) return JWA_PS384;
    else if (strcmp(str, "PS512") == 0) return JWA_PS512;
    else return JWA_UNKNOWN;
}

dword jose::create_jws(const char * jws_header, const char * jws_protected_header, const char * jws_payload, class jwk * jwk, jwa_alg alg, char * jws)
{
    // Preconditions
    if (jws_payload == NULL || jwk == NULL) {
        return 0;
    }

    json_io jws_out(jws);
    word root = jws_out.add_object(JSON_ID_ROOT, 0);
    char * jws_header_encoded = NULL;

    // Add uprotected header
    if (jws_header != NULL) {
        jws_header_encoded = (char *)alloca(4 * strlen(jws_header) / 3 + 2);
        size_t jws_header_encoded_len = encode_base64url((byte*)jws_header, jws_header_encoded, (dword)strlen(jws_header));
        jws_header_encoded[jws_header_encoded_len] = '\0';
        jws_out.add_string(root, "header", jws_header_encoded);
    }

    // Add protected header
    char * jws_protected_header_encoded = NULL;
    size_t jws_protected_header_encoded_len = 0;
    if (jws_protected_header != NULL) {
        jws_protected_header_encoded = (char *)alloca(4 * strlen(jws_protected_header) / 3 + 2);
        jws_protected_header_encoded_len = encode_base64url((byte*)jws_protected_header, jws_protected_header_encoded, (dword)strlen(jws_protected_header));
        jws_protected_header_encoded[jws_protected_header_encoded_len] = '\0';
        jws_out.add_string(root, "protected", jws_protected_header_encoded);
    }

    // Add payload
    char * jws_payload_encoded = (char *)alloca(4 * strlen(jws_payload) / 3 + 2);
    size_t jws_payload_encoded_len = encode_base64url((byte*)jws_payload, jws_payload_encoded, (dword)strlen(jws_payload));
    jws_payload_encoded[jws_payload_encoded_len] = '\0';
    jws_out.add_string(root, "payload", jws_payload_encoded);

    // Assemble contents to sign
    size_t signed_content_size = jws_protected_header_encoded_len + jws_payload_encoded_len + 1;
    char * signed_content = (char *)alloca(signed_content_size + 1);
    memcpy(signed_content, jws_protected_header_encoded, jws_protected_header_encoded_len);
    memcpy(signed_content + jws_protected_header_encoded_len, ".", 1);
    memcpy(signed_content + jws_protected_header_encoded_len + 1, jws_payload_encoded, jws_payload_encoded_len);
    memcpy(signed_content + signed_content_size, "\0", 1);

    // Add signature
    char * signature = (char *)alloca(jwk->get_signature_size());
    jwk->sign(signature, signed_content, signed_content_size);
    char * signature_encoded = (char *)alloca(4 * jwk->get_signature_size() / 3 + 2);
    size_t signature_encoded_len = encode_base64url((byte*)signature, signature_encoded, jwk->get_signature_size());
    signature_encoded[signature_encoded_len] = '\0';
    jws_out.add_string(root, "signature", signature_encoded);

    //debug->iprintf("Create signature");
    //debug->hexdump(signature, jwk->get_signature_size());

    return jws_out.encode();
}

bool jose::verify_jws(const char * jws, class jwk * jwk, jwa_alg alg)
{
    json_io jws_in((char *)jws); // It's safe to remove constness here because we don't modify the underlying buffer
    jws_in.decode();
    word jws_root = jws_in.get_object(JSON_ID_ROOT, 0);

    const char * protected_header = jws_in.get_string(jws_root, "protected");
    const char * payload = jws_in.get_string(jws_root, "payload");
    const char * signature = jws_in.get_string(jws_root, "signature");

    // Decode signature
    size_t signature_len = jwk->get_signature_size();//strlen(signature);
    char * signature_bin = (char *)alloca(signature_len);
    decode_base64url_bin(signature, (byte*)signature_bin, (dword)signature_len); // Decoded signature is shorter than encoded one

    // Assemble signed content
    size_t protected_header_len = strlen(protected_header);
    size_t payload_len = strlen(payload);
    size_t signed_content_len = protected_header_len + payload_len + 1;
    char * signed_content = (char *)alloca(signed_content_len + 1);
    memcpy(signed_content, protected_header, protected_header_len);
    memcpy(signed_content + protected_header_len, ".", 1);
    memcpy(signed_content + protected_header_len + 1, payload, payload_len);
    memcpy(signed_content + signed_content_len, "\0", 1);

    //debug->iprintf("Verify signature");
    //debug->hexdump(signature_bin, jwk->get_signature_size());

    // Verify signature
    return jwk->verify(signature_bin, signed_content, signed_content_len);
}

bool jose::parse_jws(const char * jws, char * jws_header, char * jws_protected_header, char * jws_payload)
{
    json_io jws_in((char *)jws); // It's safe to remove constness here because we don't modify the underlying buffer
    jws_in.decode();
    word jws_root = jws_in.get_object(JSON_ID_ROOT, 0);

    bool header_present = false;
    const char * header = jws_in.get_string(jws_root, "header", &header_present);
    if (header_present && jws_header != NULL) {
        decode_base64url(header, (byte*)jws_header, (dword)strlen(header));
    }

    bool protected_header_present = false;
    const char * protected_header = jws_in.get_string(jws_root, "protected", &protected_header_present);
    if (protected_header_present && jws_protected_header != NULL) {
        decode_base64url(protected_header, (byte*)jws_protected_header, (dword)strlen(protected_header));
    }

    bool payload_present = false;
    const char * payload = jws_in.get_string(jws_root, "payload", &payload_present);
    if (payload_present && jws_payload != NULL) {
        decode_base64url(payload, (byte*)jws_payload, (dword)strlen(payload));
    }

    return protected_header_present && payload_present;
}


dword jose::encode_compact_jws(const char * jws, char * compact_jws)
{
    json_io jws_in((char *)jws); // It's safe to remove constness here because we don't modify the underlying buffer
    jws_in.decode();
    word jws_root = jws_in.get_object(JSON_ID_ROOT, 0);

    const char * protected_header = jws_in.get_string(jws_root, "protected");
    const char * payload = jws_in.get_string(jws_root, "payload");
    const char * signature = jws_in.get_string(jws_root, "signature");

    // Make sure required fields are present from here on
    if (protected_header == NULL || payload == NULL || signature == NULL) {
        return 0;
    }

    size_t protected_header_len = strlen(protected_header);
    size_t payload_len = strlen(payload);
    size_t signature_len = strlen(signature);

    memcpy(compact_jws, protected_header, protected_header_len);
    memcpy(compact_jws + protected_header_len, ".", 1);
    memcpy(compact_jws + protected_header_len + 1, payload, payload_len);
    memcpy(compact_jws + protected_header_len + 1 + payload_len, ".", 1);
    memcpy(compact_jws + protected_header_len + 1 + payload_len + 1, signature, signature_len);
    memcpy(compact_jws + protected_header_len + 1 + payload_len + 1 + signature_len, "\0", 1);

    return (dword)(protected_header_len + 1 + payload_len + 1 + signature_len);
}

dword jose::decode_compact_jws(const char * compact_jws, char * jws)
{
    size_t compact_jws_len = strlen(compact_jws);

    // Compute offsets
    size_t payload_offset = 0;
    size_t signature_offset = 0;
    int current_delimiter = 0;
    for (size_t i = 0; i < compact_jws_len; i++) {
        if (compact_jws[i] == '.') {
            if (current_delimiter == 0) {
                payload_offset = i + 1;
            }
            else if (current_delimiter == 1) {
                signature_offset = i + 1;
            }
            current_delimiter++;
        }
    }
    size_t protected_header_len = payload_offset - 1;
    size_t payload_len = signature_offset - payload_offset - 1;
    size_t signature_len = compact_jws_len - signature_offset;

    // Make sure there have been exactly two delimiters. Otherwise compact JWS is malformed.
    if (current_delimiter != 2) {
        return 0;
    }

    // Check invariants
    ASSERT(0 < payload_offset && payload_offset < signature_offset, "Failed to parse compact JWS");

    // Extract (protected-)header, payload and signature
    char * protected_header = (char *)alloca(protected_header_len + 1);
    memcpy(protected_header, compact_jws, protected_header_len);
    protected_header[protected_header_len] = '\0';
    char * payload = (char *)alloca(payload_len + 1);
    memcpy(payload, compact_jws + payload_offset, payload_len);
    payload[payload_len] = '\0';
    char * signature = (char *)alloca(signature_len + 1);
    memcpy(signature, compact_jws + signature_offset, signature_len);
    signature[signature_len] = '\0';

    // Create unflattened JWS
    json_io jws_out(jws);
    word jws_root = jws_out.add_object(JSON_ID_ROOT, 0);
    jws_out.add_string(jws_root, "protected", protected_header);
    jws_out.add_string(jws_root, "payload", payload);
    jws_out.add_string(jws_root, "signature", signature);
    return jws_out.encode();
}

bool jwk::is_algorithm_supported(jwa_alg alg)
{
    if (alg == JWA_ES256) {
        return true;
    }
    return false;
}

jwk::jwk()
    : kty(JWK_KTY_NONE)
    , use(JWK_USE_NONE)
    , crv(JWK_CRV_NONE)
    , alg(JWA_NONE)
{
}

jwk::~jwk()
{
}

bool jwk::decode(const char * buffer)
{
    char * buf = (char *)alloca(strlen(buffer) + 1);
    memcpy(buf, buffer, strlen(buffer) + 1);
    json_io json(buf);
    if (json.decode()) {
        return decode(&json, JSON_ID_ROOT, 0);
    }
    return false;
}

dword jwk::encode(char * buffer, bool include_optional_parameters)
{
    return encode_impl(buffer, true, true, include_optional_parameters);
}

dword jwk::encode_public_key(char * buffer, bool include_optional_parameters)
{
    return encode_impl(buffer, true, false, include_optional_parameters);
}

dword jwk::encode_private_key(char * buffer, bool include_optional_parameters)
{
    return encode_impl(buffer, false, true, include_optional_parameters);
}

bool jwk::decode(json_io* json, word base, const char * name)
{
    this->clear();
    word root = json->get_object(base, name);

    // Parse key type
    bool kty_present = false;
    const char * kty = json->get_string(root, "kty", &kty_present);
    if (!kty_present) {
        this->kty = JWK_KTY_NONE;
    }
    else if (strcmp(kty, "EC") == 0) {
        this->kty = JWK_KTY_EC;
        bool crv_present = false;
        const char * crv = json->get_string(root, "crv", &crv_present);
        if (!crv_present) {
            this->crv = JWK_CRV_NONE;
        }
        else if (strcmp(crv, "P-256") == 0) {
            this->crv = JWK_CRV_P256;
            const char * x = json->get_string(root, "x");
            decode_base64url(x, (byte*)this->params.p256.x, (dword)strlen(x));
            const char * y = json->get_string(root, "y");
            decode_base64url(y, (byte*)this->params.p256.y, (dword)strlen(y));
            const char * d = json->get_string(root, "d");
            decode_base64url(d, (byte*)this->params.p256.d, (dword)strlen(d));
        }
        else {
            this->crv = JWK_CRV_UNKNOWN;
        }
    }
    else if (strcmp(kty, "RSA") == 0) {
        this->kty = JWK_KTY_RSA;
        // TODO not implemented yet
    }
    else {
        this->kty = JWK_KTY_UNKNOWN;
    }

    // Parse algorithm
    bool alg_present = false;
    const char * alg = json->get_string(root, "alg", &alg_present);
    this->alg = jwa::alg_from_str(alg);

    return true;
}

dword jwk::encode(json_io* json, word base, const char * name, bool include_optional_parameters)
{
    return encode_impl(json, base, name, true, true, include_optional_parameters);
}

dword jwk::encode_public_key(json_io* json, word base, const char * name, bool include_optional_parameters)
{
    return encode_impl(json, base, name, true, false, include_optional_parameters);
}

dword jwk::encode_private_key(json_io* json, word base, const char * name, bool include_optional_parameters)
{
    return encode_impl(json, base, name, false, true, include_optional_parameters);
}

void jwk::clear() {
    this->kty = JWK_KTY_NONE;
    this->use = JWK_USE_NONE;
    this->crv = JWK_CRV_NONE;
    this->alg = JWA_NONE;
    memset(&this->params, 0, sizeof(this->params));
}

bool jwk::sign(char * signature, const char * msg, size_t msg_len)
{
    return sign(signature, msg, msg_len, this->alg);
}

bool jwk::sign(char * signature, const char * msg, size_t msg_len, jwa_alg alg)
{
    if (alg == JWA_ES256) {
        return sign_es256(signature, msg, msg_len);
    }
    return false;
}

bool jwk::verify(const char * signature, const char * msg, size_t msg_len)
{
    return verify(signature, msg, msg_len, this->alg);
}

bool jwk::verify(const char * signature, const char * msg, size_t msg_len, jwa_alg alg)
{
    if (alg == JWA_ES256) {
        return verify_es256(signature, msg, msg_len);
    }
    return false;
}

dword jwk::get_signature_size()
{
    if (this->alg == JWA_ES256) {
        return 64;
    }
    return false;
}

void jwk::gen_es256_keys()
{
    clear();

    byte pub_key[64];
    ecc::make_key(pub_key, (byte*)this->params.p256.d, EC_SECP256R1);
    memcpy(this->params.p256.x, pub_key, 32);
    memcpy(this->params.p256.y, pub_key + 32, 32);
    
    this->kty = JWK_KTY_EC;
    this->crv = JWK_CRV_P256;
    this->alg = JWA_ES256;
}

void jwk::set_es256_keys(const byte * params)
{
    clear();

    memcpy(&this->params, params, sizeof(this->params));

    this->kty = JWK_KTY_EC;
    this->crv = JWK_CRV_P256;
    this->alg = JWA_ES256;
}

dword jwk::encode_impl(char * buffer, bool serialize_public_key, bool serialize_private_key, bool include_optional_parameters)
{
    json_io json(buffer);
    return encode_impl(&json, JSON_ID_ROOT, 0, serialize_public_key, serialize_private_key, include_optional_parameters);
}

dword jwk::encode_impl(json_io* json, word base, const char * name, bool serialize_public_key, bool serialize_private_key, bool include_optional_parameters)
{
    char x_base64url[45];
    char y_base64url[45];
    char d_base64url[45];
    word root = json->add_object(base, name);

    // take care that keys are lexicographically ordered, see https://datatracker.ietf.org/doc/html/rfc7638#ref-JWK!
    if (this->kty == JWK_KTY_EC) {
        if (this->crv == JWK_CRV_P256) {
            json->add_string(root, "crv", "P-256");
            json->add_string(root, "kty", "EC");
            if (serialize_public_key) {
                encode_base64url((byte*)this->params.p256.x, x_base64url, 32);
                json->add_string(root, "x", x_base64url);
                encode_base64url((byte*)this->params.p256.y, y_base64url, 32);
                json->add_string(root, "y", y_base64url);
            }
            if (serialize_private_key) {
                encode_base64url((byte*)this->params.p256.d, d_base64url, 32);
                json->add_string(root, "d", d_base64url);
            }
        }
        else {
            json->add_string(root, "kty", "EC");
        }
    }
    else if (this->kty == JWK_KTY_RSA) {
        json->add_string(root, "kty", "RSA");
    }

    if (include_optional_parameters && this->alg != JWA_UNKNOWN && this->alg != JWA_NONE) {
        json->add_string(root, "alg", jwa::alg_to_str(this->alg));
    }

    return json->encode();
}

bool jwk::sign_es256(char * signature, const char * msg, size_t msg_len)
{
    // Calculate sha256 hash
    hash sha256;
    byte hash[32];
    sha256.init(HASH_SHA256);
    sha256.update(msg, msg_len);
    sha256.final(hash);

    // Create signature
    return ecc::sign((byte*)signature, hash, 32, (byte*)this->params.p256.d, EC_SECP256R1);
}

bool jwk::verify_es256(const char * signature, const char * msg, size_t msg_len)
{
    // Calculate sha256 hash
    hash sha256;
    byte hash[32];
    sha256.init(HASH_SHA256);
    sha256.update(msg, msg_len);
    sha256.final(hash);

    // Concat public key elements
    byte pub_key[64];
    memcpy(pub_key, this->params.p256.x, 32);
    memcpy(pub_key + 32, this->params.p256.y, 32);

    // Verify signature
    return ecc::verify((byte*)signature, hash, 32, pub_key, EC_SECP256R1);
}

NAMESPACE_END
