#!/usr/bin/perl -w

$num_args = $#ARGV + 1;
if (($num_args == 0) || in_array("-h",\@ARGV) || in_array("--help",\@ARGV)){
    usage();
}

# Get the members of the group "nitrcceadmin"
@group_info = getgrnam("nitrcceadmin");
if ($#group_info < 0){
    print STDERR "ERROR: nitrcceadmin group does not exist\n";
    exit;
}
@members = split(/ /,$group_info[3]);

# Get the current user name
$user = getpwuid($<);

# If the current user is not in the group "nitrcceadmin" throw an errir
if (!in_array($user,\@members)){
    print STDERR "ERROR: You are not a member of nitrcceadmin\n";
    exit;
}

# Become root:
delete @ENV{qw(IFS CDPATH ENV BASH_ENV)}; # Make %ENV safer
$ENV{"PATH"} = "/usr/bin:/usr/local/bin:/bin"; # Make ENV PATH known
$<=$>=0;
$(=$)=0;

if ($ARGV[0] eq "has_vnc_passwd"){ # Command has_vncpasswd
    if ($num_args != 2){ usage(); }
    $username = $ARGV[1]; # Clean username
    # Look for vnc passwd file for the user
    print has_vnc_password($username) . "\n";
    exit();
} elsif ($ARGV[0] eq "set_vnc_passwd"){ # Command has_vncpasswd
    if ($num_args != 2){ usage(); }
    # Get username and untaint it
    $username = untaint($ARGV[1]);
    # Get the passwd from STDIN
    $passwd = <STDIN>;
    $passwd =~ /^(.*)$/; # Mark it untainted, though we want to leaves spaces and such
    if (length($passwd) < 7){
	die ("ERROR: The password must have at least 6 characters\n");
    }
    # Execute vncpasswd
    $cmd = "sudo -u ".$username. " -n -H sh -c \"vncpasswd\" 1>/dev/null";
    open (CMD,"|$cmd");
    print CMD $passwd;
    print CMD $passwd;
    close CMD;
    exit();
} elsif ($ARGV[0] eq "has_session"){ # Command has_session
    if ($num_args != 2){ usage(); }
    # Get username and untaint it
    $username = untaint($ARGV[1]);
    # Now check for display
    $display = get_session_display($username);
    print $display."\n";
    exit();
} elsif ($ARGV[0] eq "start_session"){ # Command start_session
    if ($num_args < 2){ usage(); }
    # Get username and untaint it
    $username = untaint($ARGV[1]);
    # Check that vnc password exists
    if (!has_vnc_password($username)) { 
	die("ERROR: The user has not set a VNC password\n"); 
    }
    # Collect any command-line options for vncserver
    $vnc_options = "";
    if ($#ARGV >= 2){
	$vnc_options = join(" ",@ARGV[2..$#ARGV]);
	$vnc_options =~ /^(.*)$/; # Mark as untainted, through we wish to leave spaces
	$vnc_options = $1;
    }
    # Check for an existing session
    $display = get_session_display($username);
    if ($display > 0){
	# Return the display number and exit
	print $display."\n";
	exit();
    }
    # There isn't a current session so create one
    # Run vncserver and get the display number
    $vnc_passwd_path = "/home/".$username."/.vnc/passwd";
    $cmd = "sudo -u ".$username. " -n -H sh -c \"vncserver ".$vnc_options." 2>&1 </dev/null\"";
    $output = `$cmd`;
    if ($output =~ /desktop is [^:]+:(\d+)/s){
	$display = $1;
    } else {
	die("ERROR: Unable to get Display number\n");
    }
    # Record lock file
    $lock_file = "/var/lib/nitrcce/sessions/$username.lck";
    open (LCK, ">$lock_file") or die ("ERROR: Unable to write lock file\n");
    print LCK $display."\n";
    close LCK;
    # Double check that we ran by getting the display number
    $display = get_session_display($username);
    print $display."\n";
    exit();
} elsif ($ARGV[0] eq "end_session"){ # Command end_session
    if ($num_args != 2){ usage(); }
    # Get username and untaint it
    $username = untaint($ARGV[1]);
    # Get the current running display
    $display = untaint(get_session_display($username));
    if ($display > 0){
	# If it is running, kill it
	$output = `sudo -u $username -n -H sh -c "vncserver -kill :$display 2>&1 </dev/null"`;
    }
    # Remove the lock file if it exists
    $lock_file = "/var/lib/nitrcce/sessions/$username.lck";
    if (-e $lock_file){
	$output = `rm $lock_file`;
    }
} else {  # Unknown command
    usage();
}

#
# Look for a lock file, grab the display, and check for running process
# Return 0 if the process isn't running, else return the display number
#
sub get_session_display{
    my ($username) = @_;
    $lock_file = "/var/lib/nitrcce/sessions/$username.lck";
    if (-e $lock_file){ # Lock file exists
	open (LCK, "<$lock_file") or return(0);
	$display = <LCK>;
	close (LCK);
	chop($display);
	if (!$display){ return(0); }
	$pid = untaint(get_vnc_pid($username,$display));
	if (!$pid){ return(0); }
	$process_info = `ps $pid | grep vnc`;
	if (!$process_info){ return(0); }
	return ($display);
    } else { # Lock file does not exist
	return (0);
    }
}

#
# untaint - Return anything to the first white space
sub untaint{
    my($a) = @_;
    $a =~ /^(\S*)/;
    return($1);
}

#
# Get the pid for a running display
#
sub get_vnc_pid{
    my ($username,$display) = @_;
    # Untaint variables
    $username = untaint($username);
    $display = untaint($display);
    $pid = `cat /home/$username/.vnc/*:$display.pid 2>/dev/null`;
    chop($pid);
    return($pid);
}

#
# Check if VNC password exists
#
sub has_vnc_password {
    my ($username) = @_;
    # Look for vnc passwd file for the user
    $vnc_passwd_path = "/home/".$username."/.vnc/passwd";
    if (-e $vnc_passwd_path){
	return(1);
    } else {
	return(0);
    }
}

#
# Prints the usage
#
sub usage {
    print "\n";
    print " Usage: nitrcce-admin [command] [options]\n";
    print "\n";
    print " nitrcce-admin has_vnc_passwd username \n";
    print "    - Prints 1 if the user has a vnc password and 0 otherwise\n";
    print "\n";
    print " nitrcce-admin set_vnc_passwd username \n";
    print "    - Reads one line from STDIN and sets the vnc password for the user\n";
    print "    - The password must be at least 6 characters long\n";
    print "\n";
    print " nitrcce-admin has_session username \n";
    print "    - Returns the display number of the session (and 0 otherwise)\n";
    print "\n";
    print " nitrcce-admin start_session username [vncserver_options... ]\n";
    print "    - Starts a session and returns the display number of the session\n";
    print "    - If a session exists, it simply returns the display number\n";
    print "    - Any vnc server options are passed to the vncserver command\n";
    print "      (if starting a new session)\n";
    print "\n";
    print " nitrcce-admin end_session username \n";
    print "    - Ends a user session\n";
    exit;
}

#
# in_array($needle,\@arr) returns 1 if $needle found in \@arr and 0 otherwise
#
sub in_array
{
    my ($needle,$arr) = @_;
    foreach (@$arr)
    {
	return 1 if $_ eq $needle;
    }
    return 0;
}
