Skip to main content

bunkr_client/preprocess/
preprocess.rs

1use crate::config::config::Config;
2use anyhow::{Result, anyhow};
3use mime_guess::from_path;
4use std::path::Path;
5use std::process::Command;
6use uuid::Uuid;
7
8pub struct PreprocessResult {
9    pub files_to_upload: Vec<String>,
10    pub preprocess_id: String,
11}
12
13pub fn preprocess_file(path: &str, max_file_size: u64, config: &Config) -> Result<PreprocessResult> {
14    let p = Path::new(path);
15    let mime = from_path(p).first_or_octet_stream();
16
17    // Video preprocessing
18    if mime.type_() == mime_guess::mime::VIDEO && config.preprocess_videos.unwrap_or(true) {
19        let metadata = p.metadata()?;
20        let size = metadata.len();
21        if size > max_file_size {
22            let parts = split_video(path, max_file_size)?;
23            return Ok(PreprocessResult {
24                files_to_upload: parts,
25                preprocess_id: "split_video".to_string(),
26            });
27        }
28    }
29
30    // Default: no preprocessing
31    Ok(PreprocessResult {
32        files_to_upload: vec![path.to_string()],
33        preprocess_id: "original".to_string(),
34    })
35}
36
37pub fn cleanup_preprocess(preprocess_id: &str, _original_path: &str, files_to_upload: &[String]) {
38    match preprocess_id {
39        "original" => {
40            // Nothing to clean up
41        }
42        "split_video" => {
43            for file in files_to_upload {
44                let _ = std::fs::remove_file(file);
45            }
46            // Try to remove the temp directory if it's empty
47            if let Some(first_file) = files_to_upload.first() {
48                if let Some(parent) = Path::new(first_file).parent() {
49                    let _ = std::fs::remove_dir(parent);
50                }
51            }
52        }
53        _ => {
54            // Unknown preprocess, do nothing
55        }
56    }
57}
58
59fn split_video(path: &str, max_file_size: u64) -> Result<Vec<String>> {
60    let p = Path::new(path);
61    let stem = p.file_stem().unwrap().to_string_lossy();
62    let extension = p.extension().unwrap_or_default().to_string_lossy();
63
64    let parent_dir = p.parent().unwrap_or(Path::new("."));
65    let temp_dir = parent_dir.join(format!("bunkr_split_{}", Uuid::new_v4()));
66    std::fs::create_dir_all(&temp_dir)?;
67
68    let hwaccel = detect_hwaccel();
69    let output = Command::new("ffprobe")
70        .args(&[
71            "-v", "quiet",
72            "-show_entries", "format=duration",
73            "-of", "default=noprint_wrappers=1:nokey=1",
74            path
75        ])
76        .output()?;
77    if !output.status.success() {
78        return Err(anyhow!("Failed to get video duration: {}", String::from_utf8_lossy(&output.stderr)));
79    }
80    let duration_str = String::from_utf8(output.stdout)?;
81    let duration: f64 = duration_str.trim().parse()?;
82
83    let metadata = p.metadata()?;
84    let size = metadata.len();
85    let parts = (size as f64 / max_file_size as f64).ceil() as u32;
86    let segment_time = duration / parts as f64;
87    let output_pattern = temp_dir.join(format!("{}_%03d.{}", stem, extension)).to_string_lossy().to_string();
88
89    // Build ffmpeg args
90    let mut args = vec![];
91    if let Some(accel) = hwaccel {
92        args.push("-hwaccel".to_string());
93        args.push(accel);
94    }
95    args.push("-loglevel".to_string());
96    args.push("quiet".to_string());
97    args.push("-i".to_string());
98    args.push(path.to_string());
99    args.push("-f".to_string());
100    args.push("segment".to_string());
101    args.push("-segment_time".to_string());
102    args.push(segment_time.to_string());
103    args.push("-c".to_string());
104    args.push("copy".to_string());
105    args.push("-reset_timestamps".to_string());
106    args.push("1".to_string());
107    args.push(output_pattern);
108
109    let status = Command::new("ffmpeg")
110        .args(&args)
111        .status()?;
112    if !status.success() {
113        return Err(anyhow!("Failed to split video"));
114    }
115
116    let mut result = vec![];
117    for i in 0..parts {
118        let part_path = temp_dir.join(format!("{}_{:03}.{}", stem, i, extension)).to_string_lossy().to_string();
119        // Check if file exists and size <= max_file_size
120        if let Ok(meta) = std::fs::metadata(&part_path) {
121            if meta.len() <= max_file_size {
122                result.push(part_path);
123            } else {
124                // If still too big, perhaps further split, but for now, include anyway
125                result.push(part_path);
126            }
127        }
128    }
129
130    Ok(result)
131}
132
133fn detect_hwaccel() -> Option<String> {
134    let output = Command::new("ffmpeg").arg("-hwaccels").output();
135    match output {
136        Ok(out) if out.status.success() => {
137            let stdout = String::from_utf8_lossy(&out.stdout);
138            let lines: Vec<&str> = stdout.lines().collect();
139            if let Some(pos) = lines.iter().position(|l| l.contains("Hardware acceleration methods:")) {
140                for line in lines.iter().skip(pos + 1) {
141                    let trimmed = line.trim();
142                    if !trimmed.is_empty() && trimmed != "none" {
143                        return Some(trimmed.to_string());
144                    }
145                }
146            }
147        }
148        _ => {}
149    }
150    None
151}