use super::super::connection::{ OmikronConnection, OmikronResult, OptionalDataValueCompat, RequiredMtpFields, }; use crate::{ db::{iota_repo, user_repo}, models::{IotaId, UserId}, }; use base64::{Engine as _, engine::general_purpose::STANDARD}; use mtp::{ codec::{CommunicationType, CommunicationValue, DataType, DataValue}, crypto::PublicKeyBundle, }; use std::sync::Arc; fn connections(connection: &OmikronConnection, iota_id: i64) -> DataValue { DataValue::Array( connection .state() .presence .iota_connections(iota_id) .unwrap_or_default() .into_iter() .map(|id| DataValue::SignedNumber(id.into())) .collect(), ) } pub async fn get_user( connection: Arc, value: CommunicationValue, ) -> OmikronResult<()> { let request_id = value.require_id()?; let sender = value.require_sender_i64()?; let state = connection.state(); let user = if let Some(id) = value .get_data(DataType::UserId) .as_number() .and_then(|id| i64::try_from(id).ok()) .filter(|id| *id > 0) { user_repo::get_by_user_id(UserId::from(id)).await.ok() } else if let Some(name) = value.get_data(DataType::Username).as_str() { user_repo::get_by_username(name).await.ok() } else { None }; let Some(user) = user else { return connection .send_error_response(request_id, CommunicationType::ErrorNotFound) .await; }; let id = user.id.0; let iota_id = user.iota_id.map(|id| id.0); let username = user.username.clone(); let display = user .display .filter(|name| !name.is_empty()) .unwrap_or_else(|| username.clone()); let mut response = CommunicationValue::new(CommunicationType::GetUserData) .with_id(request_id) .add_typed_default(DataType::Username, DataValue::Str(username)) .add_typed_default( DataType::PublicKey, DataValue::Str(user.public_key.try_to_base64()?), ) .add_typed_default(DataType::UserId, DataValue::SignedNumber(id.into())) .add_typed_default(DataType::Display, DataValue::Str(display)) .add_typed_default( DataType::SubLevel, DataValue::SignedNumber(user.sub_level as i128), ) .add_typed_default( DataType::SubEnd, DataValue::SignedNumber(user.sub_end.into()), ); if let Some(status) = user.status.filter(|value| !value.is_empty()) { response = response.add_typed_default(DataType::Status, DataValue::Str(status)); } if let Some(about) = user.about.filter(|value| !value.is_empty()) { response = response.add_typed_default(DataType::About, DataValue::Str(about)); } if let Some(avatar) = user.avatar { response = response.add_typed_default(DataType::Avatar, DataValue::Str(STANDARD.encode(avatar))); } let route = state.presence.user_route(id); let private_request = sender == id; let resolved_status = if private_request { if !state .presence .load_preference(id, &user.presence_preference) { crate::log_in!( crate::util::logger::PrintType::General, "Invalid persisted presence preference for user {}, using user_online", id ); } state.presence.resolve_private_state(id) } else { iota_id .map(|iota_id| state.presence.resolve_public_state(id, iota_id)) .unwrap_or(crate::sql::connection_status::UserStatus::user_offline) }; response = response .add_typed_default( DataType::OnlineStatus, DataValue::Str(resolved_status.to_string()), ) .add_typed_default( DataType::OmikronConnections, iota_id .map(|iota_id| connections(&connection, iota_id)) .unwrap_or_else(|| DataValue::Array(Vec::new())), ); if let Some(iota_id) = iota_id { response = response.add_typed_default(DataType::IotaId, DataValue::SignedNumber(iota_id.into())); } if let Some(route) = route { response = response.add_typed_default( DataType::OmikronId, DataValue::SignedNumber(route.omikron_id.into()), ); } connection.send(&response).await } pub async fn get_iota( connection: Arc, value: CommunicationValue, ) -> OmikronResult<()> { let request_id = value.require_id()?; let found = if let Some(id) = value .get_data(DataType::IotaId) .as_number() .and_then(|id| i64::try_from(id).ok()) .filter(|id| *id > 0) { iota_repo::get_iota_by_id(IotaId::from(id)) .await .ok() .map(|iota| (iota.id.0, iota.public_key, None, None)) } else if let Some(id) = value.get_data(DataType::UserId).as_number() { if let Some(id) = i64::try_from(id).ok().filter(|id| *id > 0) && let Ok(user) = user_repo::get_by_user_id(UserId::from(id)).await { match user.iota_id { Some(iota_id) => iota_repo::get_iota_by_id(iota_id) .await .ok() .map(|iota| (iota.id.0, iota.public_key, Some(user.id.0), None)), None => None, } } else { None } } else if let Some(name) = value.get_data(DataType::Username).as_str() { if let Ok(user) = user_repo::get_by_username(name).await { match user.iota_id { Some(iota_id) => iota_repo::get_iota_by_id(iota_id).await.ok().map(|iota| { ( iota.id.0, iota.public_key, Some(user.id.0), Some(name.to_owned()), ) }), None => None, } } else { None } } else { None }; let Some((id, key, user_id, username)) = found else { return connection .send_error_response(request_id, CommunicationType::ErrorNotFound) .await; }; let mut response = CommunicationValue::new(CommunicationType::GetIotaData) .with_id(request_id) .add_typed_default(DataType::PublicKey, DataValue::Str(key.try_to_base64()?)) .add_typed_default(DataType::IotaId, DataValue::SignedNumber(id.into())) .add_typed_default(DataType::OmikronConnections, connections(&connection, id)); if let Some(user_id) = user_id { response = response.add_typed_default(DataType::UserId, DataValue::SignedNumber(user_id.into())); } if let Some(username) = username { response = response.add_typed_default(DataType::Username, DataValue::Str(username)); } connection.send(&response).await } async fn update_user( connection: Arc, value: CommunicationValue, ) -> OmikronResult<()> { let request_id = value.require_id()?; let id = UserId::from(value.require_sender_i64()?); let mut error = None; if let Some(name) = value.get_data(DataType::Username).as_str() { error = user_repo::change_username(id, name.to_owned()) .await .err() .map(|e| e.to_string()); } if error.is_none() && let Some(name) = value.get_data(DataType::Display).as_str() { error = user_repo::change_display_name(id, name.to_owned()) .await .err() .map(|e| e.to_string()); } if error.is_none() && let Some(avatar) = value.get_data(DataType::Avatar).as_str() { error = user_repo::change_avatar(id, avatar.to_owned()) .await .err() .map(|e| e.to_string()); } if error.is_none() && let Some(about) = value.get_data(DataType::About).as_str() { error = user_repo::change_about(id, about.to_owned()) .await .err() .map(|e| e.to_string()); } if error.is_none() && let Some(status) = value.get_data(DataType::Status).as_str() { error = user_repo::change_status(id, status.to_owned()) .await .err() .map(|e| e.to_string()); } if error.is_none() && let Some(key) = value .get_data(DataType::PublicKey) .as_str() .and_then(|key| PublicKeyBundle::from_base64(key).ok()) { error = user_repo::change_keys(id, key) .await .err() .map(|e| e.to_string()); } let response = match error { None => CommunicationValue::new(CommunicationType::Success), Some(error) => CommunicationValue::new(CommunicationType::ErrorInternal) .add_typed_default(DataType::ErrorType, DataValue::Str(error)), }; connection.send(&response.with_id(request_id)).await } pub async fn change_user( connection: Arc, value: CommunicationValue, ) -> OmikronResult<()> { update_user(connection, value).await } pub async fn change_iota( connection: Arc, value: CommunicationValue, ) -> OmikronResult<()> { let request_id = value.require_id()?; let reset_data = value.get_data(DataType::ResetToken); let Some(reset) = reset_data.as_str() else { return connection .send_error_response(value.require_id()?, CommunicationType::ErrorInvalidData) .await; }; let new_token_data = value.get_data(DataType::NewToken); let Some(new_token) = new_token_data.as_str() else { return connection .send_error_response(value.require_id()?, CommunicationType::ErrorInvalidData) .await; }; let user_id = UserId::from(value.require_sender_i64()?); let user = match user_repo::get_by_user_id(user_id).await { Ok(user) => user, Err(_) => { return connection .send_error_response(request_id, CommunicationType::ErrorNotFound) .await; } }; if user.token != reset { return connection .send_error_response(request_id, CommunicationType::ErrorInvalidChallenge) .await; } let result = match user_repo::change_iota_id(user_id, Some(IotaId::from(value.require_sender_i64()?))) .await { Ok(()) => user_repo::change_token(user_id, new_token.to_owned()).await, Err(error) => Err(error), }; let response = match result { Ok(()) => CommunicationValue::new(CommunicationType::Success), Err(error) => CommunicationValue::new(CommunicationType::ErrorInternal) .add_typed_default(DataType::ErrorType, DataValue::Str(error.to_string())), }; connection.send(&response.with_id(request_id)).await }