#include <unistd.h>
#include <stdlib.h>
#include <stdio.h>
#include <errno.h>
#include <pwd.h>
#include <grp.h>
#include <sys/types.h>
#include <sys/io.h>

#define USER "nobody"
#define GROUP "nogroup"

#define VGA_BASE 0x3C6

#define io_out(value, port) 	outb(value, port)
#define io_in(port)		inb(port)

#define set_mask()		io_out(0xFF, VGA_BASE)
#define set_reg_read(reg)	io_out(reg, VGA_BASE + 1)
#define set_reg_write(reg)	io_out(reg, VGA_BASE + 2)
#define set_value(value)	io_out(value, VGA_BASE + 3)
#define get_value()		io_in(VGA_BASE + 3)

enum {
	OPT_NONE = 0,
	OPT_GRAB,
	OPT_SET
};

extern int errno;

int parse_opt (int ac, const char *av[])
{
	if (ac < 2)
		return OPT_NONE;

	if (!strcmp(av[1], "-g"))
		return OPT_GRAB;
	else if (!strcmp(av[1], "-s"))
		return OPT_SET;

	return OPT_NONE;
}

void drop_perms (void)
{
	ioperm (VGA_BASE, 4, 0);
}

void acquire_perms (void)
{
	if (ioperm(VGA_BASE, 4, 1) == -1) {
		fprintf (stderr, "Can't get raw IO privileges: %m.\n");
		if (errno == EPERM) {
			fprintf (stderr, "Program needs raw access to video card I/O registers, thus it is \n");
			fprintf (stderr, "required to have root permissions or CAP_SYS_RAWIO capabilities.\n");
		}
		exit (EXIT_FAILURE);
	}

	atexit (drop_perms);
}

void drop_privs (void)
{
	struct passwd *p;
	struct group *g;

	if (getuid() != 0)
		return;

	p = getpwnam(USER);
	if (!p) {
		fprintf (stderr, "Warning: Couldn't drop permissions: ");
		if (errno == 0)
			fprintf (stderr, "User %s not found.\n", USER);
		else 
			fprintf (stderr, "%m.\n");
		return;
	}

	g = getgrnam(GROUP);
	if (!g) {
		fprintf (stderr, "Warning: Couldn't set group: ");
		if (errno == 0)
			fprintf (stderr, "Group %s not found.\n", GROUP);
		else 
			fprintf (stderr, "%m.\n");
	} else if (setgid(g->gr_gid) == -1)
		fprintf (stderr, "Warning: Couldn't set group: %m.\n");

	if (setuid(p->pw_uid) == -1)
		fprintf (stderr, "Warning: Couldn't drop permissions: %m.\n");

	return;
}

void dump_io_error (int rs)
{
	if (rs == -1)
		fprintf (stderr, "%m.\n");
	else if (!rs)
		fprintf (stderr, "EOF.\n");
	else
		fprintf (stderr, "%d lt %d.\n", rs, 16 * 3);
}

void do_grab (void)
{
	int i, rs;
	char pal[16 * 3];

	acquire_perms();
	drop_privs();
	set_mask();
	set_reg_read (0);

	for (i = 0; i < 16 * 3; i++)
		pal[i] = get_value();

	rs = write(1, pal, 16 * 3);
	if (rs != 16 * 3) {
		fprintf (stderr, "Write error: ");
		dump_io_error (rs);
		return;
	}
}

void do_set (void)
{
	int i, rs;
	char pal[16 * 3];

	acquire_perms();
	drop_privs();

	rs = read(0, pal, 16 * 3);
	if (rs != 16 * 3) {
		fprintf (stderr, "Read error: ");
		dump_io_error (rs);
		return;
	}

	set_mask();
	set_reg_write (0);

	for (i = 0; i < 16 * 3; i++)
		set_value (pal[i]);
}

int main (int ac, const char *av[])
{
	int opt = parse_opt(ac, av);

	fprintf (stderr, "VGA palette utility v0.1 (C) 2006 Adam Wysocki\n");
	fprintf (stderr, "For newest version surf to http://www.chmurka.net/\n");

	switch (opt) {
		case OPT_NONE:
			fprintf (stderr, "Syntax for grabbing: %s -g > palette.pal\n", av[0]);
			fprintf (stderr, "Syntax for setting:  %s -s < palette.pal\n", av[0]);
			break;

		case OPT_GRAB:
			do_grab();
			break;

		case OPT_SET:
			do_set();
			break;
	}

	return 0;
}