main.rs
1use std::collections::HashMap;
2use std::io::{self, Write};
3use std::sync::Arc;
4
5use anyhow::{Context, Result, bail};
6use clap::Parser;
7use reqwest::header::HeaderMap;
8use serde::{Deserialize, Serialize};
9
10use atproto_client::client::{AppPasswordAuth, post_apppassword_json};
11use atproto_client::com::atproto::server::create_session;
12use atproto_identity::key::{KeyType, identify_key};
13use atproto_identity::model::{Document, VerificationMethod};
14use atproto_identity::resolve::{HickoryDnsResolver, InnerIdentityResolver};
15use atproto_plc::{Operation as PlcOperation, VerifyingKey as PlcVerifyingKey};
16
17const PLC_DIRECTORY: &str = "plc.directory";
18const PLC_DIRECTORY_URL: &str = "https://plc.directory";
19
20#[derive(Parser, Debug)]
21#[command(name = "plc-key")]
22#[command(about = "Add a verification method to a did:plc document")]
23struct Args {
24 /// Handle or DID of the identity
25 #[arg(short, long)]
26 subject: String,
27
28 /// Account password (can also be set via ATPROTO_PASSWORD env var)
29 #[arg(short, long, env = "ATPROTO_PASSWORD")]
30 password: String,
31
32 /// did:key to add as verification method
33 #[arg(short, long)]
34 key: String,
35
36 /// Key identifier fragment (without #)
37 #[arg(short = 'i', long)]
38 key_id: String,
39
40 /// PLC operation token (skips email request if provided)
41 #[arg(short, long, env = "ATPROTO_PLC_TOKEN")]
42 token: Option<String>,
43}
44
45/// Response from com.atproto.identity.signPlcOperation
46#[derive(Debug, Deserialize)]
47struct SignPlcOperationResponse {
48 operation: serde_json::Value,
49}
50
51/// Request for com.atproto.identity.signPlcOperation
52#[derive(Debug, Serialize)]
53#[serde(rename_all = "camelCase")]
54struct SignPlcOperationRequest {
55 token: String,
56 #[serde(skip_serializing_if = "Option::is_none")]
57 rotation_keys: Option<Vec<String>>,
58 #[serde(skip_serializing_if = "Option::is_none")]
59 also_known_as: Option<Vec<String>>,
60 #[serde(skip_serializing_if = "Option::is_none")]
61 verification_methods: Option<HashMap<String, String>>,
62 #[serde(skip_serializing_if = "Option::is_none")]
63 services: Option<HashMap<String, ServiceRecord>>,
64}
65
66#[derive(Debug, Serialize)]
67struct ServiceRecord {
68 #[serde(rename = "type")]
69 service_type: String,
70 endpoint: String,
71}
72
73/// Validate that the key is a P256 or K256 public key
74fn validate_key(key: &str) -> Result<()> {
75 let key_data = identify_key(key).with_context(|| format!("Failed to parse key: {}", key))?;
76
77 match key_data.key_type() {
78 KeyType::P256Public | KeyType::K256Public => Ok(()),
79 other => bail!(
80 "Unsupported key type: {:?}. Only P256 and K256 public keys are supported.",
81 other
82 ),
83 }
84}
85
86/// Extract the multibase key from a did:key URI
87fn extract_multibase_key(did_key: &str) -> &str {
88 did_key.strip_prefix("did:key:").unwrap_or(did_key)
89}
90
91/// Get the ID from a verification method
92fn get_verification_method_id(vm: &VerificationMethod) -> Option<&str> {
93 match vm {
94 VerificationMethod::Multikey { id, .. } => Some(id.as_str()),
95 VerificationMethod::Other { .. } => None,
96 }
97}
98
99/// Get the public key multibase from a verification method
100fn get_verification_method_key(vm: &VerificationMethod) -> Option<&str> {
101 match vm {
102 VerificationMethod::Multikey {
103 public_key_multibase,
104 ..
105 } => Some(public_key_multibase.as_str()),
106 VerificationMethod::Other { .. } => None,
107 }
108}
109
110/// Validate that the key_id is not already in use
111fn validate_key_id(doc: &Document, key_id: &str) -> Result<()> {
112 let full_id = format!("{}#{}", doc.id, key_id);
113 let fragment_id = format!("#{}", key_id);
114
115 for vm in &doc.verification_method {
116 if let Some(id) = get_verification_method_id(vm) {
117 if id == full_id || id == fragment_id || id == key_id {
118 bail!("Key ID '{}' is already in use in the DID document", key_id);
119 }
120 }
121 }
122 Ok(())
123}
124
125/// Extract the PDS host from a DID document
126fn extract_pds_host(doc: &Document) -> Result<String> {
127 for service in &doc.service {
128 if service.id.ends_with("#atproto_pds")
129 || service.id == "atproto_pds"
130 || service.r#type == "AtprotoPersonalDataServer"
131 {
132 return Ok(service.service_endpoint.clone());
133 }
134 }
135 bail!("No PDS service endpoint found in DID document")
136}
137
138/// Create an authenticated session
139async fn create_auth_session(
140 http_client: &reqwest::Client,
141 pds_host: &str,
142 identifier: &str,
143 password: &str,
144) -> Result<(AppPasswordAuth, String)> {
145 let session = create_session(http_client, pds_host, identifier, password, None)
146 .await
147 .with_context(|| "Failed to create session")?;
148
149 let auth = AppPasswordAuth {
150 access_token: session.access_jwt.clone(),
151 };
152
153 Ok((auth, session.did))
154}
155
156/// Request a PLC operation signature token (sent via email)
157async fn request_plc_signature(
158 http_client: &reqwest::Client,
159 pds_host: &str,
160 auth: &AppPasswordAuth,
161) -> Result<()> {
162 let url = format!(
163 "{}/xrpc/com.atproto.identity.requestPlcOperationSignature",
164 pds_host
165 );
166
167 let mut headers = HeaderMap::default();
168 headers.insert(
169 reqwest::header::AUTHORIZATION,
170 reqwest::header::HeaderValue::from_str(&format!("Bearer {}", auth.access_token))?,
171 );
172
173 let http_response = http_client
174 .post(url)
175 .headers(headers)
176 .send()
177 .await
178 .with_context(|| "HTTP Request failed")?;
179
180 print!("HTTP response {}", http_response.status());
181
182 Ok(())
183}
184
185/// Prompt the user to enter the PLC signing token
186fn prompt_for_token() -> Result<String> {
187 print!("Enter the PLC signing token from email: ");
188 io::stdout().flush()?;
189
190 let mut token = String::new();
191 io::stdin().read_line(&mut token)?;
192
193 let token = token.trim().to_string();
194 if token.is_empty() {
195 bail!("Token cannot be empty");
196 }
197
198 Ok(token)
199}
200
201/// Sign a PLC operation with the updated verification methods
202async fn sign_plc_operation(
203 http_client: &reqwest::Client,
204 pds_host: &str,
205 auth: &AppPasswordAuth,
206 token: &str,
207 new_key: &str,
208 key_id: &str,
209 current_doc: &Document,
210) -> Result<SignPlcOperationResponse> {
211 // Build verification_methods map from current doc + new key
212 let mut verification_methods: HashMap<String, String> = HashMap::new();
213
214 // Add existing verification methods
215 for vm in ¤t_doc.verification_method {
216 if let Some(id) = get_verification_method_id(vm) {
217 if let Some(key) = get_verification_method_key(vm) {
218 // Extract just the fragment part
219 let id_fragment = id.split('#').last().unwrap_or(id);
220 verification_methods.insert(id_fragment.to_string(), key.to_string());
221 }
222 }
223 }
224
225 // Add the new key
226 verification_methods.insert(
227 key_id.to_string(),
228 extract_multibase_key(new_key).to_string(),
229 );
230
231 println!(
232 "DEBUG: verification_methods being sent: {:?}",
233 verification_methods
234 );
235
236 let request = SignPlcOperationRequest {
237 token: token.to_string(),
238 rotation_keys: None,
239 also_known_as: None,
240 verification_methods: Some(verification_methods),
241 services: None,
242 };
243
244 let request_json = serde_json::to_value(&request)?;
245 println!(
246 "DEBUG: Full request JSON: {}",
247 serde_json::to_string_pretty(&request_json)?
248 );
249
250 let url = format!("{}/xrpc/com.atproto.identity.signPlcOperation", pds_host);
251
252 let response_value = post_apppassword_json(http_client, auth, &url, request_json)
253 .await
254 .with_context(|| "Failed to sign PLC operation")?;
255
256 println!(
257 "DEBUG: Response from signPlcOperation: {}",
258 serde_json::to_string_pretty(&response_value)?
259 );
260
261 let result: SignPlcOperationResponse = serde_json::from_value(response_value)
262 .with_context(|| "Failed to parse signPlcOperation response")?;
263
264 Ok(result)
265}
266
267/// Submit a signed PLC operation via the PDS
268async fn submit_plc_operation(
269 http_client: &reqwest::Client,
270 pds_host: &str,
271 auth: &AppPasswordAuth,
272 operation: &serde_json::Value,
273) -> Result<()> {
274 let url = format!("{}/xrpc/com.atproto.identity.submitPlcOperation", pds_host);
275
276 let response = post_apppassword_json(
277 http_client,
278 auth,
279 &url,
280 serde_json::json!({ "operation": operation }),
281 )
282 .await
283 .with_context(|| "Failed to submit PLC operation")?;
284
285 // Check for error response
286 if let Some(error) = response.get("error") {
287 let error_type = error.as_str().unwrap_or("Unknown");
288 let message = response
289 .get("message")
290 .and_then(|m| m.as_str())
291 .unwrap_or("No message");
292 bail!("PLC operation failed: {error_type} - {message} - {response:?}");
293 }
294
295 Ok(())
296}
297
298/// Fetch the audit log for a DID from plc.directory
299async fn fetch_audit_log(http_client: &reqwest::Client, did: &str) -> Result<Vec<PlcOperation>> {
300 let url = format!("{}/{}/log", PLC_DIRECTORY_URL, did);
301
302 let response = http_client
303 .get(&url)
304 .send()
305 .await
306 .with_context(|| format!("Failed to fetch audit log from {}", url))?;
307
308 if !response.status().is_success() {
309 bail!("Failed to fetch audit log: HTTP {}", response.status());
310 }
311
312 let operations: Vec<PlcOperation> = response
313 .json()
314 .await
315 .with_context(|| "Failed to parse audit log response")?;
316
317 Ok(operations)
318}
319
320/// Verify a signed operation against the rotation keys from the audit log
321fn verify_signed_operation(
322 signed_operation: &serde_json::Value,
323 audit_log: &[PlcOperation],
324) -> Result<()> {
325 // Parse the signed operation
326 let operation: PlcOperation = serde_json::from_value(signed_operation.clone())
327 .with_context(|| "Failed to parse signed operation")?;
328
329 // Get rotation keys from the last operation in the audit log
330 let last_op = audit_log
331 .last()
332 .ok_or_else(|| anyhow::anyhow!("Audit log is empty"))?;
333
334 let rotation_keys = last_op
335 .rotation_keys()
336 .ok_or_else(|| anyhow::anyhow!("Last operation has no rotation keys"))?;
337
338 // Convert rotation keys to verifying keys
339 let verifying_keys: Vec<PlcVerifyingKey> = rotation_keys
340 .iter()
341 .map(|k| PlcVerifyingKey::from_did_key(k))
342 .collect::<std::result::Result<Vec<_>, _>>()
343 .with_context(|| "Failed to parse rotation keys")?;
344
345 // Verify the signature
346 operation
347 .verify(&verifying_keys)
348 .with_context(|| "Operation signature verification failed")?;
349
350 Ok(())
351}
352
353/// Build a proposed document showing what the result will look like
354fn build_proposed_document(doc: &Document, new_key: &str, key_id: &str) -> serde_json::Value {
355 let mut proposed = serde_json::to_value(doc).unwrap();
356
357 if let Some(vms) = proposed.get_mut("verificationMethod") {
358 if let Some(arr) = vms.as_array_mut() {
359 arr.push(serde_json::json!({
360 "id": format!("{}#{}", doc.id, key_id),
361 "type": "Multikey",
362 "controller": doc.id,
363 "publicKeyMultibase": extract_multibase_key(new_key)
364 }));
365 }
366 }
367
368 proposed
369}
370
371#[tokio::main]
372async fn main() -> Result<()> {
373 let args = Args::parse();
374
375 // Create HTTP client
376 let http_client = reqwest::Client::new();
377
378 // Create identity resolver
379 let resolver = InnerIdentityResolver {
380 dns_resolver: Arc::new(HickoryDnsResolver::create_resolver(&[])),
381 http_client: http_client.clone(),
382 plc_hostname: PLC_DIRECTORY.to_string(),
383 };
384
385 // 1. Resolve subject to DID document
386 println!("Resolving identity: {}", args.subject);
387 let doc = resolver
388 .resolve(&args.subject)
389 .await
390 .with_context(|| format!("Failed to resolve identity: {}", args.subject))?;
391 let did = &doc.id;
392 println!("Resolved DID: {}", did);
393 println!("\nCurrent DID Document:");
394 println!("{}", serde_json::to_string_pretty(&doc)?);
395
396 // 2. Validate the key
397 println!("\nValidating key: {}", args.key);
398 validate_key(&args.key)?;
399 println!("Key is valid (P256 or K256)");
400
401 // 3. Validate the key_id
402 validate_key_id(&doc, &args.key_id)?;
403 println!("Key ID '{}' is available", args.key_id);
404
405 // 4. Extract PDS host from document
406 let pds_host = extract_pds_host(&doc)?;
407 println!("\nPDS: {}", pds_host);
408
409 // 5. Create session
410 println!("\nCreating session...");
411 let (auth, session_did) =
412 create_auth_session(&http_client, &pds_host, did, &args.password).await?;
413 println!("Session created for: {}", session_did);
414
415 // 6. Show proposed changes
416 println!("\nProposed DID Document (with new key):");
417 let proposed = build_proposed_document(&doc, &args.key, &args.key_id);
418 println!("{}", serde_json::to_string_pretty(&proposed)?);
419
420 // 7. Get PLC operation token
421 let token = match args.token {
422 Some(token) => {
423 println!("\nUsing provided PLC operation token.");
424 token
425 }
426 None => {
427 println!("\nRequesting PLC operation signature...");
428 println!("Check your email for the confirmation token.");
429 request_plc_signature(&http_client, &pds_host, &auth).await?;
430 prompt_for_token()?
431 }
432 };
433
434 // 8. Sign PLC operation
435 println!("\nSigning PLC operation...");
436 let response = sign_plc_operation(
437 &http_client,
438 &pds_host,
439 &auth,
440 &token,
441 &args.key,
442 &args.key_id,
443 &doc,
444 )
445 .await?;
446
447 // 9. Verify the signed operation against the audit log
448 println!("\nVerifying signed operation...");
449 let audit_log = fetch_audit_log(&http_client, did).await?;
450 verify_signed_operation(&response.operation, &audit_log)?;
451 println!("Operation signature verified against rotation keys");
452
453 // 10. Submit PLC operation via PDS
454 println!("\nSubmitting PLC operation...");
455 submit_plc_operation(&http_client, &pds_host, &auth, &response.operation).await?;
456
457 // 11. Verify the change
458 println!("\nVerifying change...");
459 tokio::time::sleep(tokio::time::Duration::from_secs(2)).await;
460 let updated_doc = resolver
461 .resolve(did)
462 .await
463 .with_context(|| "Failed to verify DID document update")?;
464 println!("\nUpdated DID Document:");
465 println!("{}", serde_json::to_string_pretty(&updated_doc)?);
466
467 println!(
468 "\nSuccess! Key '{}' has been added to the DID document.",
469 args.key_id
470 );
471
472 Ok(())
473}