#
# Copyright (c) 2021 the Hugging Face team.
#
# Licensed 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.#
import base64
import datetime
from ipaddress import IPv4Address, IPv6Address
from typing import List

import pytest
from cryptography.hazmat.primitives import hashes
from cryptography.hazmat.primitives.asymmetric import padding
from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey, RSAPublicKey
from fastapi import FastAPI, status
from httpx import AsyncClient

from app.api.dependencies import crypto
from app.models.experiment import ExperimentCreate, ExperimentCreatePublic, ExperimentPublic, ExperimentUpdate
from app.models.experiment_join import ExperimentJoinInput, ExperimentJoinOutput


# decorate all tests with @pytest.mark.asyncio
pytestmark = pytest.mark.asyncio


@pytest.fixture
def new_experiment():
    return ExperimentCreatePublic(organization_name="organization_a", model_name="model-a")


def encrypt_msg(msg: str, public_key: RSAPublicKey):

    ciphertext = public_key.encrypt(
        msg.encode(), padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA256()), algorithm=hashes.SHA256(), label=None)
    )
    return ciphertext


def decrypt_msg(ciphertext: str, private_key: RSAPrivateKey):

    plaintext = private_key.decrypt(
        ciphertext, padding.OAEP(mgf=padding.MGF1(algorithm=hashes.SHA256()), algorithm=hashes.SHA256(), label=None)
    )
    return plaintext.decode()


class TestExperimentsRoutes:
    async def test_routes_exist(self, app: FastAPI, client_wt_auth_user_1: AsyncClient) -> None:
        res = await client_wt_auth_user_1.post(app.url_path_for("experiments:create-experiment"), json={})
        assert res.status_code != status.HTTP_404_NOT_FOUND

    async def test_invalid_input_raises_error(self, app: FastAPI, client_wt_auth_user_1: AsyncClient) -> None:
        res = await client_wt_auth_user_1.post(app.url_path_for("experiments:create-experiment"), json={})
        assert res.status_code == status.HTTP_422_UNPROCESSABLE_ENTITY


class TestCreateExperiment:
    @pytest.mark.parametrize(
        "new_experiment",
        (
            (ExperimentCreatePublic(organization_name="organization_a", model_name="model-a")),
            (ExperimentCreatePublic(organization_name="Organization_a", model_name="modelA")),
            (ExperimentCreatePublic(organization_name="org_1", model_name="modelA")),
        ),
    )
    async def test_valid_input(
        self,
        app: FastAPI,
        client_wt_auth_user_1: AsyncClient,
        moonlanding_user_1,
        new_experiment: ExperimentCreatePublic,
    ) -> None:
        res = await client_wt_auth_user_1.post(
            app.url_path_for("experiments:create-experiment"),
            json={"new_experiment": new_experiment.dict()},
        )
        assert res.status_code == status.HTTP_201_CREATED

        created_experiment = ExperimentPublic(**res.json())
        assert created_experiment.organization_name == new_experiment.organization_name
        assert created_experiment.model_name == new_experiment.model_name
        assert created_experiment.creator == moonlanding_user_1.username

    @pytest.mark.parametrize(
        "invalid_payload, status_code",
        (
            (None, 422),
            ({}, 422),
        ),
    )
    async def test_invalid_input_raises_error(
        self, app: FastAPI, client_wt_auth_user_1: AsyncClient, invalid_payload: dict, status_code: int
    ) -> None:
        res = await client_wt_auth_user_1.post(
            app.url_path_for("experiments:create-experiment"), json={"new_experiment": invalid_payload}
        )
        assert res.status_code == status_code

    @pytest.mark.parametrize(
        "new_experiment, status_code",
        (
            (ExperimentCreatePublic(organization_name="org_2", model_name="model-a"), 401),
            (ExperimentCreatePublic(organization_name="fake_org", model_name="model-a"), 401),
        ),
    )
    async def test_not_admin_raises_error(
        self,
        app: FastAPI,
        client_wt_auth_user_1: AsyncClient,
        new_experiment: ExperimentCreatePublic,
        status_code: int,
    ) -> None:
        res = await client_wt_auth_user_1.post(
            app.url_path_for("experiments:create-experiment"), json={"new_experiment": new_experiment.dict()}
        )
        assert res.status_code == status_code

    async def test_unauthenticated_user_unable_to_create_experiment(
        self, app: FastAPI, client: AsyncClient, new_experiment: ExperimentCreate
    ) -> None:
        res = await client.post(
            app.url_path_for("experiments:create-experiment"),
            json={"new_experiment": new_experiment.dict()},
        )
        assert res.status_code == status.HTTP_401_UNAUTHORIZED


class TestGetExperimentById:
    async def test_get_experiment_by_id_valid_query_by_user_1(
        self,
        app: FastAPI,
        client_wt_auth_user_1: AsyncClient,
        test_experiment_1_created_by_user_1: ExperimentPublic,
    ) -> None:
        res = await client_wt_auth_user_1.get(
            app.url_path_for("experiments:get-experiment-by-id", id=test_experiment_1_created_by_user_1.id)
        )
        assert res.status_code == status.HTTP_200_OK
        experiment = ExperimentPublic(**res.json())
        assert experiment == test_experiment_1_created_by_user_1

    async def test_get_experiment_by_id_unvalid_query_by_user_1(
        self,
        app: FastAPI,
        client_wt_auth_user_1: AsyncClient,
        test_experiment_1_created_by_user_2: ExperimentPublic,
    ) -> None:
        res = await client_wt_auth_user_1.get(
            app.url_path_for("experiments:get-experiment-by-id", id=test_experiment_1_created_by_user_2.id)
        )
        assert res.status_code == status.HTTP_401_UNAUTHORIZED

    @pytest.mark.parametrize(
        "id, status_code",
        (
            (500, 401),
            (-1, 401),
            (None, 422),
        ),
    )
    async def test_wrong_id_returns_error(
        self, app: FastAPI, client_wt_auth_user_1: AsyncClient, id: int, status_code: int
    ) -> None:
        res = await client_wt_auth_user_1.get(app.url_path_for("experiments:get-experiment-by-id", id=id))
        assert res.status_code == status_code


class TestGetExperimentByOrgAndModelName:
    async def test_get_experiment_by_org_and_model_name_valid_query_by_user_1(
        self,
        app: FastAPI,
        client_wt_auth_user_1: AsyncClient,
        test_experiment_1_created_by_user_1: ExperimentPublic,
    ) -> None:
        res = await client_wt_auth_user_1.get(
            app.url_path_for("experiments:get-experiment-by-organization-and-model-name"),
            params={
                "organization_name": test_experiment_1_created_by_user_1.organization_name,
                "model_name": test_experiment_1_created_by_user_1.model_name,
            },
        )
        assert res.status_code == status.HTTP_200_OK
        experiment = ExperimentPublic(**res.json())
        assert experiment.organization_name == test_experiment_1_created_by_user_1.organization_name
        assert experiment.model_name == test_experiment_1_created_by_user_1.model_name
        assert experiment.id == test_experiment_1_created_by_user_1.id
        assert experiment.coordinator_ip == test_experiment_1_created_by_user_1.coordinator_ip
        assert experiment.coordinator_port == test_experiment_1_created_by_user_1.coordinator_port
        assert experiment.creator == test_experiment_1_created_by_user_1.creator

    async def test_get_experiment_by_id_unvalid_query_by_user_1(
        self,
        app: FastAPI,
        client_wt_auth_user_1: AsyncClient,
        test_experiment_1_created_by_user_2: ExperimentPublic,
    ) -> None:
        res = await client_wt_auth_user_1.get(
            app.url_path_for("experiments:get-experiment-by-organization-and-model-name"),
            params={
                "organization_name": test_experiment_1_created_by_user_2.organization_name,
                "model_name": test_experiment_1_created_by_user_2.model_name,
            },
        )
        assert res.status_code == status.HTTP_401_UNAUTHORIZED

    @pytest.mark.parametrize(
        "organization_name, model_name, status_code",
        (
            (500, 500, 401),
            (-1, 500, 401),
            (None, 500, 401),
            (500, 500, 401),
            (500, -1, 401),
            (500, None, 401),
        ),
    )
    async def test_wrong_id_returns_error(
        self,
        app: FastAPI,
        client_wt_auth_user_1: AsyncClient,
        organization_name: int,
        model_name: int,
        status_code: int,
    ) -> None:
        res = await client_wt_auth_user_1.get(
            app.url_path_for("experiments:get-experiment-by-organization-and-model-name"),
            params={
                "organization_name": organization_name,
                "model_name": model_name,
            },
        )
        assert res.status_code == status_code


class TestUpdateExperiment:
    @pytest.mark.parametrize(
        "attrs_to_change, values",
        (
            (["coordinator_ip"], ["192.0.2.0"]),
            (["coordinator_ip"], ["684D:1111:222:3333:4444:5555:6:77"]),
            (["coordinator_port"], [400]),
            (
                [
                    "coordinator_ip",
                    "coordinator_port",
                ],
                [
                    "0.0.0.0",
                    80,
                ],
            ),
            (
                [
                    "model_name",
                    "coordinator_ip",
                    "coordinator_port",
                ],
                [
                    "model_4",
                    "1.0.0.0",
                    890,
                ],
            ),
            (
                [
                    "organization_name",
                    "model_name",
                    "coordinator_ip",
                    "coordinator_port",
                ],
                [
                    "Organization_a",
                    "model_4",
                    "2.0.0.0",
                    880,
                ],
            ),
        ),
    )
    async def test_update_experiment_with_valid_input(
        self,
        app: FastAPI,
        client_wt_auth_user_1: AsyncClient,
        test_experiment_2_created_by_user_1_for_updates_tests: ExperimentPublic,
        attrs_to_change: List[str],
        values: List[str],
    ) -> None:
        _ = ExperimentUpdate(**{attrs_to_change[i]: values[i] for i in range(len(attrs_to_change))})
        experiment_update = {"experiment_update": {attrs_to_change[i]: values[i] for i in range(len(attrs_to_change))}}
        res = await client_wt_auth_user_1.put(
            app.url_path_for(
                "experiments:update-experiment-by-id", id=test_experiment_2_created_by_user_1_for_updates_tests.id
            ),
            json=experiment_update,
        )
        assert res.status_code == status.HTTP_200_OK
        updated_experiment = ExperimentPublic(**res.json())
        assert (
            updated_experiment.id == test_experiment_2_created_by_user_1_for_updates_tests.id
        )  # make sure it's the same experiment

        # make sure that any attribute we updated has changed to the correct value
        for attr_to_change, value in zip(attrs_to_change, values):
            # Beware, this can raise an error if someone ask to removed user not whitelisted (and it isn't an error)
            assert getattr(updated_experiment, attr_to_change) != getattr(
                test_experiment_2_created_by_user_1_for_updates_tests, attr_to_change
            )
            final_value = getattr(updated_experiment, attr_to_change)
            if isinstance(final_value, IPv4Address):
                assert final_value == IPv4Address(value)
            elif isinstance(final_value, IPv6Address):
                assert final_value == IPv6Address(value)
            else:
                assert final_value == value
        # make sure that no other attributes' values have changed
        for attr, value in updated_experiment.dict().items():
            if attr not in attrs_to_change and attr != "updated_at":
                final_value = getattr(test_experiment_2_created_by_user_1_for_updates_tests, attr)

                if isinstance(final_value, IPv4Address):
                    assert final_value == IPv4Address(value)
                elif isinstance(final_value, IPv6Address):
                    assert final_value == IPv6Address(value)
                else:
                    assert final_value == value
            if attr == "updated_at":
                assert getattr(test_experiment_2_created_by_user_1_for_updates_tests, attr) != value

    @pytest.mark.parametrize(
        "id, payload, status_code",
        (
            (-1, {"organization_name": "test"}, 422),
            (0, {"organization_name": "test2"}, 422),
            (500, {}, 401),
            (500, {"organization_name": "test3"}, 401),
        ),
    )
    async def test_update_experiment_with_invalid_input_throws_error(
        self,
        app: FastAPI,
        client_wt_auth_user_1: AsyncClient,
        id: int,
        payload: dict,
        status_code: int,
    ) -> None:
        experiment_full_update = {"experiment_update": payload}
        res = await client_wt_auth_user_1.put(
            app.url_path_for("experiments:update-experiment-by-id", id=id),
            json=experiment_full_update,
        )
        assert res.status_code == status_code


class TestDeleteExperiment:
    async def test_can_delete_experiment_successfully(
        self,
        app: FastAPI,
        client_wt_auth_user_1: AsyncClient,
        test_experiment_1_created_by_user_1: ExperimentPublic,
    ) -> None:
        # delete the experiment
        res = await client_wt_auth_user_1.delete(
            app.url_path_for("experiments:delete-experiment-by-id", id=test_experiment_1_created_by_user_1.id)
        )
        assert res.status_code == status.HTTP_200_OK
        # ensure that the experiment no longer exists
        res = await client_wt_auth_user_1.get(
            app.url_path_for("experiments:get-experiment-by-id", id=test_experiment_1_created_by_user_1.id)
        )
        assert res.status_code == status.HTTP_401_UNAUTHORIZED

    async def test_cant_delete_other_user_experiment(
        self,
        app: FastAPI,
        client_wt_auth_user_1: AsyncClient,
        test_experiment_1_created_by_user_2: ExperimentPublic,
    ) -> None:
        # delete the experiment
        res = await client_wt_auth_user_1.delete(
            app.url_path_for("experiments:delete-experiment-by-id", id=test_experiment_1_created_by_user_2.id)
        )
        assert res.status_code == status.HTTP_401_UNAUTHORIZED

    @pytest.mark.parametrize(
        "id, status_code",
        (
            (500, 401),
            (0, 422),
            (-1, 422),
            (None, 422),
        ),
    )
    async def test_can_delete_experiment_unsuccessfully(
        self,
        moonlanding_user_1,
        app: FastAPI,
        client_wt_auth_user_1: AsyncClient,
        test_experiment_1_created_by_user_1: ExperimentPublic,
        id: int,
        status_code: int,
    ) -> None:
        res = await client_wt_auth_user_1.delete(app.url_path_for("experiments:delete-experiment-by-id", id=id))
        assert res.status_code == status_code


class TestJoinExperimentById:
    async def test_can_join_experiment_successfully(
        self,
        moonlanding_user_1,
        moonlanding_user_2,
        app: FastAPI,
        client_wt_auth_user_2: AsyncClient,
        # client_wt_auth_user_2: AsyncClient,
        test_experiment_1_created_by_user_1: ExperimentPublic,
        test_experiment_join_input_1_by_user_2: ExperimentJoinInput,
        test_experiment_join_input_2_by_user_2: ExperimentJoinInput,
    ) -> None:
        join_input_1, private_key_input_1 = test_experiment_join_input_1_by_user_2
        values = join_input_1.dict()
        # Make the values JSON serializable
        values["peer_public_key"] = values["peer_public_key"].decode("utf-8")

        res = await client_wt_auth_user_2.put(
            app.url_path_for("experiments:join-experiment-by-id", id=test_experiment_1_created_by_user_1.id),
            json={"experiment_join_input": values},
        )
        assert res.status_code == status.HTTP_200_OK, res.content

        exp_pass = ExperimentJoinOutput(**res.json())
        assert getattr(exp_pass, "coordinator_ip") == test_experiment_1_created_by_user_1.coordinator_ip
        assert getattr(exp_pass, "coordinator_port") == test_experiment_1_created_by_user_1.coordinator_port

        hivemind_access = getattr(exp_pass, "hivemind_access")
        assert getattr(hivemind_access, "peer_public_key") == join_input_1.peer_public_key

        signature = base64.b64decode(getattr(hivemind_access, "signature"))

        auth_server_public_key = getattr(exp_pass, "auth_server_public_key")
        auth_server_public_key = crypto.load_public_key(auth_server_public_key)

        verif = auth_server_public_key.verify(
            signature,
            f"{hivemind_access.username} {hivemind_access.peer_public_key} {hivemind_access.expiration_time}".encode(),
            crypto.PADDING,
            crypto.HASH_ALGORITHM,
        )
        assert verif is None  # verify() returns None iff the signature is correct
        assert hivemind_access.expiration_time > datetime.datetime.utcnow()
        assert hivemind_access.username == moonlanding_user_2.username

        # Now the same user try to join a second time with another public key
        join_input_2, private_key_input_2 = test_experiment_join_input_2_by_user_2
        values = join_input_2.dict()
        # Make the values JSON serializable
        values["peer_public_key"] = values["peer_public_key"].decode("utf-8")

        res = await client_wt_auth_user_2.put(
            app.url_path_for("experiments:join-experiment-by-id", id=test_experiment_1_created_by_user_1.id),
            json={"experiment_join_input": values},
        )
        assert res.status_code == status.HTTP_200_OK, res.content

        exp_pass = ExperimentJoinOutput(**res.json())
        assert getattr(exp_pass, "coordinator_ip") == test_experiment_1_created_by_user_1.coordinator_ip
        assert getattr(exp_pass, "coordinator_port") == test_experiment_1_created_by_user_1.coordinator_port

        hivemind_access = getattr(exp_pass, "hivemind_access")
        peer_public_key = getattr(hivemind_access, "peer_public_key")
        assert peer_public_key == join_input_2.peer_public_key

        signature = base64.b64decode(getattr(hivemind_access, "signature"))

        auth_server_public_key = getattr(exp_pass, "auth_server_public_key")
        auth_server_public_key = crypto.load_public_key(auth_server_public_key)

        verif = auth_server_public_key.verify(
            signature,
            f"{hivemind_access.username} {hivemind_access.peer_public_key} {hivemind_access.expiration_time}".encode(),
            crypto.PADDING,
            crypto.HASH_ALGORITHM,
        )
        assert verif is None  # verify() returns None iff the signature is correct
        assert hivemind_access.expiration_time > datetime.datetime.utcnow()
        assert hivemind_access.username == moonlanding_user_2.username

        msg = "this is a test"
        ciphertext = encrypt_msg(msg, crypto.load_public_key(peer_public_key))
        plaintext = decrypt_msg(ciphertext, private_key_input_2)
        assert plaintext == msg

        with pytest.raises(Exception):
            plaintext = decrypt_msg(ciphertext, private_key_input_1)

    async def test_can_join_experiment_successfully_2_times_with_same_public_key(
        self,
        moonlanding_user_1,
        moonlanding_user_2,
        app: FastAPI,
        client_wt_auth_user_2: AsyncClient,
        # client_wt_auth_user_2: AsyncClient,
        test_experiment_1_created_by_user_1: ExperimentPublic,
        test_experiment_join_input_1_by_user_2: ExperimentJoinInput,
        test_experiment_join_input_2_by_user_2: ExperimentJoinInput,
    ) -> None:
        join_input_1, private_key_input_1 = test_experiment_join_input_1_by_user_2
        values = join_input_1.dict()
        # Make the values JSON serializable
        values["peer_public_key"] = values["peer_public_key"].decode("utf-8")

        res = await client_wt_auth_user_2.put(
            app.url_path_for("experiments:join-experiment-by-id", id=test_experiment_1_created_by_user_1.id),
            json={"experiment_join_input": values},
        )
        assert res.status_code == status.HTTP_200_OK, res.content

        exp_pass = ExperimentJoinOutput(**res.json())
        assert getattr(exp_pass, "coordinator_ip") == test_experiment_1_created_by_user_1.coordinator_ip
        assert getattr(exp_pass, "coordinator_port") == test_experiment_1_created_by_user_1.coordinator_port

        hivemind_access = getattr(exp_pass, "hivemind_access")
        assert getattr(hivemind_access, "peer_public_key") == join_input_1.peer_public_key

        signature = base64.b64decode(getattr(hivemind_access, "signature"))

        auth_server_public_key = getattr(exp_pass, "auth_server_public_key")
        auth_server_public_key = crypto.load_public_key(auth_server_public_key)

        verif = auth_server_public_key.verify(
            signature,
            f"{hivemind_access.username} {hivemind_access.peer_public_key} {hivemind_access.expiration_time}".encode(),
            crypto.PADDING,
            crypto.HASH_ALGORITHM,
        )
        assert verif is None  # verify() returns None iff the signature is correct
        assert hivemind_access.expiration_time > datetime.datetime.utcnow()
        assert hivemind_access.username == moonlanding_user_2.username

        # Now the same user try to join a second time with the same public key
        join_input_1, private_key_input_1 = test_experiment_join_input_1_by_user_2
        values = join_input_1.dict()
        # Make the values JSON serializable
        values["peer_public_key"] = values["peer_public_key"].decode("utf-8")

        res = await client_wt_auth_user_2.put(
            app.url_path_for("experiments:join-experiment-by-id", id=test_experiment_1_created_by_user_1.id),
            json={"experiment_join_input": values},
        )
        assert res.status_code == status.HTTP_200_OK, res.content

        exp_pass = ExperimentJoinOutput(**res.json())
        assert getattr(exp_pass, "coordinator_ip") == test_experiment_1_created_by_user_1.coordinator_ip
        assert getattr(exp_pass, "coordinator_port") == test_experiment_1_created_by_user_1.coordinator_port

        hivemind_access = getattr(exp_pass, "hivemind_access")
        assert getattr(hivemind_access, "peer_public_key") == join_input_1.peer_public_key

        signature = base64.b64decode(getattr(hivemind_access, "signature"))

        auth_server_public_key = getattr(exp_pass, "auth_server_public_key")
        auth_server_public_key = crypto.load_public_key(auth_server_public_key)

        verif = auth_server_public_key.verify(
            signature,
            f"{hivemind_access.username} {hivemind_access.peer_public_key} {hivemind_access.expiration_time}".encode(),
            crypto.PADDING,
            crypto.HASH_ALGORITHM,
        )
        assert verif is None  # verify() returns None iff the signature is correct
        assert hivemind_access.expiration_time > datetime.datetime.utcnow()
        assert hivemind_access.username == moonlanding_user_2.username

    async def test_cant_join_experiment_successfully_user_not_allowlisted(
        self,
        moonlanding_user_1,
        app: FastAPI,
        client_wt_auth_user_1: AsyncClient,
        test_experiment_1_created_by_user_2: ExperimentPublic,
        test_experiment_join_input_1_by_user_2: ExperimentJoinInput,
    ):
        join_input_1, private_key_input_1 = test_experiment_join_input_1_by_user_2
        values = join_input_1.dict()
        # Make the values JSON serializable
        values["peer_public_key"] = values["peer_public_key"].decode("utf-8")

        res = await client_wt_auth_user_1.put(
            app.url_path_for("experiments:join-experiment-by-id", id=test_experiment_1_created_by_user_2.id),
            json={"experiment_join_input": values},
        )
        assert res.status_code == status.HTTP_401_UNAUTHORIZED, res.content


class TestJoinExperimentByOrgAndModelName:
    async def test_can_join_experiment_successfully(
        self,
        moonlanding_user_2,
        app: FastAPI,
        client_wt_auth_user_2: AsyncClient,
        # client_wt_auth_user_2: AsyncClient,
        test_experiment_1_created_by_user_1: ExperimentPublic,
        test_experiment_join_input_1_by_user_2: ExperimentJoinInput,
        test_experiment_join_input_2_by_user_2: ExperimentJoinInput,
    ) -> None:
        join_input_1, private_key_input_1 = test_experiment_join_input_1_by_user_2
        values = join_input_1.dict()
        # Make the values JSON serializable
        values["peer_public_key"] = values["peer_public_key"].decode("utf-8")

        res = await client_wt_auth_user_2.put(
            app.url_path_for(
                "experiments:join-experiment-by-organization-and-model-name",
            ),
            json={"experiment_join_input": values},
            params={
                "organization_name": test_experiment_1_created_by_user_1.organization_name,
                "model_name": test_experiment_1_created_by_user_1.model_name,
            },
        )
        assert res.status_code == status.HTTP_200_OK, res.content

        exp_pass = ExperimentJoinOutput(**res.json())
        assert getattr(exp_pass, "coordinator_ip") == test_experiment_1_created_by_user_1.coordinator_ip
        assert getattr(exp_pass, "coordinator_port") == test_experiment_1_created_by_user_1.coordinator_port

        hivemind_access = getattr(exp_pass, "hivemind_access")
        assert getattr(hivemind_access, "peer_public_key") == join_input_1.peer_public_key

        signature = base64.b64decode(getattr(hivemind_access, "signature"))

        auth_server_public_key = getattr(exp_pass, "auth_server_public_key")
        auth_server_public_key = crypto.load_public_key(auth_server_public_key)

        verif = auth_server_public_key.verify(
            signature,
            f"{hivemind_access.username} {hivemind_access.peer_public_key} {hivemind_access.expiration_time}".encode(),
            crypto.PADDING,
            crypto.HASH_ALGORITHM,
        )
        assert verif is None  # verify() returns None iff the signature is correct
        assert hivemind_access.expiration_time > datetime.datetime.utcnow()
        assert hivemind_access.username == moonlanding_user_2.username

        # Now the same user try to join a second time with another public key
        join_input_2, private_key_input_2 = test_experiment_join_input_2_by_user_2
        values = join_input_2.dict()
        # Make the values JSON serializable
        values["peer_public_key"] = values["peer_public_key"].decode("utf-8")

        res = await client_wt_auth_user_2.put(
            app.url_path_for(
                "experiments:join-experiment-by-organization-and-model-name",
            ),
            json={"experiment_join_input": values},
            params={
                "organization_name": test_experiment_1_created_by_user_1.organization_name,
                "model_name": test_experiment_1_created_by_user_1.model_name,
            },
        )
        assert res.status_code == status.HTTP_200_OK, res.content

        exp_pass = ExperimentJoinOutput(**res.json())
        assert getattr(exp_pass, "coordinator_ip") == test_experiment_1_created_by_user_1.coordinator_ip
        assert getattr(exp_pass, "coordinator_port") == test_experiment_1_created_by_user_1.coordinator_port

        hivemind_access = getattr(exp_pass, "hivemind_access")
        peer_public_key = getattr(hivemind_access, "peer_public_key")
        assert peer_public_key == join_input_2.peer_public_key

        signature = base64.b64decode(getattr(hivemind_access, "signature"))

        auth_server_public_key = getattr(exp_pass, "auth_server_public_key")
        auth_server_public_key = crypto.load_public_key(auth_server_public_key)

        verif = auth_server_public_key.verify(
            signature,
            f"{hivemind_access.username} {hivemind_access.peer_public_key} {hivemind_access.expiration_time}".encode(),
            crypto.PADDING,
            crypto.HASH_ALGORITHM,
        )
        assert verif is None  # verify() returns None iff the signature is correct
        assert hivemind_access.expiration_time > datetime.datetime.utcnow()
        assert hivemind_access.username == moonlanding_user_2.username

        msg = "this is a test"
        ciphertext = encrypt_msg(msg, crypto.load_public_key(peer_public_key))
        plaintext = decrypt_msg(ciphertext, private_key_input_2)
        assert plaintext == msg

        with pytest.raises(Exception):
            plaintext = decrypt_msg(ciphertext, private_key_input_1)

    async def test_can_join_experiment_successfully_2_times_with_same_public_key(
        self,
        moonlanding_user_1,
        moonlanding_user_2,
        app: FastAPI,
        client_wt_auth_user_2: AsyncClient,
        # client_wt_auth_user_2: AsyncClient,
        test_experiment_1_created_by_user_1: ExperimentPublic,
        test_experiment_join_input_1_by_user_2: ExperimentJoinInput,
        test_experiment_join_input_2_by_user_2: ExperimentJoinInput,
    ) -> None:
        join_input_1, private_key_input_1 = test_experiment_join_input_1_by_user_2
        values = join_input_1.dict()
        # Make the values JSON serializable
        values["peer_public_key"] = values["peer_public_key"].decode("utf-8")

        res = await client_wt_auth_user_2.put(
            app.url_path_for(
                "experiments:join-experiment-by-organization-and-model-name",
            ),
            json={"experiment_join_input": values},
            params={
                "organization_name": test_experiment_1_created_by_user_1.organization_name,
                "model_name": test_experiment_1_created_by_user_1.model_name,
            },
        )
        assert res.status_code == status.HTTP_200_OK, res.content
        exp_pass = ExperimentJoinOutput(**res.json())
        assert getattr(exp_pass, "coordinator_ip") == test_experiment_1_created_by_user_1.coordinator_ip
        assert getattr(exp_pass, "coordinator_port") == test_experiment_1_created_by_user_1.coordinator_port

        hivemind_access = getattr(exp_pass, "hivemind_access")
        assert getattr(hivemind_access, "peer_public_key") == join_input_1.peer_public_key

        signature = base64.b64decode(getattr(hivemind_access, "signature"))

        auth_server_public_key = getattr(exp_pass, "auth_server_public_key")
        auth_server_public_key = crypto.load_public_key(auth_server_public_key)

        verif = auth_server_public_key.verify(
            signature,
            f"{hivemind_access.username} {hivemind_access.peer_public_key} {hivemind_access.expiration_time}".encode(),
            crypto.PADDING,
            crypto.HASH_ALGORITHM,
        )
        assert verif is None  # verify() returns None iff the signature is correct
        assert hivemind_access.expiration_time > datetime.datetime.utcnow()
        assert hivemind_access.username == moonlanding_user_2.username

        # Now the same user try to join a second time with the same public key
        join_input_1, private_key_input_1 = test_experiment_join_input_1_by_user_2
        values = join_input_1.dict()
        # Make the values JSON serializable
        values["peer_public_key"] = values["peer_public_key"].decode("utf-8")

        res = await client_wt_auth_user_2.put(
            app.url_path_for(
                "experiments:join-experiment-by-organization-and-model-name",
            ),
            json={"experiment_join_input": values},
            params={
                "organization_name": test_experiment_1_created_by_user_1.organization_name,
                "model_name": test_experiment_1_created_by_user_1.model_name,
            },
        )
        assert res.status_code == status.HTTP_200_OK, res.content

        exp_pass = ExperimentJoinOutput(**res.json())
        assert getattr(exp_pass, "coordinator_ip") == test_experiment_1_created_by_user_1.coordinator_ip
        assert getattr(exp_pass, "coordinator_port") == test_experiment_1_created_by_user_1.coordinator_port

        hivemind_access = getattr(exp_pass, "hivemind_access")
        assert getattr(hivemind_access, "peer_public_key") == join_input_1.peer_public_key

        signature = base64.b64decode(getattr(hivemind_access, "signature"))

        auth_server_public_key = getattr(exp_pass, "auth_server_public_key")
        auth_server_public_key = crypto.load_public_key(auth_server_public_key)

        verif = auth_server_public_key.verify(
            signature,
            f"{hivemind_access.username} {hivemind_access.peer_public_key} {hivemind_access.expiration_time}".encode(),
            crypto.PADDING,
            crypto.HASH_ALGORITHM,
        )
        assert verif is None  # verify() returns None iff the signature is correct
        assert hivemind_access.expiration_time > datetime.datetime.utcnow()
        assert hivemind_access.username == moonlanding_user_2.username

    async def test_cant_join_experiment_successfully_user_not_allowlisted(
        self,
        moonlanding_user_1,
        app: FastAPI,
        client_wt_auth_user_1: AsyncClient,
        test_experiment_1_created_by_user_2: ExperimentPublic,
        test_experiment_join_input_1_by_user_2: ExperimentJoinInput,
    ):
        join_input_1, private_key_input_1 = test_experiment_join_input_1_by_user_2
        values = join_input_1.dict()
        # Make the values JSON serializable
        values["peer_public_key"] = values["peer_public_key"].decode("utf-8")

        res = await client_wt_auth_user_1.put(
            app.url_path_for(
                "experiments:join-experiment-by-organization-and-model-name",
            ),
            json={"experiment_join_input": values},
            params={
                "organization_name": test_experiment_1_created_by_user_2.organization_name,
                "model_name": test_experiment_1_created_by_user_2.model_name,
            },
        )
        assert res.status_code == status.HTTP_401_UNAUTHORIZED, res.content
