/*
 * 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.sshd.client.kex;

import java.util.Collections;

import org.apache.sshd.AbstractContainerTestBase;
import org.apache.sshd.client.ClientBuilder;
import org.apache.sshd.client.SshClient;
import org.apache.sshd.client.future.AuthFuture;
import org.apache.sshd.client.session.ClientSession;
import org.apache.sshd.common.kex.BuiltinDHFactories;
import org.apache.sshd.common.keyprovider.FileKeyPairProvider;
import org.apache.sshd.common.util.security.KEM;
import org.apache.sshd.common.util.security.SecurityUtils;
import org.apache.sshd.util.test.CommonTestSupportUtils;
import org.junit.jupiter.api.Assumptions;
import org.junit.jupiter.api.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testcontainers.containers.GenericContainer;
import org.testcontainers.containers.output.Slf4jLogConsumer;
import org.testcontainers.containers.wait.strategy.Wait;
import org.testcontainers.images.builder.ImageFromDockerfile;
import org.testcontainers.junit.jupiter.Container;
import org.testcontainers.utility.MountableFile;

/**
 * Test ciphers against OpenSSH. Force resetting ciphers every time to verify that they are res-initialized correctly.
 *
 * @author <a href="mailto:dev@mina.apache.org">Apache MINA SSHD Project</a>
 */
class OpenSshMlKemTest extends AbstractContainerTestBase {

    private static final Logger LOG = LoggerFactory.getLogger(OpenSshMlKemTest.class);

    private static final String TEST_KEYS = "org/apache/sshd/testkeys/user";

    private static final String TEST_RESOURCES = "org/apache/sshd/docker";

    @Container
    static GenericContainer<?> sshdContainer = new GenericContainer<>(new ImageFromDockerfile()
            .withDockerfileFromBuilder(
                    // 3.21
                    builder -> builder.from("alpine@sha256:c3f8e73fdb79deaebaa2037150150191b9dcbfba68b4a46d70103204c53f4709") //
                            .run(discriminate()) //
                            .run("apk --update add openssh-server") // Installs OpenSSH 9.9
                            .run("ssh-keygen -A") // Generate multiple host keys
                            .run("adduser -D bob") // Add a user
                            .run("echo 'bob:passwordBob' | chpasswd") // Give it a password to unlock the user
                            .run("mkdir -p /home/bob/.ssh") // Create the SSH config directory
                            .entryPoint("/entrypoint.sh") // Sets bob as owner of anything under /home/bob and launches sshd
                            .build())) //
            .withCopyFileToContainer(MountableFile.forClasspathResource(TEST_KEYS + "/user01_ed25519.pub"),
                    "/home/bob/.ssh/authorized_keys")
            // entrypoint must be executable. Spotbugs doesn't like 0777, so use hex
            .withCopyFileToContainer(
                    MountableFile.forClasspathResource(TEST_RESOURCES + "/sshd_bob.sh", 0x1ff),
                    "/entrypoint.sh")
            .withExposedPorts(22) //
            .waitingFor(Wait.forLogMessage(".*Server listening on.*port 22.*\\n", 1)) //
            .withLogConsumer(new Slf4jLogConsumer(LOG));

    @Test
    void mlkem768x25519() throws Exception {
        Assumptions.assumeTrue(BuiltinDHFactories.mlkem768x25519.isSupported());
        LOG.info("Java {} using KEM {}", System.getProperty("java.version"), SecurityUtils.getKEM(KEM.ML_KEM_768));
        FileKeyPairProvider keyPairProvider = CommonTestSupportUtils
                .createTestKeyPairProvider(TEST_KEYS + "/user01_ed25519");
        SshClient client = setupTestClient();
        client.setKeyIdentityProvider(keyPairProvider);
        client.setKeyExchangeFactories(
                Collections.singletonList(ClientBuilder.DH2KEX.apply(BuiltinDHFactories.mlkem768x25519)));
        client.start();

        Integer actualPort = sshdContainer.getMappedPort(22);
        String actualHost = sshdContainer.getHost();
        try (ClientSession session = client.connect("bob", actualHost, actualPort).verify(CONNECT_TIMEOUT).getSession()) {
            AuthFuture authed = session.auth().verify(AUTH_TIMEOUT);
            assertTrue(authed.isDone() && authed.isSuccess());
        } finally {
            client.stop();
        }
    }
}
