#include <signal.h>
#include <errno.h>
#include <unistd.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/wait.h>

char datavar = '\xcc';

struct aslr_test {
	void *code;
	void *data;
	void *heap;
	void *ldso;
        void *libc;
	void *stack;

};

typedef struct aslr_test aslrtest_t;
typedef void (*func_t)(void);

enum {
	CODE=1,
	DATA=2,
	HEAP=4,
	LDSO=8,
	LIBC=16,
	STACK=32
};

void *
get_ldbase(char *envp[])
{
	long *p = (long *)envp;
	for (; *p; p += 2);
	for (; *p == 7; p += 2);

	return ((void *)++p);
}

void
write_aslr(char *envp[])
{
	char stackvar;
	aslrtest_t test =
	{
		.code  = write_aslr,
		.data  = &datavar,
		.heap  = malloc(1),
		.ldso  = get_ldbase(envp),
		.libc  = &errno,
		.stack = &stackvar,
	};
	write(1, &test, sizeof(test));

	exit(0);
}

int
get_aslr(char *progname)
{
	aslrtest_t a, b;
	int i, ret=0;
	int filedes[2]; pipe(filedes);

	for (i=0; i<10; i++)
	{
		if (!fork())
		{
			dup2(filedes[1], 1);
			execlp(progname, NULL);
		}

		read(filedes[0], &b, sizeof(b));
		wait(NULL);

		if (i==0)
			a = b;

		ret |= ( a.code != b.code ? CODE : 0 ) |
		       ( a.data != b.data ? DATA : 0 ) |
		       ( a.heap != b.heap ? HEAP : 0 ) |
		       ( a.libc != b.libc ? LIBC : 0 ) |
		       ( a.ldso != b.ldso ? LDSO : 0 ) |
		       ( a.stack != b.stack ? STACK : 0 );
	}

	return (ret);
}

void
depornodep(int sig)
{
    printf( "%s\n", (sig==SIGTRAP) ? "[NO DEP]" : "[DEP]" );

    exit(0);
}

void
try_depornodep(char *itsatrap)
{
	int status=0;
	func_t trap = (func_t)itsatrap;
	*itsatrap = '\xcc';

	if (!fork())
		trap();
	else
		wait(NULL);
}

int
main(int argc, char *argv[], char *envp[])
{
	char *progname = argv[0];
	const char *selfproc_path = "/proc/curproc/file";
	if ( access(selfproc_path, F_OK) == 0 )
		progname = "/proc/curproc/file";

	signal(SIGTRAP, depornodep);
	signal(SIGSEGV, depornodep);

	char stackvar;

	if (argc == 0)
		write_aslr(envp);

	int aslr = get_aslr(progname);

	/* header */
	printf("Test:   %-10s %-10s\n",  "RAND", "NX");

	printf("Code:   %-10s\n", aslr&CODE  ? "[ASLR]":"[NO ASLR]");

	printf("Data:   %-10s",  aslr&DATA  ? "[ASLR]":"[NO ASLR]"); fflush(stdout);
	try_depornodep(&datavar);

	printf("Linker: %-10s\n", aslr&LDSO  ? "[ASLR]":"[NO ASLR]");

	printf("Libc:   %-10s\n", aslr&LIBC  ? "[ASLR]":"[NO ASLR]");

	printf("Heap:   %-10s",  aslr&HEAP  ? "[ASLR]":"[NO ASLR]"); fflush(stdout);
	try_depornodep(malloc(1));

	printf("Stack:  %-10s",  aslr&STACK ? "[ASLR]":"[NO ASLR]"); fflush(stdout);
	try_depornodep(&stackvar);

	exit(EXIT_SUCCESS);
}