/*
 * Copyright (C) Tildeslash Ltd. All rights reserved.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License version 3.
 *
 * 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 General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 * In addition, as a special exception, the copyright holders give
 * permission to link the code of portions of this program with the
 * OpenSSL library under certain conditions as described in each
 * individual source file, and distribute linked combinations
 * including the two.
 *
 * You must obey the GNU Affero General Public License in all respects
 * for all of the code used other than OpenSSL.
 */

#include "config.h"

#ifdef HAVE_STRING_H
#include <string.h>
#endif

#include <stdlib.h>						// secure_getenv

#include "protocol.h"
#include "md5.h"

// libmonit
#include "exceptions/IOException.h"
#include "exceptions/ProtocolException.h"

// https://github.com/nmap/nmap/blob/master/nselib/pgsql.lua
#define PROTO_V3			3

#define PKT_Error			'E'
#define PKT_AuthRequest		'R'
#define PKT_Password		'p'
#define PKT_Terminate		'X'

#define AUTH_Plain			3
#define AUTH_MD5			5


const char * safe_getenv(const char * szName, const char * szDefault)
{
	const char * tmp;

#ifdef secure_getenv						// glibc 2.17
	tmp = secure_getenv(szName);
#else
	tmp = getenv(szName);
#endif
	if (tmp == NULL || *tmp == 0)
		return	szDefault;
	else
		return	tmp;	
}


char * str_push(char * buffer, const char * szString)
{
	if (szString == NULL)
		return	buffer;

	strcpy(buffer, szString);
	return	buffer + strlen(szString) + 1;
}


void md5(char * output, const char * data1, int len1, const char * data2, int len2)
{
	static char hex[] = "0123456789abcdef";
	md5_context_t ctx;
	md5_byte_t hash[16];
	char * tmp = output;

	md5_init(&ctx);
	md5_append(&ctx, (const md5_byte_t *) data1, len1);
	md5_append(&ctx, (const md5_byte_t *) data2, len2);
	md5_finish(&ctx, hash);
	for (int i = 0; i < 16; i++)
	{
		*tmp ++ = hex[hash[i] >> 4];
		*tmp ++ = hex[hash[i] & 0xf];
	}
}


/**
 *  PostgreSQL test.
 *
 *  @file
 */
void check_pgsql(Socket_T socket)
{
	const char * szPgUser		= safe_getenv("PGUSER", "postgres");
	const char * szPgDatabase	= safe_getenv("PGDATABASE", "postgres");
	const char * szPgPassword	= safe_getenv("PGPASSWORD", "");

	#define LOGIN_MAXLEN		255			//!!! lazy: update only one byte: requestLogin[3]
	char requestLogin[LOGIN_MAXLEN];

	#define REPLY_MINLEN		9
	#define REPLY_MAXLEN		1024
    char reply[REPLY_MAXLEN];
	int br;

	#define PKT_AuthRequestLen	9
	char responseAuthOk[PKT_AuthRequestLen] =
	{
		PKT_AuthRequest,
		0x00,                              /** Length */
		0x00,
		0x00,
		0x08,
		0x00,                              /** OK code 0 */
		0x00,
		0x00,
		0x00
	};

	char requestTerminate[5] =
	{
		PKT_Terminate,
		0x00,                              /** Length */
		0x00,
		0x00,
		0x04
	};

	size_t iLoginLen;
	char * ptr;

	iLoginLen = 8 + 4 + 1 + strlen(szPgUser) + 1 + 8 + 1 + strlen(szPgDatabase) + 1 + 1;
	if (iLoginLen > LOGIN_MAXLEN)
		THROW(ProtocolException, "PGSQL: internal error (insufficient buffer for request) -- %s", STRERROR);

	memset(requestLogin, 0, LOGIN_MAXLEN);		// set last NULL after payload
	ptr = requestLogin + 8;
	ptr = str_push(ptr, "user");
	ptr = str_push(ptr, szPgUser);
	ptr = str_push(ptr, "database");
	ptr = str_push(ptr, szPgDatabase);
	++ ptr;		// trailing NULL

	if (iLoginLen != ptr - &requestLogin[0])
		THROW(ProtocolException, "PGSQL: internal error (request packet) -- %s", STRERROR);

	requestLogin[3] = (unsigned char) iLoginLen;
	requestLogin[5] = PROTO_V3;

	ASSERT(socket);

	if (Socket_write(socket, requestLogin, iLoginLen) != iLoginLen)
		THROW(IOException, "PGSQL: error sending data -- %s", STRERROR);

	br = Socket_read(socket, reply, REPLY_MAXLEN - 1);
	if (br < REPLY_MINLEN)
		THROW(IOException, "PGSQL: error receiving data -- %s", STRERROR);

	if (*reply == PKT_Error)
		THROW(ProtocolException, "PGSQL: server replies Error -- %s", STRERROR);

// server *expects* password?
	if (*reply == PKT_AuthRequest && memcmp(reply, responseAuthOk, PKT_AuthRequestLen) != 0)
	{
		if (br == PKT_AuthRequestLen && reply[4] == br-1 && reply[8] == AUTH_Plain)
		{
			iLoginLen = 4 + strlen(szPgPassword);
			memset(requestLogin, 0, LOGIN_MAXLEN);
			requestLogin[0] = PKT_Password;
			requestLogin[4] = (unsigned char) iLoginLen;
			str_push(requestLogin + 5, szPgPassword);
			Socket_write(socket, requestLogin, iLoginLen + 1);
		}
		else
			if (br == PKT_AuthRequestLen + 4 && reply[4] == br-1 && reply[8] == AUTH_MD5)
			{
				//!!! first char is PKT_Password
				char pkt_hash[8 + 32 + 1] = "p\0\0\0\x28md5                                \0";
				md5(pkt_hash + 8, szPgPassword, strlen(szPgPassword), szPgUser, strlen(szPgUser));
				md5(pkt_hash + 8, pkt_hash + 8, 32, reply + 9, 4);
				Socket_write(socket, pkt_hash, sizeof(pkt_hash));
			}
			else
				THROW(ProtocolException, "PGSQL: unknown authentication type -- %s", STRERROR);

		br = Socket_read(socket, reply, REPLY_MAXLEN - 1);
		if (br < REPLY_MINLEN)
			THROW(IOException, "PGSQL: error receiving data -- %s", STRERROR);

//!!! PKT_Error: FATAL C28000 password authentication failed for user...
	}

// R: successful connection?
	if (br >= PKT_AuthRequestLen && memcmp(reply, responseAuthOk, PKT_AuthRequestLen) == 0)
	{
		reply[br] = 0;
		if (strstr(reply, "does not exist") != NULL)
			THROW(ProtocolException, "PGSQL: user/database not exist -- %s", STRERROR);

		Socket_write(socket, requestTerminate, sizeof(requestTerminate));
		return;
	}

	THROW(ProtocolException, "PGSQL: unknown server response -- %s", STRERROR);
}