git.delta.rocks / remowt / refs/commits / 99930b3c897c

difftreelog

source

cmds/remowt-agent/src/main.rs10.5 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_endpoints::{fs::Fs, pty::Pty, systemd::Systemd};15use remowt_link_shared::{editor::EditorEndpointsClient, Address, BifConfig};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	tracing_subscriber::fmt()229		.with_writer(std::io::stderr)230		.without_time()231		.init();232	let opts = Opts::parse();233234	let runtime = Builder::new_current_thread().enable_all().build()?;235236	match opts {237		Opts::AskPass {238			prompt,239			description,240		} => runtime.block_on(askpass::ask(&prompt, description)),241		Opts::LocalAgent => runtime.block_on(main_real()),242		Opts::Editor { path } => runtime.block_on(editor::edit(path)),243		Opts::RealAgent { path, privileged } => runtime.block_on(main_real_agent(path, privileged)),244	}245}246async fn main_real() -> anyhow::Result<()> {247	let conn = Connection::system().await?;248	let helper = SocketHelper {249		fallback: SuidHelper,250	};251	register_auth_agent(&conn, Agent::new(helper, RofiPrompter)).await?;252253	let _conn = conn;254	pending().await255}256async fn main_real_agent(path: Option<PathBuf>, privileged: bool) -> anyhow::Result<()> {257	let address = if privileged {258		Address::AgentPrivileged259	} else {260		Address::Agent261	};262	let mut rpc = Rpc::<BifConfig>::new(address);263264	Fs::new().register_endpoints(&mut rpc);265	Systemd.register_endpoints(&mut rpc);266	Pty::new().register_endpoints(&mut rpc);267268	remowt_plugin::host::serve(&mut rpc);269270	let user_prompter = PromptEndpointsClient::wrap(rpc.remote(Address::User));271	let editor_client = EditorEndpointsClient::wrap(rpc.remote(Address::User));272273	let bus = bus::spawn().await?;274	askpass::serve(&bus.conn, user_prompter.clone()).await?;275	editor::serve(&bus.conn, editor_client).await?;276277	let helpers = tempfile::Builder::new().prefix("remowt-path.").tempdir()?;278	let exe = std::env::current_exe()?;279	let askpass_helper = helpers.path().join("remowt-askpass");280	let editor_helper = helpers.path().join("remowt-editor");281	{282		let script = format!(283			"#!/bin/sh\nexec {} ask-pass \"password\" \"$1\"\n",284			sh_quote(&exe.to_string_lossy())285		);286		fs::write(&askpass_helper, script).await?;287		fs::set_permissions(&askpass_helper, Permissions::from_mode(0o755)).await?;288	}289	{290		let script = format!(291			"#!/bin/sh\nexec {} editor \"$1\"\n",292			sh_quote(&exe.to_string_lossy())293		);294		fs::write(&editor_helper, script).await?;295		fs::set_permissions(&editor_helper, Permissions::from_mode(0o755)).await?;296	}297298	// Safety: Hoping tokio own threads won't read any of those...299	unsafe {300		prepend_path(helpers.path());301		std::env::set_var("SUDO_ASKPASS", &askpass_helper);302		std::env::set_var("SSH_ASKPASS", &askpass_helper);303		std::env::set_var("SSH_ASKPASS_REQUIRE", "force");304		std::env::set_var("EDITOR", &editor_helper);305		std::env::set_var("VISUAL", &editor_helper);306		std::env::set_var("DBUS_SESSION_BUS_ADDRESS", &bus.address);307	}308309	let port = match path {310		Some(path) => from_socket(UnixStream::connect(path).await?),311		None => from_stdio(),312	};313	rpc.add_direct(Address::User, port, bifrostlink::Rtt(0));314315	let polkit_conn = if !privileged {316		// The unprivileged agent doubles as a polkit authentication agent so317		// `run0` (e.g. our own elevation) routes its prompt to the User over318		// bifrost instead of failing on a tty-less session.319		let conn = Connection::system().await?;320		let helper = SocketHelper {321			fallback: SuidHelper,322		};323		register_auth_agent(&conn, Agent::new(helper, user_prompter)).await?;324		Some(conn)325	} else {326		None327	};328329	let _keep_alive = (bus, helpers, polkit_conn);330	pending().await331}332333async fn register_auth_agent<H, P>(conn: &Connection, agent: Agent<H, P>) -> anyhow::Result<()>334where335	H: Helper + Clone + Send + Sync + 'static,336	P: Prompter + Clone + Send + Sync + 'static,337{338	let proxy = zbus_polkit::policykit1::AuthorityProxy::new(conn).await?;339	conn.object_server().at(OBJ_PATH, agent).await?;340341	let subject = auth_agent_subject()?;342	proxy343		.register_authentication_agent(&subject, "C", OBJ_PATH)344		.await?;345	debug!(kind = subject.subject_kind, "registered polkit agent");346	Ok(())347}348349fn auth_agent_subject() -> anyhow::Result<Subject> {350	let mut details = HashMap::new();351	if let Ok(session_id) = std::env::var("XDG_SESSION_ID") {352		let val: OwnedValue = Str::from(session_id).into();353		details.insert("session-id".to_string(), val);354		return Ok(Subject {355			subject_kind: "unix-session".to_string(),356			subject_details: details,357		});358	}359360	details.insert("pid".to_string(), OwnedValue::from(std::process::id()));361	Ok(Subject {362		subject_kind: "unix-process".to_string(),363		subject_details: details,364	})365}366367fn sh_quote(s: &str) -> String {368	format!("'{}'", s.replace('\'', "'\\''"))369}370371/// Prepend `dir` to the process `PATH`.372///373/// # SAFETY374///375/// Same as `set_var`376unsafe fn prepend_path(dir: &std::path::Path) {377	let value = match std::env::var_os("PATH") {378		Some(existing) => {379			let mut v = dir.as_os_str().to_owned();380			v.push(":");381			v.push(existing);382			v383		}384		None => dir.as_os_str().to_owned(),385	};386	unsafe {387		std::env::set_var("PATH", value);388	}389}