1
use std::{collections::HashMap, sync::Arc};
2

            
3
use base64::Engine;
4
use oo7::{Secret, crypto, dbus};
5
use zbus::zvariant::{ObjectPath, Optional, Value};
6

            
7
use crate::{
8
    gnome::{
9
        prompter::{PromptType, Properties, Reply},
10
        secret_exchange,
11
    },
12
    service::Service,
13
};
14

            
15
/// Helper to create a peer-to-peer connection pair using Unix socket
16
7
async fn create_p2p_connection()
17
-> Result<(zbus::Connection, zbus::Connection), Box<dyn std::error::Error>> {
18
8
    let guid = zbus::Guid::generate();
19
15
    let (p0, p1) = tokio::net::UnixStream::pair()?;
20

            
21
24
    let (client_conn, server_conn) = tokio::try_join!(
22
        // Client
23
12
        zbus::connection::Builder::unix_stream(p0).p2p().build(),
24
        // Server
25
16
        zbus::connection::Builder::unix_stream(p1)
26
8
            .server(guid)?
27
8
            .p2p()
28
8
            .build(),
29
    )?;
30

            
31
2
    Ok((server_conn, client_conn))
32
}
33

            
34
pub(crate) struct TestServiceSetup {
35
    pub server: Service,
36
    pub client_conn: zbus::Connection,
37
    pub service_api: dbus::api::Service,
38
    pub session: Arc<dbus::api::Session>,
39
    pub collections: Vec<dbus::api::Collection>,
40
    pub server_public_key: Option<oo7::Key>,
41
    pub keyring_secret: Option<oo7::Secret>,
42
    pub aes_key: Option<Arc<oo7::Key>>,
43
    pub mock_prompter: MockPrompterService,
44
}
45

            
46
impl TestServiceSetup {
47
    /// Get the default/Login collection
48
2
    pub(crate) async fn default_collection(
49
        &self,
50
    ) -> Result<&dbus::api::Collection, Box<dyn std::error::Error>> {
51
8
        for collection in &self.collections {
52
8
            let label = collection.label().await?;
53
4
            if label == "Login" {
54
2
                return Ok(collection);
55
            }
56
        }
57
        Err("Default collection not found".into())
58
    }
59

            
60
6
    pub(crate) async fn plain_session(
61
        with_default_collection: bool,
62
    ) -> Result<TestServiceSetup, Box<dyn std::error::Error>> {
63
13
        let (server_conn, client_conn) = create_p2p_connection().await?;
64

            
65
11
        let secret = if with_default_collection {
66
7
            Some(Secret::from("test-password-long-enough"))
67
        } else {
68
2
            None
69
        };
70

            
71
12
        let server = Service::run_with_connection(server_conn.clone(), secret.clone()).await?;
72

            
73
        // Create and serve the mock prompter
74
4
        let mock_prompter = MockPrompterService::new();
75
11
        client_conn
76
            .object_server()
77
3
            .at("/org/gnome/keyring/Prompter", mock_prompter.clone())
78
9
            .await?;
79

            
80
        // Give the server a moment to fully initialize
81
6
        tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
82

            
83
3
        let service_api = dbus::api::Service::new(&client_conn).await?;
84

            
85
9
        let (server_public_key, session) = service_api.open_session(None).await?;
86
6
        let session = Arc::new(session);
87

            
88
8
        let collections = service_api.collections().await?;
89

            
90
3
        Ok(TestServiceSetup {
91
3
            server,
92
3
            keyring_secret: secret,
93
3
            client_conn,
94
3
            service_api,
95
3
            session,
96
            collections,
97
2
            server_public_key,
98
            aes_key: None,
99
3
            mock_prompter,
100
        })
101
    }
102

            
103
2
    pub(crate) async fn encrypted_session(
104
        with_default_collection: bool,
105
    ) -> Result<TestServiceSetup, Box<dyn std::error::Error>> {
106
6
        let (server_conn, client_conn) = create_p2p_connection().await?;
107

            
108
6
        let secret = if with_default_collection {
109
4
            Some(Secret::from("test-password-long-enough"))
110
        } else {
111
2
            None
112
        };
113

            
114
6
        let server = Service::run_with_connection(server_conn.clone(), secret.clone()).await?;
115

            
116
        // Create and serve the mock prompter
117
2
        let mock_prompter = MockPrompterService::new();
118
8
        client_conn
119
            .object_server()
120
2
            .at("/org/gnome/keyring/Prompter", mock_prompter.clone())
121
6
            .await?;
122

            
123
        // Give the server a moment to fully initialize
124
4
        tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
125

            
126
4
        let service_api = dbus::api::Service::new(&client_conn).await?;
127

            
128
        // Generate client key pair for encrypted session
129
4
        let client_private_key = oo7::Key::generate_private_key()?;
130
4
        let client_public_key = oo7::Key::generate_public_key(&client_private_key)?;
131

            
132
6
        let (server_public_key, session) =
133
            service_api.open_session(Some(client_public_key)).await?;
134
4
        let session = Arc::new(session);
135

            
136
4
        let aes_key =
137
            oo7::Key::generate_aes_key(&client_private_key, &server_public_key.as_ref().unwrap())?;
138

            
139
6
        let collections = service_api.collections().await?;
140

            
141
2
        Ok(Self {
142
2
            server,
143
2
            keyring_secret: secret,
144
2
            client_conn,
145
2
            service_api,
146
2
            session,
147
2
            collections,
148
2
            server_public_key,
149
2
            aes_key: Some(Arc::new(aes_key)),
150
2
            mock_prompter,
151
        })
152
    }
153

            
154
    /// Create a test setup that discovers keyrings from disk
155
    /// This is useful for PAM tests that need to create keyrings on disk first
156
2
    pub(crate) async fn with_disk_keyrings(
157
        secret: Option<Secret>,
158
    ) -> Result<TestServiceSetup, Box<dyn std::error::Error>> {
159
        use zbus::proxy::Defaults;
160

            
161
6
        let (server_conn, client_conn) = create_p2p_connection().await?;
162

            
163
2
        let service = crate::Service::default();
164

            
165
8
        server_conn
166
            .object_server()
167
            .at(
168
2
                oo7::dbus::api::Service::PATH.as_deref().unwrap(),
169
2
                service.clone(),
170
            )
171
6
            .await?;
172

            
173
6
        let discovered = service.discover_keyrings(secret.clone()).await?;
174
5
        service.initialize(server_conn, discovered, false).await?;
175

            
176
2
        let mock_prompter = MockPrompterService::new();
177
8
        client_conn
178
            .object_server()
179
2
            .at("/org/gnome/keyring/Prompter", mock_prompter.clone())
180
6
            .await?;
181

            
182
4
        tokio::time::sleep(tokio::time::Duration::from_millis(50)).await;
183

            
184
2
        let service_api = dbus::api::Service::new(&client_conn).await?;
185

            
186
6
        let (server_public_key, session) = service_api.open_session(None).await?;
187
4
        let session = Arc::new(session);
188

            
189
6
        let collections = service_api.collections().await?;
190

            
191
2
        Ok(TestServiceSetup {
192
2
            server: service,
193
2
            keyring_secret: secret,
194
2
            client_conn,
195
2
            service_api,
196
2
            session,
197
            collections,
198
2
            server_public_key,
199
            aes_key: None,
200
2
            mock_prompter,
201
        })
202
    }
203
}
204

            
205
/// Mock implementation of org.gnome.keyring.internal.Prompter
206
///
207
/// This simulates the GNOME System Prompter for testing without requiring
208
/// the actual GNOME keyring prompter service to be running.
209
#[derive(Debug, Clone)]
210
pub(crate) struct MockPrompterService {
211
    /// The password to use for unlock prompts (simulates user input)
212
    unlock_password: Arc<tokio::sync::Mutex<Option<oo7::Secret>>>,
213
    /// Whether to accept (true) or dismiss (false) prompts
214
    should_accept: Arc<tokio::sync::Mutex<bool>>,
215
    /// Queue of passwords to use for for testing retry logic
216
    password_queue: Arc<tokio::sync::Mutex<Vec<oo7::Secret>>>,
217
}
218

            
219
impl MockPrompterService {
220
6
    pub fn new() -> Self {
221
        Self {
222
4
            unlock_password: Arc::new(tokio::sync::Mutex::new(Some(oo7::Secret::from(
223
                "test-password-long-enough",
224
            )))),
225
8
            should_accept: Arc::new(tokio::sync::Mutex::new(true)),
226
8
            password_queue: Arc::new(tokio::sync::Mutex::new(Vec::new())),
227
        }
228
    }
229

            
230
    /// Set whether prompts should be accepted or dismissed
231
8
    pub async fn set_accept(&self, accept: bool) {
232
4
        *self.should_accept.lock().await = accept;
233
    }
234

            
235
8
    pub async fn set_password_queue(&self, passwords: Vec<oo7::Secret>) {
236
2
        *self.password_queue.lock().await = passwords;
237
    }
238
}
239

            
240
#[zbus::interface(name = "org.gnome.keyring.internal.Prompter")]
241
impl MockPrompterService {
242
2
    async fn begin_prompting(
243
        &self,
244
        callback: ObjectPath<'_>,
245
        #[zbus(connection)] connection: &zbus::Connection,
246
    ) -> zbus::fdo::Result<()> {
247
4
        tracing::debug!("MockPrompter: begin_prompting called for {}", callback);
248
4
        let callback_path = callback.to_owned();
249
4
        let connection = connection.clone();
250

            
251
        // Spawn a task to send the initial prompt_ready call
252
8
        tokio::spawn(async move {
253
4
            tracing::debug!("MockPrompter: spawned task starting");
254
            // Small delay to ensure callback is fully registered
255
6
            tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
256

            
257
            // Call PromptReady directly without building a proxy (avoids introspection
258
            // issues in p2p)
259
2
            tracing::debug!(
260
                "MockPrompter: calling PromptReady with None on {}",
261
                callback_path
262
            );
263
2
            let properties: HashMap<String, Value> = HashMap::new();
264
4
            let empty_exchange = "";
265

            
266
8
            connection
267
2
                .call_method(
268
                    None::<()>, // No destination in p2p
269
                    &callback_path,
270
                    Some("org.gnome.keyring.internal.Prompter.Callback"),
271
                    "PromptReady",
272
4
                    &(Optional::<Reply>::from(None), properties, empty_exchange),
273
                )
274
8
                .await?;
275

            
276
2
            tracing::debug!("MockPrompter: PromptReady(None) completed");
277
2
            Ok::<_, zbus::Error>(())
278
        });
279

            
280
2
        Ok(())
281
    }
282

            
283
2
    async fn perform_prompt(
284
        &self,
285
        callback: ObjectPath<'_>,
286
        type_: PromptType,
287
        _properties: Properties,
288
        exchange: &str,
289
        #[zbus(connection)] connection: &zbus::Connection,
290
    ) -> zbus::fdo::Result<()> {
291
4
        tracing::debug!(
292
            "MockPrompter: perform_prompt called for {}, type={:?}",
293
            callback,
294
            type_
295
        );
296
        // This is called by PrompterCallback.prompter_init() with the server's exchange
297
4
        let callback_path = callback.to_owned();
298
4
        let unlock_password = self.unlock_password.clone();
299
4
        let should_accept = self.should_accept.clone();
300
4
        let password_queue = self.password_queue.clone();
301
4
        let exchange = exchange.to_owned();
302
4
        let connection = connection.clone();
303

            
304
        // Spawn a task to simulate user interaction and send final response
305
10
        tokio::spawn(async move {
306
4
            tracing::debug!("MockPrompter: perform_prompt task starting");
307
            // Small delay to simulate user interaction
308
6
            tokio::time::sleep(tokio::time::Duration::from_millis(10)).await;
309

            
310
2
            let accept = *should_accept.lock().await;
311
2
            let properties: HashMap<String, Value> = HashMap::new();
312

            
313
2
            if !accept {
314
4
                tracing::debug!("MockPrompter: dismissing prompt");
315
                // Dismiss the prompt
316
8
                connection
317
2
                    .call_method(
318
                        None::<()>, // No destination in p2p
319
                        &callback_path,
320
                        Some("org.gnome.keyring.internal.Prompter.Callback"),
321
                        "PromptReady",
322
2
                        &(Reply::No, properties, ""),
323
                    )
324
8
                    .await?;
325
2
                tracing::debug!("MockPrompter: PromptReady(no) completed");
326

            
327
2
                return Ok(());
328
4
            } else if type_ == PromptType::Password {
329
4
                tracing::debug!("MockPrompter: performing unlock (password prompt)");
330
                // Unlock prompt - perform secret exchange
331

            
332
4
                let mut queue = password_queue.lock().await;
333
8
                let password = if !queue.is_empty() {
334
4
                    let pwd = queue.remove(0);
335
4
                    tracing::debug!(
336
                        "MockPrompter: using password from queue (length: {}, queue remaining: {})",
337
                        std::str::from_utf8(pwd.as_bytes()).unwrap_or("<binary>"),
338
                        queue.len()
339
                    );
340
2
                    pwd
341
                } else {
342
4
                    let pwd = unlock_password.lock().await.clone().unwrap();
343
2
                    tracing::debug!(
344
                        "MockPrompter: using default password (length: {})",
345
                        std::str::from_utf8(pwd.as_bytes()).unwrap_or("<binary>")
346
                    );
347
2
                    pwd
348
                };
349
2
                drop(queue);
350

            
351
                // Generate our own key pair
352
2
                let private_key = oo7::Key::generate_private_key().unwrap();
353
4
                let public_key = crate::gnome::crypto::generate_public_key(&private_key).unwrap();
354

            
355
                // Handshake with server's exchange to get AES key
356
4
                let aes_key = secret_exchange::handshake(&private_key, &exchange).unwrap();
357

            
358
                // Encrypt the password
359
4
                let iv = crypto::generate_iv().unwrap();
360
4
                let encrypted = crypto::encrypt(password.as_bytes(), &aes_key, &iv).unwrap();
361

            
362
                // Create final exchange with encrypted secret
363
2
                let final_exchange = format!(
364
                    "[sx-aes-1]\npublic={}\nsecret={}\niv={}",
365
4
                    base64::prelude::BASE64_STANDARD.encode(public_key.as_ref()),
366
2
                    base64::prelude::BASE64_STANDARD.encode(&encrypted),
367
2
                    base64::prelude::BASE64_STANDARD.encode(&iv)
368
                );
369

            
370
4
                tracing::debug!("MockPrompter: calling PromptReady with yes");
371
8
                connection
372
2
                    .call_method(
373
                        None::<()>, // No destination in p2p
374
                        &callback_path,
375
                        Some("org.gnome.keyring.internal.Prompter.Callback"),
376
                        "PromptReady",
377
2
                        &(Reply::Yes, properties, final_exchange.as_str()),
378
                    )
379
8
                    .await?;
380
2
                tracing::debug!("MockPrompter: PromptReady(yes) with secret exchange completed");
381
            } else {
382
                tracing::debug!("MockPrompter: accepting confirm prompt");
383
                // Lock/confirm prompt - just accept
384
                connection
385
                    .call_method(
386
                        None::<()>, // No destination in p2p
387
                        &callback_path,
388
                        Some("org.gnome.keyring.internal.Prompter.Callback"),
389
                        "PromptReady",
390
                        &(Reply::Yes, properties, ""),
391
                    )
392
                    .await?;
393
                tracing::debug!("MockPrompter: PromptReady(yes) completed");
394
            }
395

            
396
2
            Ok::<_, zbus::Error>(())
397
        });
398

            
399
2
        Ok(())
400
    }
401

            
402
2
    async fn stop_prompting(
403
        &self,
404
        callback: ObjectPath<'_>,
405
        #[zbus(connection)] connection: &zbus::Connection,
406
    ) -> zbus::fdo::Result<()> {
407
4
        tracing::debug!("MockPrompter: stop_prompting called for {}", callback);
408
4
        let callback_path = callback.to_owned();
409
4
        let connection = connection.clone();
410

            
411
6
        tokio::spawn(async move {
412
4
            tracing::debug!("MockPrompter: calling PromptDone for {}", callback_path);
413
6
            let result = connection
414
2
                .call_method(
415
                    None::<()>,
416
                    &callback_path,
417
                    Some("org.gnome.keyring.internal.Prompter.Callback"),
418
                    "PromptDone",
419
                    &(),
420
                )
421
8
                .await;
422

            
423
4
            if let Err(err) = result {
424
                tracing::debug!("MockPrompter: PromptDone failed: {}", err);
425
            } else {
426
4
                tracing::debug!("MockPrompter: PromptDone completed for {}", callback_path);
427
            }
428
        });
429

            
430
2
        Ok(())
431
    }
432
}