#ifndef _GNU_SOURCE
# define _GNU_SOURCE
#endif

#include <sys/types.h>
#include <sys/stat.h>
#include <stdint.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <unistd.h>
#include <dlfcn.h>

// Peek the first few bytes of a file (i.e., #!/bin/sh etc)
static int peek_file(const char *filename, char *buf, int bufSize)
{
	FILE *fp = fopen(filename, "r");
	int rv;

	if (!fp)
		return -1;

	rv = fread(buf, 1, bufSize, fp);

	fclose(fp);

	return rv;
}

static char *kcovBash;
static int kcovUseDebugTrap;

static int (*orig_execve)(const char *filename, char * const argv[], char * const envp[]);
int execve(const char *filename, char * const argv[], char * const envp[])
{
	char startBytes[128];
	int sizeRead;
	const char bashArg[3] =
	{ '-', 'x', '\0' };

	if (!orig_execve)
		orig_execve = dlsym(RTLD_NEXT, "execve");

	sizeRead = peek_file(filename, startBytes, sizeof(startBytes));

	// Make sure it's NULL-terminated
	startBytes[sizeof(startBytes) - 1] = '\0';
	if (sizeRead >= 0)
		startBytes[sizeRead] = '\0';

	// Shell script being executed?
	if (kcovBash && sizeRead >= 6
			&& (strstr(startBytes, "/bin/sh") != NULL || strstr(startBytes, "/bin/bash") != NULL
					|| strstr(startBytes, "/bin/ash") != NULL || strstr(startBytes, "/bin/dash") != NULL ))
	{
		char **replacementArgv;
		unsigned i = 0;

		// Find end
		for (i = 0; argv[i]; i++)
			;

		unsigned arg = 0;

		replacementArgv = malloc(sizeof(char*) * (i + 4));
		replacementArgv[arg++] = strdup(kcovBash);
		if (!kcovUseDebugTrap)
			replacementArgv[arg++] = strdup(bashArg);
		replacementArgv[arg++] = strdup(filename);
		for (i = 0; argv[i]; i++)
			replacementArgv[i + arg] = strdup(argv[i]);
		replacementArgv[i + arg] = NULL;

		filename = strdup(kcovBash);

		argv = replacementArgv;
	}

	return orig_execve(filename, argv, envp);
}

void __attribute__((constructor)) kcov_bash_execve_at_startup(void)
{
	kcovBash = getenv("KCOV_BASH_COMMAND");
	if (getenv("KCOV_BASH_USE_DEBUG_TRAP"))
		kcovUseDebugTrap = 1;
}
