/*
 * TEA Total Copyright (c) Alex Holden <alex@linuxhacker.org> 2000, 2001.
 *
 * This code is public domain; you can do whatever you want with it, though I
 * would prefer you to contribute any bug fixes back to me if possible.
 * This software comes with NO WARRANTY WHATSOEVER, not even the implied
 * warranties of merchantability and fitness for a particular purpose.
 *
 * teauntea.c: The tea and untea applets.
 */

#include <fcntl.h>
#include <unistd.h>
#include <string.h>

#include "teatotal.h"
#include "teauntea.h"
#include "readkey.h"
#include "teaprot.h"
#include "getarg.h"

/*
 * Initialise the tea and untea applets.
 */
void init_tea_untea(teastate *state)
{
	int i;
	char *p;

	/* Check for the -h or --help arguments */
	if(getflag(state->ga, 'h', "help")) usage(state);

#ifdef BUILD_BASE64
	/* Check if the -a argument is present to enable Base64 coding */
	if(getflag(state->ga, 'a', NULL)) {
		/* It is, set the flag */
		state->base64 = 1;
	} else state->base64 = 0;
#endif
#ifdef BUILD_HUFFMAN
	/* Check if the -c argument is present to enable compression */
	if(getflag(state->ga, 'c', NULL)) {
		/* It is, so enable compression */
		state->compress = 1;
		/* Initialise the Huffman engine */
		if(!(state->huff = huff_init(MAXPKTLEN + 4)))
			die("Out of memory");
	} else state->compress = 0;
#endif 
#ifdef BUILD_PPKEYS
	/* Get the key password */
	state->inkeypass = getarg(state->ga, 'p', NULL);
#endif

	/* Allocate the packet structure */
	state->pkt = safe_malloc(sizeof(teapacket));

#ifdef BUILD_HUFFMAN
	if(state->compress) state->huff->inbuf = (u8 *) &state->pkt->seq;
#endif

	/* Choose the starting sequence number */
	choose_start_seq(state);
	state->first_block = 1;

	/* The first argument not associated with a flag should be the name of
	 * the key file */
	if(!(state->keyfile = getarg(state->ga, 0, NULL))) usage(state);

	/* Get the second argument not associated with a flag, which if it
	 * exists is the input filename or "-" for stdin. If it is not given,
	 * we use stdin and stdout. */
	if(!(state->inputfile = getarg(state->ga, 0, NULL))) {
		state->outputfile = NULL;
		return;
	}
	if(!strcmp(state->inputfile, "-")) state->inputfile = NULL;

	/* Get the second argument not associated with a flag, which if it
	 * exists is the output filename or "-" for stdout. If it does not
	 * exist, if the input is stdin we also use stdout, otherwise we
	 * need to construct the output filename by adding or trying to
	 * remove .tea from the end of the input filename. */
	if(!(state->outputfile = getarg(state->ga, 0, NULL))) {
#ifdef BUILD_BASE64
		state->fnspecified = 0;
#endif
		/* If input file is stdin, leave output as stdout */
		if(!state->inputfile) return;
		if(state->untea) {
			/* Decrypting, so remove the .tea from the end */
			if(!(p = state->outputfile =
					safe_strdup(state->inputfile)))
				die("Out of memory");
			while(1) {
				if(!*p) die("Unknown file extension");
				if((p[0] == '.') && (p[1] == 't') &&
						(p[2] == 'e') && (p[3] == 'a')
						&& (p[4] == 0)) {
					*p = 0;
					return;
				}
				p++;
			}
		} else {
			/* Encrypting, so add a .tea to the end */
			i = strlen(state->inputfile);
			p = safe_malloc(i + sizeof(".tea"));
			memcpy(p, state->inputfile, i);
			memcpy(&p[i], ".tea", sizeof(".tea"));
			state->outputfile = p;
		}
	} else {
#ifdef BUILD_BASE64
		state->fnspecified = 1;
#endif
		if(!strcmp(state->outputfile, "-")) state->outputfile = NULL;
	}
}

#ifdef BUILD_BASE64
/*
 * The tea applet with Base64 encoding.
 */
static int do_base64_tea(teastate *state)
{
	int i;
	char *p, *pp;

	/* Open the output file */
	state->outfd = openwrite(state->outputfile);

	/* Initialise the Base64 engine */
	if(!(state->b64 = base64_init(state->outfd, NULL, NULL)))
		die("Out of memory");

	/* Write the envelope start line */
	safe_write(state->outfd, ENVELOPE_START, sizeof(ENVELOPE_START) - 1);
#ifdef DOS_LINE_ENDINGS
	safe_write(state->outfd, "\r\n", 2);
#else
	safe_write(state->outfd, "\n", 1);
#endif

	/* If there is an input file name */
	if(state->inputfile) {
		/* Strip the path off */
		pp = p = state->inputfile;
		while(*p) if(*p++ == '/') pp = p;
	/* Otherwise, specify the filename as "-" */
	} else pp = "-";

	/* Write the name out */
	safe_write(state->outfd, ENVELOPE_NAME, (sizeof(ENVELOPE_NAME) - 1));
	safe_write(state->outfd, pp, strlen(pp));
#ifdef DOS_LINE_ENDINGS
	safe_write(state->outfd, "\r\n", 2);
#else
	safe_write(state->outfd, "\n", 1);
#endif

	/* Encode the file */
	while(!(i = encode_block(state)));
	if(i < 0) return i;

	/* Flush the output buffer */
	if(base64_flush(state->b64) == -1) return -2;

	/* Write the envelope start line */
	safe_write(state->outfd, ENVELOPE_END, sizeof(ENVELOPE_END) - 1);
#ifdef DOS_LINE_ENDINGS
	safe_write(state->outfd, "\r\n", 2);
#else
	safe_write(state->outfd, "\n", 1);
#endif

	return 1; /* Success */
}

/*
 * Base64 decoder callback function used to handle envelope markers.
 */
static int base64_marker(void *arg, char *line, size_t len)
{
	int i;
	char *p, *pp;
	teastate *state = arg;

	/* Check for the envelope start marker */
	if(len >= (sizeof(ENVELOPE_START) - 1))
		if(!memcmp(line, ENVELOPE_START, (sizeof(ENVELOPE_START) - 1)))
			return 1;

	/* Check for the envelope end marker */
	if(len >= (sizeof(ENVELOPE_END) - 1)) {
		if(!memcmp(line, ENVELOPE_END, (sizeof(ENVELOPE_END) - 1))){
			/* If a filename was specified on the command line,
			 * only decode the first block */
			if(state->fnspecified) return 3;
			else return 2;
		}
	}

	/* If in a Base64 data block, check for a file name marker */
	if(state->b64->decoding && (len >= (sizeof(ENVELOPE_NAME) - 1))) {
		if(!memcmp(line, ENVELOPE_NAME, (sizeof(ENVELOPE_NAME) - 1))) {

			/* If a filename was specified on the command line,
			 * open that. Otherwise, use the marker name */
			if(state->fnspecified) p = state->outputfile;
			else {
				/* Strip any path data off and remove the
				 * newline */
				i = 0;
				p = pp = &line[sizeof(ENVELOPE_NAME) - 1];
				while(i < (len - sizeof(ENVELOPE_NAME))) {
					if(!pp[i]) break;
					if(pp[i++] == '/') p = &pp[i];
					if((pp[i] == '\r') || (pp[i] == '\n')) {
						pp[i] = 0;
						break;
					}
				}
			}

			/* Close the output file if it isn't stdout */
			if(state->outfd != STDOUT_FILENO)
				close(state->outfd);

			if(p) {
				/* Handle "-" == stdout */
				if(!strcmp(p, "-")) p = NULL;
				/* Handle blank names */
				if(*p == 0) p = NULL;
			} 

			/* Remember the new file name */
			state->outputfile = p;

			/* Reread the sequence number */
			state->first_block = 1;

			/* Open the new output file */
			if((state->outfd = openwrite(p)) == -1)
				return -1;

			return 0; /* success */
		}
	}	

	/* It didn't match a known marker, so just ignore it */
	return 0;
}

/*
 * The untea applet with Base64 decoding.
 */
static int do_base64_untea(teastate *state)
{
	int i;

	/* Start off with a output fd set to STDOUT */
	state->outfd = STDOUT_FILENO;

	/* Initialise the Base64 engine */
	if(!(state->b64 = base64_init(state->infd, base64_marker, state)))
		die("Out of memory");

	/* Decode the file */
	while(!(i = decode_block(state)));

	return i;
}
#endif /* BUILD_BASE64 */

/*
 * The main loop for the tea and untea applets
 */
void do_tea_untea(teastate *state)
{
	int i;

	/* Read the private key in from the key file */
	read_key(state);

	/* Open the input file if one was specified */
	if(state->inputfile) {
		if(!(state->infd = open(state->inputfile, O_RDONLY)))
			die("Input file open failed");
	/* Otherwise, set it to the fd of the standard input stream */
	} else state->infd = STDIN_FILENO;

#ifdef BUILD_BASE64
	/* If doing Base64 encoding */
	if(state->base64) {
		/* Do the encoding or decoding */
		if(state->untea) i = do_base64_untea(state);
		else i = do_base64_tea(state);
		/* Destroy the Base64 engine */
		base64_destroy(state->b64);
	} else {
#endif
	/* Open the output file */
	state->outfd = openwrite(state->outputfile);

	/* Do the decryption or encryption */
	if(state->untea) while(!(i = decode_block(state)));
	else while(!(i = encode_block(state)));

#ifdef BUILD_BASE64
	}
#endif

	/* Determine whether the encryption stopped because of an error */
	if(i == -1) die("Read error");
	else if(i == -2) die("Write error");

	/* Close the input and output files */
	if(state->inputfile) close(state->infd);
	if(state->outputfile) close(state->outfd);
}