git.delta.rocks / remowt / refs/commits / 5cb6be498b69

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        // FIXME: block_in_place prevents to moving to current_thread runtime46        // There should be a blocking way to remove ObjectServer listener.47        // As far as I can see, it is only async because of async RwLock, shouldn't it be48        // just a sync lock?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    /// BeginAuthentication method97    #[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                // let connection = Connection::system().await?;175                // let helper = PolkitHelperProxy::new(&connection).await?;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                            // TODO: Let user choose185                            identity: identities[choosen_identity as usize].clone(),186                        }, // cookie.to_owned(), HashMap::new(), prompter.path.clone()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    /// CancelAuthentication method211    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        // debug!("Authentication cancled ! {cookie}");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}