polonium/src/main.c

// TODO: if user requires, only damage contents of PNG, BMP, WAV
//       PNG: work only with pixels
//       BMP: work only with pixels
//       WAV: work only with samples

#include <ctype.h>
#include <inttypes.h>
#include <stdbool.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>

#include "sys_id.h"
#ifdef SYS_UNIX_GENERAL
# include <sys/types.h>
# include <sys/stat.h>
#endif

#include "common.h"

#include "MTPRNG.h"
#include "atoumax_base10.h"
#include "corrupter.h"
#include "file_type.h"

/* enums */
enum arg_destinations {
	ARG_NO_DEST,
	ARG_DEST_POSSIBILITY,
	ARG_DEST_THRESHOLD,
	ARG_DEST_PASSES,
	ARG_DEST_SEED
};

/* constant values: arguments */
const char* ARGS_HELP[] =         { "-h", "--help" };
const char* ARGS_PROBABILITY[] =  { "-p", "--probability" };
const char* ARGS_THRESHOLD[] =    { "-t", "--threshold" };
const char* ARGS_PASSES[] =       { "-n", "--passes" };
const char* ARGS_CONTENTS[] =     { "-c", "--contents" };
const char* ARGS_LINE_ENDINGS[] = { "-l", "--line-endings" };
const char* ARGS_PRINTABLE[] =    { "-P", "--printable" };
const char* ARGS_SEED[] =         { "-s", "--seed" };

/* global variables */
uint32_t PRNG_seed_value;
uint8_t config = C_CONFIRM;

/* macros: procedures */
#define DOES_NOT_EXIST                                                        \
    if (errno == ENOENT) {                                                    \
        FPRINTF_MACRO(stderr, "Error: File %s doesn't exist\n", file_path);   \
        return EXIT_FAILURE;                                                  \
    }

/* macros: lambdas */
#define ARG_MATCH(s, args2) (!strcmp(s, args2[0]) || !strcmp(s, args2[1]))

/* default values */
static uint16_t probability = 100;
static uint8_t threshold = 2;
static size_t passes = 1;

/* function definitions */
static char* correct_slashes(const char* path);
static char* my_basename(const char* raw_path);
static void parse_value(uint8_t destination, const char* arg);

/* function implementations */
static char* correct_slashes(const char* path) {
	char* new_path = strdup(path);
	if (new_path == NULL)
		return NULL;

	char* ptr = new_path;
	while (*ptr != '\0') {
		if (*ptr == '\\')
			*ptr = '/';
		ptr++;
	}

	return new_path;
}

static char* my_basename(const char* raw_path) {
	char* path = correct_slashes(raw_path);
	if (path == NULL)
		return NULL;

	char* base = malloc(FILENAME_MAX * sizeof(char));
	if (base == NULL) {
		free(path);
		PERROR_MACRO("malloc");
		exit(EXIT_FAILURE);
		return NULL;
	}

	size_t fname_len = strlen(path);

	const char* last_slash = strrchr(path, '/');
	if (last_slash != NULL)
		fname_len = strlen(last_slash + 1);

	memcpy(base, last_slash + 1, fname_len);
	base[fname_len] = '\0';

	return base;
}

static void parse_value(uint8_t destination, const char* arg) {
	if (destination == ARG_DEST_SEED && !strcmp(arg, "random"))
		return;

	uintmax_t result = atoumax_base10(arg);
	switch (destination) {
	case ARG_DEST_POSSIBILITY:
		if (result == 0 || result > UINT8_MAX) {
			FATAL_ERROR("Chance value should be in range [1..%" PRIu8 "]",
			            UINT8_MAX);
			return;
		}

		probability = result;
		break;
	case ARG_DEST_THRESHOLD:
		if (result < 1 || result > CHAR_BIT) {
			FATAL_ERROR("Damage value should be in range [1.."
			            INT2STR(CHAR_BIT) "]");
			return;
		}

		threshold = result;
		break;
	case ARG_DEST_PASSES:
		if (result == 0) {
			FATAL_ERROR("The number of passes can't be less than 1");
			return;
		}

		passes = result;
		break;
	case ARG_DEST_SEED:
		SET_CONFIG(C_CUSTOM_SEED);
		PRNG_seed_value = result;
		break;
	default:
		FATAL_ERROR("Unknown argument destination (value = %" PRIu16 ")",
		           destination);
		return;
	}
}

int main(int argc, char** argv) {
	puts("Polonium: a file corrupter\n");

	if (argc < 2 || ARG_MATCH(argv[argc - 1], ARGS_HELP)) {
		char* program_name = my_basename(argv[0]);
		printf(
			"Usage: %s <file to corrupt> [parameters] [options]\n"
			"\n"
			"Both parameters and options are optional.\n"
			"\n"
			"[parameters]:\n"
			"  --probability : Determines the likelihood of a neutron "
			"radiation event\n"
			"                  occurring. The higher the value, the more "
			"likely the bits are\n"
			"                  to be corrupted.\n"
			"                  Value range: [1..%" PRIu16 "]\n"
			"                  Default value: %" PRIu16 "\n"
			"\n"
			"  --threshold   : Controls the amount of change allowed for each "
			"byte. The\n"
			"                  higher the value, the more extensive the "
			"modifications will\n"
			"                  be.\n"
			"                  Value range: [1.." INT2STR(CHAR_BIT) "]\n"
			"                  Default value: %" PRIu8 "\n"
			"\n"
			"  --passes      : This parameter determines the number of times "
			"the file will\n"
			"                  be processed. The higher the value, the longer "
			"the overall\n"
			"                  processing time will be.\n"
			"                  Value range: [1..%" PRIuMAX "]\n"
			"                  Default value: %" PRIuMAX "\n"
			"\n"
			"[options]:\n"
			"  --noconfirm    : (Caution!) Do not ask user to confirm "
			"deletion\n"
			"\n"
			"  -c, --contents : Only corrupt contents of files (PNG, BMP, WAV)\n"
			"\n"
			"  -l,\n"
			"  --line-endings : Preserve line endings\n"
			"\n"
			"  -p,\n"
			"  --printable    : Work only with printable ASCII characters\n"
			"\n"
			"  --seed         : Specify a 32-bit seed for the PRNG. If you "
			"want to keep it\n"
			"                   random, set the option to `random`.\n",
			program_name, UINT16_MAX, probability, threshold,
			(uintmax_t) SIZE_MAX, passes);
		free(program_name);
		return EXIT_FAILURE;
	}

	const char* file_path = argv[1];
	file_type_t file_type = FILE_TYPE_AUTO;

	if (argc > 2) {
		uint8_t arg_destination = ARG_NO_DEST;
		bool read_off_value = false;

		for (int i = 2; i < argc; ++i) {
			bool last_arg = i == (argc - 1);

			const char* arg = argv[i];

			if (read_off_value) {
				parse_value(arg_destination, arg);
				read_off_value = false;
				continue;
			}

			if (ARG_MATCH(arg, ARGS_PROBABILITY))
				arg_destination = ARG_DEST_POSSIBILITY;
			else if (ARG_MATCH(arg, ARGS_THRESHOLD))
				arg_destination = ARG_DEST_THRESHOLD;
			else if (ARG_MATCH(arg, ARGS_PASSES))
				arg_destination = ARG_DEST_PASSES;
			else if (ARG_MATCH(arg, ARGS_SEED))
				arg_destination = ARG_DEST_SEED;
			else if (!strcmp(arg, "--noconfirm")) {
				LOOP_ACTION_CONFIG(CLEAR_CONFIG, C_CONFIRM);
			} else if (ARG_MATCH(arg, ARGS_CONTENTS)) {
				LOOP_ACTION_CONFIG(SET_CONFIG, C_CONTENTS);
			} else if (ARG_MATCH(arg, ARGS_LINE_ENDINGS)) {
				LOOP_ACTION_CONFIG(SET_CONFIG, C_LINE_ENDINGS);
			} else if (ARG_MATCH(arg, ARGS_PRINTABLE)) {
				LOOP_ACTION_CONFIG(SET_CONFIG, C_PRINTABLE);
			} else
				FATAL_ERROR_RET("Unknown command line parameter '%s'.\n"
				                "Please invoke the program with argument %s or "
				                "%s to display help.\n",
				                arg, ARGS_HELP[0], ARGS_HELP[1]);

			if (last_arg)
				FATAL_ERROR_RET("Please provide the value.");
			else {
				read_off_value = true;
				continue;
			}
		}
	}

#ifdef SYS_UNIX_GENERAL
	// check if the file is a directory
	struct stat path_stat;
	if (stat(file_path, &path_stat) != 0) {
		if (errno == ENOENT) {
			DOES_NOT_EXIST;
		}
		PERROR_MACRO("stat");
		return EXIT_FAILURE;
	}

	if (S_ISDIR(path_stat.st_mode)) {
		FPRINTF_MACRO(stderr, "Error: %s is a directory\n", file_path);
		return EXIT_FAILURE;
	}
#endif

	// open the file
	FILE* file = fopen(file_path, "rb+");
	if (file == NULL) {
#ifndef SYS_UNIX_GENERAL
		DOES_NOT_EXIST;
#endif
		PERROR_MACRO("fopen");
		return EXIT_FAILURE;
	}

	if (READ_CONFIG(C_CONTENTS))
		// FIXME: interprets *.txt files as binary
		file_type = determine_file_type(file, file_path);

	printf("Configuration:\n"
	       "> Only damaging file contents (PNG, BMP, WAV): %s\n"
	       "> Preserving line endings: %s\n"
	       "> Operating on printable ASCII characters exclusively: %s\n"
	       "> Custom seed: %s\n"
	       "> File type: %s\n"
	       "\n",
	       YES_NO(READ_CONFIG(C_CONTENTS)),
	       YES_NO(READ_CONFIG(C_LINE_ENDINGS)),
	       YES_NO(READ_CONFIG(C_PRINTABLE)),
	       YES_NO(READ_CONFIG(C_CUSTOM_SEED)),
	       file_type_to_string(file_type));

	printf("Parameters:\n"
	       "> Probability: %" PRIu8 "\n"
	       "> Threshold: %" PRIu8 "\n"
	       "> Passes: %" PRIuMAX "\n"
	       "\n", probability, threshold, (uintmax_t) passes);

	if (READ_CONFIG(C_CONFIRM)) {
		printf("Are you sure? (Y/n): ");
		fflush(stdout);
		if ((getchar()) != (int) 'Y') {
			printf("File corruption aborted.\n");
			return EXIT_FAILURE;
		}
	}

	fflush(stdout);

	Corrupter_Param param = {
		.file = file,

		// parameters
		.probability = probability,
		.threshold = threshold,
		.passes = passes,

		.config = config,

		// configuration
		.seed = READ_CONFIG(C_CUSTOM_SEED) ? &PRNG_seed_value : NULL,
		.type = file_type
	};

	Corrupter_Result* result = corrupt_file(&param);
	if (result == NULL) {
		PERROR_MACRO("corrupt_file memory allocation");
		return EXIT_FAILURE;
	} else if (result->error) {
		free(result);
		PERROR_MACRO("corrupt_file");
		return EXIT_FAILURE;
	}

	putchar('\n');

	if (result->damaged_bytes) {
		size_t dmg = result->damaged_bytes, fsize = result->file_size,
		       passes = param.passes;
		printf("Byte hit counter: %" PRIuMAX " / %" PRIuMAX " = %.3Lf%%\n",
		       (uintmax_t) dmg, (uintmax_t) (fsize * passes),
		       ((long double) dmg * 100.l) / (long double) (fsize * passes));
	} else {
		puts("No bytes were damaged");
	}

	free(result);

	// finish working with the file
	fclose(file);

	puts("Done");

	return EXIT_SUCCESS;
}