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 polkit_shared::{emphasize, BackendRequest, Identity, PidDisplay};15use remowt_link_shared::editor::EditorEndpointsClient;16use remowt_link_shared::{Address, BifConfig, Fs, Pty, Systemd};17use tokio::fs;18use tokio::net::UnixStream;19use tokio::runtime::Builder;20use tokio::task::AbortHandle;21use tracing::{info, trace};22use ui_prompt::bifrost::PromptEndpointsClient;23use ui_prompt::{PrependSourcePrompter, Prompter, Source};24use zbus::fdo;25use zbus::zvariant::{OwnedValue, Str};26use zbus::{interface, proxy, Connection};27use zbus_polkit::policykit1::Subject;2829use self::helper::{Helper, SocketHelper, SuidHelper};3031pub mod askpass;32pub mod bus;33pub mod editor;34pub mod helper;3536struct CancelTaskOnDrop {37 tasks: Arc<Mutex<HashMap<String, AbortHandle>>>,38 handle: String,39}40impl Drop for CancelTaskOnDrop {41 fn drop(&mut self) {42 info!("cancel on drop");43 if let Some(task) = self44 .tasks45 .lock()46 .expect("not poisoned")47 .remove(&self.handle)48 {49 task.abort();50 }51 }52}5354struct Agent<H, P> {55 tasks: Arc<Mutex<HashMap<String, AbortHandle>>>,56 helper: H,57 prompter: P,58}59impl<H, P> Agent<H, P> {60 fn new(helper: H, prompter: P) -> Self {61 Agent {62 tasks: Arc::new(Mutex::new(HashMap::new())),63 helper,64 prompter,65 }66 }67}6869#[interface(name = "org.freedesktop.PolicyKit1.AuthenticationAgent")]70impl<H, P> Agent<H, P>71where72 H: Helper + Clone + Send + Sync + 'static,73 P: Prompter + Clone + Send + Sync + 'static,74{75 76 #[allow(clippy::too_many_arguments)]77 async fn begin_authentication(78 &self,79 action_id: String,80 message: String,81 _icon_name: String,82 mut details: BTreeMap<String, String>,83 cookie: String,84 identities: Vec<Identity>,85 ) -> zbus::fdo::Result<()> {86 use std::fmt::Write;87 info!("begin auth");88 let _cancel_guard = Arc::new(OnceLock::new());89 let task = {90 let helper = self.helper.clone();91 let prompter = self.prompter.clone();92 let cookie = cookie.clone();93 let _cancel_guard = _cancel_guard.clone();94 tokio::task::spawn(async move {95 let _cancel_guard = _cancel_guard.clone();96 trace!("conversation task");97 let mut description = format!("{message}\n\n<b>Action id:</b> {action_id}",);98 if let Some(subject) = details.remove("polkit.caller-pid") {99 let _ = write!(description, "\n<b>Caller:</b> ");100 if let Ok(pid) = subject.parse::<u32>() {101 let _ = write!(description, "{}", PidDisplay(pid));102 } else {103 let _ = write!(description, "{}", emphasize("invalid pid"));104 }105 }106 if let Some(subject) = details.remove("polkit.subject-pid") {107 let _ = write!(description, "\n<b>Subject:</b> ");108 if let Ok(pid) = subject.parse::<u32>() {109 let _ = write!(description, "{}", PidDisplay(pid));110 } else {111 let _ = write!(description, "{}", emphasize("invalid pid"));112 }113 }114 let mut prompter = PrependSourcePrompter {115 source: vec![Source(Cow::Borrowed("polkit agent"))],116 description: description.clone(),117 prompter,118 };119120 let identity_displays: Vec<String> =121 identities.iter().map(|v| v.to_string()).collect();122 let identity_displays: Vec<&str> =123 identity_displays.iter().map(|v| v.as_str()).collect();124 info!("choose identity");125 let choosen_identity = match identity_displays.len() {126 0 => {127 return Err(fdo::Error::AuthFailed(128 "no identity to authenticate as".to_owned(),129 ))130 }131 1 => 0,132 _ => {133 prompter134 .prompt_enum(135 "Identity",136 "Select identity to use for polkit authorization",137 &identity_displays,138 &[],139 )140 .await?141 }142 };143 info!("identity chosen");144145 let _ = write!(146 description,147 "\n<b>Identity:</b> {}",148 identities[choosen_identity as usize]149 );150 prompter.description = description;151152 prompter.source.push(Source(Cow::Borrowed("polkit daemon")));153154 helper155 .help_me(156 &cookie,157 prompter,158 identities[choosen_identity as usize].clone(),159 )160 .await161 .map_err(|e| fdo::Error::Failed(e.to_string()))?;162 163 164165 Ok(())166 })167 };168 self.tasks169 .lock()170 .unwrap()171 .insert(cookie.clone(), task.abort_handle());172 info!("abort handle stored");173 let _ = _cancel_guard.set(CancelTaskOnDrop {174 tasks: self.tasks.clone(),175 handle: cookie.clone(),176 });177178 let _ = task.await;179180 Ok(())181 }182183 184 async fn cancel_authentication(&self, cookie: &str) -> zbus::fdo::Result<()> {185 info!("auth cancelled");186 if let Some(abort) = self.tasks.lock().unwrap().remove(cookie) {187 info!("abort handle found");188 abort.abort();189 }190 191 Ok(())192 }193}194195const OBJ_PATH: &str = "/org/freedesktop/PolicyKit1/AuthenticationAgent";196197#[proxy(198 interface = "lach.PolkitHelper",199 default_service = "lach.polkit.helper1",200 default_path = "/lach/PolkitHelper"201)]202trait PolkitHelper {203 fn init_conversation(&self, request: BackendRequest) -> zbus::Result<()>;204}205206#[derive(Parser)]207enum Opts {208 AskPass {209 prompt: String,210 description: String,211 },212 Editor {213 214 path: String,215 },216 RealAgent {217 #[arg(long)]218 path: Option<PathBuf>,219 220 #[arg(long)]221 privileged: bool,222 },223}224225fn main() -> anyhow::Result<()> {226 227 228 tracing_subscriber::fmt()229 .with_writer(std::io::stderr)230 .init();231 let opts = Opts::parse();232233 let runtime = Builder::new_current_thread().enable_all().build()?;234235 match opts {236 Opts::AskPass {237 prompt,238 description,239 } => runtime.block_on(askpass::ask(&prompt, description)),240 Opts::Editor { path } => runtime.block_on(editor::edit(path)),241 Opts::RealAgent { path, privileged } => runtime.block_on(main_real_agent(path, privileged)),242 }243}244async fn main_real_agent(path: Option<PathBuf>, privileged: bool) -> anyhow::Result<()> {245 let address = if privileged {246 Address::AgentPrivileged247 } else {248 Address::Agent249 };250 let mut rpc = Rpc::<BifConfig>::new(address);251252 Fs::new().register_endpoints(&mut rpc);253 Systemd.register_endpoints(&mut rpc);254 Pty::new().register_endpoints(&mut rpc);255256 remowt_plugin::host::serve(&mut rpc);257258 let user_prompter = PromptEndpointsClient::wrap(rpc.remote(Address::User));259 let editor_client = EditorEndpointsClient::wrap(rpc.remote(Address::User));260261 let bus = bus::spawn().await?;262 askpass::serve(&bus.conn, user_prompter.clone()).await?;263 editor::serve(&bus.conn, editor_client).await?;264265 let helpers = tempfile::Builder::new().prefix("remowt-path.").tempdir()?;266 let exe = std::env::current_exe()?;267 let askpass_helper = helpers.path().join("remowt-askpass");268 let editor_helper = helpers.path().join("remowt-editor");269 {270 let script = format!(271 "#!/bin/sh\nexec {} ask-pass \"password\" \"$1\"\n",272 sh_quote(&exe.to_string_lossy())273 );274 fs::write(&askpass_helper, script).await?;275 fs::set_permissions(&askpass_helper, Permissions::from_mode(0o755)).await?;276 }277 {278 let script = format!(279 "#!/bin/sh\nexec {} editor \"$1\"\n",280 sh_quote(&exe.to_string_lossy())281 );282 fs::write(&editor_helper, script).await?;283 fs::set_permissions(&editor_helper, Permissions::from_mode(0o755)).await?;284 }285286 287 unsafe {288 prepend_path(helpers.path());289 std::env::set_var("SUDO_ASKPASS", &askpass_helper);290 std::env::set_var("SSH_ASKPASS", &askpass_helper);291 std::env::set_var("SSH_ASKPASS_REQUIRE", "force");292 std::env::set_var("EDITOR", &editor_helper);293 std::env::set_var("VISUAL", &editor_helper);294 std::env::set_var("DBUS_SESSION_BUS_ADDRESS", &bus.address);295 }296297 let port = match path {298 Some(path) => from_socket(UnixStream::connect(path).await?),299 None => from_stdio(),300 };301 rpc.add_direct(Address::User, port, bifrostlink::Rtt(0));302303 let polkit_conn = if !privileged {304 305 306 307 let conn = Connection::system().await?;308 let helper = SocketHelper {309 fallback: SuidHelper,310 };311 register_auth_agent(&conn, Agent::new(helper, user_prompter)).await?;312 Some(conn)313 } else {314 None315 };316317 let _keep_alive = (bus, helpers, polkit_conn);318 pending().await319}320321async fn register_auth_agent<H, P>(conn: &Connection, agent: Agent<H, P>) -> anyhow::Result<()>322where323 H: Helper + Clone + Send + Sync + 'static,324 P: Prompter + Clone + Send + Sync + 'static,325{326 let proxy = zbus_polkit::policykit1::AuthorityProxy::new(conn).await?;327 conn.object_server().at(OBJ_PATH, agent).await?;328329 let subject = auth_agent_subject()?;330 proxy331 .register_authentication_agent(&subject, "C", OBJ_PATH)332 .await?;333 info!(kind = subject.subject_kind, "registered polkit agent");334 Ok(())335}336337fn auth_agent_subject() -> anyhow::Result<Subject> {338 let mut details = HashMap::new();339 if let Ok(session_id) = std::env::var("XDG_SESSION_ID") {340 let val: OwnedValue = Str::from(session_id).into();341 details.insert("session-id".to_string(), val);342 return Ok(Subject {343 subject_kind: "unix-session".to_string(),344 subject_details: details,345 });346 }347348 details.insert("pid".to_string(), OwnedValue::from(std::process::id()));349 Ok(Subject {350 subject_kind: "unix-process".to_string(),351 subject_details: details,352 })353}354355fn sh_quote(s: &str) -> String {356 format!("'{}'", s.replace('\'', "'\\''"))357}358359360361362363364unsafe fn prepend_path(dir: &std::path::Path) {365 let value = match std::env::var_os("PATH") {366 Some(existing) => {367 let mut v = dir.as_os_str().to_owned();368 v.push(":");369 v.push(existing);370 v371 }372 None => dir.as_os_str().to_owned(),373 };374 unsafe {375 std::env::set_var("PATH", value);376 }377}