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

#include <socket.h>
#include <netdb.h>
#include <in.h>

#include <krb5.h>

#include "krb_common.h"

#if __INITIAL_POINTER_SIZE == 64
#pragma __required_pointer_size __save
#pragma __required_pointer_size 32
#endif

typedef char ** ArgArray;

#if __INITIAL_POINTER_SIZE == 64
#pragma __required_pointer_size __restore
#endif

/*
** Option Data Structure
*/
struct _OptData {
    char 		*prog;
    unsigned short	port;
    char		*serv;
    char		*host;
    char		*msg;
    };

/*
** Context Data Structure
*/
struct _CtxData {
    krb5_context	context;
    krb5_auth_context	auth_context;
    };

/*
** Local Routine Prototypes
*/
void parse_command_line (
    int,
    ArgArray,
    struct _OptData *);

void usage (
    char *);

int connect_to_server (
    char *,
    unsigned short);

int initialize_client_context (
    struct _OptData *,
    struct _CtxData *,
    int);

int dialog_with_server (
    struct _OptData *,
    struct _CtxData *,
    int);

int terminate_client_context (
    struct _OptData *,
    struct _CtxData *,
    int);

/*
**
**  main - Main processing routine for the Kerberos client
**
**  Functional Description:
**
**      This routine controls the Kerberos client processing.
**
**  Usage:
**
**      main argc, argv, envp
**
**  Formal parameters:
**
**      argc 		- (IN) argument count
**      argv         	- (IN) address of an argument array 
**      envp         	- (IN) address of an environment string 
**
**  Implicit Parameters:
**
**      None
**
**  Routine Value:
**
**      FAILURE		- Completed with errors
**      SUCCESS		- Completed successfully
**
**  Side Effects:
**
**      None
**
*/
main (
    int			argc,
    ArgArray            argv,
    char                *envp[]
    )
{
    struct _OptData OptData = {NULL, 0, NULL, NULL};
    struct _CtxData CtxData = {NULL, NULL};
    int status = SUCCESS,
        sock = 0;

    /*
    ** Process the command line options
    */
    parse_command_line (argc, argv, &OptData);

    /*
    ** Connect to the requested server
    */
    sock = connect_to_server (OptData.host, OptData.port);
    if (sock < 0)
        status = FAILURE;

    /*
    ** Initialize the client context
    */
    if (status == SUCCESS)
        status = initialize_client_context (&OptData, &CtxData, sock);

    /*
    ** Dialog with the server 
    */
    if (status == SUCCESS)
        status = dialog_with_server (&OptData, &CtxData, sock);

    /*
    ** Terminate the client context
    */
    terminate_client_context (&OptData, &CtxData, sock);

    /*
    ** Exit
    */
    exit (status);
}

/*
**
**  parse_command_line - Parse the command line options
**
**  Functional Description:
**
**      This routine parses the command line options.
**
**  Usage:
**
**      parse_command_line argc, argv, OptData
**
**  Formal parameters:
**
**      argc 		- (IN) argument count
**      argv         	- (IN) address of an argument array 
**      OptData		- (OUT) address of command option data structure 
**			  which will contain the parsed input.
**
**  Implicit Parameters:
**
**      None
**
**  Routine Value:
**
**      None
**
**  Side Effects:
**
**      None
**
*/
void
parse_command_line
    (
    int			argc,
    ArgArray            argv,
    struct _OptData	*OptData
    )
{
    char host_name[255 + 1];
    struct servent *sp;
    int option,
        status;

    /*
    ** Parse the program name from the first parameter
    */
    parse_prog_name (argv[0], &OptData->prog);

    /*
    ** Process the command line options
    */
    while ((option = getopt (argc, argv, "p:?")) != EOF) 
    {
        switch (option) 
    	{
        	/* 
        	** Port Number ?
        	*/
        	case 'p':
        	    OptData->port = atoi (optarg);
        	    break;
        
        	/* 
        	** Invalid argument ?
        	*/
        	case '?':
        	default:
        	    usage (OptData->prog);
        	    exit (1);
        	    break;
    	}
    }

    /*
    ** Are there extra parameters on the command line ?
    */
    if (argc - optind > 3)
    {
        usage (OptData->prog);
        exit (1);
    }

    /*
    ** Process the Message Text
    */
    if (argc - optind > 0)
        OptData->msg = argv[optind + 0];
    else
        OptData->msg = "<NULL>";

    /*
    ** Process the Host Name
    */
    if (argc - optind > 1)
        OptData->host = argv[optind + 1];
    else
    {
        status = gethostname (host_name, sizeof (host_name) - 1);
        if (status != 0)
    	{
        	perror ("gethostname");
        	exit (1);
    	}
        OptData->host = malloc (strlen (host_name) + 1);
        strcpy (OptData->host, host_name);
    }    

    /*
    ** Process the Service Name
    */
    if (argc - optind > 2)
        OptData->serv = argv[optind + 2];
    else
        OptData->serv = KRB_SERVICE;

    /*
    ** If no port was specified at the command line, then let's determine
    ** an appropriate default.
    */
    if (! OptData->port)
    {
        sp = getservbyname (OptData->serv, "tcp");
        if (sp)
            OptData->port = ntohs (sp->s_port);
        else
    	OptData->port = KRB_PORT;
    }

}

/*
**
**  usage - Display the acceptable unix style command usage
**
**  Functional Description:
**
**      This routine displays to standard output the appropriate unix style 
**	command usage.
**
**  Usage:
**
**      usage prog_name
**
**  Formal parameters:
**
**      prog_name	- (IN) address of the program name string
**
**  Implicit Parameters:
**
**      None
**
**  Routine Value:
**
**      None
**
**  Side Effects:
**
**      None
**
*/
void
usage 
    (
    char		*prog_name
    )
{
    fprintf (stderr, "Usage: %s [-p port] [message] [host] [service]\n", prog_name);
}

/*
**
**  connect_to_server - Connect to a given server
**
**  Functional Description:
**
**      This routine makes the appropriate TCP/IP socket calls to connect to 
**	the specified host name on the given port.
**
**  Usage:
**
**      parse_command_line argc, argv, envp
**
**  Formal parameters:
**
**      host 		- (IN) address of host name string
**      port         	- (IN) port number
**
**  Implicit Parameters:
**
**      None
**
**  Routine Value:
**
**      -1		- socket creation failure has occured
**	n		- The established socket number
**
**  Side Effects:
**
**      None
**
*/
int
connect_to_server
    ( 
    char		*host,
    unsigned short	port
    )
{
    struct sockaddr_in SockBuf;
    struct hostent *hp;
    int sock;

    /*
    ** Get the host entry by name
    */     
    if ((hp = gethostbyname (host)) == NULL)
    {
        fprintf (stderr, "Unknown host: %s\n", host);
        return (-1);
    }

    /*
    ** Setup the socket information
    */
    SockBuf.sin_family = hp->h_addrtype;
    memcpy ((char *) &SockBuf.sin_addr, hp->h_addr, sizeof (SockBuf.sin_addr));
    SockBuf.sin_port = htons (port);

    /*
    ** Create a socket
    */
    if ((sock = socket (AF_INET, SOCK_STREAM, 0)) < 0)
    {
        perror ("socket");
        return (-1);
    }

    /*
    ** Connect the socket 
    */
    if (connect (sock, (struct sockaddr *) &SockBuf, sizeof (SockBuf)) < 0)
    {
        perror ("connect");
        return (-1);
    }

    /*
    ** Return the socket
    */    
    return (sock);

}

/*
**
**  initialize_client_context - Initialize the client's context 
**
**  Functional Description:
**
**      This routine initializes the necessary client context information.
**	The authorization is sent to the server and if it is accepted the
**	authorization context is modified with the components necessary to
**	encrypt and decrypt messages to and from the server.
**
**  Usage:
**
**      initialize_client_context OptData, CtxData, sock
**
**  Formal parameters:
**
**      OptData		- (IN) address of the option data structure
**      CtxData        	- (OUT) address of the context data structure
**      sock		- (IN) socket number
**
**  Implicit Parameters:
**
**      None
**
**  Routine Value:
**
**      SUCCESS		- client context initialization succeeded
**	FAILURE		- client context initialization failed
**
**  Side Effects:
**
**      None
**
*/
int
initialize_client_context
    ( 
    struct _OptData	*OptData,
    struct _CtxData	*CtxData,
    int			sock
    )
{
    krb5_ap_rep_enc_part *reply_return;
    krb5_error *error_return;
    krb5_error_code status;
    krb5_principal client;
    krb5_principal server;
    krb5_data rcache_name;
    krb5_rcache rcache;
    krb5_ccache ccache;
    krb5_address kaddr;

    /*
    ** Initialize the Kerberos context
    */
    status = krb5_init_context (&CtxData->context);
    if (status)
    {
        krb5_put_message (OptData->prog, status, "while initializing krb5");
        return (FAILURE);
    }

    /*
    ** Get the default credential cache
    */
    status = krb5_cc_default (CtxData->context, &ccache);
    if (status)
    {
        krb5_put_message (OptData->prog, status,
    		      "while getting default ccache");
        return (FAILURE);
    }

    /*
    ** Get the principal name from the default credential cache
    */
    status = krb5_cc_get_principal (CtxData->context, ccache, &client);
    if (status)
    {
        krb5_put_message (OptData->prog, status,
    		      "while getting client principal name");
        return (FAILURE);
    }

    /*
    ** Convert the service name into a principal name
    */
    status = krb5_sname_to_principal (CtxData->context, OptData->host, 
    				  OptData->serv, KRB5_NT_SRV_HST,
    				  &server);
    if (status)
    {
        krb5_put_message (OptData->prog, status, 
    		      "while creating server name for %s/%s",
    		      OptData->host, OptData->serv);
        krb5_free_principal (CtxData->context, client);
        return (FAILURE);
    }

    /*
    ** Send the authorization to the server
    */
    status = krb5_sendauth (CtxData->context, &CtxData->auth_context, (krb5_pointer) &sock,
    			KRB_VERSION, client, server, 
    			AP_OPTS_MUTUAL_REQUIRED, NULL, NULL, 
    			ccache, &error_return, &reply_return, 0);
    
    /*
    ** Free the client and server principals
    */
    krb5_free_principal (CtxData->context, client);
    krb5_free_principal (CtxData->context, server);

    /*
    ** Process any authorization errors
    */
    if (status)
    {
        if (status == KRB5_SENDAUTH_REJECTED)
    	printf ("%s: krb5_sendauth rejected with '%*s'\n",
    		OptData->prog, error_return->text.length, error_return->text.data);
        else
        	krb5_put_message (OptData->prog, status, "while using sendauth");
        return (FAILURE);
    }

    /*
    ** Generate local & remote addresses
    */
    status = krb5_auth_con_genaddrs (CtxData->context, CtxData->auth_context, sock, 
    	 KRB5_AUTH_CONTEXT_GENERATE_LOCAL_FULL_ADDR | KRB5_AUTH_CONTEXT_GENERATE_REMOTE_FULL_ADDR);
    if (status)
    {
        krb5_put_message (OptData->prog, status, "while generating addresses");
        return (FAILURE);
    }

    /*
    ** Create a temporary replay cache
    */
    rcache_name.data = "rcache.tmp";
    rcache_name.length = strlen ("rcache.tmp");
    status = krb5_get_server_rcache (CtxData->context, &rcache_name, &rcache);
    if (status)
    {
        krb5_put_message (OptData->prog, status, "while getting server rcache");
        return (FAILURE);
    }

    /*
    ** Set the authorization replay cache
    */
    status = krb5_auth_con_setrcache (CtxData->context, CtxData->auth_context, rcache);
    if (status)
    {
        krb5_put_message (OptData->prog, status, "while setting server rcache");
        return (FAILURE);
    }

    /*
    ** If no reply was returned, then we've failed
    */
    if (reply_return) 
        return (SUCCESS);
    else
        return (FAILURE);
}

/*
**
**  dialog_with_server - Dialog with the given server
**
**  Functional Description:
**
**      This routine sends an encrypted message to the server and then
**	descripts the server's reply.
**
**  Usage:
**
**      dialog_with_server OptData, CtxData, sock
**
**  Formal parameters:
**
**      OptData		- (IN) address of the option data structure
**      CtxData        	- (IN) address of the context data structure
**      sock		- (IN) socket number
**
**  Implicit Parameters:
**
**      None
**
**  Routine Value:
**
**      SUCCESS		- dialog with server succeeded
**	FAILURE		- dialog with server failed
**
**  Side Effects:
**
**      None
**
*/
int
dialog_with_server 
    ( 
    struct _OptData	*OptData,
    struct _CtxData	*CtxData,
    int			sock
    )
{
    krb5_error_code status;
    krb5_data send_data,
              recv_data,
              msg_data;

    /*
    ** Initialize the message data descriptor
    */
    msg_data.data = OptData->msg;
    msg_data.length = strlen (OptData->msg);

    /*
    ** Encrypt the send data
    */
    status = krb5_mk_priv (CtxData->context, CtxData->auth_context, &msg_data, &send_data, NULL);
    if (status)
    {
        krb5_put_message (OptData->prog, status, "while encrypting message");
        return (FAILURE);
    }

    /*
    ** Write the send data length
    */
    status = net_write (sock, (char *) &send_data.length, sizeof (send_data.length));
    if (status <= 0)
    {   
        krb5_put_message (OptData->prog, 0, "while writing data length to the server");
        if (send_data.data)
        	free (send_data.data);
        return (FAILURE);
    }

    /*
    ** Write the send data
    */
    status = net_write (sock, (char *) send_data.data, send_data.length);
    if (status <= 0)
    {   
        krb5_put_message (OptData->prog, 0, "while writing data to the server");
        if (send_data.data)
        	free (send_data.data);
        return (FAILURE);
    }

    /*
    ** Log the sent message
    */
    printf ("%s: sent message '%-*.*s'\n",
    	OptData->prog, msg_data.length, msg_data.length, msg_data.data);

    /*
    ** Free the send data
    */
    if (send_data.data)
        free (send_data.data);

    /*
    ** Read the receive data length
    */
    status = net_read (sock, (char *) &recv_data.length, sizeof (recv_data.length));
    if (status <= 0)
    {
        krb5_put_message (OptData->prog, 0, "while reading data from server");
        return (FAILURE);
    }

    /*
    ** Allocate a buffer to hold the receive data
    */
    recv_data.data = malloc (recv_data.length);
    if (! recv_data.data)
    {
        krb5_put_message (OptData->prog, 0, "unable to allocate message buffer");
        return (FAILURE);
    }

    /*
    ** Read the receive data
    */
    status = net_read (sock, (char *) recv_data.data, recv_data.length);
    if (status <= 0)
    {
        krb5_put_message (OptData->prog, 0, "while reading data from server");
        if (recv_data.data)
            free (recv_data.data);
        return (FAILURE);
    }
    
    /*
    ** Decrypt the receive data
    */
    status = krb5_rd_priv (CtxData->context, CtxData->auth_context, &recv_data, &msg_data, NULL);
    if (status)
    {
        krb5_put_message (OptData->prog, status, "while decrypting message");
        return (FAILURE);
    }

    /*
    ** Free the receive data
    */
    if (recv_data.data)
        free (recv_data.data);

    /*
    ** Log the received message
    */
    printf ("%s: received message '%-*.*s'\n",
    	OptData->prog, 
    	msg_data.length, msg_data.length, msg_data.data);

    /*
    ** Free the message data
    */
    if (msg_data.data)
        free (msg_data.data);

    /*
    ** Return success
    */
    return (SUCCESS);
}

/*
**
**  terminate_client_context - Terminate the client context
**
**  Functional Description:
**
**      This routine terminates the client's context information and
**	closes it's socket.
**
**  Usage:
**
**      terminate_client_context OptData, CtxData, sock
**
**  Formal parameters:
**
**      OptData		- (IN) address of the option data structure
**      CtxData        	- (IN) address of the context data structure
**      sock		- (IN) socket number
**
**  Implicit Parameters:
**
**      None
**
**  Routine Value:
**
**      SUCCESS		- terminate client context succeeded
**	FAILURE		- terminate client context failed
**
**  Side Effects:
**
**      None
**
*/
int
terminate_client_context
    ( 
    struct _OptData	*OptData,
    struct _CtxData	*CtxData,
    int			sock
    )
{
    krb5_error_code status;

    /*
    ** Free the host name
    */
    if (OptData->host)
        free (OptData->host);

    /*
    ** Free the Kerberos context
    */
    if (CtxData->context)
    {
        krb5_free_context (CtxData->context);
        CtxData->context = 0;
    }

    /*
    ** Initialize the authenication context
    */
    if (CtxData->auth_context)
        CtxData->auth_context = 0;

    /*
    ** Shutdown and close the socket
    */
    if (sock > 0)
    {
        shutdown (sock, 2);
        close (sock);
    }

    /*
    ** Return success
    */
    return (SUCCESS);

}