1use std::borrow::Cow;2use std::collections::{BTreeMap, HashMap};3use std::future;4use std::io::{stdout, Write};5use std::path::PathBuf;6use std::sync::{Arc, Mutex, OnceLock};78use bifrostlink::{AddressT, Rpc};9use bifrostlink_ports::unix_socket::from_socket;10use clap::Parser;11use polkit_shared::{emphasize, BackendRequest, Identity, PidDisplay};12use remowt_link_shared::Address;13use tokio::io::{AsyncReadExt, AsyncWriteExt};14use tokio::net::UnixStream;15use tokio::runtime::Runtime;16use tokio::task::AbortHandle;17use tracing::{info, trace};18use ui_prompt::rofi::RofiPrompter;19use ui_prompt::{PrependSourcePrompter, Prompter, Source};20use zbus::fdo;21use zbus::zvariant::{OwnedValue, Str};22use zbus::{interface, proxy, Connection};23use zbus_polkit::policykit1::Subject;2425use self::helper::{Helper, SuidHelper};2627pub mod helper;2829struct CancelTaskOnDrop {30 tasks: Arc<Mutex<HashMap<String, AbortHandle>>>,31 handle: String,32}33impl Drop for CancelTaskOnDrop {34 fn drop(&mut self) {35 info!("cancel on drop");36 if let Some(task) = self37 .tasks38 .lock()39 .expect("not poisoned")40 .remove(&self.handle)41 {42 task.abort();43 }44 }45}4647struct Agent<H> {48 tasks: Arc<Mutex<HashMap<String, AbortHandle>>>,49 helper: H,50}51impl<H> Agent<H> {52 fn new(helper: H) -> Self {53 Agent {54 tasks: Arc::new(Mutex::new(HashMap::new())),55 helper,56 }57 }58}5960#[interface(name = "org.freedesktop.PolicyKit1.AuthenticationAgent")]61impl<H> Agent<H>62where63 H: Helper + Clone + Send + Sync + 'static,64{65 66 #[allow(clippy::too_many_arguments)]67 async fn begin_authentication(68 &self,69 action_id: String,70 message: String,71 icon_name: String,72 mut details: BTreeMap<String, String>,73 cookie: String,74 identities: Vec<Identity>,75 ) -> zbus::fdo::Result<()> {76 use std::fmt::Write;77 info!("begin auth");78 let _cancel_guard = Arc::new(OnceLock::new());79 let task = {80 let helper = self.helper.clone();81 let cookie = cookie.clone();82 let _cancel_guard = _cancel_guard.clone();83 tokio::task::spawn(async move {84 let _cancel_guard = _cancel_guard.clone();85 trace!("conversation task");86 let mut description = format!("{message}\n\n<b>Action id:</b> {action_id}",);87 if let Some(subject) = details.remove("polkit.caller-pid") {88 let _ = write!(description, "\n<b>Caller:</b> ");89 if let Ok(pid) = subject.parse::<u32>() {90 let _ = write!(description, "{}", PidDisplay(pid));91 } else {92 let _ = write!(description, "{}", emphasize("invalid pid"));93 }94 }95 if let Some(subject) = details.remove("polkit.subject-pid") {96 let _ = write!(description, "\n<b>Subject:</b> ");97 if let Ok(pid) = subject.parse::<u32>() {98 let _ = write!(description, "{}", PidDisplay(pid));99 } else {100 let _ = write!(description, "{}", emphasize("invalid pid"));101 }102 }103 let mut prompter = PrependSourcePrompter {104 source: vec![Source(Cow::Borrowed("polkit agent"))],105 description: description.clone(),106 prompter: RofiPrompter,107 };108109 let identity_displays: Vec<String> =110 identities.iter().map(|v| v.to_string()).collect();111 let identity_displays: Vec<&str> =112 identity_displays.iter().map(|v| v.as_str()).collect();113 info!("choose identity");114 let choosen_identity = match identity_displays.len() {115 0 => {116 return Err(fdo::Error::AuthFailed(117 "no identity to authenticate as".to_owned(),118 ))119 }120 1 => 0,121 _ => {122 prompter123 .prompt_enum(124 "Identity",125 "Select identity to use for polkit authorization",126 &identity_displays,127 &[],128 )129 .await?130 }131 };132 info!("identity chosen");133134 let _ = write!(135 description,136 "\n<b>Identity:</b> {}",137 identities[choosen_identity as usize]138 );139 prompter.description = description;140141 prompter.source.push(Source(Cow::Borrowed("polkit daemon")));142143 helper144 .help_me(145 &cookie,146 prompter,147 identities[choosen_identity as usize].clone(),148 )149 .await150 .map_err(|e| fdo::Error::Failed(e.to_string()))?;151 152 153154 Ok(())155 })156 };157 self.tasks158 .lock()159 .unwrap()160 .insert(cookie.clone(), task.abort_handle());161 info!("abort handle stored");162 let _ = _cancel_guard.set(CancelTaskOnDrop {163 tasks: self.tasks.clone(),164 handle: cookie.clone(),165 });166167 let _ = task.await;168169 Ok(())170 }171172 173 async fn cancel_authentication(&self, cookie: &str) -> zbus::fdo::Result<()> {174 info!("auth cancelled");175 if let Some(abort) = self.tasks.lock().unwrap().remove(cookie) {176 info!("abort handle found");177 abort.abort();178 }179 180 Ok(())181 }182}183184const OBJ_PATH: &str = "/org/freedesktop/PolicyKit1/AuthenticationAgent";185186#[proxy(187 interface = "lach.PolkitHelper",188 default_service = "lach.polkit.helper1",189 default_path = "/lach/PolkitHelper"190)]191trait PolkitHelper {192 fn init_conversation(&self, request: BackendRequest) -> zbus::Result<()>;193}194195#[derive(Parser)]196enum Opts {197 Agent,198 AskPass {199 description: String,200 },201 RealAgent {202 #[arg(long)]203 path: PathBuf,204 },205}206207fn main() -> anyhow::Result<()> {208 tracing_subscriber::fmt::init();209 let opts = Opts::parse();210211 let runtime = Runtime::new()?;212213 match opts {214 Opts::Agent => {215 216 runtime.block_on(main_agent())217 }218 Opts::AskPass { description } => runtime.block_on(main_askpass(description)),219 Opts::RealAgent { path } => runtime.block_on(main_real_agent(path)),220 }221}222async fn main_real_agent(path: PathBuf) -> anyhow::Result<()> {223 let mut stream = UnixStream::connect(path).await?;224 stream.write_all(b"REMOWT_HELLO\0").await?;225 let mut buf = [0u8; 12];226 stream.read_exact(&mut buf).await?;227 assert_eq!(&buf, b"REMOWT_EHLO\0");228 let port = from_socket(stream);229 let rpc = Rpc::<Address, remowt_link_shared::Error>::new(Address::Agent);230 rpc.add_direct(Address::User, port, bifrostlink::Rtt(0));231 Ok(())232}233234async fn main_agent() -> anyhow::Result<()> {235 trace!("started");236 let conn = Connection::system().await?;237238 let proxy = zbus_polkit::policykit1::AuthorityProxy::new(&conn).await?;239 conn.object_server()240 .at(OBJ_PATH, Agent::new(SuidHelper))241 .await?;242243 let session_id = std::env::var("XDG_SESSION_ID")?;244 let mut details = HashMap::new();245 let val: OwnedValue = {246 let wrapped: Str<'_> = session_id.into();247 wrapped.into()248 };249 details.insert("session-id".to_string(), val);250 proxy251 .register_authentication_agent(252 &Subject {253 subject_kind: "unix-session".to_string(),254 subject_details: details,255 },256 "C",257 OBJ_PATH,258 )259 .await?;260 future::pending().await261}262263async fn main_askpass(description: String) -> anyhow::Result<()> {264 let password = RofiPrompter265 .prompt_text(false, &description, "SSH password request", &[])266 .await?;267 stdout().lock().write_all(password.as_bytes())?;268 future::pending().await269}