nekrochan/src/files.rs

use actix_multipart::form::tempfile::TempFile;
use anyhow::Error;
use chrono::Utc;
use glob::glob;
use image::io::Reader as ImageReader;
use std::{collections::HashSet, process::Command};
use tokio::{
    fs::{remove_file, rename},
    task::spawn_blocking,
};

use crate::{
    cfg::Cfg,
    ctx::Ctx,
    db::models::{Board, File, Post},
    error::NekrochanError,
};

impl File {
    pub async fn new(
        cfg: &Cfg,
        temp_file: TempFile,
        spoiler: bool,
        thumb: bool,
    ) -> Result<Self, NekrochanError> {
        let original_name = temp_file.file_name.unwrap_or_else(|| "unknown".into());

        let mime = temp_file
            .content_type
            .ok_or(NekrochanError::FileError(
                original_name.clone(),
                "žádný mime typ",
            ))?
            .to_string();

        let (video, format) = match mime.as_str() {
            "image/jpeg" => (false, "jpg"),
            "image/pjpeg" => (false, "jpg"),
            "image/png" => (false, "png"),
            "image/bmp" => (false, "bmp"),
            "image/gif" => (false, "gif"),
            "image/webp" => (false, "webp"),
            "image/apng" => (false, "apng"),
            "video/mpeg" => (true, "mpeg"),
            "video/quicktime" => (true, "mov"),
            "video/mp4" => (true, "mp4"),
            "video/webm" => (true, "webm"),
            "video/x-matroska" => (true, "mkv"),
            "video/ogg" => (true, "ogg"),
            _ => {
                return Err(NekrochanError::FileError(
                    original_name,
                    "nepodporovaný formát",
                ))
            }
        };

        if video && !cfg.files.videos {
            return Err(NekrochanError::FileError(
                original_name,
                "videa nejsou podporovaná",
            ));
        }

        let size = temp_file.size;

        if size / 1_000_000 > cfg.files.max_size_mb {
            return Err(NekrochanError::FileError(
                original_name,
                "soubor je příliš velký",
            ));
        }

        let timestamp = Utc::now().timestamp_micros();
        let format = format.to_owned();

        let new_name = format!("{timestamp}.{format}");

        let (thumb_format, thumb_name) = if thumb {
            let format = if video { "png".into() } else { format.clone() };

            (Some(format.clone()), Some(format!("{timestamp}.{format}")))
        } else {
            (None, None)
        };

        rename(temp_file.file.path(), format!("/tmp/{new_name}")).await?;

        let (width, height) = if video {
            process_video(cfg, original_name.clone(), new_name.clone(), thumb_name).await?
        } else {
            process_image(cfg, original_name.clone(), new_name.clone(), thumb_name).await?
        };

        rename(format!("/tmp/{new_name}"), format!("uploads/{new_name}")).await?;

        let file = File {
            original_name,
            format,
            thumb_format,
            spoiler,
            width,
            height,
            timestamp,
            size,
        };

        Ok(file)
    }

    pub async fn delete(&self) {
        remove_file(format!("./uploads/{}.{}", self.timestamp, self.format))
            .await
            .ok();

        if let Some(thumb_format) = &self.thumb_format {
            remove_file(format!(
                "./uploads/thumb/{}.{}",
                self.timestamp, thumb_format
            ))
            .await
            .ok();
        }
    }

    pub fn file_url(&self) -> String {
        format!("/uploads/{}.{}", self.timestamp, self.format)
    }

    pub fn thumb_url(&self) -> String {
        if self.spoiler {
            "/static/spoiler.png".into()
        } else if let Some(thumb_format) = &self.thumb_format {
            format!("/uploads/thumb/{}.{}", self.timestamp, thumb_format)
        } else {
            self.file_url()
        }
    }
}

async fn process_image(
    cfg: &Cfg,
    original_name: String,
    new_name: String,
    thumb_name: Option<String>,
) -> Result<(u32, u32), NekrochanError> {
    let original_name_ = original_name.clone();

    let img = spawn_blocking(move || {
        ImageReader::open(format!("/tmp/{new_name}"))?
            .decode()
            .map_err(|_| {
                NekrochanError::FileError(original_name_, "nepodařilo se dekódovat obrázek")
            })
    })
    .await??;

    let (width, height) = (img.width(), img.height());

    if width > cfg.files.max_width || height > cfg.files.max_height {
        return Err(NekrochanError::FileError(
            original_name,
            "rozměry obrázku jsou příliš velké",
        ));
    }

    let thumb_name = match thumb_name {
        Some(thumb_name) => thumb_name,
        None => return Ok((width, height)),
    };

    let thumb_w = if width > cfg.files.thumb_size {
        cfg.files.thumb_size
    } else {
        width
    };

    let thumb_h = if height > cfg.files.thumb_size {
        cfg.files.thumb_size
    } else {
        height
    };

    spawn_blocking(move || {
        let thumb = img.thumbnail(thumb_w, thumb_h);

        thumb
            .save(format!("./uploads/thumb/{thumb_name}"))
            .map_err(|_| {
                NekrochanError::FileError(original_name, "nepodařilo se vytvořit náhled obrázku")
            })
    })
    .await??;

    Ok((width, height))
}

async fn process_video(
    cfg: &Cfg,
    original_name: String,
    new_name: String,
    thumb_name: Option<String>,
) -> Result<(u32, u32), NekrochanError> {
    let new_name_ = new_name.clone();

    let ffprobe_out = spawn_blocking(move || {
        Command::new("ffprobe")
            .args([
                "-v",
                "error",
                "-select_streams",
                "v:0",
                "-show_entries",
                "stream=width,height",
                "-of",
                "csv=s=x:p=0",
                &format!("/tmp/{new_name_}"),
            ])
            .output()
    })
    .await??;

    if !ffprobe_out.status.success() {
        return Err(NekrochanError::FileError(
            original_name,
            "nepodařilo se získat rozměry videa",
        ));
    }

    let invalid_dimensions = "ffprobe vrátil neplatné rozměry";

    let out_string = String::from_utf8_lossy(&ffprobe_out.stdout);

    let (width, height) = out_string
        .trim()
        .split_once('x')
        .ok_or(NekrochanError::FileError(
            original_name.clone(),
            invalid_dimensions,
        ))?;

    let (width, height) = (
        width
            .parse()
            .map_err(|_| NekrochanError::FileError(original_name.clone(), invalid_dimensions))?,
        height
            .parse()
            .map_err(|_| NekrochanError::FileError(original_name.clone(), invalid_dimensions))?,
    );

    if width > cfg.files.max_width || height > cfg.files.max_height {
        return Err(NekrochanError::FileError(
            original_name,
            "rozměry videa jsou příliš velké",
        ));
    }

    let thumb_name = match thumb_name {
        Some(thumb_name) => thumb_name,
        None => return Ok((width, height)),
    };

    let thumb_size = cfg.files.thumb_size;

    let output = spawn_blocking(move || {
        Command::new("ffmpeg")
            .args([
                "-i",
                &format!("/tmp/{new_name}"),
                "-ss",
                "00:00:00.50",
                "-vframes",
                "1",
                "-vf",
                &format!(
                    "scale={}",
                    if width > height {
                        format!("{thumb_size}:-2")
                    } else {
                        format!("-2:{thumb_size}")
                    }
                ),
                &format!("./uploads/thumb/{thumb_name}"),
            ])
            .output()
    })
    .await??;

    if !output.status.success() {
        return Err(NekrochanError::FileError(
            original_name,
            "nepodařilo se vytvořit náhled videa",
        ));
    }

    Ok((width, height))
}

pub async fn cleanup_files(ctx: &Ctx) -> Result<(), Error> {
    let mut keep = HashSet::new();
    let mut keep_thumbs = HashSet::new();

    let boards = Board::read_all(ctx).await?;

    for board in boards.into_iter() {
        for file in board.banners.0 {
            keep.insert(format!("{}.{}", file.timestamp, file.format));
        }

        let posts = Post::read_all(ctx, board.id.clone()).await?;

        for post in posts.into_iter() {
            for file in post.files.0 {
                keep.insert(format!("{}.{}", file.timestamp, file.format));

                if let Some(thumb_format) = file.thumb_format {
                    keep_thumbs.insert(format!("{}.{}", file.timestamp, thumb_format));
                }
            }
        }
    }

    for file in glob("uploads/*.*")? {
        let file = file?;
        let file_name = file.file_name();

        if let Some(file_name) = file_name {
            let check = file_name.to_string_lossy().to_string();

            if !keep.contains(&check) {
                remove_file(file).await?;
            }
        }
    }

    for file in glob("uploads/thumb/*.*")? {
        let file = file?;
        let file_name = file.file_name();

        if let Some(file_name) = file_name {
            let check = file_name.to_string_lossy().to_string();

            if !keep_thumbs.contains(&check) {
                remove_file(file).await?;
            }
        }
    }

    Ok(())
}