/*
 * Copyright (c) 2022, stelar7 <dudedbz@gmail.com>
 *
 * SPDX-License-Identifier: BSD-2-Clause
 */

#include <AK/Random.h>
#include <LibCrypto/Curves/Curve25519.h>
#include <LibCrypto/Curves/Ed25519.h>
#include <LibCrypto/Hash/SHA2.h>

namespace Crypto::Curves {

// https://datatracker.ietf.org/doc/html/rfc8032#section-5.1.5
ErrorOr<ByteBuffer> Ed25519::generate_private_key()
{
    // The private key is 32 octets (256 bits, corresponding to b) of
    // cryptographically secure random data.  See [RFC4086] for a discussion
    // about randomness.

    auto buffer = TRY(ByteBuffer::create_uninitialized(key_size()));
    fill_with_random(buffer);
    return buffer;
};

// https://datatracker.ietf.org/doc/html/rfc8032#section-5.1.5
ErrorOr<ByteBuffer> Ed25519::generate_public_key(ReadonlyBytes private_key)
{
    // The 32-byte public key is generated by the following steps.

    // 1. Hash the 32-byte private key using SHA-512, storing the digest in a 64-octet large buffer, denoted h.
    // Only the lower 32 bytes are used for generating the public key.
    auto digest = Crypto::Hash::SHA512::hash(private_key);

    // NOTE: we do these steps in the opposite order (since its easier to modify s)
    // 3. Interpret the buffer as the little-endian integer, forming a secret scalar s.
    memcpy(s, digest.data, 32);

    // 2. Prune the buffer:
    // The lowest three bits of the first octet are cleared,
    s[0] &= 0xF8;
    // the highest bit of the last octet is cleared,
    s[31] &= 0x7F;
    // and the second highest bit of the last octet is set.
    s[31] |= 0x40;

    // Perform a fixed-base scalar multiplication [s]B.
    point_multiply_scalar(&sb, s, &BASE_POINT);

    // 4.  The public key A is the encoding of the point [s]B.
    // First, encode the y-coordinate (in the range 0 <= y < p) as a little-endian string of 32 octets.
    // The most significant bit of the final octet is always zero.
    // To form the encoding of the point [s]B, copy the least significant bit of the x coordinate
    // to the most significant bit of the final octet.  The result is the public key.
    auto public_key = TRY(ByteBuffer::create_uninitialized(key_size()));
    encode_point(&sb, public_key.data());

    return public_key;
}

// https://datatracker.ietf.org/doc/html/rfc8032#section-5.1.6
ErrorOr<ByteBuffer> Ed25519::sign(ReadonlyBytes public_key, ReadonlyBytes private_key, ReadonlyBytes message)
{
    // 1. Hash the private key, 32 octets, using SHA-512.
    // Let h denote the resulting digest.
    auto h = Crypto::Hash::SHA512::hash(private_key);

    // Construct the secret scalar s from the first half of the digest,
    memcpy(s, h.data, 32);
    // NOTE: This is done later in step 4, but we can also do it here.
    s[0] &= 0xF8;
    s[31] &= 0x7F;
    s[31] |= 0x40;

    // and the corresponding public key A, as described in the previous section.
    // NOTE: The public key A is taken as input to this function.

    // Let prefix denote the second half of the hash digest, h[32],...,h[63].
    memcpy(p, h.data + 32, 32);

    // 2. Compute SHA-512(dom2(F, C) || p || PH(M)), where M is the message to be signed.
    Crypto::Hash::SHA512 hash;
    // NOTE: dom2(F, C) is a blank octet string when signing or verifying Ed25519
    hash.update(p, 32);
    // NOTE: PH(M) = M
    hash.update(message.data(), message.size());

    // Interpret the 64-octet digest as a little-endian integer r.
    // For efficiency, do this by first reducing r modulo L, the group order of B.
    auto digest = hash.digest();
    barrett_reduce(r, digest.data);

    // 3.  Compute the point [r]B.
    point_multiply_scalar(&rb, r, &BASE_POINT);

    auto R = TRY(ByteBuffer::create_uninitialized(32));
    // Let the string R be the encoding of this point
    encode_point(&rb, R.data());

    // 4. Compute SHA512(dom2(F, C) || R || A || PH(M)),
    // NOTE: We can reuse hash here, since digest() calls reset()
    // NOTE: dom2(F, C) is a blank octet string when signing or verifying Ed25519
    hash.update(R.data(), R.size());
    // NOTE: A == public_key
    hash.update(public_key.data(), public_key.size());
    // NOTE: PH(M) = M
    hash.update(message.data(), message.size());

    digest = hash.digest();
    // and interpret the 64-octet digest as a little-endian integer k.
    memcpy(k, digest.data, 64);

    // 5. Compute S = (r + k * s) mod L. For efficiency, again reduce k modulo L first.
    barrett_reduce(p, k);
    multiply(k, k + 32, p, s, 32);
    barrett_reduce(p, k);
    add(s, p, r, 32);

    // modular reduction
    auto reduced_s = TRY(ByteBuffer::create_uninitialized(32));
    auto is_negative = subtract(p, s, Curve25519::BASE_POINT_L_ORDER, 32);
    select(reduced_s.data(), p, s, is_negative, 32);

    // 6.  Form the signature of the concatenation of R (32 octets) and the little-endian encoding of S
    // (32 octets; the three most significant bits of the final octet are always zero).
    auto signature = TRY(ByteBuffer::copy(R));
    signature.append(reduced_s);

    return signature;
}

// https://datatracker.ietf.org/doc/html/rfc8032#section-5.1.7
bool Ed25519::verify(ReadonlyBytes public_key, ReadonlyBytes signature, ReadonlyBytes message)
{
    auto not_valid = false;

    // 1. To verify a signature on a message M using public key A,
    // with F being 0 for Ed25519ctx, 1 for Ed25519ph, and if Ed25519ctx or Ed25519ph is being used, C being the context
    // first split the signature into two 32-octet halves.
    // If any of the decodings fail (including S being out of range), the signature is invalid.

    // NOTE: We dont care about F, since we dont implement Ed25519ctx or Ed25519PH
    // NOTE: C is the internal state, so its not a parameter

    auto half_signature_size = signature_size() / 2;

    // Decode the first half as a point R
    memcpy(r, signature.data(), half_signature_size);

    // and the second half as an integer S, in the range  0 <= s < L.
    memcpy(s, signature.data() + half_signature_size, half_signature_size);

    // NOTE: Ed25519 and Ed448 signatures are not malleable due to the verification check that decoded S is smaller than l.
    // Without this check, one can add a multiple of l into a scalar part and still pass signature verification,
    // resulting in malleable signatures.
    auto is_negative = subtract(p, s, Curve25519::BASE_POINT_L_ORDER, half_signature_size);
    not_valid |= 1 ^ is_negative;

    // Decode the public key A as point A'.
    not_valid |= decode_point(&ka, public_key.data());

    // 2. Compute SHA512(dom2(F, C) || R || A || PH(M)), and interpret the 64-octet digest as a little-endian integer k.
    Crypto::Hash::SHA512 hash;
    // NOTE: dom2(F, C) is a blank octet string when signing or verifying Ed25519
    hash.update(r, half_signature_size);
    // NOTE: A == public_key
    hash.update(public_key.data(), key_size());
    // NOTE: PH(M) = M
    hash.update(message.data(), message.size());

    auto digest = hash.digest();
    auto k = digest.data;

    // 3.  Check the group equation [8][S]B = [8]R + [8][k]A'.
    // It's sufficient, but not required, to instead check [S]B = R + [k]A'.
    // NOTE: For efficiency, do this by first reducing k modulo L.
    barrett_reduce(k, k);

    // NOTE: We check [S]B - [k]A' == R
    Curve25519::modular_subtract(ka.x, Curve25519::ZERO, ka.x);
    Curve25519::modular_subtract(ka.t, Curve25519::ZERO, ka.t);
    point_multiply_scalar(&sb, s, &BASE_POINT);
    point_multiply_scalar(&ka, k, &ka);
    point_add(&ka, &sb, &ka);
    encode_point(&ka, p);

    not_valid |= compare(p, r, half_signature_size);

    return !not_valid;
}

void Ed25519::point_double(Ed25519Point* result, Ed25519Point const* point)
{
    Curve25519::modular_square(a, point->x);
    Curve25519::modular_square(b, point->y);
    Curve25519::modular_square(c, point->z);
    Curve25519::modular_add(c, c, c);
    Curve25519::modular_add(e, a, b);
    Curve25519::modular_add(f, point->x, point->y);
    Curve25519::modular_square(f, f);
    Curve25519::modular_subtract(f, e, f);
    Curve25519::modular_subtract(g, a, b);
    Curve25519::modular_add(h, c, g);
    Curve25519::modular_multiply(result->x, f, h);
    Curve25519::modular_multiply(result->y, e, g);
    Curve25519::modular_multiply(result->z, g, h);
    Curve25519::modular_multiply(result->t, e, f);
}

void Ed25519::point_multiply_scalar(Ed25519Point* result, u8 const* scalar, Ed25519Point const* point)
{
    // Set U to the neutral element (0, 1, 1, 0)
    Curve25519::set(u.x, 0);
    Curve25519::set(u.y, 1);
    Curve25519::set(u.z, 1);
    Curve25519::set(u.t, 0);

    for (i32 i = Curve25519::BITS - 1; i >= 0; i--) {
        u8 b = (scalar[i / 8] >> (i % 8)) & 1;

        // Compute U = 2 * U
        point_double(&u, &u);
        // Compute V = U + P
        point_add(&v, &u, point);

        // If b is set, then U = V
        Curve25519::select(u.x, u.x, v.x, b);
        Curve25519::select(u.y, u.y, v.y, b);
        Curve25519::select(u.z, u.z, v.z, b);
        Curve25519::select(u.t, u.t, v.t, b);
    }

    Curve25519::copy(result->x, u.x);
    Curve25519::copy(result->y, u.y);
    Curve25519::copy(result->z, u.z);
    Curve25519::copy(result->t, u.t);
}

// https://datatracker.ietf.org/doc/html/rfc8032#section-5.1.2
void Ed25519::encode_point(Ed25519Point* point, u8* data)
{
    // Retrieve affine representation
    Curve25519::modular_multiply_inverse(point->z, point->z);
    Curve25519::modular_multiply(point->x, point->x, point->z);
    Curve25519::modular_multiply(point->y, point->y, point->z);
    Curve25519::set(point->z, 1);
    Curve25519::modular_multiply(point->t, point->x, point->y);

    // First, encode the y-coordinate (in the range 0 <= y < p) as a little-endian string of 32 octets.
    // The most significant bit of the final octet is always zero.
    Curve25519::export_state(point->y, data);

    // To form the encoding of the point [s]B,
    // copy the least significant bit of the x coordinate to the most significant bit of the final octet.
    data[31] |= (point->x[0] & 1) << 7;
}

void Ed25519::barrett_reduce(u8* result, u8 const* input)
{
    // Barrett reduction b = 2^8 && k = 32
    u8 is_negative;
    u8 u[33];
    u8 v[33];

    multiply(NULL, u, input + 31, Curve25519::BARRETT_REDUCTION_QUOTIENT, 33);
    multiply(v, NULL, u, Curve25519::BASE_POINT_L_ORDER, 33);

    subtract(u, input, v, 33);

    is_negative = subtract(v, u, Curve25519::BASE_POINT_L_ORDER, 33);
    select(u, v, u, is_negative, 33);
    is_negative = subtract(v, u, Curve25519::BASE_POINT_L_ORDER, 33);
    select(u, v, u, is_negative, 33);

    copy(result, u, 32);
}

void Ed25519::multiply(u8* result_low, u8* result_high, u8 const* a, u8 const* b, u8 n)
{
    // Comba's algorithm
    u32 temp = 0;
    for (u32 i = 0; i < n; i++) {
        for (u32 j = 0; j <= i; j++) {
            temp += (uint16_t)a[j] * b[i - j];
        }

        if (result_low != NULL) {
            result_low[i] = temp & 0xFF;
        }

        temp >>= 8;
    }

    if (result_high != NULL) {
        for (u32 i = n; i < (2 * n); i++) {
            for (u32 j = i + 1 - n; j < n; j++) {
                temp += (uint16_t)a[j] * b[i - j];
            }

            result_high[i - n] = temp & 0xFF;
            temp >>= 8;
        }
    }
}

void Ed25519::add(u8* result, u8 const* a, u8 const* b, u8 n)
{
    // Compute R = A + B
    u16 temp = 0;
    for (u8 i = 0; i < n; i++) {
        temp += a[i];
        temp += b[i];
        result[i] = temp & 0xFF;
        temp >>= 8;
    }
}

u8 Ed25519::subtract(u8* result, u8 const* a, u8 const* b, u8 n)
{
    i16 temp = 0;

    // Compute R = A - B
    for (i8 i = 0; i < n; i++) {
        temp += a[i];
        temp -= b[i];
        result[i] = temp & 0xFF;
        temp >>= 8;
    }

    // Return 1 if the result of the subtraction is negative
    return temp & 1;
}

void Ed25519::select(u8* r, u8 const* a, u8 const* b, u8 c, u8 n)
{
    u8 mask = c - 1;
    for (u8 i = 0; i < n; i++) {
        r[i] = (a[i] & mask) | (b[i] & ~mask);
    }
}

// https://datatracker.ietf.org/doc/html/rfc8032#section-5.1.3
u32 Ed25519::decode_point(Ed25519Point* point, u8 const* data)
{
    u32 u[8];
    u32 v[8];
    u32 ret;
    u64 temp = 19;

    // 1. First, interpret the string as an integer in little-endian representation.
    // Bit 255 of this number is the least significant bit of the x-coordinate and denote this value x_0.
    u8 x0 = data[31] >> 7;

    // The y-coordinate is recovered simply by clearing this bit.
    Curve25519::import_state(point->y, data);
    point->y[7] &= 0x7FFFFFFF;

    // Compute U = Y + 19
    for (u32 i = 0; i < 8; i++) {
        temp += point->y[i];
        u[i] = temp & 0xFFFFFFFF;
        temp >>= 32;
    }

    // If the resulting value is >= p, decoding fails.
    ret = (u[7] >> 31) & 1;

    // 2. To recover the x-coordinate, the curve equation implies x^2 = (y^2 - 1) / (d y^2 + 1) (mod p).
    // The denominator is always non-zero mod p.
    // Let u = y^2 - 1 and v = d * y^2 + 1
    Curve25519::modular_square(v, point->y);
    Curve25519::modular_subtract_single(u, v, 1);
    Curve25519::modular_multiply(v, v, Curve25519::CURVE_D);
    Curve25519::modular_add_single(v, v, 1);

    // 3. Compute u = sqrt(u / v)
    ret |= Curve25519::modular_square_root(u, u, v);

    // If x = 0, and x_0 = 1, decoding fails.
    ret |= (Curve25519::compare(u, Curve25519::ZERO) ^ 1) & x0;

    // 4.  Finally, use the x_0 bit to select the right square root.
    Curve25519::modular_subtract(v, Curve25519::ZERO, u);
    Curve25519::select(point->x, u, v, (x0 ^ u[0]) & 1);
    Curve25519::set(point->z, 1);
    Curve25519::modular_multiply(point->t, point->x, point->y);

    // Return 0 if the point has been successfully decoded, else 1
    return ret;
}

void Ed25519::point_add(Ed25519Point* result, Ed25519Point const* p, Ed25519Point const* q)
{
    // Compute R = P + Q
    Curve25519::modular_add(c, p->y, p->x);
    Curve25519::modular_add(d, q->y, q->x);
    Curve25519::modular_multiply(a, c, d);
    Curve25519::modular_subtract(c, p->y, p->x);
    Curve25519::modular_subtract(d, q->y, q->x);
    Curve25519::modular_multiply(b, c, d);
    Curve25519::modular_multiply(c, p->z, q->z);
    Curve25519::modular_add(c, c, c);
    Curve25519::modular_multiply(d, p->t, q->t);
    Curve25519::modular_multiply(d, d, Curve25519::CURVE_D_2);
    Curve25519::modular_add(e, a, b);
    Curve25519::modular_subtract(f, a, b);
    Curve25519::modular_add(g, c, d);
    Curve25519::modular_subtract(h, c, d);
    Curve25519::modular_multiply(result->x, f, h);
    Curve25519::modular_multiply(result->y, e, g);
    Curve25519::modular_multiply(result->z, g, h);
    Curve25519::modular_multiply(result->t, e, f);
}

u8 Ed25519::compare(u8 const* a, u8 const* b, u8 n)
{
    u8 mask = 0;
    for (u32 i = 0; i < n; i++) {
        mask |= a[i] ^ b[i];
    }

    // Return 0 if A = B, else 1
    return ((u8)(mask | (~mask + 1))) >> 7;
}

void Ed25519::copy(u8* a, u8 const* b, u32 n)
{
    for (u32 i = 0; i < n; i++) {
        a[i] = b[i];
    }
}

}