git.delta.rocks / remowt / refs/commits / d819a0103251

1use std::borrow::Cow;2use std::collections::{BTreeMap, HashMap};3use std::io::{stdout, Write};4use std::marker::PhantomData;5use std::sync::{Mutex, RwLock};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::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 Agent {62    helper: PolkitHelperProxy<'static>,63    tasks: Mutex<HashMap<String, AbortHandle>>,64    connection: Connection,65}66impl Agent {67    async fn new(connection: Connection) -> anyhow::Result<Self> {68        Ok(Self {69            helper: PolkitHelperProxy::new(&connection).await?,70            tasks: Mutex::new(HashMap::new()),71            connection,72        })73    }74}7576#[interface(name = "org.freedesktop.PolicyKit1.AuthenticationAgent")]77impl Agent {78    /// BeginAuthentication method79    #[allow(clippy::too_many_arguments)]80    async fn begin_authentication(81        &mut self,82        action_id: String,83        message: String,84        icon_name: String,85        mut details: BTreeMap<String, String>,86        cookie: String,87        identities: Vec<Identity>,88    ) -> zbus::fdo::Result<()> {89        use std::fmt::Write;90        trace!("begin auth");91        let task = {92            let connection = self.connection.clone();93            let helper = self.helper.clone();94            let cookie = cookie.clone();95            tokio::task::spawn(async move {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: RofiPrompter,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                let choosen_identity = match identity_displays.len() {125                    0 => {126                        return Err(fdo::Error::AuthFailed(127                            "no identity to authenticate as".to_owned(),128                        ))129                    }130                    1 => 0,131                    _ => {132                        prompter133                            .prompt_enum(134                                "Identity",135                                "Select identity to use for polkit authorization",136                                &identity_displays,137                                &[],138                            )139                            .await?140                    }141                };142143                let _ = write!(144                    description,145                    "\n<b>Identity:</b> {}",146                    identities[choosen_identity as usize]147                );148                prompter.description = description;149150                prompter.source.push(Source(Cow::Borrowed("polkit daemon")));151                let prompter = TemporaryPrompterInterface::new(connection, prompter).await;152                helper153                    .init_conversation(154                        BackendRequest {155                            cookie: cookie.to_owned(),156                            environment: HashMap::new(),157                            prompter_path: prompter.path.clone(),158                            // TODO: Let user choose159                            identity: identities[choosen_identity as usize].clone(),160                        }, // cookie.to_owned(), HashMap::new(), prompter.path.clone()161                    )162                    .await?;163                println!("ASKED");164                dbg!(action_id, message, icon_name, details, cookie, identities);165166                Ok(())167            })168        };169170        self.tasks171            .lock()172            .unwrap()173            .insert(cookie.clone(), task.abort_handle());174        let result = task.await.expect("join error");175        // The only way to no reach this line, is to either panic in previous line, or if authorization cancelled,176        // while cancellation will remove task by itself.177        // TODO: But still it would be better to have abort guard, which will remove it from HashMap178        self.tasks.lock().unwrap().remove(&cookie);179180        result181    }182183    /// CancelAuthentication method184    async fn cancel_authentication(&self, cookie: &str) -> zbus::fdo::Result<()> {185        trace!("cancel auth");186        if let Some(abort) = self.tasks.lock().unwrap().remove(cookie) {187            abort.abort();188        }189        // debug!("Authentication cancled ! {cookie}");190        Ok(())191    }192}193194const OBJ_PATH: &str = "/org/freedesktop/PolicyKit1/AuthenticationAgent";195196#[proxy(197    interface = "lach.PolkitHelper",198    default_service = "lach.polkit.helper1",199    default_path = "/lach/PolkitHelper"200)]201trait PolkitHelper {202    fn init_conversation(&self, request: BackendRequest) -> zbus::Result<()>;203}204205#[derive(Parser)]206enum Opts {207    Agent,208    AskPass { description: String },209}210211#[tokio::main]212async fn main() -> anyhow::Result<()> {213    tracing_subscriber::fmt::init();214    let opts = Opts::parse();215216    match opts {217        Opts::Agent => {218            trace!("started");219            let conn = Connection::system().await?;220221            let proxy = zbus_polkit::policykit1::AuthorityProxy::new(&conn).await?;222            conn.object_server()223                .at(OBJ_PATH, Agent::new(conn.clone()).await?)224                .await?;225226            let session_id = std::env::var("XDG_SESSION_ID")?;227            let mut details = HashMap::new();228            let val: OwnedValue = {229                let wrapped: Str<'_> = session_id.into();230                wrapped.into()231            };232            details.insert("session-id".to_string(), val);233            proxy234                .register_authentication_agent(235                    &Subject {236                        subject_kind: "unix-session".to_string(),237                        subject_details: details,238                    },239                    "C",240                    OBJ_PATH,241                )242                .await?;243        }244        Opts::AskPass { description } => {245            let password = RofiPrompter246                .prompt_text(false, &description, "SSH password request", &[])247                .await?;248            stdout().lock().write_all(password.as_bytes())?;249        }250    }251252    future::pending().await253}