package org.bouncycastle.crypto.agreement.owl;

import java.math.BigInteger;
import java.security.SecureRandom;

import org.bouncycastle.crypto.CryptoException;
import org.bouncycastle.crypto.CryptoServicesRegistrar;
import org.bouncycastle.crypto.Digest;
import org.bouncycastle.crypto.digests.SHA256Digest;
import org.bouncycastle.math.ec.ECCurve;
import org.bouncycastle.math.ec.ECPoint;
import org.bouncycastle.util.Arrays;
import org.bouncycastle.util.Exceptions;

/**
 * A client in the Elliptic Curve Owl key exchange protocol.
 * <p>
 * Owl is an augmented password-authenticated key exchange (PAKE) protocol,
 * defined by Feng Hao, Samiran Bag, Liqun Chen, and Paul C. van Oorschot
 * in the paper <a href="https://eprint.iacr.org/2023/768.pdf">"Owl: An Augmented Password-Authenticated Key Exchange Scheme"</a>.
 * Owl builds on the same idea as J-PAKE (using Schnorr zero-knowledge proofs to enforce participants to follow the protocol specification honestly),
 * but it is augmented to provide additional protection against server compromise.
 * While J-PAKE is symmetric, Owl is asymmetric. Like J-PAKE, Owl can be implemented in either elliptic curve (EC) or finite field (FF) settings.
 * This implementation is done in the elliptic curve setting.
 * Any elliptic curve that is suitable for cryptography can be used for Owl (same for J-PAKE).
 * <p>
 * <b>Owl is an academic protocol — not (yet) an IETF standard.</b> Callers who need a standardised
 * PAKE should prefer the existing J-PAKE in {@link org.bouncycastle.crypto.agreement.jpake} /
 * {@code org.bouncycastle.crypto.agreement.ecjpake}.
 * <p>
 * In Owl, there is one client and one server communicating between each other.
 * An instance of {@link OwlServer} represents one server, and
 * an instance of {@link OwlClient} represents one client.
 * These together make up the main machine through which the protocol is facilitated.
 * <p>
 * There are two distinct phases that can be taken in Owl: user registration - where the client registers
 * as a new user on the server; login - where an existing user (client) attempts to log in and establish a shared session key with the server
 * based on password authentication. Using the session key, the user (client) can perform further actions (e.g., password update) over a secure channel,
 * but these actions are outside the scope of this key exchange program.
 * <p>
 * The user registration phase involves only one pass of communication from the client to the server. It's assumed that the following communication is done over a secure channel (e.g., using an out-of-band method).
 * <ol>
 * <li>{@link OwlClientRegistration#initiateUserRegistration()} - The client sends payload to the server. 
 * </ol>
 * After receiving the payload, the server calls {@link OwlServerRegistration#registerUseronServer(OwlInitialRegistration)} to 
 * generate a user credential file {@link OwlFinishRegistration} and safely store it on the server.
 * 
 * The login phase involves three passes of communication between the client and the server. 
 * The following communications do not have to be sent via secure channels.
 * Call the following methods in this order, with the client initiating every key exchange process. 
 * <ol>
 * <li> {@link OwlClient#authenticationInitiate()} - The client sends payload to the server </li>
 * <li> {@link OwlServer#authenticationServerResponse(OwlAuthenticationInitiate, OwlFinishRegistration)} - The server validates the payload received from the client, retrieves the user credentials stored in the server from user registration, and sends the payload generated to the client.</li>
 * <li> {@link OwlClient#authenticationFinish(OwlAuthenticationServerResponse)} - The client validates the payload from the server, and sends next payload generated by this function back to the server. </li>
 * </ol>
 * <p>After the third pass, the client has completed client-to-server authentication. At this point, both the client and the server can compute a shared key. They call the following methods to compute a raw keying material and do key confirmation.  
 * <ul>
 * <li> {@link OwlServer#authenticationServerEnd(OwlAuthenticationFinish)} - The server validates the payload from the client in the third pass.</li>
 * <li> {@link #calculateKeyingMaterial()} - Both the client and the server execute this to generate a common key material. </li>
 * <li>{@link #initiateKeyConfirmation(BigInteger)} - Either the client or the server sends the payload to the other participant for explicit key confirmation in this exchange. This is called by both client and server.</li>
 * <li>{@link validateKeyConfirmation(OwlKeyConfirmation, BigInteger)} - Either the client or the server validates the payload received from the other participant to check if the same key has been derived. This is called by both client and server.</li>
 * </ul>
 * <p>
 * Each side should derive a session key from the keying material returned by {@link #calculateKeyingMaterial()}.
 * The caller is responsible for deriving the session key using a secure key derivation function (KDF).
 * <p>
 * The explicit key confirmation process is optional but highly recommended. It does not affect the round efficiency and adds a negligible computational cost. The client-to-server key confirmation string
 * can be piggybacked in the third pass along with {@link OwlAuthenticationFinish}. The server-to-client key confirmation string can be sent in the next pass together with encrypted data.
 * If you do not execute key confirmation, then there is no assurance that both client and server have actually derived the same key, and the ensuing secure communication may fail.
 * If the key confirmation succeeds, then the keys are guaranteed to be the same on both sides.
 * <p>
 * The key confirmation process is implemented as specified in
 * <a href="https://nvlpubs.nist.gov/nistpubs/SpecialPublications/NIST.SP.800-56Ar3.pdf">NIST SP 800-56A Revision 3</a>,
 * Section 5.9.1 Unilateral Key Confirmation for Key Agreement Schemes.
 * <p>
 * This class is stateful and NOT threadsafe.
 * Each instance should only be used for ONE complete Owl key exchange
 * (i.e. a new {@link OwlServer} and {@link OwlClient} should be constructed for each new Owl key exchange).
 */
public class OwlClient
{
    /*
     * Possible internal states.  Used for state checking.
     */
    public static final int STATE_INITIALISED = 0;
    public static final int STATE_LOGIN_INITIALISED = 10;
    public static final int STATE_LOGIN_FINISHED = 20;
    public static final int STATE_KEY_CALCULATED = 30;
    public static final int STATE_KC_INITIALISED = 40;
    public static final int STATE_KC_VALIDATED = 50;

    /**
     * Unique identifier of this client.
     * The client and server in the exchange must NOT share the same id.
     * clientId is the same as the username in this PAKE
     */
    private final String clientId;

    /**
     * Shared secret.  This only contains the secret between construction
     * and the call to {@link #calculateKeyingMaterial()}.
     * <p>
     * i.e. When {@link #calculateKeyingMaterial()} is called, this buffer overwritten with 0's,
     * and the field is set to null.
     * </p>
     */
    private char[] password;

    /**
     * Digest to use during calculations.
     */
    private final Digest digest;

    /**
     * Source of secure random data.
     */
    private final SecureRandom random;

    /**
     * The serverId of the server in this exchange.
     */
    private String serverId; 

    private ECCurve.AbstractFp ecCurve;
    private BigInteger q;
    private BigInteger h;
    private BigInteger n;
    private ECPoint g;
    /**
     * Client's x1
     */
    private BigInteger x1;
    /**
     * Client's x2.
     */
    private BigInteger x2;
    /**
     Client's gx1.
     */
    private ECPoint gx1;
    /**
     * Client's gx2.
     */
    private ECPoint gx2;
    /**
     * Server's gx3.
     */
    private ECPoint gx3;
    /**
     * Server's gx4.
     */
    private ECPoint gx4;
    /**
     * Client's user specified secret t = H(username||password) mod n
     */
    private BigInteger t;
    /**
     * Shared secret used for authentication pi = H(t) mod n
     */
    private BigInteger pi;
    /**
     * The current state.
     * See the <code>STATE_*</code> constants for possible values.
     */
    private int state;
    /**
     *  The raw key K used to calculate a session key.
     */
    private ECPoint rawKey;
    /**
     * Client's alpha
     */
    private ECPoint alpha;
    /**
     * ECSchnorrZKP knowledge proof for x1, using {@link ECSchnorrZKP}
     */
    private ECSchnorrZKP knowledgeProofForX1;
    /**
     * ECSchnorrZKP knowledge proof for x2, using {@link ECSchnorrZKP}
     */
    private ECSchnorrZKP knowledgeProofForX2;
    /**
     * ECSchnorrZKP knowledge proof for x3, using {@link ECSchnorrZKP}
     */
    private ECSchnorrZKP knowledgeProofForX3;
    /**
     * ECSchnorrZKP knowledge proof for x4, using {@link ECSchnorrZKP}
     */
    private ECSchnorrZKP knowledgeProofForX4;

    /**
     * Convenience constructor for a new {@link OwlClient} that uses
     * the {@link OwlCurves#NIST_P256} elliptic curve,
     * a SHA-256 digest, and a default {@link SecureRandom} implementation.
     * <p>
     * After construction, the {@link #getState() state} will be  {@link #STATE_INITIALISED}.
     *
     * @param clientId      unique identifier of this client.
     *                      The server and client in the exchange must NOT share the same id.
     * @param password      shared secret.
     *                      A defensive copy of this array is made (and cleared once {@link #calculateKeyingMaterial()} is called).
     *                      Caller should clear the input password as soon as possible.
     * @throws NullPointerException     if any argument is null
     * @throws IllegalArgumentException if password is empty
     */
    public OwlClient(
        String clientId,
        char[] password)
    {
        this(
            clientId,
            password,
            OwlCurves.NIST_P256);
    }

    /**
     * Convenience constructor for a new {@link OwlClient} that uses
     * a SHA-256 digest and a default {@link SecureRandom} implementation.
     * <p>
     * After construction, the {@link #getState() state} will be  {@link #STATE_INITIALISED}.
     *
     * @param clientId      unique identifier of this client.
     *                      The server and client in the exchange must NOT share the same id.     
     * @param password      shared secret.
     *                      A defensive copy of this array is made (and cleared once {@link #calculateKeyingMaterial()} is called).
     *                      Caller should clear the input password as soon as possible.
     * @param curve         elliptic curve
     *                      See {@link OwlCurves} for standard curves.
     * @throws NullPointerException     if any argument is null
     * @throws IllegalArgumentException if password is empty
     */
    public OwlClient(
        String clientId,
        char[] password,
        OwlCurve curve)
    {
        this(
            clientId,
            password,
            curve,
            SHA256Digest.newInstance(),
            CryptoServicesRegistrar.getSecureRandom());
    }

    /**
     * Construct a new {@link OwlClient}.
     * <p>
     * After construction, the {@link #getState() state} will be  {@link #STATE_INITIALISED}.
     *
     * @param clientId      unique identifier of this client.
     *                      The server and client in the exchange must NOT share the same id.
     * @param password      shared secret.
     *                      A defensive copy of this array is made (and cleared once {@link #calculateKeyingMaterial()} is called).
     *                      Caller should clear the input password as soon as possible.
     * @param curve         elliptic curve.
     *                      See {@link OwlCurves} for standard curves
     * @param digest        digest to use during zero knowledge proofs and key confirmation (SHA-256 or stronger preferred)
     * @param random        source of secure random data for x1 and x2, and for the zero knowledge proofs
     * @throws NullPointerException     if any argument is null
     * @throws IllegalArgumentException if password is empty
     */
    public OwlClient(
        String clientId,
        char[] password,
        OwlCurve curve,
        Digest digest,
        SecureRandom random)
    {
        OwlUtil.validateNotNull(clientId, "clientId");
        OwlUtil.validateNotNull(password, "password");
        OwlUtil.validateNotNull(curve, "curve params");
        OwlUtil.validateNotNull(digest, "digest");
        OwlUtil.validateNotNull(random, "random");
        if (password.length == 0)
        {
            throw new IllegalArgumentException("Password must not be empty.");
        }

        this.clientId = clientId;

        /*
         * Create a defensive copy so as to fully encapsulate the password.
         *
         * This array will contain the password for the lifetime of this
         * client BEFORE {@link #calculateKeyingMaterial()} is called.
         *
         * i.e. When {@link #calculateKeyingMaterial()} is called, the array will be cleared
         * in order to remove the password from memory.
         *
         * The caller is responsible for clearing the original password array
         * given as input to this constructor.
         */
        this.password = Arrays.copyOf(password, password.length);

        this.ecCurve = curve.getCurve();
        this.g = curve.getG();
        this.h = curve.getH();
        this.n = curve.getN();
        this.q = curve.getQ();

        this.digest = digest;
        this.random = random;

        this.state = STATE_INITIALISED;
    }

    /**
     * Gets the current state of this client.
     * See the <code>STATE_*</code> constants for possible values.
     * 
     * @return the state of the client
     */
    public int getState()
    {
        return this.state;
    }
    /**
     * Creates and returns the payload to send to the server as part of the first pass of the protocol.
     * <p>
     * Must be called prior to {@link #authenticationFinish(OwlAuthenticationServerResponse)}
     * After execution, the {@link #getState() state} will be  {@link #STATE_LOGIN_INITIALISED}.
     * 
     * @return {@link OwlAuthenticationInitiate}
     * @throws IllegalStateException if called multiple times.
     */
    public OwlAuthenticationInitiate authenticationInitiate()
    {
        if (this.state >= STATE_LOGIN_INITIALISED)
        {
            throw new IllegalStateException("Login already initiated by " + clientId);
        }
        this.t = calculateT();

        this.pi = calculatePi();

        this.x1 = OwlUtil.generateX1(n, random);
        this.x2 = OwlUtil.generateX1(n, random);

        this.gx1 = OwlUtil.calculateGx(g, x1);
        this.gx2 = OwlUtil.calculateGx(g, x2);

        this.knowledgeProofForX1 = OwlUtil.calculateZeroknowledgeProof(g, n, x1, gx1, digest, clientId, random);
        this.knowledgeProofForX2 = OwlUtil.calculateZeroknowledgeProof(g, n, x2, gx2, digest, clientId, random);

        this.state = STATE_LOGIN_INITIALISED;

        return new OwlAuthenticationInitiate(clientId, gx1, gx2, knowledgeProofForX1, knowledgeProofForX2);
    }

    /**
     * Finalises the login authentication protocol by creating and sending the final payload to the server.
     * Validates the payload sent by the {@link OwlServer#authenticationServerResponse(OwlAuthenticationInitiate, OwlFinishRegistration)} after login initilisation.
     * <p>
     * Must be called prior to {@link #calculateKeyingMaterial()}.
     * <p>
     * After execution, the {@link #getState() state} will be  {@link #STATE_LOGIN_FINISHED}.
     * 
     * @param authenticationServerResponse  The payload sent by {@link OwlServer#authenticationServerResponse(OwlAuthenticationInitiate, OwlFinishRegistration)} and to be validated.
     * 
     * @return {@link OwlAuthenticationFinish}
     * @throws CryptoException       if validation fails.
     * @throws IllegalStateException if called prior to {@link #authenticationInitiate()} or called multiple times.
     */
    public OwlAuthenticationFinish authenticationFinish(OwlAuthenticationServerResponse authenticationServerResponse) 
        throws CryptoException
    {
        if (this.state >= STATE_LOGIN_FINISHED)
        {
            throw new IllegalStateException("Login authentication already finished by: " + clientId);
        }
        if (this.state < STATE_LOGIN_INITIALISED)
        {
            throw new IllegalStateException("Must initialise login authentication before calling authentication finish for: " + clientId);
        }
        this.serverId = authenticationServerResponse.getServerId();
        this.gx3 = authenticationServerResponse.getGx3();
        this.gx4 = authenticationServerResponse.getGx4();
        ECPoint beta = authenticationServerResponse.getBeta();
        ECSchnorrZKP knowledgeProofForX3 = authenticationServerResponse.getKnowledgeProofForX3();
        ECSchnorrZKP knowledgeProofForX4 = authenticationServerResponse.getKnowledgeProofForX4();
        ECSchnorrZKP knowledgeProofForBeta = authenticationServerResponse.getKnowledgeProofForBeta();

        ECPoint betaG = OwlUtil.calculateGA(gx1, gx2, gx3);

        OwlUtil.validateParticipantIdsDiffer(clientId, authenticationServerResponse.getServerId());
        OwlUtil.validateZeroknowledgeProof(g, gx3, knowledgeProofForX3, q, n, ecCurve, h, authenticationServerResponse.getServerId(), digest);
        OwlUtil.validateZeroknowledgeProof(g, gx4, knowledgeProofForX4, q, n, ecCurve, h, authenticationServerResponse.getServerId(), digest);
        OwlUtil.validateZeroknowledgeProof(betaG, beta, knowledgeProofForBeta, q, n, ecCurve, h, authenticationServerResponse.getServerId(), digest);

        ECPoint alphaG = OwlUtil.calculateGA(gx1, gx3, gx4);
        BigInteger x2pi = OwlUtil.calculateX2s(n, x2, pi); 
        ECPoint alpha = OwlUtil.calculateA(alphaG, x2pi);

        ECSchnorrZKP knowledgeProofForAlpha  = OwlUtil.calculateZeroknowledgeProof(alphaG, n, x2pi, alpha, digest, clientId, random);

        this.rawKey = OwlUtil.calculateKeyingMaterial(gx4, x2, x2pi, beta);

        BigInteger hTranscript = OwlUtil.calculateTranscript(rawKey, clientId, gx1, gx2, knowledgeProofForX1, knowledgeProofForX2, serverId, gx3, gx4, 
            knowledgeProofForX3, knowledgeProofForX4, beta, knowledgeProofForBeta, alpha, knowledgeProofForAlpha, digest);

        BigInteger r = OwlUtil.calculateR(x1, t, hTranscript, n);

        this.state = STATE_LOGIN_FINISHED;

        return new OwlAuthenticationFinish(clientId, alpha, knowledgeProofForAlpha, r);
    }


    /**
     * Calculates and returns the key material.
     * A session key must be derived from this key material using a secure key derivation function (KDF).
     * The KDF used to derive the key is handled externally (i.e. not by {@link OwlClient}).
     * <p>
     * The keying material will be identical for client and server if and only if
     * the login password is the same as the password stored by the server.  i.e. If the client and  
     * server do not share the same password, then each  will derive a different key.
     * Therefore, if you immediately start using a key derived from
     * the keying material, then you must handle detection of incorrect keys.
     * If you want to handle this detection explicitly, you can perform explicit
     * key confirmation.  See {@link OwlClient} for details on how to execute
     * key confirmation.
     * <p>
     * If the passwords used for registration and login are different then this will be caught
     * when validating r during {@link OwlServer#authenticationServerEnd(OwlAuthenticationFinish)}.
     * <p>
     * {@link #authenticationFinish(OwlAuthenticationServerResponse)} must be called prior to this method.
     * <p>
     * As a side effect, the internal {@link #password} array is cleared, since it is no longer needed.
     * <p>
     * After execution, the {@link #getState() state} will be  {@link #STATE_KEY_CALCULATED}.
     *
     * @return raw key material
     * @throws IllegalStateException if called prior to {@link #authenticationFinish(OwlAuthenticationServerResponse)},
     *                               or if called multiple times.
     */
    public BigInteger calculateKeyingMaterial()
    {
        if (this.state >= STATE_KEY_CALCULATED)
        {
            throw new IllegalStateException("Key already calculated for " + clientId);
        }
        if (this.state < STATE_LOGIN_FINISHED)
        {
            throw new IllegalStateException("Login authentication must be finished prior to creating key for " + clientId);
        }

        /*
         * Clear the password array from memory, since we don't need it anymore.
         *
         * Also set the field to null as a flag to indicate that the key has already been calculated.
         */
        Arrays.fill(password, (char)0);
        this.password = null;
        
        BigInteger keyingMaterial = rawKey.normalize().getAffineXCoord().toBigInteger();
        /*
         * Clear the ephemeral private key fields as well.
         * Note that we're relying on the garbage collector to do its job to clean these up.
         * The old objects will hang around in memory until the garbage collector destroys them.
         *
         * If the ephemeral private keys x1 and x2 are leaked,
         * the attacker might be able to brute-force the password.
         */
        this.x1 = null;
        this.x2 = null;
        this.t = null;
        this.pi = null;
        this.rawKey = null;

        /*
         * Do not clear gx* yet, since those are needed by key confirmation.
         */
        this.state = STATE_KEY_CALCULATED;

        return keyingMaterial;
    }

    /**
     * Creates and returns the payload to send to the server as part of Key Confirmation.
     * <p>
     * See {@link OwlClient} for more details on Key Confirmation.
     * <p>
     * After execution, the {@link #getState() state} will be {@link #STATE_KC_INITIALISED}.
     *
     * @return {@link OwlKeyConfirmation}
     * @param keyingMaterial            The keying material as returned from {@link #calculateKeyingMaterial()}.
     * @throws IllegalStateException    if called prior to {@link #calculateKeyingMaterial()}, or multiple times
     */
    public OwlKeyConfirmation initiateKeyConfirmation(BigInteger keyingMaterial)
    {
        if (this.state >= STATE_KC_INITIALISED)
        {
            throw new IllegalStateException("Key Confirmation already initiated for " + this.clientId);
        }
        if (this.state < STATE_KEY_CALCULATED)
        {
            throw new IllegalStateException("Keying material must be calculated prior to initialising key confirmation for " + this.clientId);
        }

        BigInteger macTag = OwlUtil.calculateMacTag(
            this.clientId,
            this.serverId,
            this.gx1,
            this.gx2,
            this.gx3,
            this.gx4,
            keyingMaterial,
            this.digest);

        this.state = STATE_KC_INITIALISED;

        return new OwlKeyConfirmation(clientId, macTag);
    }

    /**
     * Validates the key confirmation payload received by the server.
     * <p>
     * See {@link OwlClient} for more details on Key Confirmation.
     * <p>
     * After execution, the {@link #getState() state} will be {@link #STATE_KC_VALIDATED}.
     *
     * @param keyConfirmationPayload    The key confirmation payload received from the other client.
     * @param keyingMaterial            The keying material as returned from {@link #calculateKeyingMaterial()}.
     * @throws CryptoException          if validation fails.
     * @throws IllegalStateException    if called prior to {@link #calculateKeyingMaterial()}, or multiple times
     */
    public void validateKeyConfirmation(OwlKeyConfirmation keyConfirmationPayload, BigInteger keyingMaterial)
        throws CryptoException
    {
        if (this.state >= STATE_KC_VALIDATED)
        {
            throw new IllegalStateException("Validation already attempted for this payload for" + clientId);
        }
        if (this.state < STATE_KEY_CALCULATED)
        {
            throw new IllegalStateException("Keying material must be calculated prior to validating this payload for " + this.clientId);
        }
        OwlUtil.validateParticipantIdsDiffer(clientId, keyConfirmationPayload.getId());
        OwlUtil.validateParticipantIdsEqual(this.serverId, keyConfirmationPayload.getId());

        OwlUtil.validateMacTag(
            this.clientId,
            this.serverId,
            this.gx1,
            this.gx2,
            this.gx3,
            this.gx4,
            keyingMaterial,
            this.digest,
            keyConfirmationPayload.getMacTag());


        /*
         * Clear the rest of the fields.
         */
        this.gx1 = null;
        this.gx2 = null;
        this.gx3 = null;
        this.gx4 = null;

        this.state = STATE_KC_VALIDATED;
    }

    private BigInteger calculateT()
    {
        try 
        {
            // t = H(username||password). Prepend each item with its byte length (int) to set clear boundary
            return OwlUtil.calculateT(n, 
                    String.valueOf(clientId.getBytes().length) + clientId + 
                    String.valueOf(password.length) + new String(password), digest);
        } 
        catch (CryptoException e)
        {
            throw Exceptions.illegalStateException(e.getMessage(), e);
        }
    }

    private BigInteger calculatePi()
    {
        try 
        {
            return OwlUtil.calculatePi(n, t, digest);
        }
        catch (CryptoException e)
        {
            throw Exceptions.illegalStateException(e.getMessage(), e);
        }
    }
}
