git.delta.rocks / remowt / refs/commits / 875f55ab11c2

difftreelog

source

cmds/remowt-agent/src/main.rs10.6 KiBsourcehistory
1use std::borrow::Cow;2use std::collections::{BTreeMap, HashMap};3use std::fs::Permissions;4use std::future::pending;5use std::os::unix::fs::PermissionsExt as _;6use std::path::PathBuf;7use std::sync::{Arc, Mutex, OnceLock};89use bifrostlink::declarative::RemoteEndpoints;10use bifrostlink::Rpc;11use bifrostlink_ports::stdio::from_stdio;12use bifrostlink_ports::unix_socket::from_socket;13use clap::Parser;14use remowt_link_shared::editor::EditorEndpointsClient;15use remowt_link_shared::{Address, BifConfig, Fs, Pty, Systemd};16use remowt_polkit_shared::{emphasize, BackendRequest, Identity, PidDisplay};17use remowt_ui_prompt::bifrost::PromptEndpointsClient;18use remowt_ui_prompt::rofi::RofiPrompter;19use remowt_ui_prompt::{PrependSourcePrompter, Prompter, Source};20use tokio::fs;21use tokio::net::UnixStream;22use tokio::runtime::Builder;23use tokio::task::AbortHandle;24use tracing::{debug, info, trace};25use zbus::fdo;26use zbus::zvariant::{OwnedValue, Str};27use zbus::{interface, proxy, Connection};28use zbus_polkit::policykit1::Subject;2930use self::helper::{Helper, SocketHelper, SuidHelper};3132pub mod askpass;33pub mod bus;34pub mod editor;35pub mod helper;3637struct CancelTaskOnDrop {38	tasks: Arc<Mutex<HashMap<String, AbortHandle>>>,39	handle: String,40}41impl Drop for CancelTaskOnDrop {42	fn drop(&mut self) {43		debug!("cancel on drop");44		if let Some(task) = self45			.tasks46			.lock()47			.expect("not poisoned")48			.remove(&self.handle)49		{50			task.abort();51		}52	}53}5455struct Agent<H, P> {56	tasks: Arc<Mutex<HashMap<String, AbortHandle>>>,57	helper: H,58	prompter: P,59}60impl<H, P> Agent<H, P> {61	fn new(helper: H, prompter: P) -> Self {62		Agent {63			tasks: Arc::new(Mutex::new(HashMap::new())),64			helper,65			prompter,66		}67	}68}6970#[interface(name = "org.freedesktop.PolicyKit1.AuthenticationAgent")]71impl<H, P> Agent<H, P>72where73	H: Helper + Clone + Send + Sync + 'static,74	P: Prompter + Clone + Send + Sync + 'static,75{76	/// BeginAuthentication method77	#[allow(clippy::too_many_arguments)]78	async fn begin_authentication(79		&self,80		action_id: String,81		message: String,82		_icon_name: String,83		mut details: BTreeMap<String, String>,84		cookie: String,85		identities: Vec<Identity>,86	) -> zbus::fdo::Result<()> {87		use std::fmt::Write;88		info!("begin auth");89		let _cancel_guard = Arc::new(OnceLock::new());90		let task = {91			let helper = self.helper.clone();92			let prompter = self.prompter.clone();93			let cookie = cookie.clone();94			let _cancel_guard = _cancel_guard.clone();95			tokio::task::spawn(async move {96				let _cancel_guard = _cancel_guard.clone();97				trace!("conversation task");98				let mut description = format!("{message}\n\n<b>Action id:</b> {action_id}",);99				if let Some(subject) = details.remove("polkit.caller-pid") {100					let _ = write!(description, "\n<b>Caller:</b> ");101					if let Ok(pid) = subject.parse::<u32>() {102						let _ = write!(description, "{}", PidDisplay(pid));103					} else {104						let _ = write!(description, "{}", emphasize("invalid pid"));105					}106				}107				if let Some(subject) = details.remove("polkit.subject-pid") {108					let _ = write!(description, "\n<b>Subject:</b> ");109					if let Ok(pid) = subject.parse::<u32>() {110						let _ = write!(description, "{}", PidDisplay(pid));111					} else {112						let _ = write!(description, "{}", emphasize("invalid pid"));113					}114				}115				let mut prompter = PrependSourcePrompter {116					source: vec![Source(Cow::Borrowed("polkit agent"))],117					description: description.clone(),118					prompter,119				};120121				let identity_displays: Vec<String> =122					identities.iter().map(|v| v.to_string()).collect();123				let identity_displays: Vec<&str> =124					identity_displays.iter().map(|v| v.as_str()).collect();125				debug!("choose identity");126				let choosen_identity = match identity_displays.len() {127					0 => {128						return Err(fdo::Error::AuthFailed(129							"no identity to authenticate as".to_owned(),130						))131					}132					1 => 0,133					_ => {134						prompter135							.prompt_enum(136								"Identity",137								"Select identity to use for polkit authorization",138								&identity_displays,139								&[],140							)141							.await?142					}143				};144				debug!("identity chosen");145146				let _ = write!(147					description,148					"\n<b>Identity:</b> {}",149					identities[choosen_identity as usize]150				);151				prompter.description = description;152153				prompter.source.push(Source(Cow::Borrowed("polkit daemon")));154155				helper156					.help_me(157						&cookie,158						prompter,159						identities[choosen_identity as usize].clone(),160					)161					.await162					.map_err(|e| fdo::Error::Failed(e.to_string()))?;163				// let connection = Connection::system().await?;164				// let helper = PolkitHelperProxy::new(&connection).await?;165166				Ok(())167			})168		};169		self.tasks170			.lock()171			.unwrap()172			.insert(cookie.clone(), task.abort_handle());173		debug!("abort handle stored");174		let _ = _cancel_guard.set(CancelTaskOnDrop {175			tasks: self.tasks.clone(),176			handle: cookie.clone(),177		});178179		let _ = task.await;180181		Ok(())182	}183184	/// CancelAuthentication method185	async fn cancel_authentication(&self, cookie: &str) -> zbus::fdo::Result<()> {186		debug!("auth cancelled");187		if let Some(abort) = self.tasks.lock().unwrap().remove(cookie) {188			debug!("abort handle found");189			abort.abort();190		}191		// debug!("Authentication cancled ! {cookie}");192		Ok(())193	}194}195196const OBJ_PATH: &str = "/org/freedesktop/PolicyKit1/AuthenticationAgent";197198#[proxy(199	interface = "lach.PolkitHelper",200	default_service = "lach.polkit.helper1",201	default_path = "/lach/PolkitHelper"202)]203trait PolkitHelper {204	fn init_conversation(&self, request: BackendRequest) -> zbus::Result<()>;205}206207#[derive(Parser)]208enum Opts {209	AskPass {210		prompt: String,211		description: String,212	},213	Editor {214		/// Argument to nvim215		path: String,216	},217	RealAgent {218		#[arg(long)]219		path: Option<PathBuf>,220		/// Expect own address to be AgentPrivileged, skip installing polkit agent221		#[arg(long)]222		privileged: bool,223	},224	LocalAgent,225}226227fn main() -> anyhow::Result<()> {228	// Log to stderr: `privileged-agent` uses stdout as the bifrost transport,229	// so anything written there would corrupt the stream.230	tracing_subscriber::fmt()231		.with_writer(std::io::stderr)232		.without_time()233		.init();234	let opts = Opts::parse();235236	let runtime = Builder::new_current_thread().enable_all().build()?;237238	match opts {239		Opts::AskPass {240			prompt,241			description,242		} => runtime.block_on(askpass::ask(&prompt, description)),243		Opts::LocalAgent => runtime.block_on(main_real()),244		Opts::Editor { path } => runtime.block_on(editor::edit(path)),245		Opts::RealAgent { path, privileged } => runtime.block_on(main_real_agent(path, privileged)),246	}247}248async fn main_real() -> anyhow::Result<()> {249	let conn = Connection::system().await?;250	let helper = SocketHelper {251		fallback: SuidHelper,252	};253	register_auth_agent(&conn, Agent::new(helper, RofiPrompter)).await?;254255	let _conn = conn;256	pending().await257}258async fn main_real_agent(path: Option<PathBuf>, privileged: bool) -> anyhow::Result<()> {259	let address = if privileged {260		Address::AgentPrivileged261	} else {262		Address::Agent263	};264	let mut rpc = Rpc::<BifConfig>::new(address);265266	Fs::new().register_endpoints(&mut rpc);267	Systemd.register_endpoints(&mut rpc);268	Pty::new().register_endpoints(&mut rpc);269270	remowt_plugin::host::serve(&mut rpc);271272	let user_prompter = PromptEndpointsClient::wrap(rpc.remote(Address::User));273	let editor_client = EditorEndpointsClient::wrap(rpc.remote(Address::User));274275	let bus = bus::spawn().await?;276	askpass::serve(&bus.conn, user_prompter.clone()).await?;277	editor::serve(&bus.conn, editor_client).await?;278279	let helpers = tempfile::Builder::new().prefix("remowt-path.").tempdir()?;280	let exe = std::env::current_exe()?;281	let askpass_helper = helpers.path().join("remowt-askpass");282	let editor_helper = helpers.path().join("remowt-editor");283	{284		let script = format!(285			"#!/bin/sh\nexec {} ask-pass \"password\" \"$1\"\n",286			sh_quote(&exe.to_string_lossy())287		);288		fs::write(&askpass_helper, script).await?;289		fs::set_permissions(&askpass_helper, Permissions::from_mode(0o755)).await?;290	}291	{292		let script = format!(293			"#!/bin/sh\nexec {} editor \"$1\"\n",294			sh_quote(&exe.to_string_lossy())295		);296		fs::write(&editor_helper, script).await?;297		fs::set_permissions(&editor_helper, Permissions::from_mode(0o755)).await?;298	}299300	// Safety: Hoping tokio own threads won't read any of those...301	unsafe {302		prepend_path(helpers.path());303		std::env::set_var("SUDO_ASKPASS", &askpass_helper);304		std::env::set_var("SSH_ASKPASS", &askpass_helper);305		std::env::set_var("SSH_ASKPASS_REQUIRE", "force");306		std::env::set_var("EDITOR", &editor_helper);307		std::env::set_var("VISUAL", &editor_helper);308		std::env::set_var("DBUS_SESSION_BUS_ADDRESS", &bus.address);309	}310311	let port = match path {312		Some(path) => from_socket(UnixStream::connect(path).await?),313		None => from_stdio(),314	};315	rpc.add_direct(Address::User, port, bifrostlink::Rtt(0));316317	let polkit_conn = if !privileged {318		// The unprivileged agent doubles as a polkit authentication agent so319		// `run0` (e.g. our own elevation) routes its prompt to the User over320		// bifrost instead of failing on a tty-less session.321		let conn = Connection::system().await?;322		let helper = SocketHelper {323			fallback: SuidHelper,324		};325		register_auth_agent(&conn, Agent::new(helper, user_prompter)).await?;326		Some(conn)327	} else {328		None329	};330331	let _keep_alive = (bus, helpers, polkit_conn);332	pending().await333}334335async fn register_auth_agent<H, P>(conn: &Connection, agent: Agent<H, P>) -> anyhow::Result<()>336where337	H: Helper + Clone + Send + Sync + 'static,338	P: Prompter + Clone + Send + Sync + 'static,339{340	let proxy = zbus_polkit::policykit1::AuthorityProxy::new(conn).await?;341	conn.object_server().at(OBJ_PATH, agent).await?;342343	let subject = auth_agent_subject()?;344	proxy345		.register_authentication_agent(&subject, "C", OBJ_PATH)346		.await?;347	debug!(kind = subject.subject_kind, "registered polkit agent");348	Ok(())349}350351fn auth_agent_subject() -> anyhow::Result<Subject> {352	let mut details = HashMap::new();353	if let Ok(session_id) = std::env::var("XDG_SESSION_ID") {354		let val: OwnedValue = Str::from(session_id).into();355		details.insert("session-id".to_string(), val);356		return Ok(Subject {357			subject_kind: "unix-session".to_string(),358			subject_details: details,359		});360	}361362	details.insert("pid".to_string(), OwnedValue::from(std::process::id()));363	Ok(Subject {364		subject_kind: "unix-process".to_string(),365		subject_details: details,366	})367}368369fn sh_quote(s: &str) -> String {370	format!("'{}'", s.replace('\'', "'\\''"))371}372373/// Prepend `dir` to the process `PATH`.374///375/// # SAFETY376///377/// Same as `set_var`378unsafe fn prepend_path(dir: &std::path::Path) {379	let value = match std::env::var_os("PATH") {380		Some(existing) => {381			let mut v = dir.as_os_str().to_owned();382			v.push(":");383			v.push(existing);384			v385		}386		None => dir.as_os_str().to_owned(),387	};388	unsafe {389		std::env::set_var("PATH", value);390	}391}