From b2d8da489c92be8d5560165ccd4f63c0ec4baa68 Mon Sep 17 00:00:00 2001
From: Jason Volk <jason@zemos.net>
Date: Sun, 11 Aug 2024 23:07:58 +0000
Subject: [PATCH] add configurable automatic admin command execution after
 startup

Signed-off-by: Jason Volk <jason@zemos.net>
---
 src/core/config/mod.rs       |  3 +++
 src/main/clap.rs             |  9 ++++++++
 src/service/admin/console.rs |  6 ++++++
 src/service/admin/mod.rs     | 40 ++++++++++++++++++++++++++++++++++--
 4 files changed, 56 insertions(+), 2 deletions(-)

diff --git a/src/core/config/mod.rs b/src/core/config/mod.rs
index bec7c4ea..edca607e 100644
--- a/src/core/config/mod.rs
+++ b/src/core/config/mod.rs
@@ -334,6 +334,8 @@ pub struct Config {
 	pub admin_escape_commands: bool,
 	#[serde(default)]
 	pub admin_console_automatic: bool,
+	#[serde(default)]
+	pub admin_execute: Vec<String>,
 
 	#[serde(default)]
 	pub sentry: bool,
@@ -592,6 +594,7 @@ impl fmt::Display for Config {
 			"Activate admin console after startup",
 			&self.admin_console_automatic.to_string(),
 		);
+		line("Execute admin commands after startup", &self.admin_execute.join(", "));
 		line("Allow outgoing federated typing", &self.allow_outgoing_typing.to_string());
 		line("Allow incoming federated typing", &self.allow_incoming_typing.to_string());
 		line(
diff --git a/src/main/clap.rs b/src/main/clap.rs
index 92bd73c1..56d614a4 100644
--- a/src/main/clap.rs
+++ b/src/main/clap.rs
@@ -24,6 +24,10 @@ pub(crate) struct Args {
 	/// Activate admin command console automatically after startup.
 	#[arg(long, num_args(0))]
 	pub(crate) console: bool,
+
+	/// Execute console command automatically after startup.
+	#[arg(long)]
+	pub(crate) execute: Vec<String>,
 }
 
 /// Parse commandline arguments into structured data
@@ -39,6 +43,11 @@ pub(crate) fn update(mut config: Figment, args: &Args) -> Result<Figment> {
 		config = config.join(("admin_console_automatic", true));
 	}
 
+	// Execute commands after any commands listed in configuration file
+	for command in &args.execute {
+		config = config.adjoin(("admin_execute", [command]));
+	}
+
 	// All other individual overrides can go last in case we have options which
 	// set multiple conf items at once and the user still needs granular overrides.
 	for option in &args.option {
diff --git a/src/service/admin/console.rs b/src/service/admin/console.rs
index c9a288d9..7f3a6fae 100644
--- a/src/service/admin/console.rs
+++ b/src/service/admin/console.rs
@@ -194,6 +194,12 @@ impl Console {
 	}
 }
 
+/// Standalone/static markdown printer.
+pub fn print(markdown: &str) {
+	let output = configure_output(MadSkin::default_dark());
+	output.print_text(markdown);
+}
+
 fn configure_output(mut output: MadSkin) -> MadSkin {
 	use termimad::{crossterm::style::Color, Alignment, CompoundStyle, LineStyle};
 
diff --git a/src/service/admin/mod.rs b/src/service/admin/mod.rs
index 962d31dd..2719c142 100644
--- a/src/service/admin/mod.rs
+++ b/src/service/admin/mod.rs
@@ -9,7 +9,7 @@ use std::{
 };
 
 use async_trait::async_trait;
-use conduit::{debug, error, error::default_log, pdu::PduBuilder, Err, Error, PduEvent, Result, Server};
+use conduit::{debug, error, error::default_log, info, pdu::PduBuilder, Err, Error, PduEvent, Result, Server};
 pub use create::create_admin_room;
 use loole::{Receiver, Sender};
 use ruma::{
@@ -20,7 +20,10 @@ use ruma::{
 	OwnedEventId, OwnedRoomId, RoomId, UserId,
 };
 use serde_json::value::to_raw_value;
-use tokio::sync::{Mutex, RwLock};
+use tokio::{
+	sync::{Mutex, RwLock},
+	time::{sleep, Duration},
+};
 
 use crate::{globals, rooms, rooms::state::RoomMutexGuard, Dep};
 
@@ -84,6 +87,8 @@ impl crate::Service for Service {
 	async fn worker(self: Arc<Self>) -> Result<()> {
 		let receiver = self.receiver.lock().await;
 		let mut signals = self.services.server.signal.subscribe();
+
+		self.startup_execute().await;
 		self.console_auto_start().await;
 
 		loop {
@@ -356,6 +361,37 @@ impl Service {
 		}
 	}
 
+	/// Execute admin commands after startup
+	async fn startup_execute(&self) {
+		sleep(Duration::from_millis(500)).await; //TODO: remove this after run-states are broadcast
+		for (i, command) in self.services.server.config.admin_execute.iter().enumerate() {
+			self.startup_execute_command(i, command.clone()).await;
+			tokio::task::yield_now().await;
+		}
+	}
+
+	/// Execute one admin command after startup
+	async fn startup_execute_command(&self, i: usize, command: String) {
+		debug!("Startup command #{i}: executing {command:?}");
+
+		match self.command_in_place(command, None).await {
+			Err(e) => error!("Startup command #{i} failed: {e:?}"),
+			Ok(None) => info!("Startup command #{i} completed (no output)."),
+			Ok(Some(output)) => Self::startup_command_output(i, &output),
+		}
+	}
+
+	#[cfg(feature = "console")]
+	fn startup_command_output(i: usize, content: &RoomMessageEventContent) {
+		info!("Startup command #{i} completed:");
+		console::print(content.body());
+	}
+
+	#[cfg(not(feature = "console"))]
+	fn startup_command_output(i: usize, content: &RoomMessageEventContent) {
+		info!("Startup command #{i} completed:\n{:#?}", content.body());
+	}
+
 	/// Possibly spawn the terminal console at startup if configured.
 	async fn console_auto_start(&self) {
 		#[cfg(feature = "console")]