#include <stdlib.h>
#include <stdint.h>
#include <stdio.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/stat.h>
#include <time.h>
#include <unistd.h>
#include <elf.h>
#include <limits.h>

/* Config */
#define NAME_PREFIX        "poly_" /* Copies' names start with */
#define NAME_PREFIX_LENGTH 5       /* Length of prefix */
#define NAME_AFFIX_LENGTH  6       /* Number of random chars after prefix */
#define KEY_SIZE           8       /* Size of encryption key */

#define POLYMORPHIC __attribute__((noinline, section(".polymorphic"), used))

#if ( __WORDSIZE == 64 )
# define EHDR Elf64_Ehdr
# define SHDR Elf64_Shdr
#else
# define EHDR Elf32_Ehdr
# define SHDR Elf32_Shdr
#endif

/* Keys */
const uint8_t act[KEY_SIZE] = {}, /* Operators */
              key[KEY_SIZE] = {}; /* Operands */
uintptr_t polymorphicStart, polymorphicEnd;

/* Copy a file */
POLYMORPHIC
void
copy_file(char *src, char *dst)
{
	FILE *parent, *child;
	char data[1024];
	size_t n, m;

	/* Open files */
	parent = fopen(src, "rb");
	child = fopen(dst, "wb");
	if (parent == NULL || child == NULL)
		exit(EXIT_FAILURE);

	/* Copy contents across */
	do {
		n = fread(data, 1, sizeof(data), parent);
		if (n)
			m = fwrite(data, 1, n, child);
		else
			m = 0;
	} while ((n > 0) && (n == m));
	if (m)
		exit(EXIT_FAILURE);

	/* Close files */
	fclose(parent);
	fclose(child);
}

/* Polymorphic main function */
POLYMORPHIC
int
polymorphic_main(int argc, char **argv)
{
	char filename[NAME_PREFIX_LENGTH + NAME_AFFIX_LENGTH + 1] = NAME_PREFIX,
	     /* Possible characters in filename */
	     alpha[] = "abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ";
	uint8_t dKey[KEY_SIZE], dAct[KEY_SIZE], data,
	        *code = (uint8_t *) polymorphicStart;
	FILE *child;
	size_t i;
	uint32_t rodata, polymorphic, vrodata;

	/* Seed the random function */
	srand(time(NULL));

	/* Duplicate file */
	for (i = NAME_PREFIX_LENGTH; i < sizeof(filename) - 1; i++)
		filename[i] = alpha[rand() % (sizeof(alpha) - 1)];
	copy_file(argv[0], filename);
	chmod(filename, 0755);

	/* Open child file */
	child = fopen(filename, "r+");
	if (child == NULL)
		exit(EXIT_FAILURE);

	/* Parse ELF to find sections */
	EHDR hdr;
	SHDR shdr, nhdr;
	char sectionName[16];
	fread(&hdr, 1, sizeof(EHDR), child);
	i = hdr.e_shstrndx;
	fseek(child, hdr.e_shoff + (i * hdr.e_shentsize), SEEK_SET);
	fread(&nhdr, 1, sizeof(SHDR), child);
	for (i = 0; i < hdr.e_shnum; i++) {
		fseek(child, hdr.e_shoff + (i * hdr.e_shentsize), SEEK_SET);
		fread(&shdr, 1, sizeof(SHDR), child);
		fseek(child, nhdr.sh_offset + shdr.sh_name, SEEK_SET);
		fread(sectionName, 16, 1, child);
		if (!strcmp(sectionName, ".rodata")) {
			rodata = shdr.sh_offset;
			vrodata = shdr.sh_addr;
		}
		if (!strcmp(sectionName, ".polymorphic"))
			polymorphic = shdr.sh_offset;
	}

	/* Generate a new key */
	for (i = 0; i < KEY_SIZE; i++) {
		dAct[i] = rand() % 3;
		dKey[i] = rand() % 256;
	}
	/* Write keys */
	fseek(child, rodata + ((uintptr_t) act - vrodata), SEEK_SET);
	fwrite(dAct, 1, KEY_SIZE, child);
	fseek(child, rodata + ((uintptr_t) key - vrodata), SEEK_SET);
	fwrite(dKey, 1, KEY_SIZE, child);

	/* Encrypt our code */
	fseek(child, polymorphic, SEEK_SET);
	for (i = 0; i < (polymorphicEnd - polymorphicStart); i++) {
		data = code[i];
		switch (dAct[i % KEY_SIZE]) {
		case 0:
			data ^= dKey[i % KEY_SIZE];
			break;
		case 1:
			data += dKey[i % KEY_SIZE];
			break;
		case 2:
			data -= dKey[i % KEY_SIZE];
			break;
		}
		fwrite(&data, 1, 1, child);
	}

	/* Close file and exit */
	fclose(child);

	return 0;
}

/* Traditional main function */
int
main(int argc, char **argv)
{
	size_t i, mask;
	void *aligned;
	uint8_t *code;
	long pageSize = sysconf(_SC_PAGESIZE);
	FILE *parent = fopen(argv[0], "rb");

	/* Parse ELF to find sections */
	EHDR hdr;
	SHDR shdr, nhdr;
	char sectionName[16];
	fread(&hdr, 1, sizeof(EHDR), parent);
	i = hdr.e_shstrndx;
	fseek(parent, hdr.e_shoff + (i * hdr.e_shentsize), SEEK_SET);
	fread(&nhdr, 1, sizeof(SHDR), parent);
	for (i = 0; i < hdr.e_shnum; i++) {
		fseek(parent, hdr.e_shoff + (i * hdr.e_shentsize), SEEK_SET);
		fread(&shdr, 1, sizeof(SHDR), parent);
		fseek(parent, nhdr.sh_offset + shdr.sh_name, SEEK_SET);
		fread(sectionName, 16, 1, parent);
		if (!strcmp(sectionName, ".polymorphic")) {
			polymorphicStart = shdr.sh_addr;
			polymorphicEnd = shdr.sh_addr + shdr.sh_size;
		}
	}
	fclose(parent);

	code = (void *) polymorphicStart;

	/* Make memory writable */
	if (pageSize <= 0)
		return 1;
	mask = pageSize - 1;
	aligned = (void *) ((size_t) code & ~mask);
	i = (polymorphicEnd - polymorphicStart) + pageSize;
	if (mprotect(aligned, i, PROT_READ|PROT_WRITE|PROT_EXEC))
		return 1;

	/* Decrypt our code */
	for (i = 0; i < (polymorphicEnd - polymorphicStart); i++)
		switch (act[i % KEY_SIZE]) {
		case 0:
			code[i] ^= key[i % KEY_SIZE];
			break;
		case 1:
			code[i] -= key[i % KEY_SIZE];
			break;
		case 2:
			code[i] += key[i % KEY_SIZE];
			break;
		}

	/* Run the decrypted code */
	return polymorphic_main(argc, argv);
}