#!/usr/bin/perl -XU

# This file is part of the  X2Go Project - https://www.x2go.org
# Copyright (C) 2011-2020 by Oleksandr Shneyder <oleksandr.shneyder@obviously-nice.de>
# Copyright (C) 2012-2020 by Mike Gabriel <mike.gabriel@das-netzwerkteam.de>
#
# X2Go Session Broker is free software; you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation; either version 3 of the License, or
# (at your option) any later version.
#
# X2Go Session Broker is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program; if not, write to the
# Free Software Foundation, Inc.,
# 51 Franklin St, Fifth Floor, Boston, MA 02110-1301, USA.

use strict;

use File::Basename;
use File::Which;
use POSIX;

# are we running via SSH's ForceCommand?
if ($ENV{"SSH_ORIGINAL_COMMAND"} =~ m/.*\/usr\/.*\/x2go\/x2gobroker-agent\ .*/ ) {
	my $ssh_original_command = $ENV{'SSH_ORIGINAL_COMMAND'};
	my $ssh_original_command = (split(/;/, $ssh_original_command))[0];
	$ssh_original_command =~ s/^sh -c (\/usr\/.*\/x2go\/x2gobroker-agent\ .*)/\1/;
	delete $ENV{"SSH_ORIGINAL_COMMAND"};
	exit(system($ssh_original_command));
}

my $username=shift or die;
my $mode=shift or die;

my @available_tasks = (
    "availabletasks",
    "addauthkey",
    "delauthkey",
);

if ( system("x2goversion x2goserver 1>/dev/null") == 0 )
{
	push @available_tasks, (
	    "listsessions",
	    "findbusyservers",
	    "findbusyservers_by_sessionstats",
	    "checkload",
	    "getservers",
	    "suspendsession",
	    "terminatesession",
	);
}

sub InitX2GoUser
{
	my ($user, $uidNumber, $gidNumber, $home)=@_;

	if ( -f "/etc/x2go/x2gosql/sql" )
	{
		# if we use the PostgreSQL session db backend we may have to add the
		# user to the session database...
		open(F,"/etc/x2go/x2gosql/sql");
		my @buf = <F>;
		close(F);
		if ( grep (/^backend=sqlite.*/, @buf) ) {
			#if ( ! -e "$home/.x2go/sqlpass" )
			###
			### FIXME: make the below code robust if homes are on NFS
			###
			###
			#{
			#	open my $save_out, ">&STDOUT";
			#	close (STDOUT);
			#	system "x2godbadmin", "--adduser", $user;
			#	open STDOUT, ">&", $save_out;
			#}
		}
	}
	if (($ENV{"SUDO_USER"}) && ("$ENV{'SUDO_USER'}" ne "$username")) {
		die "You cannot execute x2gobroker-agent for any other user except you!";
	}
}

sub AddAuthKey
{
	my ($uid, $uidNumber, $gidNumber, $home, $pubkey, $authkeyfile)=@_;

	# rewrite %%, %u, %U, %G and %h in authkeyfile string
	$authkeyfile =~ s/%u/$uid/;
	$authkeyfile =~ s/%U/$uidNumber/;
	$authkeyfile =~ s/%G/$gidNumber/;
	$authkeyfile =~ s/%h/$home/;
	$authkeyfile =~ s/%%/%/;

	my $authkeydir = dirname($authkeyfile);
	if ( ! $authkeyfile =~ m/\/.*/ )
	{
		$authkeyfile = "$home/$authkeyfile";
	}

	# make sure dir and file for authorized_keys do exist
	system ("sudo", "-u", "$uid", "--", "mkdir", "-p", "$authkeydir");
	system ("sudo", "-u", "$uid", "--", "touch", "$authkeyfile");
	my $authorized_keys = `sudo -u $uid -- cat "$authkeyfile"`;
	if ( ! ( $authorized_keys =~ m/\Q$pubkey\E$/ ) )
	{
		open my $saveout, ">&STDOUT";
		open STDOUT, '>', "/dev/null";
		open(TEE_STDIN, "| sudo -u $uid -- tee -a $authkeyfile");
		print TEE_STDIN "$pubkey\n";
		close(TEE_STDIN);
		open STDOUT, ">&", $saveout;
	}
}

sub DelAuthKey
{
	my ($uid, $uidNumber, $gidNumber, $home, $pubkey, $authkeyfile)=@_;

	# rewrite %%, %u, %U, %G and %h in authkeyfile string
	$authkeyfile =~ s/%u/$uid/;
	$authkeyfile =~ s/%U/$uidNumber/;
	$authkeyfile =~ s/%G/$gidNumber/;
	$authkeyfile =~ s/%h/$home/;
	$authkeyfile =~ s/%%/%/;

	if ( ! $authkeyfile =~ m/\/.*/ )
	{
		$authkeyfile = "$home/$authkeyfile";
	}
	open my $saveout, ">&STDOUT";
	open STDOUT, '>', "/dev/null";
	open my $saveerr, ">&STDERR";
	open STDERR, '>', "/dev/null";
	system("sudo", "-u", "$uid", "--", "sed", "-e", "s!^$pubkey\$!!", "-e", "/^\$/d", "-i", "$authkeyfile");
	open STDOUT, ">&", $saveout;
	open STDERR, ">&", $saveerr;
}

$< = $>;
delete @ENV{qw(IFS CDPATH ENV BASH_ENV)};
$ENV{'PATH'} = '/bin:/usr/bin';

if($mode eq 'ping')
{
	print "OK\n";
	exit;
}

if ( $mode eq 'checkload' ) {
	print "OK\n";

	# Read the values from /proc/loadavg and combine all three values in a linear
	# way and make a percent value out of the result:
	open FILE, "< /proc/loadavg" or die return ("Cannot open /proc/loadavg: $!");
	my ($avg1, $avg5, $avg15, undef, undef) = split / /, <FILE>;
	close FILE;
	my $loadavgXX = ( $avg1 + $avg5 + $avg15 ) * 100/3 + 1;
	if ( $loadavgXX lt 1 ) {
		# if we get here, then our avgXX values were negative,
		# we should never get here...
		$loadavgXX = 1;
	}

	# calculate total memory vs. free memory
	my $memTotal;
	my $memFree;
	my $memAvail;
	my $pagesActiveFile;
	my $pagesInactiveFile;
	my $slabReclaimable;
	open FILE, "< /proc/meminfo" or die return ("Cannot open /proc/meminfo: $!");
	foreach(<FILE>) {
		if ( m/^MemTotal:\s+(\S+)/ ) {
			$memTotal = $1;
		}
		if ( m/^MemFree:\s+(\S+)/ ) {
			$memFree = $1;
		}
		if ( m/^MemAvailable:\s+(\S+)/ ) {
			$memAvail = $1/1000;
		}
		if ( m/^Active\(file\):\s+(\S+)/ ) {
			$pagesActiveFile = $1;
		}
		if ( m/^Inactive\(file\):\s+(\S+)/ ) {
			$pagesInactiveFile = $1;
		}
		if ( m/^SReclaimable:\s+(\S+)/ ) {
			$slabReclaimable = $1;
		}
	}
	close(FILE);

	my $myMemAvail = 0;
	if ($myMemAvail == 0) {
                # taken from Linux kernel code in fs/proc/meminfo.c (since kernel version 3.14)...
		open FILE, "< /proc/sys/vm/min_free_kbytes" or die return ("Cannot open /proc/sys/vm/min_free_kbytes: $!");
		my $memLowWatermark = <FILE>;
		$memLowWatermark =~ s/\n//g;
		close(FILE);
		$myMemAvail += $memFree - $memLowWatermark;
		my $pagecache = $pagesActiveFile + $pagesInactiveFile;
		$pagecache -= ($pagecache/2, $memLowWatermark) [$pagecache/2 > $memLowWatermark];
		$myMemAvail += $pagecache;
		$myMemAvail += $slabReclaimable - ( ($slabReclaimable/2, $memLowWatermark) [$slabReclaimable/2 > $memLowWatermark]);
		if ($myMemAvail < 0) {
			$myMemAvail = 0;
		}
		$myMemAvail = $myMemAvail / 1000;
	}

	# check number and type of available CPU cores
	my $numCPU = 0;
	my $typeCPU;
	open FILE, "< /proc/cpuinfo" or die return ("Cannot open /proc/cpuinfo: $!");
	foreach(<FILE>) {
		if ( m/model name.*CPU.*@\ ([\d\.]+)/ ) {
			$typeCPU = $1*1000;
			$numCPU += 1;
		}
	}
	close(FILE);

	# in virtual hosts, we need to obtain the CPU frequency from the cpu MHz: field.
	if ($typeCPU == 0) {
		open FILE, "< /proc/cpuinfo" or die return ("Cannot open /proc/cpuinfo: $!");
		foreach(<FILE>) {
			if ( m/^cpu\ MHz\s+:\s+(\S+)/ ) {
				$typeCPU = ceil($1);
				$numCPU += 1;
			}
		}
		close(FILE);
	}

	print sprintf 'loadavgXX:%1$d', $loadavgXX;
	print "\n";
	print sprintf 'memAvail:%1$d', $memAvail;
	print "\n";
	print sprintf 'myMemAvail:%1$d', $myMemAvail;
	print "\n";
	print sprintf 'numCPU:%1$d', $numCPU;
	print "\n";
	print sprintf 'typeCPU:%1$d', $typeCPU;
	print "\n";
	exit;
}

if($mode eq 'availabletasks')
{
	print "OK\n";
	my $available_task;
	foreach $available_task (@available_tasks) {
		print "$available_task\n";
	}
	exit;
}

my  ($uid, $passwd, $uidNumber, $gidNumber, $quota, $comment, $gcos, $home, $shell, $expire) = getpwnam($username);

if(!defined $uidNumber)
{
	die "no such user on system: $username";
}
elsif($uidNumber < 1000)
{
	die "operation on system user: $username (with UID number: $uidNumber)";
}

if($mode eq 'listsessions' && which('x2golistsessions'))
{
	InitX2GoUser($uid, $uidNumber, $gidNumber, $home);
	print "OK\n";
	system("sudo", "-u", "$uid", "--", "x2golistsessions", "--all-servers");
}

if( (($mode eq 'findbusyservers_by_sessionstats') || ($mode eq 'findbusyservers')) && which('x2gogetservers') )
{

	# Normally the session broker setup knows about all servers,
	# make sure your configuration of X2Go Session Broker is correct and
	# lists all available servers.

	# The findbusyservers algorithm only returns servers that are currently
	# in use (i.e. have running or suspended sessions on them). So the
	# result may be empty or contain a server list not containing all
	# available servers.

	# The logic of findbusyservers is this, then:
	#   1. if no server is returned, any of the configured servers is best server
	#   2. if some servers are returned, a best server is one that is not returned
	#   3. if all configured servers are returned, than evaluate the usage value
	#      (e.g. 90:server1, 20:server2, 10:server3 -> best server is server3)

	# The above interpretation has to be handled by the broker implementation
	# calling »x2gobroker-agent findbusyservers«.

	InitX2GoUser($uid, $uidNumber, $gidNumber, $home);
	print "OK\n";
	my $busy_servers = `sudo -u $uid -- x2gogetservers`;

	my %server_usage = ();
	my $num_sessions = 0;
	foreach (split('\n', $busy_servers))
	{
		my ($hostname, $num_users) = split(' ', $_);
		$server_usage{$hostname} = $num_users;
		$num_sessions += $num_users;
	}

	# render the output result
	my @result;
	for my $hostname ( keys %server_usage ) {
		my $available = $server_usage{$hostname}/$num_sessions*100;
		push @result, sprintf '%1$d:%2$s', $available, $hostname;
	}
	print join("\n", sort @result);
	print "\n";
}

if( $mode eq 'getservers' && which('x2gogetservers') )
{
	InitX2GoUser($uid, $uidNumber, $gidNumber, $home);
	print "OK\n";
	exec ("sudo", "-u", "$uid", "--", "x2gogetservers");
}

if($mode eq 'addauthkey')
{
	my $pubkey = shift or die;
	my $authkeyfile = shift or die;
	InitX2GoUser($uid, $uidNumber, $gidNumber, $home);
	print "OK\n";
	AddAuthKey($uid, $uidNumber, $gidNumber, $home, $pubkey, $authkeyfile);
}

if($mode eq 'delauthkey')
{
	my $pubkey = shift or die;
	my $authkeyfile = shift or die;
	InitX2GoUser($uid, $uidNumber, $gidNumber, $home);
	print "OK\n";
	DelAuthKey($uid, $uidNumber, $gidNumber, $home, $pubkey, $authkeyfile);
}

# use x2gosuspend-session to test if we can really do this...
if( $mode eq 'suspendsession' && which('x2gosuspend-session') )
{
	InitX2GoUser($uid, $uidNumber, $gidNumber, $home);
	print "OK\n";
	my $sid=shift;
	my $x2go_lib_path=`x2gopath lib`;
	exec ("sudo", "-u", "$uid", "--", "$x2go_lib_path/x2gochangestatus", "S",  "$sid");
}

# use x2goterminate-session to test if we can really do this...
if( $mode eq 'terminatesession' && which('x2goterminate-session') )
{
	InitX2GoUser($uid, $uidNumber, $gidNumber, $home);
	print "OK\n";
	my $sid=shift;
	my $x2go_lib_path=`x2gopath lib`;
	exec ("sudo", "-u", "$uid", "--", "$x2go_lib_path/x2gochangestatus", "T",  "$sid");
}
