/*-
 * Copyright (c) 2005 Eirik A. Nygaard.  All rights reserved.
 * 
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 * 
 * 1. Redistributions of source code must retain the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions in binary form must reproduce the above copyright
 *    notice, this list of conditions and the following disclaimer in
 *    the documentation and/or other materials provided with the
 *    distribution.
 * 3. The name of the author may not be used to endorse or promote products
 *    derived from this software without specific prior written permission.
 * 
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 * FOR A PARTICULAR PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE
 * COPYRIGHT HOLDERS OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY OR CONSEQUENTIAL DAMAGES (INCLUDING,
 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
 * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
 * AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT
 * OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF
 * SUCH DAMAGE.
 */

#include <sys/types.h>
#include <sys/param.h>

#include <stdio.h>
#include <unistd.h>
#include <string.h>
#include <stdlib.h>
#include <err.h>
#include <pwd.h>

#include <libpq-fe.h>

#include "../psqluseradm/psql.config.h"
#include "../psqluseradm/psqladm.h"

#define ACTION_ADD 1
#define ACTION_DELETE 2
#define ACTION_EDIT 3
#define ACTION_LIST 4

PGconn *pgconn;

void usage(void);
int createdatabase(char *dbname, char *owner, uid_t uid, char *encoding);
int verify_owner(char *dbname, uid_t uid);
void listdatabases(uid_t uid);
int dropdatabase(char *dbname);
char *first_dbuser(uid_t uid);
void create_schema(char *name, char *owner, PGconn *conn);

int
main(int argc, char **argv)
{
	uid_t uid;
	int action = 0, c;
	char *dbname = NULL, *owner = NULL;
	char _dbname[1024], _owner[1024];
	char *encoding = NULL, _encoding[1024];

	while ((c = getopt(argc, argv, "hadlu:o:E:")) != -1) {
		switch (c) {
		case 'h':
			usage();
			exit(1);
		case 'a':
			action = ACTION_ADD;
			break;
		case 'd':
			action = ACTION_DELETE;
			break;
		case 'e':
			action = ACTION_EDIT;
			break;
		case 'l':
			action = ACTION_LIST;
			break;
		case 'u':
			dbname = optarg;
			break;
		case 'o':
			owner = optarg;
			break;
		case 'E':
			encoding = optarg;
			break;
		default:
			fprintf(stderr, "Unknown argument: '%c'\n", c);
			usage();
			exit(1);
		}
	}
	argc -= optind;
	argv += optind;
	
	if (!dbname && argc > 0) {
		dbname = argv[0];
		argc--;
		argv++;
	}
	if (!owner && argc > 0) {
		owner = argv[0];
		argc--;
		argv++;
	}
	
	if (!dbname &&
	    (action == ACTION_ADD || action == ACTION_EDIT || action == ACTION_DELETE)) {
		usage();
		errx(1, "Specify dbname!");
	}
	if (action == 0) {
		usage();
		errx(1, "No action specified!");
	}

	uid = get_cred();
	pgconn = db_connect(PSQL_USER, PSQL_PASSWORD, PSQL_DBNAME, PSQL_HOST);

	/* verify input */
	if (dbname && !verifystr(dbname, STR_TYPE_NAME))
		errx(1, "Invalid dbname input");
	if (owner && !verifystr(owner, STR_TYPE_NAME))
		errx(1, "Invalid owner input");
	if (encoding && !verifystr(encoding, STR_TYPE_NAME))
		errx(1, "Invalid encoding input");

	if (dbname) {
		struct passwd *passwd;
		if (strlen(dbname) > 1000)
			errx(1, "Too long dbname");
		PQescapeString(_dbname, dbname, MIN(strlen(dbname), sizeof(_dbname)));
		passwd = getpwuid(uid);
		if (!passwd)
			errx(1, "Unable to get login info");
		if (!verify_startwith(dbname, passwd->pw_name))
			errx(1, "Your database must start with your username");
	}
	if (!owner) {
		if ((owner = first_dbuser(uid)) == NULL)
			errx(1, "You must create a database user first");
		printf("Defaulting to '%s' as dbowner!\n", owner);
	}
	if (encoding) {
		PQescapeString(_encoding, encoding, MIN(strlen(encoding), sizeof(_encoding)));
		encoding = _encoding;
	}
	if (strlen(owner) > 1000)
		errx(1, "Too long owner name");
	PQescapeString(_owner , owner, MIN(strlen(owner), sizeof(_owner)));
	
	switch (action) {
	case ACTION_ADD:
		fprintf(stderr, "Creating database %s... ", _dbname);
		if (createdatabase(_dbname, _owner, uid, encoding) != 0)
			errx(1, "Unable to add new user");
		break;
	case ACTION_DELETE:
		if (!verify_owner(_dbname, uid))
			errx(1, "You are not the owner for that dbuser.");
		fprintf(stderr, "Deleting user... ");
		dropdatabase(dbname);
		break;
	case ACTION_EDIT:
		if (!verify_owner(_dbname, uid))
			errx(1, "You are not the owner for that dbuser.");
		fprintf(stderr, "Editing user... nothing to do here, sorry!\n");
		break;
	case ACTION_LIST:
		listdatabases(uid);
	}
	
	return(0);
}

void
usage(void)
{
	fprintf(stderr, "psql-dbadm [-h] [-a] [-d] [-l] [-E encoding] [dbname [owner]]\n");
	fprintf(stderr, "\t-h: This help screen\n");
	fprintf(stderr, "\t-a: Create a new database\n");
	fprintf(stderr, "\t-d: Drop a database\n");
	fprintf(stderr, "\t-l: List your databases\n");
	fprintf(stderr, "\t-E encoding: Specify encoding when you create a database\n");
}

int
createdatabase(char *dbname, char *owner, uid_t uid, char *encoding)
{
	char query[1024];
	int nfields;
	PGresult *res;
	PGconn *newpgconn;

	snprintf(query, sizeof(query), "SELECT id FROM psqldbadm WHERE dbname = '%s'", dbname);
	res = PQexec(pgconn, query);
	if (PQresultStatus(res) != PGRES_TUPLES_OK)
		errx(1, "Could not execute query: '%s': %s", query, PQerrorMessage(pgconn));
	nfields = PQntuples(res);
	if (nfields > 0)
		errx(1, "A database with that name already exists");
	snprintf(query, sizeof(query), "CREATE DATABASE %s OWNER = %s ",
	         dbname, owner);
	if (encoding) {
		char enc[1024];
		snprintf(enc, sizeof(enc), " ENCODING '%s'", encoding);
		strncat(query, enc, sizeof(query) - strlen(query));
	}
	res = PQexec(pgconn, query);
	if (PQresultStatus(res) != PGRES_COMMAND_OK)
		errx(1, "unable to add info db: %s", PQerrorMessage(pgconn));
	snprintf(query, sizeof(query), "INSERT INTO psqldbadm (uid, dbname, owner) " \
	                               "VALUES(%d, '%s', '%s')",
	         uid, dbname, owner);
	res = PQexec(pgconn, query);
	if (PQresultStatus(res) != PGRES_COMMAND_OK)
		errx(1, "unable to add database: %s", PQerrorMessage(pgconn));
	snprintf(query, sizeof(query), "GRANT ALL ON DATABASE %s TO %s", dbname, owner);
	res = PQexec(pgconn, query);
	if (PQresultStatus(res) != PGRES_COMMAND_OK)
		errx(1, "unable to grant owner access to database: %s", PQerrorMessage(pgconn));
	printf("done!\n");
	newpgconn = db_connect(PSQL_USER, PSQL_PASSWORD, (const char *)dbname, PSQL_HOST);
	create_schema(owner, owner, newpgconn);
	return(0);
}

void
create_schema(char *name, char *owner, PGconn *conn)
{
	char query[1024];
	PGresult *res;

	snprintf(query, sizeof(query), "CREATE SCHEMA %s AUTHORIZATION %s", name, owner);
	res = PQexec(conn, query);
	if (PQresultStatus(res) != PGRES_COMMAND_OK)
		errx(1, "unable to create user schema: %s", PQerrorMessage(conn));
}

void
listdatabases(uid_t uid)
{
	char query[1024];
	PGresult *res;
	int i, nfields;

	snprintf(query, sizeof(query), "SELECT dbname, owner, ts FROM psqldbadm WHERE uid = %d", uid);
	res = PQexec(pgconn, query);
	if (PQresultStatus(res) != PGRES_TUPLES_OK)
		errx(1, "unable to list databases: %s: %s", query, PQerrorMessage(pgconn));
	nfields = PQnfields(res);
	printf("%-15s %-15s %-15s\n", "DB", "DBowner", "Created");
	for (i = 0; i < PQntuples(res); i++) {
		/* XXX: Format timestamp... */
		printf("%-15s %-15s %-15s\n", PQgetvalue(res, i, 0), PQgetvalue(res, i, 1),
		       PQgetvalue(res, i, 2));
	}
	if (PQntuples(res) == 0)
		printf("You don't have any databases...\n");
}

int
dropdatabase(char *dbname)
{
	char query[1024];
	PGresult *res;
	
	snprintf(query, sizeof(query), "DROP DATABASE %s", dbname);
	res = PQexec(pgconn, query);
	if (PQresultStatus(res) != PGRES_COMMAND_OK)
		errx(1, "unable to drop database: %s: %s", query, PQerrorMessage(pgconn));
	snprintf(query, sizeof(query), "DELETE FROM psqldbadm WHERE dbname = '%s'", dbname);
	res = PQexec(pgconn, query);
	if (PQresultStatus(res) != PGRES_COMMAND_OK)
		errx(1, "unable to delete info database: %s: %s", query, PQerrorMessage(pgconn));
	printf("done!\n");
	return(0);
}

int
verify_owner(char *dbname, uid_t uid)
{
	char query[1024];
	PGresult *res;

	snprintf(query, sizeof(query), "SELECT id FROM psqldbadm " \
	                                "WHERE dbname = '%s' AND uid = %d",
	         dbname, uid);
	res = PQexec(pgconn, query);
	if (PQresultStatus(res) != PGRES_TUPLES_OK)
		errx(1, "verify query failed: %s: %s", query, PQerrorMessage(pgconn));
	if (PQntuples(res) > 0)
		return(1);
	return(0);
}


char *
first_dbuser(uid_t uid)
{
	char query[1024];
	char *owner;
	PGresult *res;

	snprintf(query, sizeof(query), "SELECT username FROM psqluseradm WHERE uid = %d LIMIT 1", uid);
	res = PQexec(pgconn, query);
	if (PQresultStatus(res) != PGRES_TUPLES_OK)
		errx(1, "could not find first dbuser: %s: %s", query, PQerrorMessage(pgconn));
	if (PQntuples(res) == 0)
		return(NULL);
	owner = PQgetvalue(res, 0, 0);
	return(owner);
}