/*
 * aes-ctr: AES-CTR Encryption and Decryption Utility
 * Copyright (c) 2024 Ali Polatel <alip@chesswob.org>
 * SPDX-License-Identifier: GPL-3.0
 *
 * This program uses OpenSSL to encrypt and decrypt data using AES-CTR
 * mode. It supports both encryption and decryption operations, with
 * key and IV provided as hexadecimal strings.
 *
 * Compile:
 *        cc aes-ctr.c -o aes-ctr -lssl -lcrypto
 *
 * Usage: aes-ctr [-hV] -e|-d -k <hex-encoded key> -i <hex-encoded iv>
 *
 * Options:
 *   -h        Print this help message and exit
 *   -V        Print version information and exit
 *   -e        Encrypt the input data
 *   -d        Decrypt the input data
 *   -k <key>  Hex-encoded key (256 bits for AES-CTR)
 *   -i <iv>   Hex-encoded IV (128 bits for AES-CTR)
 *
 * Examples:
 *   Encrypt: echo -n "data" | ./aes-ctr -e -k <key> -i <iv>
 *   Decrypt: ./aes-ctr -d -k <key> -i <iv> < encrypted_data.bin
 */

#include <openssl/evp.h>
#include <openssl/err.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <getopt.h>

#define VERSION "0.1.0"
#define KEY_LEN 32
#define IV_LEN 16
#define BUFFER_SIZE 4096

void handleErrors(void)
{
	ERR_print_errors_fp(stderr);
	exit(127);
}

/**
 * Memory allocation with error checking.
 * @param size The size of memory to allocate.
 * @return Pointer to the allocated memory.
 */
void *xmalloc(size_t size)
{
	void *ptr = malloc(size);
	if (!ptr) {
		fprintf(stderr, "Memory allocation failed\n");
		exit(127);
	}
	return ptr;
}

/**
 * Memory reallocation with error checking.
 * @param ptr The original memory pointer.
 * @param size The new size of memory to allocate.
 * @return Pointer to the reallocated memory.
 */
void *xrealloc(void *ptr, size_t size)
{
	ptr = realloc(ptr, size);
	if (!ptr) {
		fprintf(stderr, "Memory reallocation failed\n");
		exit(127);
	}
	return ptr;
}

/**
 * Convert a hexadecimal string to a byte array.
 * @param hex The hexadecimal string.
 * @param bytes The output byte array.
 * @param bytes_len The length of the output byte array.
 */
void hex_to_bytes(const char *hex, unsigned char *bytes, size_t bytes_len)
{
	for (size_t i = 0; i < bytes_len; i++) {
		sscanf(hex + 2 * i, "%2hhx", &bytes[i]);
	}
}

/**
 * Encrypt or decrypt data using AES-CTR.
 * @param encrypt Set to 1 for encryption, 0 for decryption.
 * @param key The AES key.
 * @param iv The IV (tweak) value.
 */
void process_data(int encrypt, const unsigned char *key,
                  const unsigned char *iv)
{
	EVP_CIPHER_CTX *ctx = EVP_CIPHER_CTX_new();
	if (!ctx) handleErrors();

	if (encrypt) {
		if (1 != EVP_EncryptInit_ex(ctx, EVP_aes_256_ctr(), NULL, key,
		                            iv)) handleErrors();
	} else {
		if (1 != EVP_DecryptInit_ex(ctx, EVP_aes_256_ctr(), NULL, key,
		                            iv)) handleErrors();
	}

	unsigned char inbuf[BUFFER_SIZE];
	unsigned char outbuf[BUFFER_SIZE + EVP_CIPHER_block_size(EVP_aes_256_ctr())];
	int inlen, outlen;

	while ((inlen = fread(inbuf, 1, BUFFER_SIZE, stdin)) > 0) {
		if (encrypt) {
			if (1 != EVP_EncryptUpdate(ctx, outbuf, &outlen, inbuf, inlen))
				handleErrors();
		} else {
			if (1 != EVP_DecryptUpdate(ctx, outbuf, &outlen, inbuf, inlen))
				handleErrors();
		}
		fwrite(outbuf, 1, outlen, stdout);
	}

	if (encrypt) {
		if (1 != EVP_EncryptFinal_ex(ctx, outbuf, &outlen))
			handleErrors();
	} else {
		if (1 != EVP_DecryptFinal_ex(ctx, outbuf, &outlen))
			handleErrors();
	}
	fwrite(outbuf, 1, outlen, stdout);

	EVP_CIPHER_CTX_free(ctx);
}

int main(int argc, char *argv[])
{
	int opt;
	int encrypt = -1;
	char *key_hex = NULL;
	char *iv_hex = NULL;

	while ((opt = getopt(argc, argv, "hVdek:i:")) != -1) {
		switch (opt) {
		case 'h':
			printf("Usage: aes-ctr [-hV] -e|-d -k <hex-encoded key> -i <hex-encoded iv>\n");
			printf("  -h        Print this help message and exit\n");
			printf("  -V        Print version information and exit\n");
			printf("  -e        Encrypt the input data\n");
			printf("  -d        Decrypt the input data\n");
			printf("  -k <key>  Hex-encoded key (256 bits for AES-CTR)\n");
			printf("  -i <iv>   Hex-encoded IV (128 bits for AES-CTR)\n");
			return EXIT_SUCCESS;
		case 'V':
			printf("%s v%s\n", argv[0], VERSION);
			return EXIT_SUCCESS;
		case 'e':
			encrypt = 1;
			break;
		case 'd':
			encrypt = 0;
			break;
		case 'k':
			key_hex = optarg;
			break;
		case 'i':
			iv_hex = optarg;
			break;
		default:
			fprintf(stderr,
			        "Usage: aes-ctr [-hV] -e|-d -k <hex-encoded key> -i <hex-encoded iv>\n");
			return EXIT_FAILURE;
		}
	}

	if (encrypt == -1 || key_hex == NULL || iv_hex == NULL) {
		fprintf(stderr, "Error: -e or -d and both -k and -i options are required.\n");
		fprintf(stderr,
		        "Usage: aes-ctr [-hV] -e|-d -k <hex-encoded key> -i <hex-encoded iv>\n");
		return EXIT_FAILURE;
	}

	if (strlen(key_hex) != 2 * KEY_LEN) {
		fprintf(stderr, "Error: Key must be 256 bits (32 bytes) in length.\n");
		return 1;
	}

	if (strlen(iv_hex) != 2 * IV_LEN) {
		fprintf(stderr, "Error: IV must be 128 bits (16 bytes) in length.\n");
		return 1;
	}

	// AES-256-CTR requires a 256-bit key (32 bytes)
	unsigned char key[KEY_LEN];
	// CTR mode uses a 128-bit tweak (16 bytes)
	unsigned char iv[IV_LEN];
	hex_to_bytes(key_hex, key, sizeof(key));
	hex_to_bytes(iv_hex, iv, IV_LEN);

	process_data(encrypt, key, iv);

	return EXIT_SUCCESS;
}
