/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements. See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License. You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.kafka.common.security.scram.internals;

import org.apache.kafka.common.security.scram.internals.ScramMessages.AbstractScramMessage;
import org.apache.kafka.common.security.scram.internals.ScramMessages.ClientFinalMessage;
import org.apache.kafka.common.security.scram.internals.ScramMessages.ClientFirstMessage;
import org.apache.kafka.common.security.scram.internals.ScramMessages.ServerFinalMessage;
import org.apache.kafka.common.security.scram.internals.ScramMessages.ServerFirstMessage;

import org.junit.jupiter.api.BeforeEach;
import org.junit.jupiter.api.Test;

import java.nio.charset.StandardCharsets;
import java.util.Base64;
import java.util.Collections;

import javax.security.sasl.SaslException;

import static org.junit.jupiter.api.Assertions.assertArrayEquals;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertNull;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.junit.jupiter.api.Assertions.fail;

public class ScramMessagesTest {

    private static final String[] VALID_EXTENSIONS = {
        "ext=val1",
        "anotherext=name1=value1 name2=another test value \"\'!$[]()",
        "first=val1,second=name1 = value ,third=123"
    };
    private static final String[] INVALID_EXTENSIONS = {
        "ext1=value",
        "ext",
        "ext=value1,value2",
        "ext=,",
        "ext =value"
    };

    private static final String[] VALID_RESERVED = {
        "m=reserved-value",
        "m=name1=value1 name2=another test value \"\'!$[]()"
    };
    private static final String[] INVALID_RESERVED = {
        "m",
        "m=name,value",
        "m=,"
    };

    private ScramFormatter formatter;

    @BeforeEach
    public void setUp() throws Exception {
        formatter  = new ScramFormatter(ScramMechanism.SCRAM_SHA_256);
    }

    @Test
    public void validClientFirstMessage() throws SaslException {
        String nonce = formatter.secureRandomString();
        ClientFirstMessage m = new ClientFirstMessage("someuser", nonce, Collections.emptyMap());
        checkClientFirstMessage(m, "someuser", nonce, "");

        // Default format used by Kafka client: only user and nonce are specified
        String str = String.format("n,,n=testuser,r=%s", nonce);
        m = createScramMessage(ClientFirstMessage.class, str);
        checkClientFirstMessage(m, "testuser", nonce, "");
        m = new ClientFirstMessage(m.toBytes());
        checkClientFirstMessage(m, "testuser", nonce, "");

        // Username containing comma, encoded as =2C
        str = String.format("n,,n=test=2Cuser,r=%s", nonce);
        m = createScramMessage(ClientFirstMessage.class, str);
        checkClientFirstMessage(m, "test=2Cuser", nonce, "");
        assertEquals("test,user", ScramFormatter.username(m.saslName()));

        // Username containing equals, encoded as =3D
        str = String.format("n,,n=test=3Duser,r=%s", nonce);
        m = createScramMessage(ClientFirstMessage.class, str);
        checkClientFirstMessage(m, "test=3Duser", nonce, "");
        assertEquals("test=user", ScramFormatter.username(m.saslName()));

        // Optional authorization id specified
        str = String.format("n,a=testauthzid,n=testuser,r=%s", nonce);
        checkClientFirstMessage(createScramMessage(ClientFirstMessage.class, str), "testuser", nonce, "testauthzid");

        // Optional reserved value specified
        for (String reserved : VALID_RESERVED) {
            str = String.format("n,,%s,n=testuser,r=%s", reserved, nonce);
            checkClientFirstMessage(createScramMessage(ClientFirstMessage.class, str), "testuser", nonce, "");
        }

        // Optional extension specified
        for (String extension : VALID_EXTENSIONS) {
            str = String.format("n,,n=testuser,r=%s,%s", nonce, extension);
            checkClientFirstMessage(createScramMessage(ClientFirstMessage.class, str), "testuser", nonce, "");
        }

        //optional tokenauth specified as extensions
        str = String.format("n,,n=testuser,r=%s,%s", nonce, "tokenauth=true");
        m = createScramMessage(ClientFirstMessage.class, str);
        assertTrue(m.extensions().tokenAuthenticated(), "Token authentication not set from extensions");
    }

    @Test
    public void invalidClientFirstMessage() {
        String nonce = formatter.secureRandomString();
        // Invalid entry in gs2-header
        String invalid = String.format("n,x=something,n=testuser,r=%s", nonce);
        checkInvalidScramMessage(ClientFirstMessage.class, invalid);

        // Invalid reserved entry
        for (String reserved : INVALID_RESERVED) {
            invalid = String.format("n,,%s,n=testuser,r=%s", reserved, nonce);
            checkInvalidScramMessage(ClientFirstMessage.class, invalid);
        }

        // Invalid extension
        for (String extension : INVALID_EXTENSIONS) {
            invalid = String.format("n,,n=testuser,r=%s,%s", nonce, extension);
            checkInvalidScramMessage(ClientFirstMessage.class, invalid);
        }
    }

    @Test
    public void validServerFirstMessage() throws SaslException {
        String clientNonce = formatter.secureRandomString();
        String serverNonce = formatter.secureRandomString();
        String nonce = clientNonce + serverNonce;
        String salt = randomBytesAsString();

        ServerFirstMessage m = new ServerFirstMessage(clientNonce, serverNonce, toBytes(salt), 8192);
        checkServerFirstMessage(m, nonce, salt, 8192);

        // Default format used by Kafka clients, only nonce, salt and iterations are specified
        String str = String.format("r=%s,s=%s,i=4096", nonce, salt);
        m = createScramMessage(ServerFirstMessage.class, str);
        checkServerFirstMessage(m, nonce, salt, 4096);
        m = new ServerFirstMessage(m.toBytes());
        checkServerFirstMessage(m, nonce, salt, 4096);

        // Optional reserved value
        for (String reserved : VALID_RESERVED) {
            str = String.format("%s,r=%s,s=%s,i=4096", reserved, nonce, salt);
            checkServerFirstMessage(createScramMessage(ServerFirstMessage.class, str), nonce, salt, 4096);
        }

        // Optional extension
        for (String extension : VALID_EXTENSIONS) {
            str = String.format("r=%s,s=%s,i=4096,%s", nonce, salt, extension);
            checkServerFirstMessage(createScramMessage(ServerFirstMessage.class, str), nonce, salt, 4096);
        }
    }

    @Test
    public void invalidServerFirstMessage() {
        String nonce = formatter.secureRandomString();
        String salt = randomBytesAsString();

        // Invalid iterations
        String invalid = String.format("r=%s,s=%s,i=0", nonce, salt);
        checkInvalidScramMessage(ServerFirstMessage.class, invalid);

        // Invalid salt
        invalid = String.format("r=%s,s=%s,i=4096", nonce, "=123");
        checkInvalidScramMessage(ServerFirstMessage.class, invalid);

        // Invalid format
        invalid = String.format("r=%s,invalid,s=%s,i=4096", nonce, salt);
        checkInvalidScramMessage(ServerFirstMessage.class, invalid);

        // Invalid reserved entry
        for (String reserved : INVALID_RESERVED) {
            invalid = String.format("%s,r=%s,s=%s,i=4096", reserved, nonce, salt);
            checkInvalidScramMessage(ServerFirstMessage.class, invalid);
        }

        // Invalid extension
        for (String extension : INVALID_EXTENSIONS) {
            invalid = String.format("r=%s,s=%s,i=4096,%s", nonce, salt, extension);
            checkInvalidScramMessage(ServerFirstMessage.class, invalid);
        }
    }

    @Test
    public void validClientFinalMessage() throws SaslException {
        String nonce = formatter.secureRandomString();
        String channelBinding = randomBytesAsString();
        String proof = randomBytesAsString();

        ClientFinalMessage m = new ClientFinalMessage(toBytes(channelBinding), nonce);
        assertNull(m.proof(), "Invalid proof");
        m.proof(toBytes(proof));
        checkClientFinalMessage(m, channelBinding, nonce, proof);

        // Default format used by Kafka client: channel-binding, nonce and proof are specified
        String str = String.format("c=%s,r=%s,p=%s", channelBinding, nonce, proof);
        m = createScramMessage(ClientFinalMessage.class, str);
        checkClientFinalMessage(m, channelBinding, nonce, proof);
        m = new ClientFinalMessage(m.toBytes());
        checkClientFinalMessage(m, channelBinding, nonce, proof);

        // Optional extension specified
        for (String extension : VALID_EXTENSIONS) {
            str = String.format("c=%s,r=%s,%s,p=%s", channelBinding, nonce, extension, proof);
            checkClientFinalMessage(createScramMessage(ClientFinalMessage.class, str), channelBinding, nonce, proof);
        }
    }

    @Test
    public void invalidClientFinalMessage() {
        String nonce = formatter.secureRandomString();
        String channelBinding = randomBytesAsString();
        String proof = randomBytesAsString();

        // Invalid channel binding
        String invalid = String.format("c=ab,r=%s,p=%s", nonce, proof);
        checkInvalidScramMessage(ClientFirstMessage.class, invalid);

        // Invalid proof
        invalid = String.format("c=%s,r=%s,p=123", channelBinding, nonce);
        checkInvalidScramMessage(ClientFirstMessage.class, invalid);

        // Invalid extensions
        for (String extension : INVALID_EXTENSIONS) {
            invalid = String.format("c=%s,r=%s,%s,p=%s", channelBinding, nonce, extension, proof);
            checkInvalidScramMessage(ClientFinalMessage.class, invalid);
        }
    }

    @Test
    public void validServerFinalMessage() throws SaslException {
        String serverSignature = randomBytesAsString();

        ServerFinalMessage m = new ServerFinalMessage("unknown-user", null);
        checkServerFinalMessage(m, "unknown-user", null);
        m = new ServerFinalMessage(null, toBytes(serverSignature));
        checkServerFinalMessage(m, null, serverSignature);

        // Default format used by Kafka clients for successful final message
        String str = String.format("v=%s", serverSignature);
        m = createScramMessage(ServerFinalMessage.class, str);
        checkServerFinalMessage(m, null, serverSignature);
        m = new ServerFinalMessage(m.toBytes());
        checkServerFinalMessage(m, null, serverSignature);

        // Default format used by Kafka clients for final message with error
        str = "e=other-error";
        m = createScramMessage(ServerFinalMessage.class, str);
        checkServerFinalMessage(m, "other-error", null);
        m = new ServerFinalMessage(m.toBytes());
        checkServerFinalMessage(m, "other-error", null);

        // Optional extension
        for (String extension : VALID_EXTENSIONS) {
            str = String.format("v=%s,%s", serverSignature, extension);
            checkServerFinalMessage(createScramMessage(ServerFinalMessage.class, str), null, serverSignature);
        }
    }

    @Test
    public void invalidServerFinalMessage() {
        String serverSignature = randomBytesAsString();

        // Invalid error
        String invalid = "e=error1,error2";
        checkInvalidScramMessage(ServerFinalMessage.class, invalid);

        // Invalid server signature
        invalid = "v=1=23";
        checkInvalidScramMessage(ServerFinalMessage.class, invalid);

        // Invalid extensions
        for (String extension : INVALID_EXTENSIONS) {
            invalid = String.format("v=%s,%s", serverSignature, extension);
            checkInvalidScramMessage(ServerFinalMessage.class, invalid);

            invalid = String.format("e=unknown-user,%s", extension);
            checkInvalidScramMessage(ServerFinalMessage.class, invalid);
        }
    }

    private String randomBytesAsString() {
        return Base64.getEncoder().encodeToString(formatter.secureRandomBytes());
    }

    private byte[] toBytes(String base64Str) {
        return Base64.getDecoder().decode(base64Str);
    }

    private void checkClientFirstMessage(ClientFirstMessage message, String saslName, String nonce, String authzid) {
        assertEquals(saslName, message.saslName());
        assertEquals(nonce, message.nonce());
        assertEquals(authzid, message.authorizationId());
    }

    private void checkServerFirstMessage(ServerFirstMessage message, String nonce, String salt, int iterations) {
        assertEquals(nonce, message.nonce());
        assertArrayEquals(Base64.getDecoder().decode(salt), message.salt());
        assertEquals(iterations, message.iterations());
    }

    private void checkClientFinalMessage(ClientFinalMessage message, String channelBinding, String nonce, String proof) {
        assertArrayEquals(Base64.getDecoder().decode(channelBinding), message.channelBinding());
        assertEquals(nonce, message.nonce());
        assertArrayEquals(Base64.getDecoder().decode(proof), message.proof());
    }

    private void checkServerFinalMessage(ServerFinalMessage message, String error, String serverSignature) {
        assertEquals(error, message.error());
        if (serverSignature == null)
            assertNull(message.serverSignature(), "Unexpected server signature");
        else
            assertArrayEquals(Base64.getDecoder().decode(serverSignature), message.serverSignature());
    }

    @SuppressWarnings("unchecked")
    private <T extends AbstractScramMessage> T createScramMessage(Class<T> clazz, String message) throws SaslException {
        byte[] bytes = message.getBytes(StandardCharsets.UTF_8);
        if (clazz == ClientFirstMessage.class)
            return (T) new ClientFirstMessage(bytes);
        else if (clazz == ServerFirstMessage.class)
            return (T) new ServerFirstMessage(bytes);
        else if (clazz == ClientFinalMessage.class)
            return (T) new ClientFinalMessage(bytes);
        else if (clazz == ServerFinalMessage.class)
            return (T) new ServerFinalMessage(bytes);
        else
            throw new IllegalArgumentException("Unknown message type: " + clazz);
    }

    private <T extends AbstractScramMessage> void checkInvalidScramMessage(Class<T> clazz, String message) {
        try {
            createScramMessage(clazz, message);
            fail("Exception not throws for invalid message of type " + clazz + " : " + message);
        } catch (SaslException e) {
            // Expected exception
        }
    }
}
