From 1014299d38c9c8b4492d051721ca558a2804d702 Mon Sep 17 00:00:00 2001
From: Tsu-ba-me <ynho.li.aa.e@gmail.com>
Date: Thu, 5 Aug 2021 18:46:42 -0400
Subject: [PATCH] fix(tools): enable anvil-get-server-screenshot to be a job

---
 share/words.xml                   |  4 ++
 tools/anvil-get-server-screenshot | 99 +++++++++++++++++++++++++------
 2 files changed, 84 insertions(+), 19 deletions(-)

diff --git a/share/words.xml b/share/words.xml
index 5e521616..3e2f71ff 100644
--- a/share/words.xml
+++ b/share/words.xml
@@ -2256,6 +2256,10 @@ Are you sure that you want to delete the server: [#!variable!server_name!#]? [Ty
 		<key name="message_0260">Finished [#!variable!operation!#] VNC pipe for server UUID [#!variable!server_uuid!#] from host UUID [#!variable!host_uuid!#].</key>
 		<key name="message_0261">Finished dropping VNC pipes table.</key>
 		<key name="message_0262">Finished managing VNC pipes; no operations happened because requirements not met.</key>
+		<key name="message_0263">Preparing to get server VM screenshot.</key>
+		<key name="message_0264">Finished getting server VM screenshot.</key>
+		<key name="message_0265">Failed to get server VM screenshot; got non-zero return code.</key>
+		<key name="message_0266">Finished attempting to get server VM screenshot; no operations happened because requirements not met.</key>
 		
 		<!-- Success messages shown to the user --> 
 		<key name="ok_0001">Saved the mail server information successfully!</key>
diff --git a/tools/anvil-get-server-screenshot b/tools/anvil-get-server-screenshot
index d7f5e1e9..587edfe2 100755
--- a/tools/anvil-get-server-screenshot
+++ b/tools/anvil-get-server-screenshot
@@ -18,17 +18,6 @@ if (($running_directory =~ /^\./) && ($ENV{PWD}))
 
 my $anvil = Anvil::Tools->new();
 
-$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});
-}
-
 sub system_call
 {
 	my $parameters = shift;
@@ -67,22 +56,94 @@ sub insert_server_screenshot
 	});
 }
 
-my $server_uuid   = $anvil->data->{switches}{'server-uuid'};
-my $is_stdout     = $anvil->data->{switches}{'stdout'};
+$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});
+}
+
+# Try to get a job UUID if not given.
+if (not $anvil->data->{switches}{'job-uuid'})
+{
+	$anvil->data->{switches}{'job-uuid'} = $anvil->Job->get_job_uuid({ program => $THIS_FILE });
+	$anvil->Log->variables({ source => $THIS_FILE, line => __LINE__, level => 2, list => {
+		"switches::job-uuid" => $anvil->data->{switches}{'job-uuid'}
+	} });
+}
+
+# Handle this script as a job when job UUID is provided.
+if ($anvil->data->{switches}{'job-uuid'})
+{
+	$anvil->Job->clear();
+	$anvil->Job->get_job_details();
+	$anvil->Job->update_progress({
+		progress         => 1,
+		job_picked_up_by => $$,
+		job_picked_up_at => time,
+		message          => "message_0263"
+	});
+
+	foreach my $line (split/\n/, $anvil->data->{jobs}{job_data})
+	{
+		if ($line =~ /server-uuid=(.*?)$/)
+		{
+			$anvil->data->{switches}{'server-uuid'} = $1;
+			$anvil->Log->variables({ source => $THIS_FILE, line => __LINE__, level => 2, list => {
+				'switches::server-uuid' => $anvil->data->{switches}{'server-uuid'}
+			} });
+		}
+
+		if ($line =~ /stdout=(.*?)$/)
+		{
+			$anvil->data->{switches}{'stdout'} = $1;
+			$anvil->Log->variables({ source => $THIS_FILE, line => __LINE__, level => 2, list => {
+				'switches::stdout' => $anvil->data->{switches}{'stdout'}
+			} });
+		}
+	}
+}
+
+my $server_uuid = $anvil->data->{switches}{'server-uuid'};
+my $is_stdout   = $anvil->data->{switches}{'stdout'};
+my $job_uuid    = $anvil->data->{switches}{'job-uuid'};
+
+$anvil->Log->variables({ source => $THIS_FILE, line => __LINE__, level => 2, list => {
+	server_uuid => $server_uuid,
+	is_stdout   => $is_stdout,
+	job_uuid    => $job_uuid
+} });
 
 if ($server_uuid)
 {
 	my $encoded_image = get_server_screenshot({ server_uuid => $server_uuid });
 
-	if ($is_stdout)
+	if (defined $encoded_image)
 	{
-		print($encoded_image);
+		if ($is_stdout)
+		{
+			print($encoded_image);
+		}
+		else
+		{
+			insert_server_screenshot({
+				server_uuid   => $server_uuid,
+				encoded_image => $encoded_image
+			});
+		}
+
+		$anvil->Job->update_progress({ progress => 100, message => "message_0264" });
 	}
 	else
 	{
-		insert_server_screenshot({
-			server_uuid   => $server_uuid,
-			encoded_image => $encoded_image
-		});
+		$anvil->Job->update_progress({ progress => 100, message => "message_0265" });
 	}
 }
+else
+{
+	$anvil->Job->update_progress({ progress => 100, message => "message_0266" });
+}