Something went wrong. Try again.
Personal ATProto tools.
Something went wrong. Try again.
17 kB · 340 lines
Rust
at main
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340//! Requests to the atproto API.use std::env;use std::fs;use std::sync::Arc;
use base64::Engine;use reqwest::Client;use tokio::sync::Mutex;
use crate::types::{LabelsVecWithSeq, RetrievedLabelResponse, SignatureBytes};enum ApiEndpoint { Authorized, Public,}/// Agent for interactions with atproto.#[derive(Clone)]pub struct Agent { /// The access JWT. pub access_jwt: Arc<Mutex<String>>, /// The refresh JWT. pub refresh_jwt: Arc<Mutex<String>>, /// The reqwest client. pub client: Client, /// The DID of the labeler. pub self_did: Arc<String>,}impl Default for Agent { fn default() -> Self { drop(dotenvy::dotenv().expect("Failed to load .env file")); Self { access_jwt: Arc::new(Mutex::new(env::var("ACCESS_JWT").expect("ACCESS_JWT must be set"))), refresh_jwt: Arc::new(Mutex::new(env::var("REFRESH_JWT").expect("REFRESH_JWT must be set"))), client: Client::new(), self_did: Arc::new(env::var("SELF_DID").expect("SELF_DID must be set")), } }}impl Agent { /// The base URL of the atproto API's XRPC endpoint. /// Rate limit: 3_000 per 5 minutes const AUTH_URL: &'static str = "https://bsky.social/xrpc/"; const PUBLIC_URL: &'static str = "https://public.api.bsky.app/xrpc/"; async fn client_get( &self, path: &str, parameters: &[(&str, &str)], api_endpoint: &ApiEndpoint, ) -> reqwest::Response { self.client .get(format!("{}{}", match api_endpoint { ApiEndpoint::Authorized => Self::AUTH_URL, ApiEndpoint::Public => Self::PUBLIC_URL, }, &path)) .header("Content-Type", "application/json") .header("Authorization", format!("Bearer {}", self.access_jwt.lock().await)) .header("atproto-accept-labelers", self.self_did.as_str()) .query(parameters) .send() .await.expect("Expected to be able to send request, but failed.") } async fn client_refresh(&self) { tracing::warn!("Token expired, refreshing"); let response = self.client .post(format!( "{}{}", Self::AUTH_URL, "com.atproto.server.refreshSession" )) .header("Content-Type", "application/json") .header("Authorization", format!("Bearer {}", self.refresh_jwt.lock().await)) .header("atproto-accept-labelers", self.self_did.as_str()) .send() .await.expect("Expected to be able to send request, but failed."); let json = response.json::<serde_json::Value>().await.expect("Expected to be able to read response as JSON, but failed."); if let Some(error) = json["error"].as_str() { match error { "InvalidRequest" => { tracing::warn!("Invalid request"); return; }, "ExpiredToken" => { tracing::warn!("Token expired"); return; }, "AccountDeactivated" => { tracing::warn!("Account deactivated"); return; }, "AccountTakedown" => { tracing::warn!("Account has been suspended (Takedown)"); return; }, _ => { tracing::warn!("Unknown error from HTTP response: {:?}", json); return; } } } *self.refresh_jwt.lock().await = json["refreshJwt"].as_str().expect("Expected to be able to read refreshJwt as str, but failed.").to_owned(); *self.access_jwt.lock().await = json["accessJwt"].as_str().expect("Expected to be able to read accessJwt as str, but failed.").to_owned(); let new_env = format!( "ACCESS_JWT={}\nREFRESH_JWT={}\n", self.access_jwt.lock().await, self.refresh_jwt.lock().await ); fs::write(".env", new_env).expect("Failed to write to .env"); tracing::info!("Token refreshed"); } /// Get a JSON response from the atproto API. Used internal to this struct. async fn get( &self, path: &str, parameters: &[(&str, &str)], api_endpoint: ApiEndpoint, ) -> Result<serde_json::Value, Box<dyn std::error::Error + Send + Sync>> { let response = self.client_get(path, parameters, &api_endpoint).await; if response.status() == reqwest::StatusCode::TOO_MANY_REQUESTS { tracing::warn!("Rate limited, sleeping for 5 minutes"); tracing::warn!("We were working on {} with parameters {:?}", path, parameters); tokio::time::sleep(std::time::Duration::from_secs(305)).await; // 5 minutes and 5 seconds let response = self.client_get(path, parameters, &api_endpoint).await; return Ok(response.json::<serde_json::Value>().await.expect("Expected to be able to read response as JSON, but failed.")); } if response.status() == reqwest::StatusCode::BAD_REQUEST { let json = &response.json::<serde_json::Value>().await.expect("Expected to be able to read response as JSON, but failed."); match json["error"].as_str().expect("Expected to be able to read error as str, but failed.") { "ExpiredToken" => { self.client_refresh().await; let response = self.client_get(path, parameters, &api_endpoint).await; return Ok(response.json::<serde_json::Value>().await.expect("Expected to be able to read response as JSON, but failed.")); }, "AccountDeactivated" => { tracing::warn!("Account deactivated"); return Err(Box::new(std::io::Error::new( std::io::ErrorKind::Other, "Account deactivated", ))); }, "AccountTakedown" => { tracing::warn!("Account has been suspended (Takedown)"); return Err(Box::new(std::io::Error::new( std::io::ErrorKind::Other, "Account deactivated", ))); }, "InvalidRequest" => { // Check if the message is "Profile not found" if json["message"].as_str().expect("Expected to be able to read message as str, but failed.") == "Profile not found" { tracing::warn!("Profile not found"); return Err(Box::new(std::io::Error::new( std::io::ErrorKind::NotFound, "Profile not found", ))); } tracing::warn!("Unknown invalid request: {:?}", json); return Err(Box::new(std::io::Error::new( std::io::ErrorKind::Other, "Unknown invalid request", ))); }, _ => { tracing::warn!("Unknown error from HTTP response: {:?}", json); return Err(Box::new(std::io::Error::new( std::io::ErrorKind::Other, "Unknown bad request", ))); } }; } if response.status() != reqwest::StatusCode::OK { return Err(Box::new(std::io::Error::new( std::io::ErrorKind::Other, "Unknown HTTP error", ))); } let json = response.json::<serde_json::Value>().await.expect("Expected to be able to read response as JSON, but failed."); Ok(json) } /// Get a profile from the atproto API. pub async fn get_profile( &mut self, profile_id: &str, ) -> Result<serde_json::Value, Box<dyn std::error::Error + Send + Sync>> { let path = "app.bsky.actor.getProfile"; let parameters = [("actor", profile_id)]; self.get(path, ¶meters, ApiEndpoint::Public).await } /// Get multiple profiles. pub async fn get_profiles( &mut self, profile_ids: &[String], ) -> Result<serde_json::Value, Box<dyn std::error::Error + Send + Sync>> { let path = "app.bsky.actor.getProfiles"; let mut parameters = Vec::new(); for profile_id in profile_ids { parameters.push(("actors", profile_id.as_str())); } self.get(path, parameters.as_slice(), ApiEndpoint::Authorized).await } /// Check if a list of profiles has a label from us. pub async fn check_profiles( &mut self, profile_ids: &[(String, i64)], ) -> Result<Vec<(bool, (String, i64))>, Box<dyn std::error::Error + Send + Sync>> { let mut found_labels: Vec<(bool, (String, i64))> = Vec::new(); let profile_ids_uris = profile_ids.iter().map(|(profile_id, _)| profile_id.clone()).collect::<Vec<String>>(); let profile_ids_seqs = profile_ids.iter().map(|(_, seq)| seq).collect::<Vec<&i64>>(); let profiles = self.get_profiles(profile_ids_uris.as_slice()).await?; let profiles_array = profiles["profiles"].as_array(); if profiles_array.is_none() { tracing::warn!("No profiles json found for profiles: {:?}", profiles); return Ok(vec![]); } for profile in profiles_array.unwrap_or_else(|| panic!("Expected to be able to read profiles as array, but failed. Profiles: {:?}", profiles)) { let labels = &profile["labels"]; let mut found = false; let label_array = labels.as_array(); if label_array.is_none() { tracing::warn!("No labels json found for profile: {:?}", profile); continue; } let did = profile["did"].as_str().expect("Expected to be able to read did as str, but failed."); let seq = profile_ids_seqs[profile_ids_uris.iter().position(|x| x == did).expect("Expected to be able to find the index of the uri.")]; for label in label_array.unwrap_or_else(|| panic!("Expected to be able to read labels as array, but failed. Profile: {:?}", profile)) { if label["src"].as_str().expect("Expected to be able to read src as str, but failed.") == self.self_did.as_str() { found = true; break; } } found_labels.push((found, (did.to_owned(), *seq))); } Ok(found_labels) } /// After getting a profile, check the labels on it, and see if one from us ("src:") is there. pub async fn check_profile( &mut self, profile_did: &str, ) -> Result<bool, Box<dyn std::error::Error + Send + Sync>> { let profile = self.get_profile(profile_did).await?; let labels = &profile["labels"]; let label_array = labels.as_array(); if label_array.is_none() { tracing::warn!("No labels json found for profile: {:?}", profile); return Ok(false); } for label in label_array.unwrap_or_else(|| panic!("Expected to be able to read labels as array, but failed. Profile: {:?}", profile)) { if label["src"].as_str().expect("Expected to be able to read src as str, but failed.") == self.self_did.as_str() { return Ok(true); } } Ok(false) } /// Get a label from the provided URL, then validate the signature. pub async fn get_label_and_validate( &self, url: &str, ) -> Result<(), Box<dyn std::error::Error>> { tracing::debug!("Getting label from {}", url); let response = reqwest::get(url).await.expect("Expected to be able to get response, but failed."); tracing::debug!("Response: {:?}", response); let response_json = response.json::<serde_json::Value>().await.expect("Expected to be able to read response as JSON, but failed."); tracing::debug!("Response JSON: {:?}", response_json); let sig = &response_json["labels"][0]["sig"]; tracing::debug!("Signature: {:?}", sig); let retrieved_label = RetrievedLabelResponse { // id: response_json["labels"][0]["id"].as_u64().unwrap(), cts: response_json["labels"][0]["cts"].as_str().expect("Expected to be able to read cts as str, but failed.").to_owned(), neg: response_json["labels"][0]["neg"].as_str() == Some("true"), src: response_json["labels"][0]["src"].as_str().expect("Expected to be able to read src as str, but failed.").to_owned().parse().expect("Failed to parse DID"), uri: response_json["labels"][0]["uri"].as_str().expect("Expected to be able to read uri as str, but failed.").to_owned().parse().expect("Failed to parse URI"), val: response_json["labels"][0]["val"].as_str().expect("Expected to be able to read val as str, but failed.").to_owned(), ver: response_json["labels"][0]["ver"].as_u64().expect("Expected to be able to read ver as u64, but failed."), }; let crypto = crate::crypto::Crypto::new(); let pub_key = "zQ3shreqyXEdouQeEQSFKfoSEN5eig74BXuqQyTaiE9uzADqZ"; let sig_string = sig["$bytes"].as_str().expect("Expected to be able to read sig as str, but failed."); if crypto.validate(retrieved_label, sig_string, pub_key) { tracing::info!("Valid signature"); Ok(()) } else { tracing::info!("Invalid signature"); Err(Box::new(std::io::Error::new( std::io::ErrorKind::Other, "Invalid signature", ))) } } /// Get a label from a websocket URL, then validate the signature. /// Similar to what's done in webserve.rs, but in reverse, we'll need to decode the message. pub async fn get_label_and_validate_ws( &self, // url: &str, ) -> Result<(), Box<dyn std::error::Error>> { // For now, use this mock response, represented in base64: let response = "omF0ZyNsYWJlbHNib3ABomNzZXEYG2ZsYWJlbHOBp2NjdHN4GzIwMjUtMDItMDlUMDM6MjU6MjcuOTI4MDIzWmNuZWf0Y3NpZ1hAXLIRXAG5mF5bCWWCwEhbYvC8YYVP9fWwbVVL6IBXXlIrZ6sr6MQ4DfNdpGhwRWawA4Mq44HlEDsJ7OvcGsDCDWNzcmN4IGRpZDpwbGM6bTZhZHB0bjYyZGNhaGZhcTM0dGNlM2o1Y3VyaXggZGlkOnBsYzptNmFkcHRuNjJkY2FoZmFxMzR0Y2UzajVjdmFsbmpvaW5lZC0yMDI1LTAyY3ZlcgE="; tracing::debug!("Response: {:?}", response); let response_bytes = base64::engine::GeneralPurpose::new( &base64::alphabet::STANDARD, base64::engine::general_purpose::PAD).decode(response).expect("Expected to be able to decode base64 response."); tracing::debug!("Response bytes: {:?}", response_bytes); let reponse_bytes_in_hex = hex::encode(&response_bytes); tracing::debug!("Response bytes in hex: {:?}", reponse_bytes_in_hex); let response_0 = &response_bytes[0..response_bytes.iter().position(|&r| r == 0x01).expect("Expected to find 0x01 in response bytes.")]; let response_1 = &response_bytes[response_bytes.iter().position(|&r| r == 0x01).expect("Expected to find 0x01 in response bytes.") + 1..]; tracing::debug!("Response 0: {:?}", hex::encode(response_0)); tracing::debug!("Response 1: {:?}", hex::encode(response_1)); let response_cbor: LabelsVecWithSeq = serde_cbor::from_slice(response_1).expect("Expected to be able to deserialize response 1 as LabelsVecWithSeq, but failed."); tracing::debug!("Response CBOR: {:?}", response_cbor); let unsigned_response = RetrievedLabelResponse { cts: response_cbor.labels[0].cts.clone(), neg: response_cbor.labels[0].neg, src: response_cbor.labels[0].src.clone(), uri: response_cbor.labels[0].uri.clone(), val: response_cbor.labels[0].val.clone(), ver: response_cbor.labels[0].ver, }; let sig_base64 = SignatureBytes::from_bytes(response_cbor.labels[0].sig).as_base64(); tracing::debug!("Retrieved label: {:?}", response_cbor); let crypto = crate::crypto::Crypto::new(); let public_key = "zQ3shreqyXEdouQeEQSFKfoSEN5eig74BXuqQyTaiE9uzADqZ"; if crypto.validate(unsigned_response, &sig_base64, public_key) { tracing::info!("Valid signature"); Ok(()) } else { tracing::info!("Invalid signature"); Err(Box::new(std::io::Error::new( std::io::ErrorKind::Other, "Invalid signature", ))) } } /// getLikes pub async fn get_likes( &mut self, uri: &str, ) -> Result<Vec<serde_json::Value>, Box<dyn std::error::Error + Send + Sync>> { let path = "app.bsky.feed.getLikes"; let parameters = [("uri", uri)]; self.get(path, ¶meters, ApiEndpoint::Public).await.map(|response| response["likes"].as_array().expect("Expected to be able to read likes as array, but failed.").to_owned()) }}