pkrs_fork/

client.rs

use std::fmt::Display;

use reqwest::{Client, RequestBuilder, Response, StatusCode};
use serde::{Deserialize, Serialize};
use thiserror::Error;
use time::format_description::well_known::Rfc3339;
use time::OffsetDateTime;
use uuid::Uuid;

#[cfg(feature = "metrics")]
use std::time::Instant;
#[cfg(feature = "metrics")]
use url::Url;

use crate::model::*;

const BASE_URL: &str = "https://api.pluralkit.me/v2/";

pub struct PkClient {
    pub client: Client,
    pub token: String,
    pub user_agent: String,
}

impl Default for PkClient {
    fn default() -> Self {
        PkClient {
            client: Client::builder()
                .build()
                .expect("failed to create reqwest client"),
            token: "".to_string(),
            user_agent: "pk + rust project".to_string(), // todo: this is kind of bad
        }
    }
}

impl PkClient {
    /// return a new PkClient based on the current one, but with the specified token
    pub fn with_token(&self, token: String) -> Self {
        Self {
            token,
            client: self.client.clone(),
            user_agent: self.user_agent.clone(),
        }
    }

    pub async fn get_system(&self, system_id: &PkId) -> Result<System, PluralKitError> {
        self.get(format!("systems/{}", system_id)).await
    }

    pub async fn update_system(&self, system: &System) -> Result<System, PluralKitError> {
        self.patch_get_json(format!("systems/{}", system.id), system)
            .await
    }

    pub async fn get_system_settings(
        &self,
        system_id: &PkId,
    ) -> Result<SystemSettings, PluralKitError> {
        self.get(format!("systems/{}/settings", system_id)).await
    }

    pub async fn update_system_settings(
        &self,
        system_id: &PkId,
        settings: &SystemSettings,
    ) -> Result<SystemSettings, PluralKitError> {
        self.patch_get_json(format!("systems/{}/settings", system_id), settings)
            .await
    }

    pub async fn get_system_guild_settings(
        &self,
        system_id: &PkId,
        guild_id: &str,
    ) -> Result<SystemGuildSettings, PluralKitError> {
        let req = format!("systems/{}/settings/guilds/{}", system_id, guild_id);
        self.get(req).await
    }

    pub async fn update_system_guild_settings(
        &self,
        system_id: &PkId,
        guild_id: &str,
        settings: &SystemGuildSettings,
    ) -> Result<SystemGuildSettings, PluralKitError> {
        let req = format!("systems/{}/settings/guilds/{}", system_id, guild_id);
        self.patch_get_json(req, settings).await
    }

    pub async fn get_system_autoproxy_settings(
        &self,
        system_id: &PkId,
        guild_id: &str,
    ) -> Result<AutoProxySettings, PluralKitError> {
        let req = format!("systems/{}/autoproxy", system_id);
        self.get_query_get_json(req, &[("guild_id", guild_id)])
            .await
    }

    pub async fn update_system_autoproxy_settings(
        &self,
        system_id: &PkId,
        guild_id: &str,
        settings: &AutoProxySettings,
    ) -> Result<AutoProxySettings, PluralKitError> {
        let req = format!("systems/{}/autoproxy", system_id);
        self.patch_query_get_json(req, settings, &[("guild_id", guild_id)])
            .await
    }

    pub async fn get_system_members(
        &self,
        system_id: &PkId,
    ) -> Result<Vec<Member>, PluralKitError> {
        self.get(format!("systems/{}/members", system_id)).await
    }

    pub async fn create_member(&self, member: &Member) -> Result<Member, PluralKitError> {
        self.post_get_json("members".to_string(), member).await
    }

    pub async fn get_member(&self, member_id: &PkId) -> Result<Member, PluralKitError> {
        self.get(format!("members/{}", member_id)).await
    }

    pub async fn update_member(&self, member: &Member) -> Result<Member, PluralKitError> {
        self.patch_get_json(format!("members/{}", member.id), member)
            .await
    }

    pub async fn delete_member(&self, member_id: &PkId) -> Result<(), PluralKitError> {
        expect_no_content(&self.delete(format!("members/{}", member_id)).await?)
    }

    pub async fn get_member_groups(&self, member_id: &PkId) -> Result<Vec<Group>, PluralKitError> {
        self.get(format!("members/{}/groups", member_id)).await
    }

    async fn member_groups(
        &self,
        action: &str,
        member_id: &PkId,
        group_ids: &[&PkId],
    ) -> Result<(), PluralKitError> {
        expect_no_content(
            &self
                .request(
                    self.client
                        .post(format!(
                            "{}members/{}/groups/{}",
                            BASE_URL, member_id, action
                        ))
                        .json(group_ids),
                )
                .await?,
        )
    }

    pub async fn add_member_groups(
        &self,
        member_id: &PkId,
        group_ids: &[&PkId],
    ) -> Result<(), PluralKitError> {
        self.member_groups("add", member_id, group_ids).await
    }

    pub async fn remove_member_groups(
        &self,
        member_id: &PkId,
        group_ids: &[&PkId],
    ) -> Result<(), PluralKitError> {
        self.member_groups("remove", member_id, group_ids).await
    }

    pub async fn overwrite_member_groups(
        &self,
        member_id: &PkId,
        group_ids: &[&PkId],
    ) -> Result<(), PluralKitError> {
        self.member_groups("overwrite", member_id, group_ids).await
    }

    pub async fn get_member_guild_settings(
        &self,
        member_id: &PkId,
        guild_id: &str,
    ) -> Result<Vec<MemberGuildSettings>, PluralKitError> {
        self.get(format!("members/{}/guilds/{}", member_id, guild_id))
            .await
    }

    pub async fn update_member_guild_settings(
        &self,
        member_id: &PkId,
        guild_id: &str,
        settings: &MemberGuildSettings,
    ) -> Result<MemberGuildSettings, PluralKitError> {
        self.patch_get_json(
            format!("members/{}/guilds/{}", member_id, guild_id),
            settings,
        )
        .await
    }

    pub async fn get_system_groups(&self, system_id: &PkId) -> Result<Vec<Group>, PluralKitError> {
        self.get(format!("systems/{}/groups", system_id)).await
    }

    pub async fn create_group(&self, group: &Group) -> Result<Group, PluralKitError> {
        self.post_get_json("groups".to_string(), group).await
    }

    pub async fn get_group(&self, group_id: &PkId) -> Result<Group, PluralKitError> {
        self.get(format!("groups/{}", group_id)).await
    }

    pub async fn update_group(&self, group: &Group) -> Result<Group, PluralKitError> {
        self.patch_get_json(format!("groups/{}", group.id), group)
            .await
    }

    pub async fn delete_group(&self, group_id: &PkId) -> Result<(), PluralKitError> {
        expect_no_content(&self.delete(format!("groups/{}", group_id)).await?)
    }

    pub async fn get_group_members(&self, group_id: &PkId) -> Result<Vec<Member>, PluralKitError> {
        self.get(format!("groups/{}/members", group_id)).await
    }

    async fn group_members(
        &self,
        action: &str,
        group_id: &PkId,
        member_ids: &[&PkId],
    ) -> Result<(), PluralKitError> {
        expect_no_content(
            &self
                .request(
                    self.client
                        .post(format!(
                            "{}groups/{}/members/{}",
                            BASE_URL, group_id, action
                        ))
                        .json(member_ids),
                )
                .await?,
        )
    }

    pub async fn add_group_members(
        &self,
        group_id: &PkId,
        member_ids: &[&PkId],
    ) -> Result<(), PluralKitError> {
        self.group_members("add", group_id, member_ids).await
    }

    pub async fn remove_group_members(
        &self,
        group_id: &PkId,
        member_ids: &[&PkId],
    ) -> Result<(), PluralKitError> {
        self.group_members("remove", group_id, member_ids).await
    }

    pub async fn overwrite_group_members(
        &self,
        group_id: &PkId,
        member_ids: &[&PkId],
    ) -> Result<(), PluralKitError> {
        self.group_members("overwrite", group_id, member_ids).await
    }

    pub async fn get_system_switches(
        &self,
        system_id: &PkId,
        before: &OffsetDateTime,
        limit: &i32,
    ) -> Result<Vec<Switch>, PluralKitError> {
        self.get_query_get_json(
            format!("systems/{}/switches", system_id),
            &[
                ("before", before.format(&Rfc3339).unwrap().as_str()),
                ("limit", limit.to_string().as_str()),
            ],
        )
        .await
    }

    pub async fn get_system_fronters(
        &self,
        system_id: &PkId,
    ) -> Result<Option<Switch>, PluralKitError> {
        let resp = self
            .request(
                self.client
                    .get(format!("{BASE_URL}systems/{system_id}/fronters")),
            )
            .await?;

        if resp.status() == StatusCode::NO_CONTENT {
            return Ok(None);
        }

        Ok(Some(resp.json().await?))
    }

    pub async fn create_switch(
        &self,
        system_id: &PkId,
        member_ids: Vec<PkId>,
        time: Option<OffsetDateTime>,
    ) -> Result<Response, PluralKitError> {
        #[derive(Serialize, Deserialize, Debug)]
        struct SwitchCreate {
            #[serde(with = "time::serde::rfc3339::option")]
            timestamp: Option<OffsetDateTime>,
            members: Vec<PkId>,
        }
        self.post(
            format!("systems/{}/switches", system_id),
            &SwitchCreate {
                timestamp: time,
                members: member_ids,
            },
        )
        .await
    }

    pub async fn get_switch(
        &self,
        system_id: &PkId,
        switch_id: &Uuid,
    ) -> Result<Vec<Switch>, PluralKitError> {
        self.get(format!("systems/{}/switches/{}", system_id, switch_id))
            .await
    }

    pub async fn update_switch(
        &self,
        system_id: &PkId,
        switch_id: &Uuid,
        time: OffsetDateTime,
    ) -> Result<Switch, PluralKitError> {
        let req = format!("systems/{}/switches/{}", system_id, switch_id);
        #[derive(Serialize, Deserialize, Debug)]
        struct SwitchTimeUpdate {
            #[serde(with = "time::serde::rfc3339")]
            timestamp: OffsetDateTime,
        }
        self.get_response_json(
            self.client
                .patch(BASE_URL.to_string() + req.as_str())
                .json(&SwitchTimeUpdate { timestamp: time }),
        )
        .await
    }

    pub async fn update_switch_members(
        &self,
        system_id: &PkId,
        switch_id: &Uuid,
        members: &[&PkId],
    ) -> Result<Switch, PluralKitError> {
        let req = format!("systems/{}/switches/{}/members", system_id, switch_id);
        self.get_response_json(
            self.client
                .patch(BASE_URL.to_string() + req.as_str())
                .json(members),
        )
        .await
    }

    pub async fn delete_switch(
        &self,
        system_id: &PkId,
        switch_id: &Uuid,
    ) -> Result<(), PluralKitError> {
        expect_no_content(
            &self
                .delete(format!("systems/{}/switches/{}", system_id, switch_id))
                .await?,
        )
    }

    pub async fn get_message(&self, id: &str) -> Result<Message, PluralKitError> {
        self.get(format!("messages/{}", id)).await
    }

    // all
    async fn get<T: for<'a> Deserialize<'a>>(&self, endpoint: String) -> Result<T, PluralKitError> {
        self.get_response_json(self.client.get(BASE_URL.to_string() + &*endpoint))
            .await
    }

    // of this
    async fn get_query_get_json<T: for<'a> Deserialize<'a>>(
        &self,
        endpoint: String,
        query: &[(&str, &str)],
    ) -> Result<T, PluralKitError> {
        self.get_response_json(
            self.client
                .get(BASE_URL.to_string() + &*endpoint)
                .query(query),
        )
        .await
    }

    // duplication
    async fn patch_get_json<T>(&self, endpoint: String, body: &T) -> Result<T, PluralKitError>
    where
        T: Serialize + for<'a> Deserialize<'a>,
    {
        self.get_response_json(
            self.client
                .patch(BASE_URL.to_string() + &*endpoint)
                .json(body),
        )
        .await
    }

    // feels
    async fn patch_query_get_json<T>(
        &self,
        endpoint: String,
        body: &T,
        query: &[(&str, &str)],
    ) -> Result<T, PluralKitError>
    where
        T: Serialize + for<'a> Deserialize<'a>,
    {
        self.get_response_json(
            self.client
                .patch(BASE_URL.to_string() + &*endpoint)
                .query(query)
                .json(body),
        )
        .await
    }

    // extremely
    async fn post_get_json<T>(&self, endpoint: String, body: &T) -> Result<T, PluralKitError>
    where
        T: Serialize + for<'a> Deserialize<'a>,
    {
        self.get_response_json(
            self.client
                .post(BASE_URL.to_string() + &*endpoint)
                .json(body),
        )
        .await
    }

    // unclean
    async fn post<T: Serialize>(
        &self,
        endpoint: String,
        body: &T,
    ) -> Result<Response, PluralKitError> {
        self.request(
            self.client
                .post(BASE_URL.to_string() + &*endpoint)
                .json(body),
        )
        .await
    }

    async fn get_response_json<T: for<'a> Deserialize<'a>>(
        &self,
        builder: RequestBuilder,
    ) -> Result<T, PluralKitError> {
        Ok(self.request(builder).await?.json::<T>().await?)
    }

    async fn delete(&self, endpoint: String) -> Result<Response, PluralKitError> {
        self.request(self.client.delete(BASE_URL.to_string() + &*endpoint))
            .await
    }

    async fn request(&self, builder: RequestBuilder) -> Result<Response, PluralKitError> {
        #[cfg(feature = "metrics")]
        let now = Instant::now();

        let req = builder
            .header("User-Agent", &self.user_agent)
            .header("Authorization", &self.token)
            .build()?;

        #[cfg(feature = "metrics")]
        let method = req.method().to_string();
        #[cfg(feature = "metrics")]
        let url = req.url().clone();

        let resp = self.client.execute(req).await;

        #[cfg(feature = "metrics")]
        match resp {
            Ok(ref resp) => track_request(
                method,
                &url,
                Some(resp.status()),
                now.elapsed().as_secs_f64(),
            ),
            Err(_) => track_request(method, &url, None, now.elapsed().as_secs_f64()),
        }

        let resp = match resp {
            Ok(resp) => resp,
            Err(err) => return Err(PluralKitError::Reqwest(err)),
        };

        let err_for_status = match resp.error_for_status_ref() {
            Ok(_) => return Ok(resp),
            Err(err) => err,
        };

        match resp.json::<ErrorMessage>().await {
            Ok(message) => Err(PluralKitError::Pk(
                err_for_status
                    .status()
                    .expect("error_for_status() always populates status()"),
                message,
            )),
            Err(_) => Err(PluralKitError::Reqwest(err_for_status)),
        }
    }
}

#[cfg(feature = "metrics")]
fn track_request(method: String, url: &Url, status: Option<StatusCode>, latency: f64) {
    // normalise the request path stripping ids/snowflakes/uuids
    let normalized_path: String = match url.path_segments() {
        None => "/".to_string(),
        #[rustfmt::skip]
        Some(segments) => {
            // we skip the api version, v1 should just fallthrough
            let segment_vec: Vec<&str> = segments.skip(1).collect();
            match segment_vec[..] {
                // /members
                // /groups
                [resource                                     ] => format!("/v2/{resource}"),
                // /systems/{systemRef}
                // /members/{memberRef}
                // /groups/{groupRef}
                // /messages/{message}
                [resource , _                                 ] => format!("/v2/{resource}/!"),
                // /members/{memberRef}/groups/add
                // /members/{memberRef}/groups/remove
                // /members/{memberRef}/groups/overwrite
                ["members", _, "groups"   , action            ] => format!("/v2/members/!/groups/{action}"),
                // /groups/{groupRef}/members/add
                // /groups/{groupRef}/members/remove
                // /groups/{groupRef}/members/overwrite
                ["groups" , _, "members"  , action            ] => format!("/v2/groups/!/members/{action}"),
                // /systems/{systemRef}/settings
                // /systems/@me/autoproxy
                // /systems/{systemRef}/members
                // /members/{memberRef}/groups
                // /systems/{systemRef}/groups
                // /groups/{groupRef}/members
                // /systems/{systemRef}/switches
                // /systems/{systemRef}/fronters
                [resource , _, subresource                    ] => format!("/v2/{resource}/!/{subresource}"),
                // /systems/@me/guilds/{guild_id}
                // /members/{memberRef}/guilds/{guild_id}
                // /systems/{systemRef}/switches/{switchRef}
                [resource , _, subresource, _                 ] => format!("/v2/{resource}/!/{subresource}/!"),
                // /systems/{systemRef}/switches/{switchRef}/members
                [resource , _, subresource, _, subsubresource ] => format!("/v2/{resource}/!/{subresource}/!/{subsubresource}"),
                _ => url.path().to_string(),
            }
        }
    };

    metrics::histogram!(
        "pkrs_request_latency",
        "path" => normalized_path.clone(),
        "method" => method.clone(),
    )
    .record(latency);
    metrics::counter!(
        "pkrs_request",
        "path" => normalized_path,
        "method" => method,
        "status" => status.map(|s| s.to_string()).unwrap_or_else(|| "Unknown".to_string()),
    )
    .increment(1);
}

/// convenience function for handling endpoints that return `204 No Content`
/// on success
fn expect_no_content(response: &Response) -> Result<(), PluralKitError> {
    if response.status() == StatusCode::NO_CONTENT {
        Ok(())
    } else {
        Err(PluralKitError::Other(
            format!(
                "received status {} expected {}",
                response.status(),
                StatusCode::NO_CONTENT
            )
            .into(),
        ))
    }
}

#[derive(Error, Debug)]
pub enum PluralKitError {
    #[error("pluralkit api error, status: {0}, error: {1}")]
    Pk(StatusCode, ErrorMessage),
    #[error(transparent)]
    Reqwest(#[from] reqwest::Error),
    #[error(transparent)]
    Other(#[from] Box<dyn std::error::Error + Send + Sync>),
}

#[derive(Deserialize, Serialize, Debug)]
pub struct ErrorMessage {
    pub code: u32,
    pub message: String,
    pub errors: Option<Vec<ErrorMessageObject>>,
    pub retry_after: Option<u32>,
}

impl Display for ErrorMessage {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "{}: {}", self.code, self.message)
    }
}

#[derive(Deserialize, Serialize, Debug)]
pub struct ErrorMessageObject {
    pub message: String,
    pub max_length: Option<u32>,
    pub actual_length: Option<u32>,
}