/*
 * The contents of this file are subject to the Mozilla Public
 * License Version 1.1 (the "License"); you may not use this file
 * except in compliance with the License. You may obtain a copy of
 * the License at http://www.mozilla.org/MPL/
 * 
 * Software distributed under the License is distributed on an "AS
 * IS" basis, WITHOUT WARRANTY OF ANY KIND, either express or
 * implied. See the License for the specific language governing
 * rights and limitations under the License.
 * 
 * The Original Code is the Netscape security libraries.
 * 
 * The Initial Developer of the Original Code is Netscape
 * Communications Corporation.  Portions created by Netscape are 
 * Copyright (C) 1994-2000 Netscape Communications Corporation.  All
 * Rights Reserved.
 * 
 * Contributor(s):
 * 
 * Alternatively, the contents of this file may be used under the
 * terms of the GNU General Public License Version 2 or later (the
 * "GPL"), in which case the provisions of the GPL are applicable 
 * instead of those above.  If you wish to allow use of your 
 * version of this file only under the terms of the GPL and not to
 * allow others to use your version of this file under the MPL,
 * indicate your decision by deleting the provisions above and
 * replace them with the notice and other provisions required by
 * the GPL.  If you do not delete the provisions above, a recipient
 * may use your version of this file under either the MPL or the
 * GPL.
 */

/*
 * SSL Sample. XXX
 *
 */

/****************************************************************************
 *  SSL client program that sets up a connection to SSL server, transmits   *
 *  some data and then reads the reply                                      *
 ****************************************************************************/

/* Platform specific header files */

#ifdef XP_UNIX
#define _XOPEN_SOURCE 1 /* to get declarations of getopt() and getpass() */

#ifdef AIX
#define _ALL_SOURCE 1	/* to get declaration for uint on AIX */
#endif
#ifdef OSF1
#define _OSF_SOURCE 1 /* to get declaration for uint on HPUX */
#endif
#ifdef SOLARIS
#define __EXTENSIONS__ 1 /* to get declaration for getpass() on solaris */
#endif
#ifdef IRIX
#define _BSD_TYPES 1	/* to get declaration for uint on IRIX 6.5 */
#endif

#include <unistd.h>	/* for getopt() + getpass() */

#ifdef SOLARIS
#include <signal.h>
#endif	
#endif	/* XP_UNIX */

#ifdef WIN32
#include <winsock.h>
#endif

/* Generic header files */

#include <stdio.h>
#include <string.h>

/* NSPR header files */
#include "prerror.h"
#include "nspr.h"
#include "prnetdb.h"

/* NSS header files */

#include "pk11func.h"
#include "secitem.h"
#include "ssl.h"
#include "certt.h"
#include "nss.h"
#include "secrng.h"
#include "secder.h"
#include "key.h"
#include "sslproto.h"

/* Custom header files */

#include "sslerror.h"

#ifdef _WINDOWS
extern int getopt(int argc, char **argv, char *optstring);
extern char *optarg;
#include <conio.h>
#else
#ifdef STRANGE_PLATFORMS
#include <getopt.h>
#endif
#endif

/* Function Prototypes */

void		usage			(char *progname);

PRFileDesc * 	setup_socket		(char *nickname);

void 		exitErr			(char *function);

SECStatus 	ownGetClientAuthData	(void *arg, PRFileDesc *socket, 
					 struct CERTDistNamesStr *caNames, 
			       		 struct CERTCertificateStr **pRetCert, 
			       		 struct SECKEYPrivateKeyStr **pRetKey);

SECStatus 	ownAuthCertificate 	(void *arg, 
			      		 PRFileDesc *socket, 
			      		 PRBool checksig, 
			      		 PRBool isServer);

SECStatus 	ownBadCertHandler 	(void *arg, PRFileDesc *socket);

char * 		ownPasswd		(PK11SlotInfo *slot,
					 PRBool required,
					 void *arg);

SECStatus 	ownHandshakeCallback	(PRFileDesc *socket, void *arg);

int 		readThisMany		(PRFileDesc *fd,
					 char *buffer,
					 int thisMany);

int 		writeThisMany		(PRFileDesc *fd,
					 char *buffer,
					 int thisMany);

/* Function implementation */

/* Function: void usage()
 *
 * Purpose: Upon a bogus invocation of the program, print the usage information
 * and exit with a nonzero status.
 */

void usage(char *progname) 
{

    printf("Usage: %s [-d <cert dir>] -s <host> -p <port> -n <nickname> [-q <query>]\n\n", progname);
    exit(1);
}

/* Function: PRFileDesc * setup_socket()
 *
 * Purpose: setup_socket takes a char *, creates a new socket, configures it
 * for SSL, and returns the address of the new socket. The nickname parameter
 * is used in the setting of callback functions.
 */

PRFileDesc * setup_socket(char *nickname) 
{

    PRFileDesc *	socket = NULL;
    PRFileDesc *	sslsocket = NULL;
    PRSocketOptionData	sockdata;

    /* Create a new socket */

    socket = PR_NewTCPSocket();
    if (!socket) {
	exitErr("PR_NewTCPSocket");
    }

    /* Ensure that the socking is blocking. The default on some platforms
     * is non blocking
     */

    sockdata.option = PR_SockOpt_Nonblocking;
    sockdata.value.non_blocking = PR_FALSE;

    if (PR_SetSocketOption(socket, &sockdata) != PR_SUCCESS ) {
	exitErr("PR_SetSocketOption");
    }

    /* Import this socket into the SSL layer */

    sslsocket = SSL_ImportFD(NULL, socket);
    if (!sslsocket) {
	exitErr("SSL_ImportFD");
    }

    /* Set some flags on this new SSL socket */

    if ((SSL_Enable(sslsocket, SSL_SECURITY, 1)) != SECSuccess ) {
	exitErr("SSL_Enable:SSL_SECURITY");
    }

    if ((SSL_Enable(sslsocket, SSL_HANDSHAKE_AS_CLIENT, 1)) != SECSuccess) {
	exitErr("SSL_Enable:SSL_HANDSHAKE_AS_CLIENT");
    }

    /* Set the desired callbacks. */

    if ((SSL_GetClientAuthDataHook(sslsocket,
	                           (SSLGetClientAuthData)ownGetClientAuthData,
	                           (void *)nickname)) != SECSuccess) {

	exitErr("SSL_GetClientAuthDataHook");
    }

    if ((SSL_AuthCertificateHook(sslsocket,
	                         (SSLAuthCertificate)ownAuthCertificate,
	                         (void *)CERT_GetDefaultCertDB()))
	!= SECSuccess ) {

	exitErr("SSL_AuthCertificateHook");
    }

    if ((SSL_BadCertHook(sslsocket,
	                 (SSLBadCertHandler)ownBadCertHandler,
			 NULL)) != SECSuccess ) {

	exitErr("SSL_BadCertHook");
    }

    if ((SSL_HandshakeCallback(sslsocket,
	                       (SSLHandshakeCallback)ownHandshakeCallback,
			       NULL)) != SECSuccess ) {

	exitErr("SSL_HandshakeCallback");
    }

    /* Return the address of the configured socket */

    return sslsocket;
}

/* Function void exitErr()
 *
 * Purpose: When a security function returns an error, get the error number
 * using PR_GetError() and find a matching error string using the custom
 * function SSL_Strerror(). Print out all relevant information, shutdown 
 * NSS, clean up NSPR, and then exit with a nonzero status.
 *
 * Note: SSL_Strerror is defined in sslerror.h.
 */

void exitErr(char *function) 
{

    PRErrorCode	errornumber = PR_GetError();
    const char *errorstring = SSL_Strerror(errornumber);

    printf("Error in function %s: %d\n - %s\n",
	    function, errornumber, errorstring);

    NSS_Shutdown();
    PR_Cleanup();
    exit(1);
}

/* Function: SECStatus ownGetClientAuthData()
 *
 * Purpose: This callback is used by SSL to pull client certificate 
 * information upon server request.
 */

SECStatus ownGetClientAuthData(	void *arg,
				PRFileDesc *socket,
				struct CERTDistNamesStr *caNames,
				struct CERTCertificateStr **pRetCert,
				struct SECKEYPrivateKeyStr **pRetKey) 
{

    CERTCertificate *		cert;
    SECKEYPrivateKey *		privKey;
    char *			chosenNickName	= (char *)arg;
    void *			proto_win 	= NULL;
    SECStatus			rv		= SECFailure;

    proto_win = SSL_RevealPinArg(socket);

    if (chosenNickName) {
	cert = PK11_FindCertFromNickname(chosenNickName, proto_win);
	if (cert) {
	    privKey = PK11_FindKeyByAnyCert(cert, proto_win);
	    if (privKey) {
		rv = SECSuccess;
	    } else {
		CERT_DestroyCertificate(cert);
	    }
	}
    } else { /* no nickname given, automatically find the right cert */
	CERTCertNicknames * 	names;
	int			i;

	names = CERT_GetCertNicknames(	CERT_GetDefaultCertDB(), 
					SEC_CERT_NICKNAMES_USER,
					proto_win);

	if (names != NULL) {
	    for(i = 0; i < names->numnicknames; i++ ) {

		cert = PK11_FindCertFromNickname(names->nicknames[i],
						 proto_win);
		if (!cert) {
		    continue;
		}

		/* Only check unexpired certs */
		if (CERT_CheckCertValidTimes(cert, PR_Now(), PR_FALSE)
		    != secCertTimeValid ) {

		    CERT_DestroyCertificate(cert);
		    continue;
		}

		rv = NSS_CmpCertChainWCANames(cert, caNames);
		if (rv == SECSuccess) {
		    privKey = PK11_FindKeyByAnyCert(cert, proto_win);
		    if (privKey) {
			break;
		    }
		    rv = SECFailure;
		    break;
		}
		CERT_FreeNicknames(names);
	    } /* for loop */
	}
    }

    if (rv == SECSuccess) {
	*pRetCert = cert;
	*pRetKey  = privKey;
    }

    return rv;
}

/* Function: SECStatus ownAuthCertificate()
 *
 * Purpose: This function checks the validity of an incoming certificate.
 */

SECStatus ownAuthCertificate( 	void *arg,
				PRFileDesc *socket,
				PRBool checksig,
				PRBool isServer) 
{

    SECCertUsage	certUsage;
    CERTCertificate *	cert;
    void *		pinArg;
    char *		hostname;
    SECStatus		rv 		= SECFailure;

    if (!arg || !socket) {
	return rv;
    }

    certUsage = isServer ? certUsageSSLClient : certUsageSSLServer;

    cert = SSL_PeerCertificate(socket);
	
    pinArg = SSL_RevealPinArg(socket);

    rv = CERT_VerifyCertNow((CERTCertDBHandle*)arg,
			     cert,
			     checksig,
			     certUsage,
			     pinArg);

    if ( (rv != SECSuccess) || isServer ) {
	return rv;
    }
	
    /* Certificate is OK. Since this is the client side of an SSL
     * connection, we need to verify that the name field in the cert
     * matches the desired hostname. This is our defense against
     * man-in-the-middle attacks.
     */

    /* SSL_RevealURL returns a hostname, not a URL. */

    hostname = SSL_RevealURL(socket);

    if (hostname && hostname[0]) {
	rv = CERT_VerifyCertName(cert, hostname);
    } else {
	PR_SetError(SSL_ERROR_BAD_CERT_DOMAIN, 0);
	rv = SECFailure;
    }
    if (hostname)
    	PR_Free(hostname);

    return rv;
}

/* Function: SECStatus ownBadCertHandler()
 *
 * Purpose: This callback is used whent the incoming cert is not valid. 
 * It should return SECSuccess to accept the cert anyway, SECFailure
 * to reject.
 */

SECStatus ownBadCertHandler(void *arg, PRFileDesc *socket) 
{

    /* Can log invalid cert here. */
    printf("Bad server certificate: %d, %s\n", PR_GetError(),
           SSL_Strerror(PR_GetError()));

    /* Note here that this particular callback implementation always
     * rejects a bad cert.
     */
    return SECFailure;
}

/* Function: char * ownPasswd()
 *
 * Purpose: Called by SSL when retriving own cert from database. Returns
 * pointer to string with a password for the database. (In this case, cert
 * database is not password protected.)
 */

char * ownPasswd(PK11SlotInfo *slot, PRBool retry, void *arg) 
{

    char *	passwd = NULL;
#ifdef _WINDOWS
    char	buf[512];
#endif
    static int	tries  = 0;

    if (!retry) 
    	tries = 0; /* starting over */

    if ( ++tries < 3 ) {
	char * tmp = NULL;

	/* The following call to _cgets is insecure, since
	 * _cgets does not suspend tty echoing. Production
	 * code should take the necessary precautions.
	 */	

	while ( !tmp && arg ) {
#ifdef _WINDOWS
	    printf("%s",(char *)arg);
	    fflush(stdout);
	    tmp = _cgets(buf);
#else
	    tmp = getpass((char *)arg);
#endif
	}
	
	passwd = PL_strdup(tmp);
	memset(tmp, 0, strlen(tmp));
    }

    return passwd;
}

/* Function: SECStatus ownHandshakeCallback()
 *
 * Purpose: Called by SSL to inform application that the handshake is
 * complete. This function is mostly used on the server side of an SSL
 * connection, although it is provided for a client as well.
 */

SECStatus ownHandshakeCallback(PRFileDesc *socket, void *arg) 
{
    printf("Handshake has completed, ready to send data securely.\n");
    return SECSuccess;
}

/* Function: int readThisMany()
 *
 * Purpose: This is a wrapper function around PR_Read that will read exactly
 * the number of bytes requested and return the number of bytes read.
 */

int readThisMany(PRFileDesc *fd, char *buffer, int thisMany) 
{
    int	total = 0;

    while ( total < thisMany ) {
	int got = PR_Read(fd, buffer+total, thisMany-total);

	if ( got < 0 ) {
	    if ( PR_GetError() != PR_WOULD_BLOCK_ERROR ) {
		break;
	    }
	    continue;
	}

	total += got;
    }

    return total;
}

/* Function: int writeThisMany()
 *
 * Purpose: This is a wrapper function around PR_Write that will write exactly
 * the number of bytes requested and return the number of bytes written.
 */

int writeThisMany(PRFileDesc *fd, char *buffer, int thisMany) 
{
    int	total = 0;

    while ( total < thisMany ) {
	int got = PR_Write(fd, buffer+total, thisMany-total);

	if ( got < 0 ) {
	    if ( PR_GetError() != PR_WOULD_BLOCK_ERROR ) {
		break;
	    }
	    continue;
	}

	total += got;
    }

    return total;
}

/* Function: int main()
 *
 * Purpose: This is the client's main function.
 */

int main (int argc, char **argv) 
{

    int 	option, querylen;
    char *	progName;
    int 	port = 0;
    char * 	serverhost = NULL;
    char * 	nickname = NULL;
    char * 	certdir = NULL;
    char * 	query = NULL;

    PRFileDesc *sock;
    PRNetAddr 	addr;
    PRInt32 	numbytes;
    PRHostEnt 	hostentry;
  
    char 	buffer[256];
    SECStatus 	secstatus = SECFailure;
    PRStatus  	prstatus  = PR_FAILURE;
 
#ifdef SOLARIS
    sigignore(SIGPIPE);
#endif

    PR_Init(PR_USER_THREAD, PR_PRIORITY_NORMAL, 1);      
    progName = PL_strdup(argv[0]);
  
    while ((option = getopt(argc, argv, "d:p:n:s:q:")) != -1) {
	switch (option) {
	case 'd': certdir    = optarg; 		break;
	case 'p': port       = atoi(optarg); 	break;
	case 'n': nickname   = optarg; 		break;
	case 's': serverhost = optarg; 		break;
	case 'q': query      = optarg; 		break;
	default: usage(progName); 		break;
	}
    }
  
    if (nickname == NULL || port == 0) {
	usage(progName);
    }
  
    if (certdir == NULL) {
	certdir = PR_smprintf("%s/.netscape", getenv("HOME"));
    }
  
    if (serverhost == NULL) {
	usage(progName); 
    }

    if (query == NULL) {
	query = PR_smprintf("GET");
    }

    querylen = strlen(query);
    
     /* Do the libsec initialization */
    PK11_SetPasswordFunc(ownPasswd);  /* set passwd callback */

    secstatus = NSS_Init(certdir);
    if (secstatus != SECSuccess) {
	exitErr("NSS_Init");
    }

    /*
     * All the ciphers except SSL_RSA_WITH_NULL_MD5 are on by default.
     * Enable encryption, enable NULL cipher. 
     */

    secstatus = NSS_SetDomesticPolicy();
    if (secstatus != SECSuccess ) {
	exitErr("NSS_SetDomesticPolicy");
    }

    SSL_EnableCipher(SSL_RSA_WITH_NULL_MD5, SSL_ALLOWED);
    
    /* set up SSL secure socket */
    sock = setup_socket(nickname);
    if (sock == NULL ) {
	exitErr("setup_socket");
    }

    secstatus = SSL_SetPKCS11PinArg(sock, "Enter database password: ");
    if (secstatus != SECSuccess) {
	exitErr("SSL_SetPKCS11PinArg");
    }
      
    secstatus = SSL_SetURL(sock, serverhost);
    if (secstatus != SECSuccess) {
	exitErr("SSL_SetURL");
    }

    /* prepare and setup network connection */
    prstatus = PR_GetHostByName( serverhost,  buffer, 256, &hostentry);
    if (prstatus != PR_SUCCESS) {
	exitErr("PR_GetHostByName");
    }
	
    if (PR_EnumerateHostEnt(0, &hostentry, port, &addr)<0) {
	exitErr("PR_EnumerateHostEnt");
    }

    prstatus = PR_Connect(sock, &addr, PR_INTERVAL_NO_TIMEOUT);
    if (prstatus == PR_FAILURE) {
	exitErr("PR_Connect");  
    }
  
    /* established SSL connection, ready to send data */
    secstatus = SSL_ForceHandshake(sock); 
    if (secstatus != SECSuccess) {
	exitErr("SSL_ForceHandshake");
    }
  
    printf("Sending query: *%s*\n", query);

    querylen = PR_htonl(querylen);  
    numbytes = writeThisMany(sock, (void *)&querylen, sizeof(querylen));
    if (numbytes != sizeof(querylen)) {
	exitErr("PR_Write querylen");
    }

    querylen = PR_ntohl(querylen);  
    numbytes = writeThisMany(sock, (void *)query, querylen); 
    if (numbytes != querylen) {
	exitErr("PR_Write query");
    }
    
    numbytes = readThisMany(sock, (void *)&querylen, sizeof(querylen));
    if (numbytes != sizeof(querylen)) {
	exitErr("PR_Read querylen"); 
    }

    querylen = PR_ntohl(querylen);
  
    numbytes = readThisMany(sock, (void *)buffer, querylen);
    if (numbytes <querylen) {
	exitErr("PR_Read server response");
    }

    buffer[querylen] = 0;

    printf("Received : *%s*\n", buffer);

    printf("\nDone, exiting... \n");

    /* now, close down SSL */
    PR_Close(sock);
    NSS_Shutdown();
    PR_Cleanup();
    return 0;
}

