/*
 * Encryption Functions Module
 *
 * This module provides commands that can be used to encrypt or decrypt data.
 *
 * Copyright (C) 2023 Brian O'Hagan
 *
 */

#include "tlsInt.h"
#include "tclOpts.h"
#include <tcl.h>
#include <stdio.h>
#include <string.h>
#include <openssl/evp.h>
#if OPENSSL_VERSION_NUMBER >= 0x30000000L
#include <openssl/params.h>
#endif

/* Macros */
#define BUFFER_SIZE	32768

/* Encryption functions */
#define TYPE_MD		0x010
#define TYPE_HMAC	0x020
#define TYPE_CMAC	0x040
#define TYPE_MAC	0x080
#define TYPE_ENCRYPT	0x100
#define TYPE_DECRYPT	0x200
#define TYPE_SIGN	0x400
#define TYPE_VERIFY	0x800

/*******************************************************************/

/*
 * This structure defines the per-instance state of a encrypt operation.
 */
typedef struct EncryptState {
	int type;		/* Operation type */

	Tcl_Interp *interp;	/* Current interpreter */
	EVP_CIPHER_CTX *ctx;	/* Cipher Context */
	Tcl_Command token;	/* Command token */
} EncryptState;


/*
 *-------------------------------------------------------------------
 *
 * EncryptStateNew --
 *
 *	This function creates a per-instance state data structure
 *
 * Returns:
 *	State structure pointer
 *
 * Side effects:
 *	Creates structure
 *
 *-------------------------------------------------------------------
 */
EncryptState *EncryptStateNew(Tcl_Interp *interp, int type) {
    EncryptState *statePtr = (EncryptState *) ckalloc((unsigned) sizeof(EncryptState));

    if (statePtr != NULL) {
	memset(statePtr, 0, sizeof(EncryptState));
	statePtr->type = type;		/* Operation type */
	statePtr->interp = interp;	/* Current interpreter */
	statePtr->ctx = NULL;		/* Cipher Context */
	statePtr->token = NULL;		/* Command token */
    }
    return statePtr;
}

/*
 *-------------------------------------------------------------------
 *
 * EncryptStateFree --
 *
 *	This function deletes a state data structure
 *
 * Returns:
 *	Nothing
 *
 * Side effects:
 *	Removes structure
 *
 *-------------------------------------------------------------------
 */
void EncryptStateFree(EncryptState *statePtr) {
    if (statePtr == (EncryptState *) NULL) {
	return;
    }

    /* Free context structures */
    if (statePtr->ctx != (EVP_CIPHER_CTX *) NULL) {
	EVP_CIPHER_CTX_free(statePtr->ctx);
    }
    ckfree(statePtr);
}

/*******************************************************************/

/*
 *-------------------------------------------------------------------
 *
 * EncryptInitialize --
 *
 *	Initialize an encryption function
 *
 * Returns:
 *	TCL_OK if successful or TCL_ERROR for failure with result set
 *	to error message.
 *
 * Side effects:
 *	No result or error message
 *
 *-------------------------------------------------------------------
 */
int EncryptInitialize(Tcl_Interp *interp, int type, EVP_CIPHER_CTX **ctx,
	Tcl_Obj *cipherObj, Tcl_Obj *keyObj, Tcl_Obj *ivObj) {
    const EVP_CIPHER *cipher;
    char *cipherName =  NULL, *key = NULL, *iv = NULL;
    int cipher_len = 0, data_len = 0, key_len = 0, iv_len = 0, res;

    dprintf("Called");

    /* Get encryption parameters */
    if (cipherObj != NULL) {
	cipherName = Tcl_GetStringFromObj(cipherObj, &cipher_len);
    }
    if (keyObj != NULL) {
	key = Tcl_GetStringFromObj(keyObj, &key_len);
    }
    if (ivObj != NULL) {
	iv = Tcl_GetStringFromObj(ivObj, &iv_len);
    }

    /* Get cipher name */
#if OPENSSL_VERSION_NUMBER < 0x30000000L
    cipher = EVP_get_cipherbyname(cipherName);
#else
    cipher = EVP_CIPHER_fetch(NULL, cipherName, NULL);
#endif
    if (cipher == NULL) {
	Tcl_AppendResult(interp, "Invalid cipher: \"", cipherName, "\"", NULL);
	return TCL_ERROR;
    }

    /* Create and initialize the context */
    if((*ctx = EVP_CIPHER_CTX_new()) == NULL) {
	Tcl_AppendResult(interp, "Memory allocation error", (char *) NULL);
	return TCL_ERROR;
    }

    /* Initialize the operation. Need appropriate key and iv size. */
#if OPENSSL_VERSION_NUMBER < 0x30000000L
    if (type == TYPE_ENCRYPT) {
	res = EVP_EncryptInit_ex(*ctx, cipher, NULL, key, iv);
    } else {
	res = EVP_DecryptInit_ex(*ctx, cipher, NULL, key, iv);
    }
#else
	OSSL_PARAM params[2];
	int index = 0;

	if (iv != NULL) {
	    params[index++] = OSSL_PARAM_construct_octet_string(OSSL_CIPHER_PARAM_IV, (void *) iv, (size_t) iv_len);
	}
	params[index] = OSSL_PARAM_construct_end();

    if (type == TYPE_ENCRYPT) {
	res = EVP_EncryptInit_ex2(ctx, cipher, key, iv, params);
    } else {
	res = EVP_DecryptInit_ex2(ctx, cipher, key, iv, params);
    }
#endif

    if(!res) {
	Tcl_AppendResult(interp, "Initialize failed: ", REASON(), NULL);
	return TCL_ERROR;
    }
    return TCL_OK;
}

/*
 *-------------------------------------------------------------------
 *
 * EncryptUpdate --
 *
 *	Update an encryption function with data
 *
 * Returns:
 *	1 if successful or 0 for failure
 *
 * Side effects:
 *	Adds encrypted data to buffer or sets result to error message
 *
 *-------------------------------------------------------------------
 */
int EncryptUpdate(Tcl_Interp *interp, int type, EVP_CIPHER_CTX *ctx, unsigned char *outbuf,
	int *out_len, unsigned char *data, int data_len) {
    int res, len = 0;

    dprintf("Called");

    /* Encrypt/decrypt data */
    if (type == TYPE_ENCRYPT) {
	res = EVP_EncryptUpdate(ctx, outbuf, out_len, data, data_len);
    } else {
	res = EVP_DecryptUpdate(ctx, outbuf, out_len, data, data_len);
    }

    if (res) {
	return TCL_OK;
    } else {
	Tcl_AppendResult(interp, "Update failed: ", REASON(), NULL);
	return TCL_ERROR;
    }
}

/*
 *-------------------------------------------------------------------
 *
 * EncryptFinalize --
 *
 *	Finalize an encryption function
 *
 * Returns:
 *	TCL_OK if successful or TCL_ERROR for failure with result set
 *	to error message.
 *
 * Side effects:
 *	Adds encrypted data to buffer or sets result to error message
 *
 *-------------------------------------------------------------------
 */
int EncryptFinalize(Tcl_Interp *interp, int type, EVP_CIPHER_CTX *ctx, unsigned char *outbuf,
	int *out_len) {
    int res, len = 0;

    dprintf("Called");

    /* Finalize data */
    if (type == TYPE_ENCRYPT) {
	res = EVP_EncryptFinal_ex(ctx, outbuf, out_len);
    } else {
	res = EVP_DecryptFinal_ex(ctx, outbuf, out_len);
    }

    if (res) {
	return TCL_OK;
    } else {
	Tcl_AppendResult(interp, "Finalize failed: ", REASON(), NULL);
	return TCL_ERROR;
    }
}

/*******************************************************************/

/*
 *-------------------------------------------------------------------
 *
 * EncryptInstanceObjCmd --
 *
 *	Handler for encrypt/decrypt command instances. Used to update
 *	and finalize data for encrypt/decrypt function.
 *
 * Returns:
 *	TCL_OK or TCL_ERROR
 *
 * Side effects:
 *	Adds data to encrypt/decrypt function
 *
 *-------------------------------------------------------------------
 */
int EncryptInstanceObjCmd(ClientData clientData, Tcl_Interp *interp, int objc, Tcl_Obj *const objv[]) {
    EncryptState *statePtr = (EncryptState *) clientData;
    int fn, data_len = 0, out_len;
    char *data = NULL;
    Tcl_Obj *resultObj;
    unsigned char *outbuf;
    static const char *instance_fns [] = { "finalize", "update", NULL };

    dprintf("Called");

    /* Validate arg count */
    if (objc < 2 || objc > 3) {
	Tcl_WrongNumArgs(interp, 1, objv, "function ?data?");
	return TCL_ERROR;
    }

    /* Get function */
    if (Tcl_GetIndexFromObj(interp, objv[1], instance_fns, "function", 0, &fn) != TCL_OK) {
	return TCL_ERROR;
    }

    /* Allocate storage for result. Size should be data size + block size. */
    resultObj = Tcl_NewObj();
    outbuf = Tcl_SetByteArrayLength(resultObj, data_len+EVP_MAX_BLOCK_LENGTH);
    if (resultObj == NULL || outbuf == NULL) {
	Tcl_AppendResult(interp, "Memory allocation error", (char *) NULL);
	Tcl_DecrRefCount(resultObj);
	return TCL_ERROR;
    }

    /* Do function */
    if (fn) {
	/* Get data or return error if none */
	if (objc == 3) {
	    data = Tcl_GetByteArrayFromObj(objv[2], &data_len);
	} else {
	    Tcl_WrongNumArgs(interp, 1, objv, "update data");
	    Tcl_DecrRefCount(resultObj);
	    return TCL_ERROR;
	}

	/* Update function */
	if (EncryptUpdate(interp, statePtr->type, statePtr->ctx, outbuf, &out_len, data, data_len) == TCL_OK) {
	    outbuf = Tcl_SetByteArrayLength(resultObj, out_len);
	    Tcl_SetObjResult(interp, resultObj);
	} else {
	    Tcl_DecrRefCount(resultObj);
	    return TCL_ERROR;
	}

    } else {
	/* Finalize function */
	if (EncryptFinalize(interp, statePtr->type, statePtr->ctx, outbuf, &out_len) == TCL_OK) {
	    outbuf = Tcl_SetByteArrayLength(resultObj, out_len);
	    Tcl_SetObjResult(interp, resultObj);
	} else {
	    Tcl_DecrRefCount(resultObj);
	    return TCL_ERROR;
	}

	/* Clean-up */
	Tcl_DeleteCommandFromToken(interp, statePtr->token);
    }
    return TCL_OK;
}

/*
 *-------------------------------------------------------------------
 *
 * EncryptCommandDeleteHandler --
 *
 *	 Callback to clean-up when encrypt/decrypt command is deleted.
 *
 * Returns:
 *	Nothing
 *
 * Side effects:
 *	Destroys state info structure
 *
 *-------------------------------------------------------------------
 */
void EncryptCommandDeleteHandler(ClientData clientData) {
    EncryptState *statePtr = (EncryptState *) clientData;

    /* Clean-up */
    EncryptStateFree(statePtr);
}

/*
 *-------------------------------------------------------------------
 *
 * EncryptCommandHandler --
 *
 *	 Create command to add data to encrypt/decrypt function.
 *
 * Returns:
 *	TCL_OK or TCL_ERROR
 *
 * Side effects:
 *	Creates command or error message
 *
 *-------------------------------------------------------------------
 */
int EncryptCommandHandler(Tcl_Interp *interp, int type, Tcl_Obj *cmdObj,
	Tcl_Obj *cipherObj, Tcl_Obj *digestObj, Tcl_Obj *keyObj, Tcl_Obj *ivObj) {
    EncryptState *statePtr;
    char *cmdName = Tcl_GetStringFromObj(cmdObj, NULL);

    dprintf("Called");

    if ((statePtr = EncryptStateNew(interp, type)) == NULL) {
	Tcl_AppendResult(interp, "Memory allocation error", (char *) NULL);
	return TCL_ERROR;
    }

    /* Initialize function */
    if (EncryptInitialize(interp, type, &statePtr->ctx, cipherObj, keyObj, ivObj) != TCL_OK) {
	EncryptStateFree(statePtr);
	return TCL_ERROR;
    }

    /* Create instance command */
    statePtr->token = Tcl_CreateObjCommand(interp, cmdName, EncryptInstanceObjCmd,
	(ClientData) statePtr, EncryptCommandDeleteHandler);

    /* Return command name */
    Tcl_SetObjResult(interp, cmdObj);
    return TCL_OK;
}

/*******************************************************************/

/*
 *-------------------------------------------------------------------
 *
 * EncryptDataHandler --
 *
 *	Perform encryption function on a block of data and return result.
 *
 * Returns:
 *	TCL_OK or TCL_ERROR
 *
 * Side effects:
 *	Sets result or error message
 *
 *-------------------------------------------------------------------
 */
int EncryptDataHandler(Tcl_Interp *interp, int type, Tcl_Obj *dataObj, Tcl_Obj *cipherObj,
	Tcl_Obj *digestObj, Tcl_Obj *keyObj, Tcl_Obj *ivObj) {
    EVP_CIPHER_CTX *ctx = NULL;
    int data_len = 0, out_len = 0, len = 0, res = TCL_OK;
    unsigned char *data, *outbuf;
    Tcl_Obj *resultObj;

    dprintf("Called");

    /* Get data */
    if (dataObj != NULL) {
	data = Tcl_GetByteArrayFromObj(dataObj, &data_len);
    } else {
	Tcl_AppendResult(interp, "No data", NULL);
	return TCL_ERROR;
    }

    /* Allocate storage for result. Size should be data size + block size. */
    resultObj = Tcl_NewObj();
    outbuf = Tcl_SetByteArrayLength(resultObj, data_len+EVP_MAX_BLOCK_LENGTH);
    if (resultObj == NULL || outbuf == NULL) {
	Tcl_AppendResult(interp, "Memory allocation error", (char *) NULL);
	return TCL_ERROR;
    }

    /* Perform operation */
    if (EncryptInitialize(interp, type, &ctx, cipherObj, keyObj, ivObj) != TCL_OK ||
	EncryptUpdate(interp, type, ctx, outbuf, &out_len, data, data_len) != TCL_OK ||
	EncryptFinalize(interp, type, ctx, outbuf+out_len, &len) != TCL_OK) {
	res = TCL_ERROR;
	goto done;
    }
    out_len += len;

done:
    /* Set output result */
    if (res == TCL_OK) {
	outbuf = Tcl_SetByteArrayLength(resultObj, out_len);
	Tcl_SetObjResult(interp, resultObj);
    } else {
	Tcl_DecrRefCount(resultObj);
	/* Result is error message */
    }

    /* Clean up */
    if (ctx != NULL) {
	EVP_CIPHER_CTX_free(ctx);
    }
    return res;
}

/*******************************************************************/

/*
 *-------------------------------------------------------------------
 *
 * EncryptFileHandler --
 *
 *	Perform encryption function on a block of data and return result.
 *
 * Returns:
 *	TCL_OK or TCL_ERROR
 *
 * Side effects:
 *	Encrypts or decrypts inFile data to outFile and sets result to
 *	size of outFile, or an error message.
 *
 *-------------------------------------------------------------------
 */
int EncryptFileHandler(Tcl_Interp *interp, int type, Tcl_Obj *inFileObj, Tcl_Obj *outFileObj,
	Tcl_Obj *cipherObj, Tcl_Obj *digestObj, Tcl_Obj *keyObj, Tcl_Obj *ivObj) {
    EVP_CIPHER_CTX *ctx = NULL;
    int total = 0, res, out_len = 0, len;
    Tcl_Channel in = NULL, out = NULL;
    unsigned char in_buf[BUFFER_SIZE];
    unsigned char out_buf[BUFFER_SIZE+EVP_MAX_BLOCK_LENGTH];

    dprintf("Called");

    /* Open input file */
    if ((in = Tcl_FSOpenFileChannel(interp, inFileObj, "rb", 0444)) == (Tcl_Channel) NULL) {
	return TCL_ERROR;
    }

    /* Open output file */
    if ((out = Tcl_FSOpenFileChannel(interp, outFileObj, "wb", 0644)) == (Tcl_Channel) NULL) {
	Tcl_Close(interp, in);
	return TCL_ERROR;
    }

    /* Initialize operation */
    if ((res = EncryptInitialize(interp, type, &ctx, cipherObj, keyObj, ivObj)) != TCL_OK) {
	goto done;
    }

    /* Read file data from inFile, encrypt/decrypt it, then output to outFile */
    while (!Tcl_Eof(in)) {
	int read = Tcl_ReadRaw(in, (char *) in_buf, BUFFER_SIZE);
	if (read > 0) {
	    if ((res = EncryptUpdate(interp, type, ctx, out_buf, &out_len, in_buf, read)) == TCL_OK) {
		if (out_len > 0) {
		    len = Tcl_WriteRaw(out, (const char *) out_buf, out_len);
		    if (len >= 0) {
			total += len;
		    } else {
			Tcl_AppendResult(interp, "Write error: ", Tcl_ErrnoMsg(Tcl_GetErrno()), (char *) NULL);
			res = TCL_ERROR;
			goto done;
		    }
		}
	    } else {
		goto done;
	    }
	} else if (read < 0) {
	    Tcl_AppendResult(interp, "Read error: ", Tcl_ErrnoMsg(Tcl_GetErrno()), (char *) NULL);
	    res = TCL_ERROR;
	    goto done;
	}
    }

    /* Finalize data and write any remaining data in block */
    if ((res = EncryptFinalize(interp, type, ctx, out_buf, &out_len)) == TCL_OK) {
	if (out_len > 0) {
	    len = Tcl_WriteRaw(out, (const char *) out_buf, out_len);
	    if (len >= 0) {
		total += len;
	    } else {
		Tcl_AppendResult(interp, "Write error: ", Tcl_ErrnoMsg(Tcl_GetErrno()), (char *) NULL);
		res = TCL_ERROR;
		goto done;
	    }
	}
	Tcl_SetObjResult(interp, Tcl_NewIntObj(total));
    } else {
	goto done;
    }

done:
    /* Clean up */
    if (in != NULL) {
	Tcl_Close(interp, in);
    }
    if (out != NULL) {
	Tcl_Close(interp, out);
    }
    if (ctx != NULL) {
	EVP_CIPHER_CTX_free(ctx);
    }
    return res;
}

/*******************************************************************/

/*
 *-------------------------------------------------------------------
 *
 * EncryptMain --
 *
 *	Perform encryption function and return result.
 *
 * Returns:
 *	TCL_OK or TCL_ERROR
 *
 * Side effects:
 *	Sets result or error message
 *
 *-------------------------------------------------------------------
 */
static int EncryptMain(int type, Tcl_Interp *interp, int objc, Tcl_Obj *const objv[]) {
    int res = TCL_OK;
    Tcl_Obj *cipherObj = NULL, *cmdObj = NULL, *dataObj = NULL, *digestObj = NULL;
    Tcl_Obj *inFileObj = NULL, *outFileObj = NULL, *keyObj = NULL, *ivObj = NULL, *macObj = NULL;
    const char *channel = NULL, *opt;
    const EVP_MD *md = NULL;
    const EVP_CIPHER *cipher = NULL;

    dprintf("Called");

    /* Clear interp result */
    Tcl_ResetResult(interp);

    /* Validate arg count */
    if (objc < 3 || objc > 12) {
	Tcl_WrongNumArgs(interp, 1, objv, "-cipher name ?-digest name? -key key ?-iv string? [-command cmdName | -infile filename -outfile filename | -data data]");
	return TCL_ERROR;
    }

    /* Get options */
    for (int idx = 1; idx < objc; idx++) {
	opt = Tcl_GetStringFromObj(objv[idx], NULL);

	if (opt[0] != '-') {
	    break;
	}

	OPTOBJ("-cipher", cipherObj);
	OPTOBJ("-command", cmdObj);
	OPTOBJ("-data", dataObj);
	OPTOBJ("-digest", digestObj);
	OPTOBJ("-infile", inFileObj);
	OPTOBJ("-outfile", outFileObj);
	OPTOBJ("-key", keyObj);
	OPTOBJ("-iv", ivObj);

	OPTBAD("option", "-cipher, -command, -data, -digest, -infile, -key, -iv, -outfile");
	return TCL_ERROR;
    }

    /* Check for required options */
    if (cipherObj == NULL) {
	Tcl_AppendResult(interp, "No cipher", NULL);
    } else if (keyObj == NULL) {
	Tcl_AppendResult(interp, "No key", NULL);
	return TCL_ERROR;
    }

    /* Perform encryption function on file, stacked channel, using instance command, or data blob */
    if (inFileObj != NULL && outFileObj != NULL) {
	res = EncryptFileHandler(interp, type, inFileObj, outFileObj, cipherObj, digestObj, keyObj, ivObj);
    } else if (cmdObj != NULL) {
	res = EncryptCommandHandler(interp, type, cmdObj, cipherObj, digestObj, keyObj, ivObj);
    } else if (dataObj != NULL) {
	res = EncryptDataHandler(interp, type, dataObj, cipherObj, digestObj, keyObj, ivObj);
    } else {
	Tcl_AppendResult(interp, "No operation specified: Use -command, -data, -infile, or -outfile option", NULL);
	res = TCL_ERROR;
    }
    return res;
}

/*
 *-------------------------------------------------------------------
 *
 * Encryption Commands --
 *
 *	Perform encryption function and return results
 *
 * Returns:
 *	TCL_OK or TCL_ERROR
 *
 * Side effects:
 *	Command dependent
 *
 *-------------------------------------------------------------------
 */
static int EncryptObjCmd(ClientData clientData, Tcl_Interp *interp, int objc, Tcl_Obj *const objv[]) {
    return EncryptMain(TYPE_ENCRYPT, interp, objc, objv);
}

static int DecryptObjCmd(ClientData clientData, Tcl_Interp *interp, int objc, Tcl_Obj *const objv[]) {
    return EncryptMain(TYPE_DECRYPT, interp, objc, objv);
}

/*
 *-------------------------------------------------------------------
 *
 * Encrypt_Initialize --
 *
 *	Create namespace, commands, and register package version
 *
 * Returns:
 *	TCL_OK or TCL_ERROR
 *
 * Side effects:
 *	Creates commands
 *
 *-------------------------------------------------------------------
 */
int Tls_EncryptCommands(Tcl_Interp *interp) {
    Tcl_CreateObjCommand(interp, "tls::encrypt", EncryptObjCmd, (ClientData) 0, (Tcl_CmdDeleteProc *) NULL);
    Tcl_CreateObjCommand(interp, "tls::decrypt", DecryptObjCmd, (ClientData) 0, (Tcl_CmdDeleteProc *) NULL);
    return TCL_OK;
}