1use std::borrow::Cow;2use std::collections::{BTreeMap, HashMap};3use std::io::{stdout, Write};4use std::marker::PhantomData;5use std::sync::{Arc, Mutex, OnceLock};6use std::{future, process};78use clap::Parser;9use polkit_shared::{emphasize, BackendRequest, Identity, PidDisplay};10use tokio::runtime::Handle;11use tokio::task::{AbortHandle, JoinHandle, LocalSet};12use tracing::{info, trace};13use ui_prompt::dbus::DbusPrompterInterface;14use ui_prompt::rofi::RofiPrompter;15use ui_prompt::{PrependSourcePrompter, Prompter, Source};16use zbus::zvariant::{OwnedValue, Str};17use zbus::{fdo, ObjectServer};18use zbus::{interface, proxy, Connection};19use zbus_polkit::policykit1::Subject;2021struct TemporaryPrompterInterface<P: Prompter + Send + Sync + 'static> {22 connection: Connection,23 path: String,24 _marker: PhantomData<P>,25}26impl<P: Prompter + Send + Sync + 'static> TemporaryPrompterInterface<P> {27 async fn new(connection: Connection, prompter: P) -> Self {28 let path = format!(29 "/remowt/prompters/{}",30 uuid::Uuid::new_v4().to_string().replace("-", "_")31 );32 let _ = connection33 .object_server()34 .at(path.clone(), DbusPrompterInterface(prompter))35 .await;36 Self {37 connection,38 path,39 _marker: PhantomData,40 }41 }42}43impl<P: Prompter + Send + Sync + 'static> Drop for TemporaryPrompterInterface<P> {44 fn drop(&mut self) {45 46 47 48 49 tokio::task::block_in_place(move || {50 Handle::current().block_on(async {51 let _ = self52 .connection53 .object_server()54 .remove::<DbusPrompterInterface<P>, String>(self.path.clone())55 .await;56 });57 });58 }59}6061struct CancelTaskOnDrop {62 tasks: Arc<Mutex<HashMap<String, AbortHandle>>>,63 handle: String,64}65impl Drop for CancelTaskOnDrop {66 fn drop(&mut self) {67 info!("cancel on drop");68 if let Some(task) = self69 .tasks70 .lock()71 .expect("not poisoned")72 .remove(&self.handle)73 {74 task.abort();75 }76 }77}7879struct Agent {80 helper: PolkitHelperProxy<'static>,81 tasks: Arc<Mutex<HashMap<String, AbortHandle>>>,82 connection: Connection,83}84impl Agent {85 async fn new(connection: Connection) -> anyhow::Result<Self> {86 Ok(Self {87 helper: PolkitHelperProxy::new(&connection).await?,88 tasks: Arc::new(Mutex::new(HashMap::new())),89 connection,90 })91 }92}9394#[interface(name = "org.freedesktop.PolicyKit1.AuthenticationAgent")]95impl Agent {96 97 #[allow(clippy::too_many_arguments)]98 async fn begin_authentication(99 &self,100 action_id: String,101 message: String,102 icon_name: String,103 mut details: BTreeMap<String, String>,104 cookie: String,105 identities: Vec<Identity>,106 ) -> zbus::fdo::Result<()> {107 use std::fmt::Write;108 info!("begin auth");109 let _cancel_guard = Arc::new(OnceLock::new());110 let task = {111 let connection = self.connection.clone();112 let helper = self.helper.clone();113 let cookie = cookie.clone();114 let _cancel_guard = _cancel_guard.clone();115 tokio::task::spawn(async move {116 let _cancel_guard = _cancel_guard.clone();117 trace!("conversation task");118 let mut description = format!("{message}\n\n<b>Action id:</b> {action_id}",);119 if let Some(subject) = details.remove("polkit.caller-pid") {120 let _ = write!(description, "\n<b>Caller:</b> ");121 if let Ok(pid) = subject.parse::<u32>() {122 let _ = write!(description, "{}", PidDisplay(pid));123 } else {124 let _ = write!(description, "{}", emphasize("invalid pid"));125 }126 }127 if let Some(subject) = details.remove("polkit.subject-pid") {128 let _ = write!(description, "\n<b>Subject:</b> ");129 if let Ok(pid) = subject.parse::<u32>() {130 let _ = write!(description, "{}", PidDisplay(pid));131 } else {132 let _ = write!(description, "{}", emphasize("invalid pid"));133 }134 }135 let mut prompter = PrependSourcePrompter {136 source: vec![Source(Cow::Borrowed("polkit agent"))],137 description: description.clone(),138 prompter: RofiPrompter,139 };140141 let identity_displays: Vec<String> =142 identities.iter().map(|v| v.to_string()).collect();143 let identity_displays: Vec<&str> =144 identity_displays.iter().map(|v| v.as_str()).collect();145 info!("choose identity");146 let choosen_identity = match identity_displays.len() {147 0 => {148 return Err(fdo::Error::AuthFailed(149 "no identity to authenticate as".to_owned(),150 ))151 }152 1 => 0,153 _ => {154 prompter155 .prompt_enum(156 "Identity",157 "Select identity to use for polkit authorization",158 &identity_displays,159 &[],160 )161 .await?162 }163 };164 info!("identity chosen");165166 let _ = write!(167 description,168 "\n<b>Identity:</b> {}",169 identities[choosen_identity as usize]170 );171 prompter.description = description;172173 prompter.source.push(Source(Cow::Borrowed("polkit daemon")));174 175 176 let prompter = TemporaryPrompterInterface::new(connection, prompter).await;177 info!("init conv");178 helper179 .init_conversation(180 BackendRequest {181 cookie: cookie.to_owned(),182 environment: HashMap::new(),183 prompter_path: prompter.path.clone(),184 185 identity: identities[choosen_identity as usize].clone(),186 }, 187 )188 .await?;189 println!("ASKED");190 dbg!(action_id, message, icon_name, details, cookie, identities);191192 Ok(())193 })194 };195 self.tasks196 .lock()197 .unwrap()198 .insert(cookie.clone(), task.abort_handle());199 info!("abort handle stored");200 let _ = _cancel_guard.set(CancelTaskOnDrop {201 tasks: self.tasks.clone(),202 handle: cookie.clone(),203 });204205 let _ = task.await;206207 Ok(())208 }209210 211 async fn cancel_authentication(&self, cookie: &str) -> zbus::fdo::Result<()> {212 info!("auth cancelled");213 if let Some(abort) = self.tasks.lock().unwrap().remove(cookie) {214 info!("abort handle found");215 abort.abort();216 }217 218 Ok(())219 }220}221222const OBJ_PATH: &str = "/org/freedesktop/PolicyKit1/AuthenticationAgent";223224#[proxy(225 interface = "lach.PolkitHelper",226 default_service = "lach.polkit.helper1",227 default_path = "/lach/PolkitHelper"228)]229trait PolkitHelper {230 fn init_conversation(&self, request: BackendRequest) -> zbus::Result<()>;231}232233#[derive(Parser)]234enum Opts {235 Agent,236 AskPass { description: String },237}238239#[tokio::main]240async fn main() -> anyhow::Result<()> {241 tracing_subscriber::fmt::init();242 let opts = Opts::parse();243244 match opts {245 Opts::Agent => {246 trace!("started");247 let conn = Connection::system().await?;248249 let proxy = zbus_polkit::policykit1::AuthorityProxy::new(&conn).await?;250 conn.object_server()251 .at(OBJ_PATH, Agent::new(conn.clone()).await?)252 .await?;253254 let session_id = std::env::var("XDG_SESSION_ID")?;255 let mut details = HashMap::new();256 let val: OwnedValue = {257 let wrapped: Str<'_> = session_id.into();258 wrapped.into()259 };260 details.insert("session-id".to_string(), val);261 proxy262 .register_authentication_agent(263 &Subject {264 subject_kind: "unix-session".to_string(),265 subject_details: details,266 },267 "C",268 OBJ_PATH,269 )270 .await?;271 }272 Opts::AskPass { description } => {273 let password = RofiPrompter274 .prompt_text(false, &description, "SSH password request", &[])275 .await?;276 stdout().lock().write_all(password.as_bytes())?;277 }278 }279280 future::pending().await281}