#include "piauth.h"
#include "picrypt.h"
#include "piliterals_string.h"

#include "gtest/gtest.h"

//! \~english Tests for the PIAuth signed key-exchange state machine.
//! \~russian Тесты машины состояний аутентификации PIAuth с подписанным обменом ключами.
//!
//! The protocol is driven in-process: two \c PIAuth endpoints exchange the same
//! \c PIByteArray packets the network transport would carry. Each \a receive()
//! updates the endpoint state and writes the reply back into the passed buffer.
//! \~russian Протокол прогоняется в одном процессе: две конечные точки \c PIAuth
//! обмениваются теми же пакетами \c PIByteArray, которые переносил бы транспорт.
//! Каждый \a receive() обновляет состояние и записывает ответ обратно в буфер.


namespace {

//! \~english Deterministic signing secret key for a human-readable seed.
//! \~russian Детерминированный секретный ключ подписи по текстовому seed.
PIByteArray signKey(const char * seed) {
	return PIAuth::generateSign(PIByteArray::fromAscii(seed));
}

//! \~english One-element authorized public key list.
//! \~russian Список доверенных открытых ключей из одного элемента.
PIVector<PIByteArray> keyList(const PIByteArray & key) {
	return PIVector<PIByteArray>{key};
}

} // namespace


class PIAuthTest: public ::testing::Test {};


//! \~english When both peers already trust each other's signing key the handshake
//! skips the password stage: probe, key exchange, connected.
//! \~russian Когда узлы уже доверяют ключам подписи друг друга, рукопожатие
//! пропускает этап пароля: probe, обмен ключом, connected.
TEST_F(PIAuthTest, AuthorizedKeysFullHandshake) {
	PIAuth client(signKey("client-seed"));
	PIAuth server(signKey("server-seed"));

	bool client_connected   = false;
	bool server_connected   = false;
	int client_disconnects  = 0;
	bool password_requested = false;
	PIString server_info;

	CONNECTL(&client, connected, [&](PIString) { client_connected = true; });
	CONNECTL(&server, connected, [&](PIString info) {
		server_connected = true;
		server_info      = info;
	});
	CONNECTL(&client, disconnected, [&](PIString) { ++client_disconnects; });
	CONNECTL(&client, passwordRequest, [&](PIString * pass) {
		password_requested = true;
		*pass              = "unused";
	});

	client.setInfoData(PIByteArray::fromAscii("hello-server"));
	client.setAuthorizedPublicKeys(keyList(server.getSignPublicKey()));
	server.setAuthorizedPublicKeys(keyList(client.getSignPublicKey()));
	ASSERT_FALSE(client.getSignPublicKey().isEmpty());
	ASSERT_FALSE(server.getSignPublicKey().isEmpty());

	client.startClient();
	PIByteArray pkt = server.startServer();
	ASSERT_FALSE(pkt.isEmpty());

	EXPECT_EQ(PIAuth::AuthReply, client.receive(pkt));
	EXPECT_EQ(PIAuth::KeyExchange, server.receive(pkt));
	EXPECT_EQ(PIAuth::Connected, client.receive(pkt));
	EXPECT_EQ(PIAuth::Connected, server.receive(pkt));

	EXPECT_TRUE(client_connected);
	EXPECT_TRUE(server_connected);
	EXPECT_EQ(0, client_disconnects);
	EXPECT_FALSE(password_requested);
	EXPECT_EQ("hello-server"_a, server_info);

	ASSERT_FALSE(client.getSecretKey().isEmpty());
	EXPECT_EQ(PICrypt::sizeKey(), client.getSecretKey().size());
	EXPECT_EQ(client.getSecretKey(), server.getSecretKey());
}


//! \~english After both peers are Connected an extra round-trip re-sends the
//! application info and keeps the same session key (keep-alive path).
//! \~russian После перехода обоих узлов в Connected дополнительный обмен
//! повторно передает прикладные данные и сохраняет тот же сеансовый ключ
//! (путь keep-alive).
TEST_F(PIAuthTest, ConnectedKeepAliveRoundTrip) {
	PIAuth client(signKey("client-seed"));
	PIAuth server(signKey("server-seed"));

	int client_connected = 0;
	int server_connected = 0;
	PIString server_info;

	client.setInfoData(PIByteArray::fromAscii("kept-info"));
	client.setAuthorizedPublicKeys(keyList(server.getSignPublicKey()));
	server.setAuthorizedPublicKeys(keyList(client.getSignPublicKey()));
	CONNECTL(&client, connected, [&](PIString) { ++client_connected; });
	CONNECTL(&server, connected, [&](PIString info) {
		++server_connected;
		server_info = info;
	});

	client.startClient();
	PIByteArray pkt = server.startServer();
	EXPECT_EQ(PIAuth::AuthReply, client.receive(pkt));
	EXPECT_EQ(PIAuth::KeyExchange, server.receive(pkt));
	EXPECT_EQ(PIAuth::Connected, client.receive(pkt));
	EXPECT_EQ(PIAuth::Connected, server.receive(pkt));

	// server.receive() left its keep-alive reply in "pkt".
	EXPECT_EQ(PIAuth::Connected, client.receive(pkt));
	EXPECT_EQ(PIAuth::Connected, server.receive(pkt));

	EXPECT_EQ(2, client_connected);
	EXPECT_EQ(2, server_connected);
	EXPECT_EQ("kept-info"_a, server_info);
	EXPECT_EQ(client.getSecretKey(), server.getSecretKey());
}


//! \~english With no mutual trust the client approves the server via the
//! \a authorize event, the server requests a password and both reach Connected
//! with the same session key.
//! \~russian Без взаимного доверия клиент подтверждает сервер через событие
//! \a authorize, сервер запрашивает пароль, и обе стороны доходят до Connected
//! с одинаковым сеансовым ключом.
TEST_F(PIAuthTest, PasswordFullHandshake) {
	PIAuth client(signKey("client-seed"));
	PIAuth server(signKey("server-seed"));

	bool authorized        = false;
	bool password_asked    = false;
	bool password_checked  = false;
	bool password_ok       = false;
	int client_disconnects = 0;
	PIString server_info;

	server.setServerPassword("s3cret");
	client.setInfoData(PIByteArray::fromAscii("my-info"));

	CONNECTL(&client, authorize, [&](PIByteArray, bool * ok) {
		authorized = true;
		*ok        = true;
	});
	CONNECTL(&client, passwordRequest, [&](PIString * pass) {
		password_asked = true;
		*pass          = "s3cret";
	});
	CONNECTL(&server, passwordCheck, [&](bool result) {
		password_checked = true;
		password_ok      = result;
	});
	CONNECTL(&client, disconnected, [&](PIString) { ++client_disconnects; });
	CONNECTL(&client, connected, [&](PIString) {});
	CONNECTL(&server, connected, [&](PIString info) { server_info = info; });

	client.startClient();
	PIByteArray pkt = server.startServer();

	EXPECT_EQ(PIAuth::AuthReply, client.receive(pkt));
	EXPECT_EQ(PIAuth::PassRequest, server.receive(pkt));
	EXPECT_EQ(PIAuth::PassRequest, client.receive(pkt));
	EXPECT_EQ(PIAuth::KeyExchange, server.receive(pkt));
	EXPECT_EQ(PIAuth::Connected, client.receive(pkt));
	EXPECT_EQ(PIAuth::Connected, server.receive(pkt));

	EXPECT_TRUE(authorized);
	EXPECT_TRUE(password_asked);
	EXPECT_TRUE(password_checked);
	EXPECT_TRUE(password_ok);
	EXPECT_EQ(0, client_disconnects);
	EXPECT_EQ("my-info"_a, server_info);
	ASSERT_FALSE(client.getSecretKey().isEmpty());
	EXPECT_EQ(client.getSecretKey(), server.getSecretKey());
}


//! \~english A client that already trusts the server key still has to pass the
//! server password check when the server does not know the client key. The
//! \a authorize event must not fire in that case.
//! \~russian Клиент, уже доверяющий ключу сервера, всё равно проходит проверку
//! пароля, если сервер не знает ключ клиента. Событие \a authorize не должно
//! срабатывать.
TEST_F(PIAuthTest, TrustedServerPasswordStage) {
	PIAuth client(signKey("client-seed"));
	PIAuth server(signKey("server-seed"));

	bool authorized     = false;
	bool password_asked = false;

	server.setServerPassword("pw");
	client.setAuthorizedPublicKeys(keyList(server.getSignPublicKey()));
	CONNECTL(&client, authorize, [&](PIByteArray, bool * ok) {
		authorized = true;
		*ok        = true;
	});
	CONNECTL(&client, passwordRequest, [&](PIString * pass) {
		password_asked = true;
		*pass          = "pw";
	});
	CONNECTL(&server, passwordCheck, [](bool) {});
	CONNECTL(&client, connected, [](PIString) {});
	CONNECTL(&server, connected, [](PIString) {});

	client.startClient();
	PIByteArray pkt = server.startServer();

	EXPECT_EQ(PIAuth::AuthReply, client.receive(pkt));
	EXPECT_EQ(PIAuth::PassRequest, server.receive(pkt));
	EXPECT_EQ(PIAuth::PassRequest, client.receive(pkt));
	EXPECT_EQ(PIAuth::KeyExchange, server.receive(pkt));
	EXPECT_EQ(PIAuth::Connected, client.receive(pkt));
	EXPECT_EQ(PIAuth::Connected, server.receive(pkt));

	EXPECT_FALSE(authorized);
	EXPECT_TRUE(password_asked);
	EXPECT_FALSE(client.getSecretKey().isEmpty());
	EXPECT_EQ(client.getSecretKey(), server.getSecretKey());
}


//! \~english A wrong password must abort the handshake on the server side and
//! report a failed password check.
//! \~russian Неверный пароль должен прервать рукопожатие на сервере и сообщить
//! о неудачной проверке пароля.
TEST_F(PIAuthTest, WrongPasswordDisconnects) {
	PIAuth client(signKey("client-seed"));
	PIAuth server(signKey("server-seed"));

	bool password_checked = false;
	bool password_ok      = true;
	PIString disconnect_reason;

	server.setServerPassword("right");
	CONNECTL(&client, authorize, [](PIByteArray, bool * ok) { *ok = true; });
	CONNECTL(&client, passwordRequest, [](PIString * pass) { *pass = "wrong"; });
	CONNECTL(&server, passwordCheck, [&](bool result) {
		password_checked = true;
		password_ok      = result;
	});
	CONNECTL(&server, disconnected, [&](PIString reason) { disconnect_reason = reason; });
	CONNECTL(&client, connected, [](PIString) {});
	CONNECTL(&server, connected, [](PIString) {});

	client.startClient();
	PIByteArray pkt = server.startServer();

	EXPECT_EQ(PIAuth::AuthReply, client.receive(pkt));
	EXPECT_EQ(PIAuth::PassRequest, server.receive(pkt));
	EXPECT_EQ(PIAuth::PassRequest, client.receive(pkt));
	EXPECT_EQ(PIAuth::NotConnected, server.receive(pkt));

	EXPECT_TRUE(password_checked);
	EXPECT_FALSE(password_ok);
	EXPECT_TRUE(disconnect_reason.contains("Invalid password"));
	EXPECT_TRUE(server.getSecretKey().isEmpty());
	EXPECT_TRUE(client.getSecretKey().isEmpty());
}


//! \~english The server refuses password authentication when no password is
//! configured, even if an empty hash is submitted.
//! \~russian Сервер отказывает в парольной аутентификации, если пароль не задан,
//! даже если передан пустой хеш.
TEST_F(PIAuthTest, EmptyServerPasswordRejects) {
	PIAuth client(signKey("client-seed"));
	PIAuth server(signKey("server-seed"));

	bool password_ok = true;
	CONNECTL(&client, authorize, [](PIByteArray, bool * ok) { *ok = true; });
	CONNECTL(&client, passwordRequest, [](PIString * pass) { *pass = "anything"; });
	CONNECTL(&server, passwordCheck, [&](bool result) { password_ok = result; });
	CONNECTL(&client, connected, [](PIString) {});
	CONNECTL(&server, connected, [](PIString) {});

	client.startClient();
	PIByteArray pkt = server.startServer();

	EXPECT_EQ(PIAuth::AuthReply, client.receive(pkt));
	EXPECT_EQ(PIAuth::PassRequest, server.receive(pkt));
	EXPECT_EQ(PIAuth::PassRequest, client.receive(pkt));
	EXPECT_EQ(PIAuth::NotConnected, server.receive(pkt));
	EXPECT_FALSE(password_ok);
}


//! \~english Rejecting an unknown server in the \a authorize callback aborts the
//! handshake on the client.
//! \~russian Отказ от неизвестного сервера в обработчике \a authorize прерывает
//! рукопожатие на клиенте.
TEST_F(PIAuthTest, AuthorizeRejectedDisconnects) {
	PIAuth client(signKey("client-seed"));
	PIAuth server(signKey("server-seed"));

	bool authorize_called = false;
	PIString disconnect_reason;

	CONNECTL(&client, authorize, [&](PIByteArray, bool * ok) {
		authorize_called = true;
		*ok              = false;
	});
	CONNECTL(&client, disconnected, [&](PIString reason) { disconnect_reason = reason; });
	CONNECTL(&client, connected, [](PIString) {});

	client.startClient();
	PIByteArray pkt = server.startServer();

	EXPECT_EQ(PIAuth::NotConnected, client.receive(pkt));
	EXPECT_TRUE(authorize_called);
	EXPECT_TRUE(disconnect_reason.contains("Unauthorised"));
	EXPECT_TRUE(client.getSecretKey().isEmpty());
}


//! \~english An empty password from the client callback cancels the handshake.
//! \~russian Пустой пароль из клиентского обработчика отменяет рукопожатие.
TEST_F(PIAuthTest, PasswordRequestCancelledDisconnects) {
	PIAuth client(signKey("client-seed"));
	PIAuth server(signKey("server-seed"));

	PIString disconnect_reason;

	server.setServerPassword("pw");
	CONNECTL(&client, authorize, [](PIByteArray, bool * ok) { *ok = true; });
	CONNECTL(&client, passwordRequest, [](PIString * pass) { pass->clear(); });
	CONNECTL(&client, disconnected, [&](PIString reason) { disconnect_reason = reason; });
	CONNECTL(&client, connected, [](PIString) {});
	CONNECTL(&server, connected, [](PIString) {});

	client.startClient();
	PIByteArray pkt = server.startServer();

	EXPECT_EQ(PIAuth::AuthReply, client.receive(pkt));
	EXPECT_EQ(PIAuth::PassRequest, server.receive(pkt));
	EXPECT_EQ(PIAuth::NotConnected, client.receive(pkt));
	EXPECT_TRUE(disconnect_reason.contains("Canceled by user"));
}


//! \~english A tampered server probe fails signature verification on the client.
//! \~russian Измененный пакет-проба сервера не проходит проверку подписи на клиенте.
TEST_F(PIAuthTest, TamperedProbeRejected) {
	PIAuth client(signKey("client-seed"));
	PIAuth server(signKey("server-seed"));

	PIString disconnect_reason;
	CONNECTL(&client, authorize, [](PIByteArray, bool * ok) { *ok = true; });
	CONNECTL(&client, disconnected, [&](PIString reason) { disconnect_reason = reason; });
	CONNECTL(&client, connected, [](PIString) {});

	client.startClient();
	PIByteArray pkt = server.startServer();
	ASSERT_TRUE(pkt.size() > 0);
	pkt[pkt.size() - 1] ^= 0xFF;

	EXPECT_EQ(PIAuth::NotConnected, client.receive(pkt));
	EXPECT_TRUE(disconnect_reason.contains("Incorrect sign"));
}


//! \~english A too-short packet is rejected without touching the session.
//! \~russian Слишком короткий пакет отклоняется без изменения сессии.
TEST_F(PIAuthTest, InvalidDataSizeRejected) {
	PIAuth client(signKey("client-seed"));

	PIString disconnect_reason;
	CONNECTL(&client, disconnected, [&](PIString reason) { disconnect_reason = reason; });

	client.startClient();
	PIByteArray empty;
	EXPECT_EQ(PIAuth::NotConnected, client.receive(empty));
	EXPECT_TRUE(disconnect_reason.contains("invalid data size"));

	PIByteArray short_pkt;
	short_pkt << (uchar)1 << (uchar)2;
	EXPECT_EQ(PIAuth::NotConnected, client.receive(short_pkt));
}


//! \~english A packet whose state does not match the current one is rejected.
//! \~russian Пакет, состояние которого не соответствует текущему, отклоняется.
TEST_F(PIAuthTest, UnexpectedStateRejected) {
	PIAuth client(signKey("client-seed"));

	PIString disconnect_reason;
	CONNECTL(&client, disconnected, [&](PIString reason) { disconnect_reason = reason; });

	client.startClient();
	PIByteArray pkt;
	pkt << (int)PIAuth::PassRequest << PIByteArray::fromAscii("junk");
	EXPECT_EQ(PIAuth::NotConnected, client.receive(pkt));
	EXPECT_TRUE(disconnect_reason.contains("invalid state"));
}


//! \~english Before the handshake no session key is available; \a stop() clears an
//! established one and the endpoint can handshake again afterwards.
//! \~russian До рукопожатия сеансовый ключ недоступен; \a stop() очищает
//! установленный ключ, после чего конечная точка может пройти рукопожатие снова.
TEST_F(PIAuthTest, SecretKeyAvailableOnlyWhenConnected) {
	PIAuth client(signKey("client-seed"));
	PIAuth server(signKey("server-seed"));

	client.setAuthorizedPublicKeys(keyList(server.getSignPublicKey()));
	server.setAuthorizedPublicKeys(keyList(client.getSignPublicKey()));
	CONNECTL(&client, connected, [](PIString) {});
	CONNECTL(&server, connected, [](PIString) {});

	client.startClient();
	EXPECT_TRUE(client.getSecretKey().isEmpty());

	PIByteArray pkt = server.startServer();
	EXPECT_EQ(PIAuth::AuthReply, client.receive(pkt));
	EXPECT_TRUE(client.getSecretKey().isEmpty());
	EXPECT_EQ(PIAuth::KeyExchange, server.receive(pkt));
	EXPECT_EQ(PIAuth::Connected, client.receive(pkt));
	ASSERT_FALSE(client.getSecretKey().isEmpty());

	client.stop();
	EXPECT_TRUE(client.getSecretKey().isEmpty());

	// A stopped client can run a fresh handshake with a new server.
	PIAuth server2(signKey("server-seed-2"));
	client.setAuthorizedPublicKeys(keyList(server2.getSignPublicKey()));
	server2.setAuthorizedPublicKeys(keyList(client.getSignPublicKey()));
	CONNECTL(&server2, connected, [](PIString) {});
	client.startClient();
	pkt = server2.startServer();
	EXPECT_EQ(PIAuth::AuthReply, client.receive(pkt));
	EXPECT_EQ(PIAuth::KeyExchange, server2.receive(pkt));
	EXPECT_EQ(PIAuth::Connected, client.receive(pkt));
	EXPECT_EQ(PIAuth::Connected, server2.receive(pkt));
	ASSERT_FALSE(client.getSecretKey().isEmpty());
	EXPECT_EQ(client.getSecretKey(), server2.getSecretKey());
}


//! \~english Two independent handshakes must derive different session keys.
//! \~russian Два независимых рукопожатия должны получить разные сеансовые ключи.
TEST_F(PIAuthTest, SessionKeysAreUnique) {
	auto handshake = [](const PIByteArray & client_key, const PIByteArray & server_key) {
		PIAuth client(client_key);
		PIAuth server(server_key);
		client.setAuthorizedPublicKeys(keyList(server.getSignPublicKey()));
		server.setAuthorizedPublicKeys(keyList(client.getSignPublicKey()));
		CONNECTL(&client, connected, [](PIString) {});
		CONNECTL(&server, connected, [](PIString) {});
		client.startClient();
		PIByteArray pkt = server.startServer();
		client.receive(pkt);
		server.receive(pkt);
		client.receive(pkt);
		server.receive(pkt);
		return client.getSecretKey();
	};

	PIByteArray first  = handshake(signKey("client-seed"), signKey("server-seed"));
	PIByteArray second = handshake(signKey("client-seed"), signKey("server-seed"));
	ASSERT_FALSE(first.isEmpty());
	ASSERT_FALSE(second.isEmpty());
	EXPECT_NE(first, second);
}


//! \~english \a generateSign() is deterministic and the endpoint derives the
//! matching public key from the secret.
//! \~russian \a generateSign() детерминирован, а конечная точка выводит
//! соответствующий открытый ключ из секретного.
TEST_F(PIAuthTest, GenerateSignAndPublicKey) {
	PIByteArray sk1 = PIAuth::generateSign(PIByteArray::fromAscii("seed"));
	PIByteArray sk2 = PIAuth::generateSign(PIByteArray::fromAscii("seed"));
	PIByteArray sk3 = PIAuth::generateSign(PIByteArray::fromAscii("other-seed"));

	EXPECT_FALSE(sk1.isEmpty());
	EXPECT_EQ(sk1, sk2);
	EXPECT_NE(sk1, sk3);

	PIAuth endpoint(sk1);
	EXPECT_EQ(PICrypt::extractSignPublicKey(sk1), endpoint.getSignPublicKey());
	EXPECT_FALSE(endpoint.getSignPublicKey().isEmpty());
}


//! \~english The authorized key list is exposed and replaced through the accessors.
//! \~russian Список доверенных ключей читается и заменяется через аксессоры.
TEST_F(PIAuthTest, AuthorizedKeysAccessors) {
	PIAuth endpoint(signKey("seed"));
	PIByteArray first  = signKey("first");
	PIByteArray second = signKey("second");

	EXPECT_TRUE(endpoint.getAuthorizedPublicKeys().isEmpty());
	endpoint.setAuthorizedPublicKeys(keyList(first));
	ASSERT_EQ(1, endpoint.getAuthorizedPublicKeys().size());
	EXPECT_EQ(first, endpoint.getAuthorizedPublicKeys()[0]);

	endpoint.setAuthorizedPublicKeys(PIVector<PIByteArray>{first, second});
	ASSERT_EQ(2, endpoint.getAuthorizedPublicKeys().size());
	EXPECT_EQ(first, endpoint.getAuthorizedPublicKeys()[0]);
	EXPECT_EQ(second, endpoint.getAuthorizedPublicKeys()[1]);

	endpoint.setAuthorizedPublicKeys(PIVector<PIByteArray>());
	EXPECT_TRUE(endpoint.getAuthorizedPublicKeys().isEmpty());
}