#!/usr/bin/perl
# 
# Manages VNC ports for server VMs that have VNC enabled.
# 

use strict;
use warnings;
use Anvil::Tools;
use Data::Dumper;
use JSON;

$| = 1;

my $THIS_FILE           =  ($0 =~ /^.*\/(.*)$/)[0];
my $running_directory   =  ($0 =~ /^(.*?)\/$THIS_FILE$/)[0];
if (($running_directory =~ /^\./) && ($ENV{PWD}))
{
	$running_directory =~ s/^\./$ENV{PWD}/;
}

my $anvil = Anvil::Tools->new();

$anvil->Log->level({ set => 2 });

sub get_server_info
{
	my $parameters  = shift;
	my $server_uuid = $parameters->{server_uuid};
	my $info_hash   = {};

	my $query = "
SELECT
	ser.server_name,
	hos.host_name,
	hos.host_uuid
FROM
	public.servers AS ser
JOIN
	public.hosts AS hos
ON
	ser.server_host_uuid = hos.host_uuid
WHERE
	server_uuid = ".$anvil->Database->quote($server_uuid)."
;";
	my $results = $anvil->Database->query({ query => $query, source => $THIS_FILE, line => __LINE__ });
	my $count   = @{$results};

	if ($count == 1)
	{
		my $row         = $results->[0];
		my $server_name = $row->[0];
		my $host_name   = $row->[1];
		my $host_uuid   = $row->[2];

		$anvil->Log->variables({ source => $THIS_FILE, line => __LINE__, level => 2, list => {
			server_name => $server_name,
			host_name   => $host_name,
			host_uuid   => $host_uuid
		} });

		$info_hash->{server_name} = $server_name;
		$info_hash->{host_name}   = $host_name;
		$info_hash->{host_uuid}   = $host_uuid;
	}

	return $info_hash;
}

sub get_vnc_info
{
	my $parameters  = shift;
	my $host_name   = $parameters->{host_name};
	my $server_name = $parameters->{server_name};
	my $port_base   = 5900;
	my $vnc_info    = { host_name => $host_name };
	# Requires root to access VM information.
	my $shell_call  = "ssh -n root@".$host_name." \"virsh vncdisplay ".$server_name."\"";

	my ($shell_output, $shell_return_code) = $anvil->System->call({ shell_call => $shell_call });
	$anvil->Log->variables({ source => $THIS_FILE, line => __LINE__, level => 2, list => {
		shell_output      => $shell_output,
		shell_return_code => $shell_return_code
	} });

	if ($shell_return_code == 0)
	{
		my ($port_offset) = $shell_output =~ /:(\d+)$/;
		$vnc_info->{port} = $port_base + int($port_offset);

		$anvil->Log->variables({ source => $THIS_FILE, line => __LINE__, level => 2, list => {
			port_offset => $port_offset,
			vnc_port    => $vnc_info->{port}
		} });
	}

	return $vnc_info;
}

sub is_websockify_process
{
	my $parameters = shift;
	my $host_name  = $parameters->{host_name};
	my $ws_pid     = $parameters->{ws_pid};
	my $shell_call = "ssh -n ".$host_name." \"ps -o comm -h -p ".$ws_pid."\"";

	my ($shell_output, $shell_return_code) = $anvil->System->call({ shell_call => $shell_call });
	$anvil->Log->variables({ source => $THIS_FILE, line => __LINE__, level => 2, list => {
		shell_output      => $shell_output,
		shell_return_code => $shell_return_code
	} });

	return $shell_output eq "websockify" ? 1 : 0;
}

sub is_ssh_process
{
	my $parameters     = shift;
	my $ssh_tunnel_pid = $parameters->{ssh_tunnel_pid};
	my $shell_call     = "ps -o comm -h -p ".$ssh_tunnel_pid;

	my ($shell_output, $shell_return_code) = $anvil->System->call({ shell_call => $shell_call });
	$anvil->Log->variables({ source => $THIS_FILE, line => __LINE__, level => 2, list => {
		shell_output      => $shell_output,
		shell_return_code => $shell_return_code
	} });

	return $shell_output eq "ssh" ? 1 : 0;
}

sub is_websockify_exists
{
	my $parameters      = shift;
	my $server_uuid     = $parameters->{server_uuid};
	my $server_vnc_port = $parameters->{server_vnc_port};

	my $query = "
SELECT
	vnc.server_vnc_port, hos.host_name, vnc.ws_pid
FROM
	public.vnc_pipes AS vnc
JOIN
	public.hosts AS hos
ON
	vnc.ws_host_uuid = hos.host_uuid
WHERE
	server_uuid = ".$anvil->Database->quote($server_uuid)."
;";

	my $results   = $anvil->Database->query({ query => $query, source => $THIS_FILE, line => __LINE__ });
	my $count     = @{$results};
	my $is_exists = 0;

	if ($count > 1)
	{
		my $row                       = $results->[0];
		my $server_vnc_port_in_record = $row->[0];
		my $host_name                 = $row->[1];
		my $ws_pid                    = $row->[2];
		my $clean_up_parameters       = { host_name => $host_name, ws_pid => $ws_pid };

		if ($server_vnc_port != $server_vnc_port_in_record)
		{
			# VNC server port mismatch/oudated; require clean up.

			stop_websockify($clean_up_parameters);

			return $is_exists;
		}

		if (not is_websockify_process($clean_up_parameters))
		{
			# Process died; require clean up.
			return $is_exists;
		}

		# Passed all tests; process considered exists.
		$is_exists = 1;
	}

	return $is_exists;
}

sub start_websockify
{
	my $parameters       = shift;
	my $server_uuid      = $parameters->{server_uuid};
	my $host_name        = $parameters->{host_name};
	my $target_port      = $parameters->{target_port};
	my $ws_info;

	if (not is_websockify_exists({ server_uuid => $server_uuid, server_vnc_port => $target_port }))
	{
		my $source_port_base = 10000;
		my $source_port      = $source_port_base + $target_port;
		my $shell_call       = "ssh -n ".$host_name." 'websockify ".$source_port." :".$target_port." & echo pid:\$!'";

		my ($shell_output, $shell_return_code) = $anvil->System->call({ shell_call => $shell_call });
		$anvil->Log->variables({ source => $THIS_FILE, line => __LINE__, level => 2, list => {
			shell_output      => $shell_output,
			shell_return_code => $shell_return_code
		} });

		if ($shell_return_code == 0)
		{
			my ($ws_pid) = $shell_output =~ /^pid:(\d+)$/;

			$ws_info                = {};
			$ws_info->{pid}         = $ws_pid;
			$ws_info->{source_port} = $source_port;

			$anvil->Log->variables({ source => $THIS_FILE, line => __LINE__, level => 2, list => { ws_pid => $ws_pid } });
		}
	}

	return $ws_info;
}

sub stop_websockify
{
	my $parameters = shift;
	my $host_name  = $parameters->{host_name};
	my $ws_pid     = $parameters->{ws_pid};

	if (is_websockify_process($parameters))
	{
		my $shell_call = "ssh -n ".$host_name." \"kill -9 ".$ws_pid."\"";

		my ($shell_output, $shell_return_code) = $anvil->System->call({ shell_call => $shell_call });
		$anvil->Log->variables({ source => $THIS_FILE, line => __LINE__, level => 2, list => {
			shell_output      => $shell_output,
			shell_return_code => $shell_return_code
		} });
	}
}

sub stop_ssh_tunnel
{
	my $parameters     = shift;
	my $ssh_tunnel_pid = $parameters->{ssh_tunnel_pid};

	if (is_ssh_process($parameters))
	{
		my $shell_call = "kill -9 ".$ssh_tunnel_pid;

		my ($shell_output, $shell_return_code) = $anvil->System->call({ shell_call => $shell_call });
		$anvil->Log->variables({ source => $THIS_FILE, line => __LINE__, level => 2, list => {
			shell_output      => $shell_output,
			shell_return_code => $shell_return_code
		} });
	}
}

sub create_vnc_pipes_table
{
	my $query = "
CREATE TABLE IF NOT EXISTS public.vnc_pipes (
	uuid			uuid				not null	primary key,
	server_uuid		uuid				not null,
	server_vnc_port		numeric				not null,
	ws_host_uuid		uuid				not null,
	ws_pid			numeric				not null,
	ssh_tunnel_host_uuid	uuid				not null,
	ssh_tunnel_pid		numeric				not null,
	modified_date		timestamp with time zone	not null
);";

	$anvil->Database->write({ query => $query, source => $THIS_FILE, line => __LINE__ });
}

sub drop_vnc_pipes_table
{
	my $query   = "DROP TABLE IF EXISTS public.vnc_pipes;";

	$anvil->Database->write({ query => $query, source => $THIS_FILE, line => __LINE__ });
}

sub insert_vnc_pipe
{
	my $parameters           = shift;
	my $server_uuid          = $parameters->{server_uuid};
	my $server_vnc_port      = $parameters->{server_vnc_port};
	my $ws_host_uuid         = $parameters->{ws_host_uuid};
	my $ws_pid               = $parameters->{ws_pid};
	my $ssh_tunnel_host_uuid = $parameters->{ssh_tunnel_host_uuid};
	my $ssh_tunnel_pid       = $parameters->{ssh_tunnel_pid};

	my $query = "
INSERT INTO public.vnc_pipes (
	server_uuid, server_vnc_port, ws_host_uuid, ws_pid, ssh_tunnel_host_uuid, ssh_tunnel_pid
) VALUES (
	".$anvil->Database->quote($server_uuid).",
	".$anvil->Database->quote($ws_host_uuid).",
	".$anvil->Database->quote($ws_pid).",
	".$anvil->Database->quote($ssh_tunnel_host_uuid).",
	".$anvil->Database->quote($ssh_tunnel_pid)."
);";

	$anvil->Database->write({ query => $query, source => $THIS_FILE, line => __LINE__ });
}

sub get_vnc_pipe
{
	my $parameters  = shift;
	my $server_uuid = $parameters->{server_uuid};
	my $host_uuid   = $parameters->{host_uuid};
	my $vnc_pipe_info;

	my $query = "
SELECT
	hos.host_name, vnc.ws_pid, vnc.ssh_tunnel_pid
FROM
	public.vnc_pipes AS vnc
JOIN
	public.hosts AS hos
ON
	vnc.ws_host_uuid = hos.host_uuid
WHERE
	server_uuid = ".$anvil->Database->quote($server_uuid)."
AND
	ssh_tunnel_host_uuid = ".$anvil->Database->quote($host_uuid)."
;";

	my $results = $anvil->Database->query({ query => $query, source => $THIS_FILE, line => __LINE__ });
	my $count   = @{$results};

	if ($count == 1)
	{
		my $row = $results->[0];

		$vnc_pipe_info                   = {};
		$vnc_pipe_info->{host_name}      = $row->[0];
		$vnc_pipe_info->{ws_pid}         = $row->[1];
		$vnc_pipe_info->{ssh_tunnel_pid} = $row->[2];

		$anvil->Log->variables({ source => $THIS_FILE, line => __LINE__, level => 2, list => {
			host_name      => $vnc_pipe_info->{host_name},
			ws_pid         => $vnc_pipe_info->{ws_pid},
			ssh_tunnel_pid => $vnc_pipe_info->{ssh_tunnel_pid}
		} });
	}

	return $vnc_pipe_info;
}

sub delete_vnc_pipe
{
	my $parameters  = shift;
	my $server_uuid = $parameters->{server_uuid};
	my $host_uuid   = $parameters->{host_uuid};

	my $query = "
DELETE FROM public.vnc_pipes
WHERE
	server_uuid = ".$anvil->Database->quote($server_uuid)."
AND
	ssh_tunnel_host_uuid = ".$anvil->Database->quote($host_uuid)."
;";

	$anvil->Database->write({ query => $query, source => $THIS_FILE, line => __LINE__ });
}

$anvil->Get->switches;

$anvil->Database->connect;
$anvil->Log->entry({ source => $THIS_FILE, line => __LINE__, level => 2, secure => 0, key => "log_0132" });
if (not $anvil->data->{sys}{database}{connections})
{
	# No databases, exit.
	$anvil->Log->entry({ source => $THIS_FILE, line => __LINE__, level => 0, 'print' => 1, priority => "err", key => "error_0003" });
	$anvil->nice_exit({ exit_code => 1 });
}

# Read in any CGI variables, if needed.
$anvil->Get->cgi();

$anvil->Database->get_hosts();
$anvil->Database->get_anvils();

print $anvil->Template->get({ file => "shared.html", name => "json_headers", show_name => 0 })."\n";

my $response_body = {};
my $request_body;

if (defined $anvil->data->{cgi}{PUTDATA}{value})
{
	my $is_decode_json_success = eval {
		$request_body = decode_json($anvil->data->{cgi}{PUTDATA}{value});
	};

	if (not $is_decode_json_success)
	{
		$anvil->Log->entry({
			source    => $THIS_FILE,
			line      => __LINE__,
			level     => 0,
			'print'   => 1,
			priority  => "err",
			key       => "error_0304",
			variables => { request_body_string => $anvil->data->{cgi}{PUTDATA}{value}, json_decode_error => $@ }
		});
	}
}

my $server_uuid = exists $request_body->{server_uuid} ? $request_body->{server_uuid} : $anvil->data->{switches}{'server-uuid'};
my $is_open     = exists $request_body->{is_open}     ? $request_body->{is_open}     : $anvil->data->{switches}{'is-open'};

$anvil->Log->variables({ source => $THIS_FILE, line => __LINE__, level => 2, list => {
	server_uuid => $server_uuid,
	is_open     => $is_open
} });

if ($server_uuid)
{
	if ($is_open)
	{
		my $server_info = get_server_info({ server_uuid => $server_uuid });
		my $vnc_info    = get_vnc_info($server_info);

		create_vnc_pipes_table();

		my $ws_info = start_websockify({ server_uuid => $server_uuid, host_name => $server_info->{host_name}, target_port => $vnc_info->{port} });
	}
	else
	{
		my $vnc_pipe_parameters = { server_uuid => $server_uuid, host_uuid => $anvil->Get->host_uuid() };
		my $vnc_pipe_info       = get_vnc_pipe($vnc_pipe_parameters);

		stop_websockify({ host_name => $vnc_pipe_info->{host_name}, ws_pid => $vnc_pipe_info->{ws_pid} });
		stop_ssh_tunnel({ ssh_tunnel_pid => $vnc_pipe_info->{ssh_tunnel_pid} });
		delete_vnc_pipe($vnc_pipe_parameters);
	}
}
elsif ($anvil->data->{switches}{'drop-table'})
{
	drop_vnc_pipes_table();
}

print JSON->new->utf8->encode($response_body)."\n";