/*
	PIP - Platform Independent Primitives
	PIP Authentication API
	Andrey Bychkov work.a.b@yandex.ru

	This program is free software: you can redistribute it and/or modify
	it under the terms of the GNU Lesser General Public License as published by
	the Free Software Foundation, either version 3 of the License, or
	(at your option) any later version.

	This program is distributed in the hope that it will be useful,
	but WITHOUT ANY WARRANTY; without even the implied warranty of
	MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
	GNU Lesser General Public License for more details.

	You should have received a copy of the GNU Lesser General Public License
	along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/

#include "piauth.h"
#define PIAUTH_NOISE_MAX_SIZE 256

PIAuth::PIAuth(const PIByteArray & sign) : PIObject() {
	setName("Client");
	role = Client;
	state = NotConnected;
	sign_sk = sign;
	sign_pk = crypt.extractSignPublicKey(sign);
}


void PIAuth::setServerPassword(const PIString & ps) {
	pass_hash = crypt.passwordHash(ps, PIByteArray::fromHex("PIAuth"));
}


void PIAuth::stop() {
	role = Client;
	state = NotConnected;
	auth_sign.clear();
	box_sk.clear();
	box_pk.clear();
	my_pk.clear();
}


void PIAuth::startClient() {
	role = Client;
	state = AuthProbe;
}


PIByteArray PIAuth::startServer() {
	setName("Server");
	role = Server;
	state = AuthProbe;
	PIByteArray ba;
	crypt.generateKeypair(my_pk, box_sk);
	PIByteArray noise = crypt.generateRandomBuff(randomi()%PIAUTH_NOISE_MAX_SIZE+128);
	ba << (int)state << custom_info << sign_pk << my_pk << noise;
	PIByteArray sign = crypt.signMessage(ba, sign_sk);
	ba << sign;
	return ba;
}


PIAuth::State PIAuth::receive(PIByteArray & ba) {
	if (ba.size() < sizeof(int)) return disconnect(ba, "invalid data size");
	State rstate;
	int s;
	ba >> s;
	rstate = (State)s;
//	if (state != rstate) return disconect(ba);

	//client side
	if (role == Client) {
		if (state == AuthProbe && rstate == AuthProbe) {
			if (ba.size() < sizeof(int)*5) return disconnect(ba, "invalid data size");
			PIByteArray rinfo;
			PIByteArray rsign;
			PIByteArray rsign_pk;
			PIByteArray noise;
			ba >> rinfo >> rsign_pk >> box_pk >> noise >> rsign;
			if (rsign_pk.isEmpty() || box_pk.isEmpty() || rsign.isEmpty()) return disconnect(ba, "invalid key size");

			PIByteArray tba;
			tba << (int)rstate << rinfo << rsign_pk << box_pk << noise;
			if (!crypt.verifySign(tba, rsign, rsign_pk)) return disconnect(ba, "Incorrect sign");
			bool auth = false;
			if (isAuthorizedKey(rsign_pk)) {
				auth = true;
			} else {
				authorize(rinfo, &auth);
				if (auth) auth_pkeys << rsign_pk;
			}
			if (!auth) return disconnect(ba, "Unauthorised");
			ba.clear();
			auth_sign = rsign_pk;
			crypt.generateKeypair(my_pk, box_sk);
			tba.clear();
			tba << sign_pk << my_pk << box_pk;
			PIByteArray sign = crypt.signMessage(tba, sign_sk);
			tba << sign;
			tba = crypt.crypt(tba, box_pk, box_sk);
			state = AuthReply;
			noise = crypt.generateRandomBuff(randomi()%PIAUTH_NOISE_MAX_SIZE);
			ba << (int)state << tba << my_pk << noise;
			sign = crypt.signMessage(ba, sign_sk);
			ba << sign;
			return state;
		}
		if (state == AuthReply && rstate == PassRequest) {
			PIByteArray ctba, tba;
			PIByteArray noise;
			PIByteArray rsign, rsign_pk;
			ba >> ctba >> rsign;
			bool ok = false;
			tba = crypt.decrypt(ctba, box_pk, box_sk, &ok);
			if (tba.isEmpty() || !ok) return disconnect(ba, "Message corrupted");
			ba.clear();
			ba << (int)rstate << ctba;
			if (!crypt.verifySign(ba, rsign, auth_sign)) return disconnect(ba, "Incorrect sign");
			ctba.clear();
			tba >> rsign_pk >> noise >> ctba;
			if (rsign_pk != auth_sign || ctba != my_pk) return disconnect(ba, "Invalid public key");
			PIString ps;
			PIByteArray ph;
			passwordRequest(&ps);
			if (ps.isEmpty()) return disconnect(ba, "Canceled by user");
			ph = crypt.passwordHash(ps, PIByteArray::fromHex("PIAuth"));
			ps.fill(0);
			tba.clear();
			tba << ph << auth_sign << sign_pk;
			tba = crypt.crypt(tba, box_pk, box_sk);
			ba.clear();
			state = PassRequest;
			ba << (int)state << tba;
			rsign = crypt.signMessage(ba, sign_sk);
			ba << rsign;
			return state;
		}
		if ((state == AuthReply && rstate == KeyExchange) || (state == PassRequest && rstate == KeyExchange)) {
			PIByteArray tba, ctba;
			PIByteArray rsign;
			ba >> ctba >> rsign;
			bool ok = false;
			tba = crypt.decrypt(ctba, box_pk, box_sk, &ok);
			if (tba.isEmpty() || !ok) return disconnect(ba, "Message corrupted");
			ba.clear();
			ba << (int)rstate << ctba;
			if (!crypt.verifySign(ba, rsign, auth_sign)) return disconnect(ba, "Incorrect sign");
			tba >> secret_key;
			if (secret_key.size() != crypt.sizeKey()) return disconnect(ba, "Invalid key");
			ba.clear();
			state = Connected;
			connected(PIString());
			ba << (int)state << crypt.crypt(custom_info, secret_key) << crypt.generateRandomBuff(randomi()%PIAUTH_NOISE_MAX_SIZE);
			return state;
		}
		if (state == Connected && rstate == Connected) {
			ba.clear();
			state = Connected;
			connected(PIString());
			ba << (int)state << crypt.crypt(custom_info, secret_key) << crypt.generateRandomBuff(randomi()%PIAUTH_NOISE_MAX_SIZE);
			return state;
		}
	}

	// server side
	if (role == Server) {
		if (state == AuthProbe && rstate == AuthReply) {
			if (ba.size() < sizeof(int)*4) return disconnect(ba, "invalid data size");
			PIByteArray ctba, tba;
			PIByteArray noise;
			PIByteArray rsign1, rsign2;
			PIByteArray rsign_pk;
			PIByteArray pk, mpk;
			ba >> ctba >> pk >> noise >> rsign1;
			bool ok = false;
			tba = crypt.decrypt(ctba, pk, box_sk, &ok);
			if (tba.isEmpty() || !ok) return disconnect(ba, "Message corrupted");
			if (tba.size() < sizeof(int)*3) return disconnect(tba, "invalid data size");
			tba >> rsign_pk >> box_pk >> mpk >> rsign2;
			if (pk != box_pk || mpk != my_pk) return disconnect(ba, "Invalid public key");
			ba.clear();
			ba << (int)rstate << ctba << box_pk << noise;
			if (!crypt.verifySign(ba, rsign1, rsign_pk)) return disconnect(ba, "Incorrect sign");
			ba.clear();
			ba << rsign_pk << box_pk << my_pk;
			if (!crypt.verifySign(ba, rsign2, rsign_pk)) return disconnect(ba, "Incorrect sign");
			auth_sign = rsign_pk;
			if (isAuthorizedKey(rsign_pk)) {
				state = KeyExchange;
				ba = createSKMessage();
				return state;
			} else {
				ba.clear();
				tba.clear();
				state = PassRequest;
				noise = crypt.generateRandomBuff(randomi()%PIAUTH_NOISE_MAX_SIZE);
				tba << sign_pk << noise << box_pk;
				tba = crypt.crypt(tba, box_pk, box_sk);
				ba << (int)state << tba;
				rsign1 = crypt.signMessage(ba, sign_sk);
				ba << rsign1;
				return state;
			}
		}
		if (state == PassRequest && rstate == PassRequest) {
			PIByteArray tba, ctba;
			PIByteArray rsign_pk, rsign, mpk;
			ba >> ctba >> rsign;
			bool ok = false;
			tba = crypt.decrypt(ctba, box_pk, box_sk, &ok);
			if (tba.isEmpty() || !ok) return disconnect(ba, "Message corrupted");
			ba.clear();
			ba << (int)rstate << ctba;
			if (!crypt.verifySign(ba, rsign, auth_sign)) return disconnect(ba, "Incorrect sign");
			ctba.clear();
			tba >> ctba >> mpk >> rsign_pk;
			if (rsign_pk != auth_sign || mpk != sign_pk) return disconnect(ba, "Invalid public key");
			bool auth = (ctba == pass_hash);
			if (ctba.isEmpty() || pass_hash.isEmpty()) auth = false;
			passwordCheck(auth);
			if (!auth) {
//				piSleep(1);
				return disconnect(ba, "Invalid password");
			}
			state = KeyExchange;
			ba = createSKMessage();
			return state;
		}
		if ((state == KeyExchange && rstate == Connected) || (state == Connected && rstate == Connected)) {
			ba.clear();
			PIByteArray rinfo;
			ba >> rinfo;
			bool ok = false;
			rinfo = crypt.decrypt(rinfo, secret_key, &ok);
			if (!ok) return disconnect(ba, "Error while exchange keys");
			state = Connected;
			connected(rinfo);
			ba << (int)state << crypt.generateRandomBuff(randomi()%PIAUTH_NOISE_MAX_SIZE);
			return state;
		}
	}

	return disconnect(ba, "invalid state " + PIString::fromNumber((int)state));
}


PIByteArray PIAuth::getSecretKey() {
	if (state == Connected) return secret_key;
	return PIByteArray();
}


PIByteArray PIAuth::generateSign(const PIByteArray & seed) {
	PIByteArray pk, sk;
	PICrypt::generateSignKeys(pk, sk, seed);
	return sk;
}


PIAuth::State PIAuth::disconnect(PIByteArray & ba, const PIString & error) {
	if (!error.isEmpty()) piCoutObj << error;
	auth_sign.clear();
	box_sk.clear();
	box_pk.clear();
	my_pk.clear();
	secret_key.clear();
	ba.clear();
	state = NotConnected;
	disconnected(error);
	return state;
}


bool PIAuth::isAuthorizedKey(const PIByteArray & pkey) {
	for (int i=0; i<auth_pkeys.size_s(); ++i) {
		if (pkey == auth_pkeys[i]) return true;
	}
	return false;
}


PIByteArray PIAuth::createSKMessage() {
	secret_key = crypt.generateKey();
	PIByteArray tba;
	PIByteArray noise = crypt.generateRandomBuff(randomi()%PIAUTH_NOISE_MAX_SIZE);
	tba << secret_key << noise;
	tba = crypt.crypt(tba, box_pk, box_sk);
	PIByteArray ret;
	ret << (int)state << tba;
	PIByteArray sign = crypt.signMessage(ret, sign_sk);
	ret << sign;
	return ret;
}