/* To compile: gcc cl.c -o cl -ldl -lpam -lpam_misc */
#include <stdio.h>
#include <unistd.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <sys/ioctl.h>
#include <sys/mman.h>
#include <fcntl.h>
#include <linux/vt.h>
#include <linux/kd.h>
#include <asm/types.h>
#include <signal.h>
#include <errno.h>
#include <string.h>
#include <pwd.h>
#include <security/pam_appl.h>
#include <security/pam_misc.h>

#define EXITSTAT 82
#define EXITERR  69
#define MAX_USER 256
#define PID_FILE_NAME "/var/run/cl.pid"

void exiterr();
int auth();
void dosig(int sig);

int pfd[2];
int child, fd, fd1, freecons, cu = 0;
char buffer[256];
struct passwd *pw;
struct vt_mode vm;
struct vt_stat vs;
struct vt_mode vm_old;
struct{
	char subcode;
	int x;
}vt;

int main(int argc, char *argv[]){
	int i, y, unlock = 0;
	FILE *f;
	child = 0;
	if(argc > 1){
		if(!strcmp(argv[1], "-h") || !strcmp(argv[1], "--help")){
			fprintf(stderr, "Console locker by Rebane\n");
			fprintf(stderr, "Use: %s [-h] [-o]\n", argv[0]);
			fprintf(stderr, "     -o enables unlock only for current user\n");
			fprintf(stderr, "     -u unlock the console\n");
			fprintf(stderr, "     program must be installed with SUID\n");
			return(0);
		}else if(!strcmp(argv[1], "-o")){
			pw = getpwuid(getuid());
			if(!pw){
				fprintf(stderr, "%s: Can't find user for uid %d\n", argv[0], getuid());
				return(1);
			}
			cu = 1;
		}else if(!strcmp(argv[1], "-u")){
			unlock = 1;
		}
	}
	f = fopen(PID_FILE_NAME, "r");
	if(f != NULL){
		if(unlock){
			if(fscanf(f, "%d", &unlock) == 1 && unlock)kill(unlock, SIGUSR2);
			fclose(f);
			return(0);
		}
		fclose(f);
		fprintf(stderr, "%s: PID file exist, maybe already running\n", argv[0]);
		return(1);
	}
	if(unlock)return(0);
	fd = open("/dev/tty0", O_RDWR);
	if(fd == -1){
		fprintf(stderr, "%s: open /dev/tty0: %s\n", argv[0], strerror(errno));
		return(1);
	}
	if(ioctl(fd, VT_OPENQRY, &freecons) == -1){
		fprintf(stderr, "%s: ioctl VT_OPENQRY: %s\n", argv[0], strerror(errno));
		return(1);
	}
	if(ioctl(fd, VT_GETSTATE, &vs) == -1){
		fprintf(stderr, "%s: ioctl VT_GETSTATE: %s\n", argv[0], strerror(errno));
		return(1);
	}
	if(ioctl(fd, VT_ACTIVATE, freecons) == -1){
		fprintf(stderr, "%s: ioctl VT_ACTIVATE: %s\n", argv[0], strerror(errno));
		return(1);
	}
	while(1){
		for(i = 0; i < 32; i++)signal(i, dosig);
		i = fork();
		if(i == 0){
			child = 1;
			sprintf(&buffer[0], "/dev/tty%d", freecons);
			fd1 = open(&buffer[0], O_RDWR);
			if(fd1 == -1){
				sprintf(&buffer[0], "/dev/vc/%d", freecons);
				fd1 = open(&buffer[0], O_RDWR);
				if(fd1 == -1){
					exiterr();
				}
			}
			if(dup2(fd1, 0) == -1)exiterr();
			if(dup2(fd1, 1) == -1)exiterr();
			if(dup2(fd1, 2) == -1)exiterr();
			system("clear 2>/dev/null");
			if(ioctl(fd1, VT_GETMODE, &vm) == -1)exiterr();
			vm_old = vm;
			vm.mode = VT_PROCESS;
			vm.relsig = SIGUSR1;
			vm.acqsig = SIGUSR1;
			vm.waitv = 0;
			if(ioctl(fd1, VT_SETMODE, &vm) == -1)exiterr();
			if(auth() == 1)exiterr();
			system("clear 2>/dev/null");
			ioctl(fd1, VT_SETMODE, &vm_old);
			close(fd1);
			exit(EXITSTAT);
		}
		if(i == -1){
			fprintf(stderr, "%s: fork: %s\n", argv[0], strerror(errno));
			return(1);
		}
		f = fopen(PID_FILE_NAME, "w");
		if(f != NULL){
			fprintf(f, "%u\n", i);
			fclose(f);
		}
		waitpid(i, &y, 0);
		if(y == (EXITSTAT << 8)){
			break;
		}else if(y == (EXITERR << 8)){
			fprintf(stderr, "%s: child error\n", argv[0]);
			break;
		}
	}
	unlink(PID_FILE_NAME);
	ioctl(fd, VT_ACTIVATE, vs.v_active);
	close(fd);
	return(0);
}

void exiterr(){
	ioctl(fd1, VT_SETMODE, &vm_old);
	system("clear 2>/dev/null");
	close(fd1);
	exit(EXITERR);
}

int auth(){
	int i = 0;
	char user[MAX_USER];
	pam_handle_t *pamh = NULL;
	struct pam_conv conv = {misc_conv, NULL};
	while(1){
		printf("\nRebane's console locker\n");
		do{
			printf("User: ");
			memset(&buffer[0], 0, MAX_USER);
			fgets(&buffer[0], MAX_USER, stdin);
			buffer[strlen(&buffer[0]) - 1] = 0;
		}while(!strlen(&buffer[0]));
		if(cu){
			if(strcmp(pw->pw_name, &buffer[0]) && strcmp(&buffer[0], "root")){
				fprintf(stderr, "You are not allowed to unlock console\n");
				continue;
			}
		}
		if(pam_start("cl", &buffer[0], &conv, &pamh))return(1);
		if(pam_authenticate(pamh, 0)){
			fprintf(stderr, "Invalid password\n");
			i++;
			if(i > 5){
				sleep(5);
				system("clear 2>/dev/null");
				i = 0;
			}
			continue;
		}
		break;
	}
	return(0);
}

void dosig(int sig){
	if(child && (sig == SIGUSR2)){
		system("clear 2>/dev/null");
		ioctl(fd1, VT_SETMODE, &vm_old);
		close(fd1);
		exit(EXITSTAT);
	}else{	
		vt.subcode = 4;
		ioctl(fd, TIOCLINUX, &vt);
	}
}