keep/src/compression_engine.rs

use anyhow::{anyhow, Context, Result};
use std::io;
use std::io::{Read, Write};
use std::path::PathBuf;
use std::process::{Command, Stdio};
use strum::IntoEnumIterator;

use log::*;

use lazy_static::lazy_static;

extern crate enum_map;
use enum_map::enum_map;
use enum_map::{Enum, EnumMap};

pub mod gzip;
pub mod lz4;
pub mod none;
pub mod program;

use crate::compression_engine::gzip::CompressionEngineGZip;
use crate::compression_engine::lz4::CompressionEngineLZ4;
use crate::compression_engine::none::CompressionEngineNone;
use crate::compression_engine::program::CompressionEngineProgram;

#[derive(Debug, Eq, PartialEq, Clone, strum::EnumIter, strum::Display, strum::EnumString, Enum)]
#[strum(ascii_case_insensitive)]
pub enum CompressionType {
    LZ4,
    GZip,
    BZip2,
    XZ,
    ZStd,
    None,
}

pub trait CompressionEngine {
    fn open(&self, file_path: PathBuf) -> Result<Box<dyn Read>>;
    fn create(&self, file_path: PathBuf) -> Result<Box<dyn Write>>;

    fn is_supported(&self) -> bool {
        true
    }

    fn copy(&self, file_path: PathBuf, writer: &mut dyn Write) -> Result<()> {
        let mut reader = self.open(file_path)?;
        io::copy(&mut reader, writer)?;
        writer.flush()?;
        Ok(())
    }

    fn cat(&self, file_path: PathBuf) -> Result<()> {
        let mut stdout = io::stdout().lock();
        self.copy(file_path, &mut stdout)
    }

    fn magic(&self, file_path: PathBuf) -> Result<String> {
        let mut buffer = [0; libc::BUFSIZ as usize];
        let mut reader = self.open(file_path)?;
        let buffer_size = reader.read(&mut buffer[..])?;

        //let cookie = magic::Cookie::open(magic::cookie::Flags::ERROR)?;

        // load the system's default database
        //let database = Default::default();
        //let cookie = cookie.load(&database).map_err(|error| { anyhow!(error.to_string())})?;
        //let magic = cookie.buffer(&buffer[0..n])?;

        let program = "file";
        let args = ["-bE", "-"];

        let process = Command::new(program)
            .args(args)
            .stdin(Stdio::piped())
            .stdout(Stdio::piped())
            .spawn()
            .context(anyhow!(
                "Unable to spawn child process: {:?} {:?}",
                program,
                args
            ))?;

        let mut stdin = process.stdin.unwrap();
        stdin.write_all(&buffer[0..buffer_size])?;
        drop(stdin);

        let mut magic = String::new();
        process.stdout.unwrap().read_to_string(&mut magic)?;
        Ok(magic)
    }

    fn size(&self, file_path: PathBuf) -> Result<usize> {
        let mut reader = self.open(file_path)?;
        let mut buffer = [0; libc::BUFSIZ as usize];
        let mut size: usize = 0;

        loop {
            let n = reader.read(&mut buffer[..libc::BUFSIZ as usize])?;
            if n == 0 {
                debug!("COMPRESSION: EOF");
                break;
            }

            size += n;
        }

        Ok(size)
    }
}

lazy_static! {
    pub static ref COMPRESSION_PROGRAMS: EnumMap<CompressionType, Option<CompressionEngineProgram>> = enum_map! {
        CompressionType::LZ4 => None,
        CompressionType::GZip => None,
        CompressionType::BZip2 => {
            let program = CompressionEngineProgram::new("bzip2", vec!["-qcf"], vec!["-dcf"]);
            if program.supported { Some(program) } else { None }
        },
        CompressionType::XZ => {
            let program = CompressionEngineProgram::new("xz", vec!["-qcf"], vec!["-dcf"]);
            if program.supported { Some(program) } else { None }
        },
        CompressionType::ZStd => {
            let program = CompressionEngineProgram::new("zstd", vec!["-qcf"], vec!["-dcf"]);
            if program.supported { Some(program) } else { None }
        },
        CompressionType::None => None
    };
}

pub fn get_compression_engine(compression_type: CompressionType) -> Result<Box<dyn CompressionEngine>> {
    match compression_type {
        CompressionType::LZ4 => Ok(Box::new(CompressionEngineLZ4::new())),
        CompressionType::GZip => Ok(Box::new(CompressionEngineGZip::new())),
        CompressionType::None => Ok(Box::new(CompressionEngineNone::new())),
        compression_type => Ok(Box::new(
            COMPRESSION_PROGRAMS[compression_type.clone()]
                .clone()
                .unwrap(),
        )),
    }
}

pub fn default_compression_type() -> CompressionType {
    let mut default = CompressionType::None;
    for compression_type in CompressionType::iter() {
        let compression_engine = get_compression_engine(compression_type.clone()).expect("Missing engine");
        if compression_engine.is_supported() {
            default = compression_type;
            break;
        }
    }
    default
}