git.delta.rocks / fleet / refs/commits / 48d120aa4aef

difftreelog

source

remowt/cmds/remowt-agent/src/main.rs11.7 KiBsourcehistory
1use std::borrow::Cow;2use std::collections::{BTreeMap, HashMap};3use std::fs::Permissions;4use std::future::pending;5use std::io;6use std::os::unix::fs::PermissionsExt as _;7use std::path::PathBuf;8use std::sync::{Arc, Mutex, OnceLock};910use bifrostlink::Rpc;11use bifrostlink::declarative::RemoteEndpoints;12use bifrostlink_ports::stdio::from_stdio;13use bifrostlink_ports::unix_socket::from_socket;14use clap::Parser;15use remowt_endpoints::{16	forward::Forward, fs::Fs, iroh_tunnel::IrohTunnel, nix_daemon::NixDaemon, pty::Pty,17	subprocess::Subprocess, systemd::Systemd,18};19use remowt_link_shared::iroh_tunnel::TunnelDialer;20use remowt_link_shared::{Address, BifConfig, gateway};21use remowt_polkit_shared::{Identity, PidDisplay, emphasize};22use remowt_ui_prompt::bifrost::{PromptEndpointsClient, serve_prompts};23use remowt_ui_prompt::rofi::RofiPrompter;24use remowt_ui_prompt::{PrependSourcePrompter, Prompter, Source};25use tokio::fs;26use tokio::net::UnixStream;27use tokio::runtime::Builder;28use tokio::task::AbortHandle;29use tracing::{debug, trace};30use zbus::fdo;31use zbus::zvariant::{OwnedValue, Str};32use zbus::{Connection, interface};33use zbus_polkit::policykit1::Subject;3435use self::helper::{Helper, SocketHelper, SuidHelper};3637pub mod askpass;38pub mod editor;39pub mod helper;4041struct CancelTaskOnDrop {42	tasks: Arc<Mutex<HashMap<String, AbortHandle>>>,43	handle: String,44}45impl Drop for CancelTaskOnDrop {46	fn drop(&mut self) {47		debug!("cancel on drop");48		if let Some(task) = self49			.tasks50			.lock()51			.expect("not poisoned")52			.remove(&self.handle)53		{54			task.abort();55		}56	}57}5859struct Agent<H, P> {60	tasks: Arc<Mutex<HashMap<String, AbortHandle>>>,61	helper: H,62	prompter: P,63}64impl<H, P> Agent<H, P> {65	fn new(helper: H, prompter: P) -> Self {66		Agent {67			tasks: Arc::new(Mutex::new(HashMap::new())),68			helper,69			prompter,70		}71	}72}7374#[interface(name = "org.freedesktop.PolicyKit1.AuthenticationAgent")]75impl<H, P> Agent<H, P>76where77	H: Helper + Clone + Send + Sync + 'static,78	P: Prompter + Clone + Send + Sync + 'static,79{80	/// BeginAuthentication method81	#[allow(clippy::too_many_arguments)]82	async fn begin_authentication(83		&self,84		action_id: String,85		message: String,86		_icon_name: String,87		mut details: BTreeMap<String, String>,88		cookie: String,89		identities: Vec<Identity>,90	) -> zbus::fdo::Result<()> {91		use std::fmt::Write;92		debug!("begin auth");93		let _cancel_guard = Arc::new(OnceLock::new());94		let task = {95			let helper = self.helper.clone();96			let prompter = self.prompter.clone();97			let cookie = cookie.clone();98			let _cancel_guard = _cancel_guard.clone();99			tokio::task::spawn(async move {100				let _cancel_guard = _cancel_guard.clone();101				trace!("conversation task");102				let mut description = format!("{message}\n\n<b>Action id:</b> {action_id}",);103				if let Some(subject) = details.remove("polkit.caller-pid") {104					let _ = write!(description, "\n<b>Caller:</b> ");105					if let Ok(pid) = subject.parse::<u32>() {106						let _ = write!(description, "{}", PidDisplay(pid));107					} else {108						let _ = write!(description, "{}", emphasize("invalid pid"));109					}110				}111				if let Some(subject) = details.remove("polkit.subject-pid") {112					let _ = write!(description, "\n<b>Subject:</b> ");113					if let Ok(pid) = subject.parse::<u32>() {114						let _ = write!(description, "{}", PidDisplay(pid));115					} else {116						let _ = write!(description, "{}", emphasize("invalid pid"));117					}118				}119				let mut prompter = PrependSourcePrompter {120					source: vec![Source(Cow::Borrowed("polkit agent"))],121					description: description.clone(),122					prompter,123				};124125				let identity_displays: Vec<String> =126					identities.iter().map(|v| v.to_string()).collect();127				let identity_displays: Vec<&str> =128					identity_displays.iter().map(|v| v.as_str()).collect();129				debug!("choose identity");130				let choosen_identity = match identity_displays.len() {131					0 => {132						return Err(fdo::Error::AuthFailed(133							"no identity to authenticate as".to_owned(),134						));135					}136					1 => 0,137					_ => prompter138						.prompt_enum(139							"Identity",140							"Select identity to use for polkit authorization",141							&identity_displays,142							&[],143						)144						.await145						.map_err(prompt_err)?,146				};147				debug!("identity chosen");148149				let _ = write!(150					description,151					"\n<b>Identity:</b> {}",152					identities[choosen_identity as usize]153				);154				prompter.description = description;155156				prompter.source.push(Source(Cow::Borrowed("polkit daemon")));157158				helper159					.help_me(160						&cookie,161						prompter,162						identities[choosen_identity as usize].clone(),163					)164					.await165					.map_err(|e| fdo::Error::Failed(e.to_string()))?;166				// let connection = Connection::system().await?;167				// let helper = PolkitHelperProxy::new(&connection).await?;168169				Ok(())170			})171		};172		self.tasks173			.lock()174			.unwrap()175			.insert(cookie.clone(), task.abort_handle());176		debug!("abort handle stored");177		let _ = _cancel_guard.set(CancelTaskOnDrop {178			tasks: self.tasks.clone(),179			handle: cookie.clone(),180		});181182		let _ = task.await;183184		Ok(())185	}186187	/// CancelAuthentication method188	async fn cancel_authentication(&self, cookie: &str) -> zbus::fdo::Result<()> {189		debug!("auth cancelled");190		if let Some(abort) = self.tasks.lock().unwrap().remove(cookie) {191			debug!("abort handle found");192			abort.abort();193		}194		// debug!("Authentication cancled ! {cookie}");195		Ok(())196	}197}198199fn prompt_err(value: remowt_ui_prompt::Error) -> fdo::Error {200	use remowt_ui_prompt::Error;201	match value {202		Error::Cancel => fdo::Error::NoReply("input was cancelled".to_owned()),203		Error::Remote(e) => fdo::Error::NoReply(format!("remote error occured: {e}")),204		Error::InputError(e) => fdo::Error::Failed(e),205	}206}207208const OBJ_PATH: &str = "/org/freedesktop/PolicyKit1/AuthenticationAgent";209210#[derive(Parser)]211enum Opts {212	AskPass {213		prompt: String,214		description: String,215	},216	Editor {217		/// Argument to nvim218		path: String,219	},220	/// Expose a remote TCP/UDP address on the local (User) machine, printing the local port.221	Forward {222		/// Protocol: `tcp` (ssh or iroh) or `udp` (iroh-only).223		proto: ForwardProto,224		/// Remote socket address to expose locally, e.g. `127.0.0.1:8080`.225		addr: String,226	},227	RealAgent {228		#[arg(long)]229		path: Option<PathBuf>,230		/// Expect own address to be AgentPrivileged, skip installing polkit agent231		#[arg(long)]232		privileged: bool,233		#[arg(long)]234		local: bool,235	},236	LocalAgent,237}238239#[derive(Clone, Copy, clap::ValueEnum)]240enum ForwardProto {241	Tcp,242	Udp,243}244245fn main() -> anyhow::Result<()> {246	tracing_subscriber::fmt()247		.with_writer(io::stderr)248		.without_time()249		.init();250	let opts = Opts::parse();251252	let runtime = Builder::new_current_thread().enable_all().build()?;253254	match opts {255		Opts::AskPass {256			prompt,257			description,258		} => runtime.block_on(askpass::ask(&prompt, description)),259		Opts::LocalAgent => runtime.block_on(main_real()),260		Opts::Editor { path } => runtime.block_on(editor::edit(path)),261		Opts::Forward { proto, addr } => {262			runtime.block_on(editor::forward(matches!(proto, ForwardProto::Udp), addr))263		}264		Opts::RealAgent {265			path,266			privileged,267			local,268		} => runtime.block_on(main_real_agent(path, privileged, local)),269	}270}271async fn main_real() -> anyhow::Result<()> {272	let system_conn = Connection::system().await?;273	let helper = SocketHelper {274		fallback: SuidHelper,275	};276	register_auth_agent(&system_conn, Agent::new(helper, RofiPrompter)).await?;277278	let mut rpc = Rpc::<BifConfig>::new(Address::User);279	serve_prompts(&mut rpc, RofiPrompter);280281	gateway::serve(rpc.clone(), &gateway::local_socket()?).await?;282283	let _keep_alive = (system_conn, rpc);284	pending().await285}286async fn main_real_agent(287	path: Option<PathBuf>,288	privileged: bool,289	local: bool,290) -> anyhow::Result<()> {291	let address = if privileged {292		Address::AgentPrivileged293	} else {294		Address::Agent295	};296	let mut rpc = Rpc::<BifConfig>::new(address);297298	let dialer = Arc::new(TunnelDialer::new());299	Fs::new().register_endpoints(&mut rpc);300	Systemd.register_endpoints(&mut rpc);301	Pty::new(dialer.clone()).register_endpoints(&mut rpc);302	Subprocess::new(dialer.clone()).register_endpoints(&mut rpc);303	NixDaemon::new(dialer.clone()).register_endpoints(&mut rpc);304	IrohTunnel::new(dialer.clone()).register_endpoints(&mut rpc);305	Forward::new(dialer.clone()).register_endpoints(&mut rpc);306307	remowt_plugin::host::serve(&mut rpc);308309	let user_prompter = PromptEndpointsClient::wrap(rpc.remote(Address::User));310311	let helpers = tempfile::Builder::new().prefix("remowt-path.").tempdir()?;312	let gateway_socket = helpers.path().join(gateway::SOCKET_NAME);313	let exe = std::env::current_exe()?;314	let askpass_helper = helpers.path().join("remowt-askpass");315	let editor_helper = helpers.path().join("remowt-editor");316	let forward_helper = helpers.path().join("remowt-forward");317	{318		let script = format!(319			"#!/bin/sh\nexec {} ask-pass \"password\" \"$1\"\n",320			sh_quote(&exe.to_string_lossy())321		);322		fs::write(&askpass_helper, script).await?;323		fs::set_permissions(&askpass_helper, Permissions::from_mode(0o755)).await?;324	}325	{326		let script = format!(327			"#!/bin/sh\nexec {} editor \"$1\"\n",328			sh_quote(&exe.to_string_lossy())329		);330		fs::write(&editor_helper, script).await?;331		fs::set_permissions(&editor_helper, Permissions::from_mode(0o755)).await?;332	}333	{334		let script = format!(335			"#!/bin/sh\nexec {} forward \"$@\"\n",336			sh_quote(&exe.to_string_lossy())337		);338		fs::write(&forward_helper, script).await?;339		fs::set_permissions(&forward_helper, Permissions::from_mode(0o755)).await?;340	}341342	// Safety: Hoping tokio own threads won't read any of those...343	unsafe {344		prepend_path(helpers.path());345		std::env::set_var("SUDO_ASKPASS", &askpass_helper);346		std::env::set_var("SSH_ASKPASS", &askpass_helper);347		std::env::set_var("SSH_ASKPASS_REQUIRE", "force");348		std::env::set_var("EDITOR", &editor_helper);349		std::env::set_var("VISUAL", &editor_helper);350		std::env::set_var(gateway::SOCKET_ENV, &gateway_socket);351	}352353	let port = match path {354		Some(path) => from_socket(UnixStream::connect(path).await?),355		None => from_stdio(),356	};357	rpc.add_direct(Address::User, port, bifrostlink::Rtt(0));358359	gateway::serve(rpc.clone(), &gateway_socket).await?;360361	let polkit_conn = if !privileged && !local {362		let conn = Connection::system().await?;363		let helper = SocketHelper {364			fallback: SuidHelper,365		};366		register_auth_agent(&conn, Agent::new(helper, user_prompter)).await?;367		Some(conn)368	} else {369		None370	};371372	let _keep_alive = (helpers, polkit_conn);373	pending().await374}375376async fn register_auth_agent<H, P>(conn: &Connection, agent: Agent<H, P>) -> anyhow::Result<()>377where378	H: Helper + Clone + Send + Sync + 'static,379	P: Prompter + Clone + Send + Sync + 'static,380{381	let proxy = zbus_polkit::policykit1::AuthorityProxy::new(conn).await?;382	conn.object_server().at(OBJ_PATH, agent).await?;383384	let subject = auth_agent_subject()?;385	proxy386		.register_authentication_agent(&subject, "C", OBJ_PATH)387		.await?;388	debug!(kind = subject.subject_kind, "registered polkit agent");389	Ok(())390}391392fn auth_agent_subject() -> anyhow::Result<Subject> {393	let mut details = HashMap::new();394	if let Ok(session_id) = std::env::var("XDG_SESSION_ID") {395		let val: OwnedValue = Str::from(session_id).into();396		details.insert("session-id".to_string(), val);397		return Ok(Subject {398			subject_kind: "unix-session".to_string(),399			subject_details: details,400		});401	}402403	details.insert("pid".to_string(), OwnedValue::from(std::process::id()));404	Ok(Subject {405		subject_kind: "unix-process".to_string(),406		subject_details: details,407	})408}409410fn sh_quote(s: &str) -> String {411	format!("'{}'", s.replace('\'', "'\\''"))412}413414/// Prepend `dir` to the process `PATH`.415///416/// # SAFETY417///418/// Same as `set_var`419unsafe fn prepend_path(dir: &std::path::Path) {420	let value = match std::env::var_os("PATH") {421		Some(existing) => {422			let mut v = dir.as_os_str().to_owned();423			v.push(":");424			v.push(existing);425			v426		}427		None => dir.as_os_str().to_owned(),428	};429	unsafe {430		std::env::set_var("PATH", value);431	}432}