git.delta.rocks / remowt / refs/commits / 2499daa8100a

1use std::collections::HashMap;2use std::io::{stdout, Write};3use std::marker::PhantomData;4use std::sync::{Mutex, RwLock};5use std::{future, process};67use clap::Parser;8use polkit_shared::{BackendRequest, Identity};9use tokio::runtime::Handle;10use tokio::task::{AbortHandle, JoinHandle, LocalSet};11use tracing::trace;12use ui_prompt::dbus::DbusPrompterInterface;13use ui_prompt::rofi::RofiPrompter;14use ui_prompt::Prompter;15use zbus::zvariant::{OwnedValue, Str};16use zbus::ObjectServer;17use zbus::{interface, proxy, Connection};18use zbus_polkit::policykit1::Subject;1920struct TemporaryPrompterInterface<P: Prompter + Send + Sync + 'static> {21    connection: Connection,22    path: String,23    _marker: PhantomData<P>,24}25impl<P: Prompter + Send + Sync + 'static> TemporaryPrompterInterface<P> {26    async fn new(connection: Connection, prompter: P) -> Self {27        let path = format!(28            "/remowt/prompters/{}",29            uuid::Uuid::new_v4().to_string().replace("-", "_")30        );31        let _ = connection32            .object_server()33            .at(path.clone(), DbusPrompterInterface(prompter))34            .await;35        Self {36            connection,37            path,38            _marker: PhantomData,39        }40    }41}42impl<P: Prompter + Send + Sync + 'static> Drop for TemporaryPrompterInterface<P> {43    fn drop(&mut self) {44        // FIXME: block_in_place prevents to moving to current_thread runtime45        // There should be a blocking way to remove ObjectServer listener.46        // As far as I can see, it is only async because of async RwLock, shouldn't it be47        // just a sync lock?48        tokio::task::block_in_place(move || {49            Handle::current().block_on(async {50                let _ = self51                    .connection52                    .object_server()53                    .remove::<DbusPrompterInterface<P>, String>(self.path.clone())54                    .await;55            });56        });57    }58}5960struct Agent {61    helper: PolkitHelperProxy<'static>,62    tasks: Mutex<HashMap<String, AbortHandle>>,63    connection: Connection,64}65impl Agent {66    async fn new(connection: Connection) -> anyhow::Result<Self> {67        Ok(Self {68            helper: PolkitHelperProxy::new(&connection).await?,69            tasks: Mutex::new(HashMap::new()),70            connection,71        })72    }73}7475#[interface(name = "org.freedesktop.PolicyKit1.AuthenticationAgent")]76impl Agent {77    /// BeginAuthentication method78    #[allow(clippy::too_many_arguments)]79    async fn begin_authentication(80        &mut self,81        action_id: String,82        message: String,83        icon_name: String,84        details: HashMap<String, String>,85        cookie: String,86        identities: Vec<Identity>,87    ) -> zbus::fdo::Result<()> {88        trace!("begin auth");89        let task = {90            let connection = self.connection.clone();91            let helper = self.helper.clone();92            let cookie = cookie.clone();93            tokio::task::spawn(async move {94                trace!("conversation task");95                let prompter = TemporaryPrompterInterface::new(connection, RofiPrompter).await;96                helper97                    .init_conversation(98                        BackendRequest {99                            cookie: cookie.to_owned(),100                            environment: HashMap::new(),101                            prompter_path: prompter.path.clone(),102                            // TODO: Let user choose103                            identity: identities.get(0).expect("first always exists").clone(),104                        }, // cookie.to_owned(), HashMap::new(), prompter.path.clone()105                    )106                    .await?;107                println!("ASKED");108                dbg!(action_id, message, icon_name, details, cookie, identities);109110                Ok(())111            })112        };113114        self.tasks115            .lock()116            .unwrap()117            .insert(cookie.clone(), task.abort_handle());118        let result = task.await.expect("join error");119        // The only way to no reach this line, is to either panic in previous line, or if authorization cancelled,120        // while cancellation will remove task by itself.121        // TODO: But still it would be better to have abort guard, which will remove it from HashMap122        self.tasks.lock().unwrap().remove(&cookie);123124        result125    }126127    /// CancelAuthentication method128    async fn cancel_authentication(&self, cookie: &str) -> zbus::fdo::Result<()> {129        trace!("cancel auth");130        if let Some(abort) = self.tasks.lock().unwrap().remove(cookie) {131            abort.abort();132        }133        // debug!("Authentication cancled ! {cookie}");134        Ok(())135    }136}137138const OBJ_PATH: &str = "/0lach/polkitAgent";139140#[proxy(141    interface = "lach.PolkitHelper",142    default_service = "lach.polkit.helper1",143    default_path = "/lach/PolkitHelper"144)]145trait PolkitHelper {146    fn init_conversation(&self, request: BackendRequest) -> zbus::Result<()>;147}148149#[derive(Parser)]150enum Opts {151    Agent,152    AskPass { description: String },153}154155#[tokio::main]156async fn main() -> anyhow::Result<()> {157    tracing_subscriber::fmt::init();158    let opts = Opts::parse();159160    match opts {161        Opts::Agent => {162            trace!("started");163            let conn = Connection::system().await?;164165            let proxy = zbus_polkit::policykit1::AuthorityProxy::new(&conn).await?;166            conn.object_server()167                .at(OBJ_PATH, Agent::new(conn.clone()).await?)168                .await?;169170            let session_id = std::env::var("XDG_SESSION_ID")?;171            let mut details = HashMap::new();172            let val: OwnedValue = {173                let wrapped: Str<'_> = session_id.into();174                wrapped.into()175            };176            details.insert("session-id".to_string(), val);177            proxy178                .register_authentication_agent(179                    &Subject {180                        subject_kind: "unix-session".to_string(),181                        subject_details: details,182                    },183                    "C",184                    OBJ_PATH,185                )186                .await?;187        }188        Opts::AskPass { description } => {189            let password = RofiPrompter190                .prompt_text(false, &description, "SSH password request", &[])191                .await?;192            stdout().lock().write_all(password.as_bytes())?;193        }194    }195196    future::pending().await197}