/* Copyright (c) 2011 Nanakos Chrysostomos <nanakos@wired-net.gr>
   Simple and lightweight Yubikey OTP-OATH/HOTP Validation Server

   yubiserver-admin is placed under the GNU General Public License, version 2 or later.

   This program is free software: you can redistribute it and/or modify
   it under the terms of the GNU 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 General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with this program.  If not, see <http://www.gnu.org/licenses/>.
*/
#include <stdio.h>
#include <stdlib.h>
#include <unistd.h>
#include <errno.h>
#include <string.h>
#include <fcntl.h>
#include <signal.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <gcrypt.h>
#include <sqlite3.h>
#include <time.h>
#include <sys/timeb.h>
#include <getopt.h>
#include <math.h>

#include <mhash.h>

#include "yubiserver.h"

int yubikey_table = 0;
int oath_table = 0;
int api_table = 0;
int ext_db = 0;
int Argc;
char **Argv;
#define YUBIKEY	0
#define OATH	1
#define API	2

static char encoding_table[] = {'A', 'B', 'C', 'D', 'E', 'F', 'G', 'H',
        'I', 'J', 'K', 'L', 'M', 'N', 'O', 'P',
        'Q', 'R', 'S', 'T', 'U', 'V', 'W', 'X',
        'Y', 'Z', 'a', 'b', 'c', 'd', 'e', 'f',
        'g', 'h', 'i', 'j', 'k', 'l', 'm', 'n',
        'o', 'p', 'q', 'r', 's', 't', 'u', 'v',
        'w', 'x', 'y', 'z', '0', '1', '2', '3',
        '4', '5', '6', '7', '8', '9', '+', '/'};

static int mod_table[] = {0, 2, 1};

char *base64_encode(const unsigned char *data,size_t input_length,size_t output_length) {
        int i,j;
        output_length = (size_t) (4.0 * ceil((double) input_length / 3.0));
        char *encoded_data = malloc(output_length);
        if (encoded_data == NULL) return NULL;
        for (i = 0, j = 0; i < input_length;) {
                uint32_t octet_a = i < input_length ? data[i++] : 0;
                uint32_t octet_b = i < input_length ? data[i++] : 0;
                uint32_t octet_c = i < input_length ? data[i++] : 0;
                uint32_t triple = (octet_a << 0x10) + (octet_b << 0x08) + octet_c;

                encoded_data[j++] = encoding_table[(triple >> 3 * 6) & 0x3F];
                encoded_data[j++] = encoding_table[(triple >> 2 * 6) & 0x3F];
                encoded_data[j++] = encoding_table[(triple >> 1 * 6) & 0x3F];
                encoded_data[j++] = encoding_table[(triple >> 0 * 6) & 0x3F];
        }

        for (i = 0; i < mod_table[input_length % 3]; i++)
                encoded_data[output_length - 1 - i] = '=';
        return encoded_data;
}

void print_table(char *table,int id)
{
	sqlite3 *handle;
	sqlite3_stmt *stmt;
	int retval,rows=0;
	char *query = NULL;

	if(id==YUBIKEY || id==OATH)
	{
		query = (char *)malloc((size_t)100);
		sprintf(query,"SELECT nickname,publicname,active FROM %s",table);
	}
	else if(id==API)
	{
		query = (char *)malloc((size_t)100);
		sprintf(query,"SELECT nickname,id FROM %s",table);
	}
	
	retval = sqlite3_open(sqlite3_dbpath,&handle);

	if(retval)
	{
		fprintf(stderr,"Database connection failed");
		exit(2);
	}

	if(query != NULL) retval = sqlite3_prepare_v2(handle,query,-1,&stmt,0);

	if(retval)
	{
		fprintf(stderr,"Selecting data from DB failed");
		exit(2);
	}

	while(1)
	{
		retval = sqlite3_step(stmt);

		if(retval==SQLITE_ROW)
		{
			if(id==YUBIKEY || id==OATH)
			{
				if(!rows) printf("%-20s %-20s %-20s\n","[Username]","[Public Token ID]","Active");
				printf("%-20s %-20s %-20s\n",sqlite3_column_text(stmt,0),sqlite3_column_text(stmt,1),sqlite3_column_text(stmt,2));
				rows++;
			}else if(id==API)
			{
				if(!rows) printf("%-20s %-20s\n","[Username]","[API ID]");
				printf("%-20s %-20s\n",sqlite3_column_text(stmt,0),sqlite3_column_text(stmt,1));
				rows++;
			}
		}
		else if(retval==SQLITE_DONE)
		{
			break;
		}
		else
		{
			fprintf(stderr,"Database error encountered");
			break;
		}
	}
	sqlite3_close(handle);
	free(query);

	if(!rows) printf("No keys in database.\n");
	printf("Total keys in database: %d\n",rows);
	exit(0);
}

void db_ops(char *table,char *user,int operation,int id)
{
	sqlite3 *handle;
        sqlite3_stmt *stmt;
        int retval;
        char *query = NULL;
	time_t t;
	struct tm *tmp;
	struct timeb tp;
	char *cdatetime = NULL;
	char *cdatetime_tmp = NULL;

	if(id == YUBIKEY || id == OATH)
	{
		if(operation == ENABLE_USER || operation == DISABLE_USER || operation == DELETE_USER)
		{
			query = (char *)malloc((size_t)100);
	                sprintf(query,"SELECT * FROM %s WHERE nickname='%s'",table,user);
			retval = sqlite3_open(sqlite3_dbpath,&handle);

			if(retval)
			{
				fprintf(stderr,"Database connection failed");
				exit(2);
			}
			
			if(query != NULL) retval = sqlite3_prepare_v2(handle,query,-1,&stmt,0);
			
			if(retval)
			{
				fprintf(stderr,"Selecting data from DB failed");
				goto out;
			}
			retval = sqlite3_step(stmt);
				
			if(retval==SQLITE_ROW)
			{
				fprintf(stdout,"User '%s' found\n",user);
				if(operation == ENABLE_USER)
				{
					sprintf(query,"SELECT * FROM %s WHERE nickname = '%s'  AND active = '1'",table,user);
					retval = sqlite3_prepare_v2(handle,query,-1,&stmt,0);
				}
				else if(operation == DISABLE_USER)
				{
					sprintf(query,"SELECT * FROM %s WHERE nickname = '%s'  AND active = '0'",table,user);
					retval = sqlite3_prepare_v2(handle,query,-1,&stmt,0);
				}
				else if(operation == DELETE_USER)
				{
					sprintf(query,"DELETE FROM %s WHERE nickname = '%s'",table,user);
					retval = sqlite3_exec(handle,query,0,0,0);
					if(retval)
				       	{
						fprintf(stderr,"Deleting user '%s' failed",user);
						goto out;
					}
					fprintf(stdout,"User '%s' deleted\n",user);
					goto out;
				}

				if(retval)
				{
					fprintf(stderr,"Selecting data from DB failed");
					goto out;
				}
				retval = sqlite3_step(stmt);
				if(retval==SQLITE_ROW)
				{
					if(operation == ENABLE_USER)
						fprintf(stdout,"User '%s' is already enabled\n",user);
					else if(operation == DISABLE_USER)
						fprintf(stdout,"User '%s' is already disabled\n",user);
				}
				else if(retval == SQLITE_DONE)
				{
					if(operation == ENABLE_USER)
					{
						fprintf(stdout,"Enabling user '%s'\n",user);
						sprintf(query,"UPDATE %s SET active = '1' WHERE nickname = '%s'",table,user);
						retval = sqlite3_exec(handle,query,0,0,0);
						if(retval)
						{
							fprintf(stderr,"Selecting data from DB failed");
							goto out;
						}
						fprintf(stdout,"User '%s' enabled\n",user);

					}
					else if(operation == DISABLE_USER)
					{
						fprintf(stdout,"Disabling user '%s'\n",user);
						sprintf(query,"UPDATE %s SET active = '0' WHERE nickname = '%s'",table,user);
						retval = sqlite3_exec(handle,query,0,0,0);
						if(retval)
						{
							fprintf(stderr,"Selecting data from DB failed");
							exit(2);
						}
						fprintf(stdout,"User '%s' disabled\n",user);
					}
				}
				else
				{
					fprintf(stderr,"Database error encountered");
					goto out;
				}
			}
			else if(retval==SQLITE_DONE)
			{
				fprintf(stdout,"User '%s' not found\n",user);
			}
			else
			{
				fprintf(stderr,"Database error encountered");
				exit(0);
			}
		}else if(operation == ADD_USER)
		{
			if(id==YUBIKEY)
			{
				if(!ext_db)
				{
					if((strlen(user) > 16 || strlen(Argv[4])!=12 || strlen(Argv[5]) != 12 || strlen(Argv[6]) != 32))
					{
						fprintf(stderr,"Parameters are not correct. Please try again.\n");
						fprintf(stderr,"Public token ID must be 12 characters long.\n");
						fprintf(stderr,"Secret token ID must be 12 characters long.\n");
						fprintf(stderr,"AES key must be 32 characters long.\n");
						exit(2);
					}
				}else if(ext_db)
				{
					if((strlen(user) > 16 || strlen(Argv[6])!=12 || strlen(Argv[7]) != 12 || strlen(Argv[8]) != 32))
					{
						fprintf(stderr,"Parameters are not correct. Please try again.\n");
						fprintf(stderr,"Public token ID must be 12 characters long.\n");
						fprintf(stderr,"Secret token ID must be 12 characters long.\n");
						fprintf(stderr,"AES key must be 32 characters long.\n");
						exit(2);
					}
				}
				query = (char *)malloc((size_t)200);
				if(!ext_db)
				{
					sprintf(query,"SELECT * FROM %s WHERE nickname='%s' OR publicname='%s'",table,user,Argv[4]);
				}else if(ext_db)
				{
					sprintf(query,"SELECT * FROM %s WHERE nickname='%s' OR publicname='%s'",table,user,Argv[6]);
				}
				retval = sqlite3_open(sqlite3_dbpath,&handle);
				if(retval)
				{
					fprintf(stderr,"Database connection failed");
					exit(2);
				}
			
				if(query != NULL) retval = sqlite3_prepare_v2(handle,query,-1,&stmt,0);
			
				if(retval)
				{
					fprintf(stderr,"Selecting data from DB failed");
					goto out;
				}
				retval = sqlite3_step(stmt);
				
				if(retval==SQLITE_ROW)
				{
					fprintf(stdout,"Username or key already exist into database. Delete it before adding the same key.\n");
					goto out;
				}else if(retval == SQLITE_DONE)
				{
					fprintf(stdout,"Username/key not found.Adding it into database.\n");
					t = time(NULL);
					tmp = localtime(&t);
					ftime(&tp);
					cdatetime = (char *)malloc((char)25);
					cdatetime_tmp = (char *)malloc((char)20);
					strftime(cdatetime_tmp,20,"%Y-%m-%dT%H:%M:%S",tmp);
					sprintf(cdatetime,"%s.%.3dZ",cdatetime_tmp,tp.millitm);
					if(!ext_db)
					{
						sprintf(query,"INSERT INTO %s VALUES('%s','%s','%s','%s','%s',1,1,1)",table,user,Argv[4],cdatetime,Argv[5],Argv[6]);
					}else if(ext_db)
					{
						sprintf(query,"INSERT INTO %s VALUES('%s','%s','%s','%s','%s',1,1,1)",table,user,Argv[6],cdatetime,Argv[7],Argv[8]);
					}
					retval = sqlite3_exec(handle,query,0,0,0);
					free(cdatetime); free(cdatetime_tmp);
					if(retval)
					{
						fprintf(stderr,"Adding user '%s' failed",user);
						goto out;
					}
					fprintf(stdout,"User '%s' added to database\n",user);
					goto out;
				}

			}else if(id==OATH)
			{
				if(!ext_db)
				{
					if((strlen(user) > 16 || strlen(Argv[4])!=12 || strlen(Argv[5]) != 40))
					{
						fprintf(stderr,"Parameters are not correct. Please try again.\n");
						fprintf(stderr,"Public token ID must be 12 characters long.\n");
						fprintf(stderr,"Secret key must be 40 characters long.\n");
						exit(2);
					}
				}
				else  if(ext_db)
				{
					if((strlen(user) > 16 || strlen(Argv[6])!=12 || strlen(Argv[7]) != 40))
					{
						fprintf(stderr,"Parameters are not correct. Please try again.\n");
						fprintf(stderr,"Public token ID must be 12 characters long.\n");
						fprintf(stderr,"Secret key must be 40 characters long.\n");
						exit(2);
					}
				}
				query = (char *)calloc(1,(size_t)150);
				if(!ext_db)
				{
					sprintf(query,"SELECT * FROM %s WHERE nickname='%s' OR publicname='%s'",table,user,Argv[4]);
				}else if(ext_db)
				{
					sprintf(query,"SELECT * FROM %s WHERE nickname='%s' OR publicname='%s'",table,user,Argv[6]);
				}
				retval = sqlite3_open(sqlite3_dbpath,&handle);
				if(retval)
				{
					fprintf(stderr,"Database connection failed");
					exit(2);
				}
			
				if(query != NULL) retval = sqlite3_prepare_v2(handle,query,-1,&stmt,0);
			
				if(retval)
				{
					fprintf(stderr,"Selecting data from DB failed");
					goto out;
				}
				retval = sqlite3_step(stmt);
				
				if(retval==SQLITE_ROW)
				{
					fprintf(stdout,"Username or key already exist into database. Delete it before adding the same key.\n");
					goto out;
				}else if(retval == SQLITE_DONE)
				{
					fprintf(stdout,"Username/key not found.Adding it into database.\n");
					t = time(NULL);
					tmp = localtime(&t);
					cdatetime = (char *)malloc((char)21);
					strftime(cdatetime,21,"%Y-%m-%dT%H:%M:%SZ",tmp);
					if(!ext_db)
					{
						sprintf(query,"INSERT INTO %s VALUES('%s','%s','%s','%s',1,1)",table,user,Argv[4],cdatetime,Argv[5]);
					}else if(ext_db)
					{
						sprintf(query,"INSERT INTO %s VALUES('%s','%s','%s','%s',1,1)",table,user,Argv[6],cdatetime,Argv[7]);
					}	
					retval = sqlite3_exec(handle,query,0,0,0);
					free(cdatetime);
					if(retval)
					{
						fprintf(stderr,"Adding user '%s' failed\n",user);
						goto out;
					}
					fprintf(stdout,"User '%s' added to database\n",user);
					goto out;
				}
			}
		}
	}else if(id == API)
	{
		if(operation == ENABLE_USER || operation == DISABLE_USER)
		{
			fprintf(stderr,"You cannot enable/disable '%s' user in API's table\n",user);
			exit(0);
		}
		else if(operation == DELETE_USER)
		{
			query = (char *)malloc((size_t)100);
			sprintf(query,"SELECT * FROM %s WHERE nickname='%s'",table,user);
			retval = sqlite3_open(sqlite3_dbpath,&handle);
			if(retval)
			{
				fprintf(stderr,"Database connection failed");
				exit(2);
			}
			
			if(query != NULL) retval = sqlite3_prepare_v2(handle,query,-1,&stmt,0);
			
			if(retval)
			{
				fprintf(stderr,"Selecting data from DB failed");
				goto out;
			}
			retval = sqlite3_step(stmt);
			
			if(retval==SQLITE_ROW)
			{
				fprintf(stdout,"User '%s' found\n",user);
			       	sprintf(query,"DELETE FROM %s WHERE nickname = '%s'",table,user);
				
			       	retval = sqlite3_exec(handle,query,0,0,0);
		       		if(retval)
				{
					 fprintf(stderr,"Deleting user '%s' failed",user);
					 goto out;
				}
				fprintf(stdout,"User '%s' deleted\n",user);
				goto out;
			}else if(retval==SQLITE_DONE)
			{
				fprintf(stdout,"User '%s' not found\n",user);
			}
			else
			{
				fprintf(stderr,"Database error encountered");
				exit(0);
			}
		}else if(operation == ADD_USER)
		{
			query = (char *)malloc((size_t)100);
			sprintf(query,"SELECT * FROM %s WHERE nickname='%s'",table,user);
			int api_id = 0;

			retval = sqlite3_open(sqlite3_dbpath,&handle);

			if(retval)
			{
				fprintf(stderr,"Database connection failed");
				exit(2);
			}

			if(query != NULL) retval = sqlite3_prepare_v2(handle,query,-1,&stmt,0);
			
			if(retval)
			{
				fprintf(stderr,"Selecting data from DB failed");
				goto out;
			}
			retval = sqlite3_step(stmt);
			
			if(retval==SQLITE_ROW)
			{
				fprintf(stdout,"API Key for username '%s' is already present. Remove it or choose another one.\n",user);
				goto out;
			}else if(retval==SQLITE_DONE)
			{
				fprintf(stdout,"User '%s' not found.Adding API Key\n",user);
				sprintf(query,"SELECT id FROM apikeys ORDER BY id DESC LIMIT 1");
				if(query != NULL) retval = sqlite3_prepare_v2(handle,query,-1,&stmt,0);
				if(retval)
				{
					fprintf(stderr,"Selecting data from DB failed");
					goto out;
				}
				retval = sqlite3_step(stmt);

				if(retval == SQLITE_ROW)
				{
					api_id = atoi((char *)sqlite3_column_text(stmt,0))+1;

				}else if(retval==SQLITE_DONE)
				{
					api_id = 1;
				}
				else{
					fprintf(stderr,"Database error encountered");
	                                exit(0);
				}
				if(!ext_db)
				{
					if(Argc!=5)
					{
						fprintf(stderr,"No API Key given.\n");
						goto out;
					}
				}else if(ext_db)
				{
					if(Argc!=7)
					{
						fprintf(stderr,"No API Key given.\n");
						goto out;
					}
				}
				if(!ext_db)
				{
					if(strlen(Argv[4])>20 || strlen(Argv[4])<20)
					{
						fprintf(stderr,"Invalid API Key. It must be 20 characters long.\n");
						goto out;
					}
					sprintf(query,"INSERT INTO apikeys VALUES ('%s','%s','%d')",user,Argv[4],api_id);
				}
				else if(ext_db)
				{
					if(strlen(Argv[6])>20 || strlen(Argv[6])<20)
					{
						fprintf(stderr,"Invalid API Key. It must be 20 characters long.\n");
						goto out;
					}
					sprintf(query,"INSERT INTO apikeys VALUES ('%s','%s','%d')",user,Argv[6],api_id);
				}
				retval = sqlite3_exec(handle,query,0,0,0);
				if(retval)
				{
					fprintf(stderr,"Adding new API KEY for '%s' failed",user);
					goto out;
				}
				if(!ext_db)
				{
					fprintf(stdout,"New API Key for '%s': %s\n",user,base64_encode((unsigned char *)Argv[4],20,20));
				}else if(ext_db)
				{
					fprintf(stdout,"New API Key for '%s': %s\n",user,base64_encode((unsigned char *)Argv[6],20,20));
				}
					fprintf(stdout,"Your API Key ID is: %d\n",api_id);
				goto out;
			}
			else
			{
				fprintf(stderr,"Database error encountered");
				exit(0);
			}
		}
	}

out:
	sqlite3_close(handle);
	free(query);
	exit(0);
}	

void show_table()
{
	if(yubikey_table && !oath_table && !api_table)
		print_table("yubikeys",YUBIKEY);
	else if(!yubikey_table && oath_table && !api_table)
		print_table("oathtokens",OATH);
	else if(!yubikey_table && !oath_table && api_table)
		print_table("apikeys",API);
	else 
		printf("Please choose a table to list.\n");
}

void user_ops(char *user,int operation)
{
	if(yubikey_table && !oath_table && !api_table)
		db_ops("yubikeys",user,operation,YUBIKEY);
	else if(!yubikey_table && oath_table && !api_table)
		db_ops("oathtokens",user,operation,OATH);
	else if(!yubikey_table && !oath_table && api_table)
		db_ops("apikeys",user,operation,API);
	else
		printf("Please choose a table first.\n");
		exit(2);
}

void usage()
{
	fprintf(stderr,"yubiserve-admin Yubikey Database Management Tool\n"
		       "Version " VERSION ". Written and copyrights by Nanakos Chrysostomos.\n"
		       "THIS SOFTWARE COMES WITH ABSOLUTELY NO WARRANTY! USE AT YOUR OWN RISK!\n"
		       "\nUsage: yubiserver-admin [[-b FILE]] [table] [options] [[attributes]] \n\n"
		       "Options supported:\n"
		       "   --version or -V			Print version information\n"
		       "   --help or -h				Print this help screen\n"
		       "   --database or -b			Use this SQLite3 database file (optional)\n"
		       "   --yubikey or -y			Choose Yubikey Token table\n"
		       "   --oath or -o				Choose OATH Token table\n"
		       "   --api or -p				Choose API Key table\n"
		       "   --add N [P S [A]] or -a N [P S [A]]	Add Yubikey OTP & HOTP/OATH token or API Key 'N' user\n"
		       "					where N is the username, P the Public Token ID,\n"
		       "					S the Secret ID and A the AES Key\n"
		       "					N must be 16 characters max,P must be 12 characters\n"
		       "					for Yubikey OTP and 12 characters max for HOTP/OATH\n"
		       "					S must be 12 characters for Yubikey OTP and 40 for HOTP/OATH\n"
		       "					and AES key must be 32 characters\n"
		       "					Adding a user to API keys requires a username\n"
		       "					and a API Key 20 characters long\n"
		       "   --delete N or -x N			Delete Yubikey OTP, HOTP/OATH token or API Key 'N' user\n"
		       "   --enable N or -e N			Enable Yubikey OTP, HOTP/OATH token 'N' user\n"
		       "   --disable N or -d N			Disable Yubikey OTP, HOTP/OATH token 'N' user\n"
		       "   --list or -l				List Yubikey OTP, HOTP/OATH token or API Key\n"
		       "This version of yubiserver-admin has been configured with '" SQLITE3_DB_PATH "' as its default\n"
                       "SQLite3 database file.\n");
}
int main(int argc, char **argv)
{

	struct option long_options[] = 
	{
		/*Takes no parameters */
		{"version",0,0,'V'},
		{"help",0,0,'h'},
		{"list",0,0,'l'},
		{"yubikey",0,0,'y'},
		{"oath",0,0,'o'},
		{"api",0,0,'p'},

		/* Takes parameters */
		{"add",1,0,'a'},
		{"delete",1,0,'x'},
		{"enable",1,0,'e'},
		{"disable",1,0,'d'},
		{"database",1,0,'b'},
		{0,0,0}
	};
	int option_index = 0,c;
	if(argc<2)
	{
		usage();
		exit(0);
	}

	while( (c=getopt_long(argc,argv,"Vhlyopa:x:e:d:b:",long_options, &option_index)) != -1)
	{
		switch(c)
		{
			case 'V':
				printf("yubiserver-admin version " VERSION ". Copyright (C) 2011 Nanakos Chrysostomos.\n"
					"This program is free software; you can redistribute it and/or modify\n"
					"it under the terms of the GNU General Public License as published by\n"
					"the Free Software Foundation; either version 2 of the License, or\n"
					"(at your option) any later version.\n" );
				exit(0);
			case 'h':
				usage();
				exit(0);
			case 'b':
				sqlite3_dbpath = strdup(optarg);
				ext_db = 1;
				break;
			case 'a':
				Argc = argc;
				Argv = argv;
				user_ops(strdup(optarg),ADD_USER);
				fprintf(stderr,"Never here\n");
				exit(2);
			case 'y':
				yubikey_table = 1;
				break;
			case 'o':
				oath_table = 1;
				break;
			case 'p':
				api_table = 1;
				break;
			case 'l':
				show_table();
				exit(0);
			case 'e':
				user_ops(strdup(optarg),ENABLE_USER);
				fprintf(stderr,"Never here\n");
				exit(2);
			case 'd':
				user_ops(strdup(optarg),DISABLE_USER);
				fprintf(stderr,"Never here\n");
				exit(2);
			case 'x':
				user_ops(strdup(optarg),DELETE_USER);
				fprintf(stderr,"Never here\n");
				exit(2);
		}
	}

	//printf("Exiting...\n");
	exit(0);

}