magnetar/src/model/processing/drive.rs

use crate::model::data::drive::PackFileBaseInput;
use crate::model::processing::PackResult;
use crate::model::{PackType, PackingContext};
use magnetar_calckey_model::ck;
use magnetar_sdk::types::drive::{DriveFileBase, PackDriveFileBase};
use magnetar_sdk::types::Id;
use magnetar_sdk::{Packed, Required};
use tracing::warn;
use url::Url;

pub struct DriveModel;

impl DriveModel {
    pub fn media_proxy_url(
        &self,
        ctx: &PackingContext,
        url: &str,
        is_thumbnail: bool,
    ) -> Option<Url> {
        if let Some(proxy) = &ctx.service.config.networking.media_proxy {
            let params = if is_thumbnail {
                vec![("url", url), ("thumbnail", "1")]
            } else {
                vec![("url", url)]
            };

            let url = Url::parse_with_params(proxy, &params);

            if let Err(e) = url {
                warn!("Url parse error: {e}");
            }

            return url.ok();
        }

        None
    }

    pub fn builtin_proxy_file(
        &self,
        ctx: &PackingContext,
        file: &ck::drive_file::Model,
        is_thumbnail: bool,
    ) -> Option<Url> {
        let key = if is_thumbnail {
            file.thumbnail_access_key.as_deref()
        } else {
            file.webpublic_access_key.as_deref()
        };

        if let Some(k) = key {
            if k != "/" {
                let url_raw = format!(
                    "{}://{}/files/{}",
                    ctx.service.config.networking.protocol.as_ref(),
                    &ctx.service.config.networking.host,
                    k
                );
                let url = Url::parse(&url_raw);

                if let Err(e) = url {
                    warn!("Url parse error: {e}");
                }

                return url.ok();
            }
        }

        None
    }

    pub fn get_public_url(
        &self,
        ctx: &PackingContext,
        file: &ck::drive_file::Model,
        is_thumbnail: bool,
    ) -> Option<Url> {
        if let Some(uri) = &file.uri {
            if file.user_host.is_none() {
                if let Some(media_proxy_url) = self.media_proxy_url(ctx, uri, is_thumbnail) {
                    return Some(media_proxy_url);
                }

                if file.is_link && ctx.service.config.networking.proxy_remote_files {
                    if let Some(proxy_url) = self.builtin_proxy_file(ctx, file, is_thumbnail) {
                        return Some(proxy_url);
                    }
                }
            }
        }

        let is_image = matches!(
            file.r#type.as_str(),
            "image/png"
                | "image/apng"
                | "image/gif"
                | "image/jpeg"
                | "image/webp"
                | "image/svg+xml"
                | "image/avif"
        );

        let url_raw = if is_thumbnail {
            file.thumbnail_url
                .as_deref()
                .or(is_image.then_some(file.webpublic_url.as_deref().unwrap_or(file.url.as_str())))
        } else {
            file.webpublic_url.as_deref().or(Some(file.url.as_str()))
        };

        if let Some(u) = url_raw {
            let url = Url::parse(u);

            if let Err(e) = url {
                warn!("Url parse error: {e}");
            }

            return url.ok();
        }

        None
    }

    pub fn pack_existing(
        &self,
        ctx: &PackingContext,
        file: &ck::drive_file::Model,
    ) -> PackDriveFileBase {
        let url = self.get_public_url(ctx, file, false);
        let thumbnail_url = self.get_public_url(ctx, file, false);

        PackDriveFileBase::pack_from((
            Required(Id::from(&file.id)),
            Required(DriveFileBase::extract(
                ctx,
                PackFileBaseInput {
                    file: &file,
                    effective_url: url.as_ref(),
                    effective_thumbnail_url: thumbnail_url.as_ref(),
                },
            )),
        ))
    }

    pub async fn get_cached_base(
        &self,
        ctx: &PackingContext,
        id: &str,
    ) -> PackResult<Option<PackDriveFileBase>> {
        let Some(file) = ctx.service.drive_file_cache.get(id).await? else {
            return Ok(None);
        };

        Ok(Some(self.pack_existing(ctx, file.as_ref())))
    }
}