continuwuity/src/admin/space/roles.rs

use std::fmt::Write;

use clap::Subcommand;
use conduwuit::{Err, Event, Result, matrix::pdu::PduBuilder};
use conduwuit_core::matrix::space_roles::{
	RoleDefinition, SPACE_ROLE_MEMBER_EVENT_TYPE, SPACE_ROLE_ROOM_EVENT_TYPE,
	SPACE_ROLES_EVENT_TYPE, SpaceRoleMemberEventContent, SpaceRoleRoomEventContent,
	SpaceRolesEventContent,
};
use futures::StreamExt;
use ruma::{OwnedRoomId, OwnedRoomOrAliasId, OwnedUserId, events::StateEventType};
use serde_json::value::to_raw_value;

use crate::{admin_command, admin_command_dispatch};

macro_rules! require_enabled {
	($self:expr) => {
		if !$self.services.rooms.roles.is_enabled() {
			return $self
				.write_str(
					"Space permission cascading is disabled. Enable it with \
					 `space_permission_cascading = true` in your config.",
				)
				.await;
		}
	};
}

macro_rules! resolve_space {
	($self:expr, $space:expr) => {{
		require_enabled!($self);
		let space_id = $self.services.rooms.alias.resolve(&$space).await?;
		if !matches!(
			$self
				.services
				.rooms
				.state_accessor
				.get_room_type(&space_id)
				.await,
			Ok(ruma::room::RoomType::Space)
		) {
			return Err!("The specified room is not a Space.");
		}
		space_id
	}};
}

macro_rules! custom_state_pdu {
	($event_type:expr, $state_key:expr, $content:expr) => {
		PduBuilder {
			event_type: $event_type.to_owned().into(),
			content: to_raw_value($content).map_err(|e| {
				conduwuit::Error::Err(
					format!("Failed to serialize custom state event content: {e}").into(),
				)
			})?,
			state_key: Some($state_key.to_owned().into()),
			..PduBuilder::default()
		}
	};
}

#[admin_command_dispatch]
#[derive(Debug, Subcommand)]
pub enum SpaceRolesCommand {
	/// List all roles defined in a space
	List {
		space: OwnedRoomOrAliasId,
	},
	/// Add a new role to a space
	Add {
		space: OwnedRoomOrAliasId,
		role_name: String,
		#[arg(long)]
		description: Option<String>,
		#[arg(long)]
		power_level: Option<i64>,
	},
	/// Remove a role from a space
	Remove {
		space: OwnedRoomOrAliasId,
		role_name: String,
	},
	/// Assign a role to a user
	Assign {
		space: OwnedRoomOrAliasId,
		user_id: OwnedUserId,
		role_name: String,
	},
	/// Revoke a role from a user
	Revoke {
		space: OwnedRoomOrAliasId,
		user_id: OwnedUserId,
		role_name: String,
	},
	/// Require a role for a room
	Require {
		space: OwnedRoomOrAliasId,
		room_id: OwnedRoomId,
		role_name: String,
	},
	/// Remove a role requirement from a room
	Unrequire {
		space: OwnedRoomOrAliasId,
		room_id: OwnedRoomId,
		role_name: String,
	},
	/// Show a user's roles in a space
	User {
		space: OwnedRoomOrAliasId,
		user_id: OwnedUserId,
	},
	/// Show a room's role requirements in a space
	Room {
		space: OwnedRoomOrAliasId,
		room_id: OwnedRoomId,
	},
}

#[admin_command]
async fn list(&self, space: OwnedRoomOrAliasId) -> Result {
	let space_id = resolve_space!(self, space);
	let roles_event_type = StateEventType::from(SPACE_ROLES_EVENT_TYPE.to_owned());

	let content: SpaceRolesEventContent = self
		.services
		.rooms
		.state_accessor
		.room_state_get_content(&space_id, &roles_event_type, "")
		.await
		.unwrap_or_default();

	if content.roles.is_empty() {
		return self.write_str("No roles defined in this space.").await;
	}

	let mut msg = format!("Roles in {space_id}:\n```\n");
	for (name, def) in &content.roles {
		let pl = def
			.power_level
			.map(|p| format!(" (power_level: {p})"))
			.unwrap_or_default();
		let _ = writeln!(msg, "- {name}: {}{pl}", def.description);
	}
	msg.push_str("```");

	self.write_str(&msg).await
}

#[admin_command]
async fn add(
	&self,
	space: OwnedRoomOrAliasId,
	role_name: String,
	description: Option<String>,
	power_level: Option<i64>,
) -> Result {
	let space_id = resolve_space!(self, space);
	let roles_event_type = StateEventType::from(SPACE_ROLES_EVENT_TYPE.to_owned());

	let mut content: SpaceRolesEventContent = self
		.services
		.rooms
		.state_accessor
		.room_state_get_content(&space_id, &roles_event_type, "")
		.await
		.unwrap_or_default();

	if content.roles.contains_key(&role_name) {
		return Err!("Role '{role_name}' already exists in this space.");
	}

	content.roles.insert(role_name.clone(), RoleDefinition {
		description: description.unwrap_or_else(|| role_name.clone()),
		power_level,
	});

	let state_lock = self.services.rooms.state.mutex.lock(&space_id).await;
	let server_user = &self.services.globals.server_user;

	self.services
		.rooms
		.timeline
		.build_and_append_pdu(
			custom_state_pdu!(SPACE_ROLES_EVENT_TYPE, "", &content),
			server_user,
			Some(&space_id),
			&state_lock,
		)
		.await?;

	self.write_str(&format!("Added role '{role_name}' to space {space_id}."))
		.await
}

#[admin_command]
async fn remove(&self, space: OwnedRoomOrAliasId, role_name: String) -> Result {
	let space_id = resolve_space!(self, space);
	let roles_event_type = StateEventType::from(SPACE_ROLES_EVENT_TYPE.to_owned());

	let mut content: SpaceRolesEventContent = self
		.services
		.rooms
		.state_accessor
		.room_state_get_content(&space_id, &roles_event_type, "")
		.await
		.unwrap_or_default();

	if content.roles.remove(&role_name).is_none() {
		return Err!("Role '{role_name}' does not exist in this space.");
	}

	let state_lock = self.services.rooms.state.mutex.lock(&space_id).await;
	let server_user = &self.services.globals.server_user;

	self.services
		.rooms
		.timeline
		.build_and_append_pdu(
			custom_state_pdu!(SPACE_ROLES_EVENT_TYPE, "", &content),
			server_user,
			Some(&space_id),
			&state_lock,
		)
		.await?;

	// Cascade: remove the role from all users' member events
	let member_event_type = StateEventType::from(SPACE_ROLE_MEMBER_EVENT_TYPE.to_owned());
	if let Ok(shortstatehash) = self
		.services
		.rooms
		.state
		.get_room_shortstatehash(&space_id)
		.await
	{
		let user_entries: Vec<(_, ruma::OwnedEventId)> = self
			.services
			.rooms
			.state_accessor
			.state_keys_with_ids(shortstatehash, &member_event_type)
			.collect()
			.await;

		for (state_key, event_id) in user_entries {
			if let Ok(pdu) = self.services.rooms.timeline.get_pdu(&event_id).await {
				if let Ok(mut member_content) = pdu.get_content::<SpaceRoleMemberEventContent>() {
					if member_content.roles.contains(&role_name) {
						member_content.roles.retain(|r| r != &role_name);
						self.services
							.rooms
							.timeline
							.build_and_append_pdu(
								custom_state_pdu!(
									SPACE_ROLE_MEMBER_EVENT_TYPE,
									&state_key,
									&member_content
								),
								server_user,
								Some(&space_id),
								&state_lock,
							)
							.await?;
					}
				}
			}
		}

		// Cascade: remove the role from all rooms' role requirement events
		let room_event_type = StateEventType::from(SPACE_ROLE_ROOM_EVENT_TYPE.to_owned());

		let room_entries: Vec<(_, ruma::OwnedEventId)> = self
			.services
			.rooms
			.state_accessor
			.state_keys_with_ids(shortstatehash, &room_event_type)
			.collect()
			.await;

		for (state_key, event_id) in room_entries {
			if let Ok(pdu) = self.services.rooms.timeline.get_pdu(&event_id).await {
				if let Ok(mut room_content) = pdu.get_content::<SpaceRoleRoomEventContent>() {
					if room_content.required_roles.contains(&role_name) {
						room_content.required_roles.retain(|r| r != &role_name);
						self.services
							.rooms
							.timeline
							.build_and_append_pdu(
								custom_state_pdu!(
									SPACE_ROLE_ROOM_EVENT_TYPE,
									&state_key,
									&room_content
								),
								server_user,
								Some(&space_id),
								&state_lock,
							)
							.await?;
					}
				}
			}
		}
	}

	self.write_str(&format!("Removed role '{role_name}' from space {space_id}."))
		.await
}

#[admin_command]
async fn assign(
	&self,
	space: OwnedRoomOrAliasId,
	user_id: OwnedUserId,
	role_name: String,
) -> Result {
	let space_id = resolve_space!(self, space);

	// Read current role definitions to validate the role name
	let roles_event_type = StateEventType::from(SPACE_ROLES_EVENT_TYPE.to_owned());
	let role_defs: SpaceRolesEventContent = self
		.services
		.rooms
		.state_accessor
		.room_state_get_content(&space_id, &roles_event_type, "")
		.await
		.unwrap_or_default();

	if !role_defs.roles.contains_key(&role_name) {
		return Err!("Role '{role_name}' does not exist in this space.");
	}

	let member_event_type = StateEventType::from(SPACE_ROLE_MEMBER_EVENT_TYPE.to_owned());

	let mut content: SpaceRoleMemberEventContent = self
		.services
		.rooms
		.state_accessor
		.room_state_get_content(&space_id, &member_event_type, user_id.as_str())
		.await
		.unwrap_or_default();

	if content.roles.contains(&role_name) {
		return Err!("User {user_id} already has role '{role_name}' in this space.");
	}

	content.roles.push(role_name.clone());

	let state_lock = self.services.rooms.state.mutex.lock(&space_id).await;
	let server_user = &self.services.globals.server_user;

	self.services
		.rooms
		.timeline
		.build_and_append_pdu(
			custom_state_pdu!(SPACE_ROLE_MEMBER_EVENT_TYPE, user_id.as_str(), &content),
			server_user,
			Some(&space_id),
			&state_lock,
		)
		.await?;

	self.write_str(&format!("Assigned role '{role_name}' to {user_id} in space {space_id}."))
		.await
}

#[admin_command]
async fn revoke(
	&self,
	space: OwnedRoomOrAliasId,
	user_id: OwnedUserId,
	role_name: String,
) -> Result {
	let space_id = resolve_space!(self, space);
	let member_event_type = StateEventType::from(SPACE_ROLE_MEMBER_EVENT_TYPE.to_owned());

	let mut content: SpaceRoleMemberEventContent = self
		.services
		.rooms
		.state_accessor
		.room_state_get_content(&space_id, &member_event_type, user_id.as_str())
		.await
		.unwrap_or_default();

	let original_len = content.roles.len();
	content.roles.retain(|r| r != &role_name);

	if content.roles.len() == original_len {
		return Err!("User {user_id} does not have role '{role_name}' in this space.");
	}

	let state_lock = self.services.rooms.state.mutex.lock(&space_id).await;
	let server_user = &self.services.globals.server_user;

	self.services
		.rooms
		.timeline
		.build_and_append_pdu(
			custom_state_pdu!(SPACE_ROLE_MEMBER_EVENT_TYPE, user_id.as_str(), &content),
			server_user,
			Some(&space_id),
			&state_lock,
		)
		.await?;

	self.write_str(&format!("Revoked role '{role_name}' from {user_id} in space {space_id}."))
		.await
}

#[admin_command]
async fn require(
	&self,
	space: OwnedRoomOrAliasId,
	room_id: OwnedRoomId,
	role_name: String,
) -> Result {
	let space_id = resolve_space!(self, space);

	// Read current role definitions to validate the role name
	let roles_event_type = StateEventType::from(SPACE_ROLES_EVENT_TYPE.to_owned());
	let role_defs: SpaceRolesEventContent = self
		.services
		.rooms
		.state_accessor
		.room_state_get_content(&space_id, &roles_event_type, "")
		.await
		.unwrap_or_default();

	if !role_defs.roles.contains_key(&role_name) {
		return Err!("Role '{role_name}' does not exist in this space.");
	}

	let room_event_type = StateEventType::from(SPACE_ROLE_ROOM_EVENT_TYPE.to_owned());

	let mut content: SpaceRoleRoomEventContent = self
		.services
		.rooms
		.state_accessor
		.room_state_get_content(&space_id, &room_event_type, room_id.as_str())
		.await
		.unwrap_or_default();

	if content.required_roles.contains(&role_name) {
		return Err!("Room {room_id} already requires role '{role_name}' in this space.");
	}

	content.required_roles.push(role_name.clone());

	let state_lock = self.services.rooms.state.mutex.lock(&space_id).await;
	let server_user = &self.services.globals.server_user;

	self.services
		.rooms
		.timeline
		.build_and_append_pdu(
			custom_state_pdu!(SPACE_ROLE_ROOM_EVENT_TYPE, room_id.as_str(), &content),
			server_user,
			Some(&space_id),
			&state_lock,
		)
		.await?;

	self.write_str(&format!(
		"Room {room_id} now requires role '{role_name}' in space {space_id}."
	))
	.await
}

#[admin_command]
async fn unrequire(
	&self,
	space: OwnedRoomOrAliasId,
	room_id: OwnedRoomId,
	role_name: String,
) -> Result {
	let space_id = resolve_space!(self, space);
	let room_event_type = StateEventType::from(SPACE_ROLE_ROOM_EVENT_TYPE.to_owned());

	let mut content: SpaceRoleRoomEventContent = self
		.services
		.rooms
		.state_accessor
		.room_state_get_content(&space_id, &room_event_type, room_id.as_str())
		.await
		.unwrap_or_default();

	let original_len = content.required_roles.len();
	content.required_roles.retain(|r| r != &role_name);

	if content.required_roles.len() == original_len {
		return Err!("Room {room_id} does not require role '{role_name}' in this space.");
	}

	let state_lock = self.services.rooms.state.mutex.lock(&space_id).await;
	let server_user = &self.services.globals.server_user;

	self.services
		.rooms
		.timeline
		.build_and_append_pdu(
			custom_state_pdu!(SPACE_ROLE_ROOM_EVENT_TYPE, room_id.as_str(), &content),
			server_user,
			Some(&space_id),
			&state_lock,
		)
		.await?;

	self.write_str(&format!(
		"Removed role requirement '{role_name}' from room {room_id} in space {space_id}."
	))
	.await
}

#[admin_command]
async fn user(&self, space: OwnedRoomOrAliasId, user_id: OwnedUserId) -> Result {
	let space_id = resolve_space!(self, space);

	let user_roles = self.services.rooms.roles.user_roles.read().await;
	let roles = user_roles
		.get(&space_id)
		.and_then(|space_users| space_users.get(&user_id));

	match roles {
		| Some(roles) if !roles.is_empty() => {
			let roles_list: Vec<&String> = roles.iter().collect();
			self.write_str(&format!(
				"Roles for {user_id} in space {space_id}:\n```\n{}\n```",
				roles_list
					.iter()
					.map(|r| format!("- {r}"))
					.collect::<Vec<_>>()
					.join("\n")
			))
			.await
		},
		| _ =>
			self.write_str(&format!("User {user_id} has no roles in space {space_id}."))
				.await,
	}
}

#[admin_command]
async fn room(&self, space: OwnedRoomOrAliasId, room_id: OwnedRoomId) -> Result {
	let space_id = resolve_space!(self, space);

	let room_reqs = self.services.rooms.roles.room_requirements.read().await;
	let requirements = room_reqs
		.get(&space_id)
		.and_then(|space_rooms| space_rooms.get(&room_id));

	match requirements {
		| Some(reqs) if !reqs.is_empty() => {
			let reqs_list: Vec<&String> = reqs.iter().collect();
			self.write_str(&format!(
				"Required roles for room {room_id} in space {space_id}:\n```\n{}\n```",
				reqs_list
					.iter()
					.map(|r| format!("- {r}"))
					.collect::<Vec<_>>()
					.join("\n")
			))
			.await
		},
		| _ =>
			self.write_str(&format!(
				"Room {room_id} has no role requirements in space {space_id}."
			))
			.await,
	}
}