drawer/src/parsing.rs

use std::ffi::OsString;
use std::io::Write;
use std::path::{Path, PathBuf};
use age::ssh::Identity;
use flate2::Compression;
use flate2::read::GzDecoder;
use flate2::write::GzEncoder;
use clap::{Parser, Subcommand};
use ssh_key::{Algorithm, LineEnding, PrivateKey};
use crate::error::DrawerError;
use crate::parsing::KeyType::*;

#[derive(Debug, Copy, Clone, PartialEq)]
pub enum OperationType {
    Open,
    Close,
    Key,
}

#[derive(Debug, PartialEq)]
pub struct DrawerOperation {
    pub operation_type: OperationType,
    pub drawer_file: PathBuf,
    pub key: KeyType,
    pub target_path: PathBuf,
    pub force: bool,
    pub target_inferred: bool,
}

#[derive(Debug, PartialEq)]
pub enum KeyType {
    EnvVar, Path(PathBuf)
}

#[derive(Parser)]
#[command(name = "drawer", about = "Encrypt and decrypt project directories as drawer files")]
struct Cli {
    #[command(subcommand)]
    command: Command,
}

#[derive(Subcommand)]
enum Command {
    /// Decrypt and expand a drawer file into a directory
    Open {
        drawer_file: PathBuf,
        target_path: Option<PathBuf>,
        #[arg(short = 'i')]
        key: Option<PathBuf>,
        #[arg(short = 'f')]
        force: bool,
    },
    /// Compress and encrypt a directory into a drawer file
    Close {
        drawer_file: PathBuf,
        target_path: Option<PathBuf>,
        #[arg(short = 'i')]
        key: Option<PathBuf>,
    },
    /// Generate a new SSH key and save it to a file
    Key {
        filename: PathBuf,
    },
}

pub fn parse_args<I, T>(args: I) -> Result<DrawerOperation, DrawerError>
where
    I: IntoIterator<Item = T>,
    T: Into<OsString> + Clone,
{
    let cli = Cli::try_parse_from(args)
        .map_err(|e| DrawerError::ParseError(e.to_string()))?;

    match cli.command {
        Command::Open { drawer_file, target_path, key, force } => {
            let key = resolve_key(key);
            let (target, target_inferred) = match target_path {
                Some(p) => (p, false),
                None => (PathBuf::from(drawer_file.file_stem().unwrap_or_default()), true),
            };
            Ok(DrawerOperation {
                operation_type: OperationType::Open,
                drawer_file,
                key,
                target_path: target,
                force,
                target_inferred,
            })
        }
        Command::Close { drawer_file, target_path, key } => {
            let key = resolve_key(key);
            let (target, target_inferred) = match target_path {
                Some(p) => (p, false),
                None => (PathBuf::from(drawer_file.file_stem().unwrap_or_default()), true),
            };
            Ok(DrawerOperation {
                operation_type: OperationType::Close,
                drawer_file,
                key,
                target_path: target,
                force: false,
                target_inferred
            })
        }
        Command::Key { filename } => {
            Ok(DrawerOperation {
                operation_type: OperationType::Key,
                drawer_file: filename,
                key: KeyType::EnvVar,
                target_path: PathBuf::from("."),
                force: false,
                target_inferred: false,
            })
        }
    }
}

fn resolve_key(key: Option<PathBuf>) -> KeyType {
    if let Some(k) = key {
        return Path(k);
    }
    EnvVar
}

fn single_file_entry(data: &[u8]) -> Option<PathBuf> {
    let cursor = std::io::Cursor::new(data);
    let gz = GzDecoder::new(cursor);
    let mut archive = tar::Archive::new(gz);
    let mut entries = archive.entries().ok()?;
    let entry = entries.next()?.ok()?;
    if !entry.header().entry_type().is_file() {
        return None;
    }
    if entries.next().is_some() {
        return None;
    }
    Some(entry.path().ok()?.into_owned())
}

impl DrawerOperation {
    pub fn perform(&self) -> Result<(), DrawerError> {
        self.validate()?;
        match self.operation_type {
            OperationType::Open => self.perform_open(),
            OperationType::Close => self.perform_close(),
            OperationType::Key => self.perform_key(),
        }
    }

    fn perform_open(&self) -> Result<(), DrawerError> {
        let encrypted = std::fs::read(&self.drawer_file)
            .map_err(|_| DrawerError::DrawerFileNotFound(self.drawer_file.clone()))?;

        let decrypted = self.decrypt(encrypted)?;

        if let Some(entry_name) = single_file_entry(&decrypted) {
            let output = if self.target_inferred {
                PathBuf::from(entry_name.file_name().unwrap_or_default())
            } else {
                self.target_path.clone()
            };
            let cursor = std::io::Cursor::new(&decrypted);
            let gz = GzDecoder::new(cursor);
            let mut archive = tar::Archive::new(gz);
            let mut entry = archive.entries()
                .map_err(|_| DrawerError::UntarFailed)?
                .next()
                .ok_or(DrawerError::UntarFailed)?
                .map_err(|_| DrawerError::UntarFailed)?;
            entry.unpack(&output).map_err(|_| DrawerError::UntarFailed)?;
        } else {
            let cursor = std::io::Cursor::new(decrypted);
            let gz = GzDecoder::new(cursor);
            let mut archive = tar::Archive::new(gz);
            archive.unpack(&self.target_path)
                .map_err(|_| DrawerError::UntarFailed)?;
        }
        Ok(())
    }

    fn decrypt(&self, encrypted: Vec<u8>) -> Result<Vec<u8>, DrawerError> {
        let (key_path, identity) = self.identity()?;

        let decryptor = age::Decryptor::new(encrypted.as_slice())
            .map_err(|_| DrawerError::DecryptFailed(self.drawer_file.clone(), key_path.clone()))?;
        let mut decrypted: Vec<u8> = Vec::new();
        let mut stream = decryptor.decrypt(std::iter::once(&identity as &dyn age::Identity))
            .map_err(|_| DrawerError::DecryptFailed(self.drawer_file.clone(), key_path.clone()))?;
        std::io::Read::read_to_end(&mut stream, &mut decrypted)
            .map_err(|_| DrawerError::DecryptFailed(self.drawer_file.clone(), key_path.clone()))?;
        Ok(decrypted)
    }

    fn identity(&self) -> Result<(PathBuf, Identity), DrawerError> {
        let key_path = self.key_path()?;
        let file = std::fs::File::open(&key_path)
            .map_err(|_| DrawerError::KeyInvalid(key_path.clone()))?;
        let reader = std::io::BufReader::new(file);
        let identity = age::ssh::Identity::from_buffer(reader, None)
            .map_err(|_| DrawerError::DecryptFailed(self.drawer_file.clone(), key_path.clone()))?;
        Ok((key_path, identity))
    }

    fn perform_close(&self) -> Result<(), DrawerError> {
        let mut buf: Vec<u8> = Vec::new();
        {
            let gz = GzEncoder::new(&mut buf, Compression::default());
            let mut archive = tar::Builder::new(gz);
            if self.target_path.is_dir() {
                archive.append_dir_all(".", &self.target_path)
                    .map_err(|_| DrawerError::TarFailed)?;
            } else {
                let name = self.target_path.file_name().unwrap_or_default();
                archive.append_path_with_name(&self.target_path, name)
                    .map_err(|_| DrawerError::TarFailed)?;
            }
            let gz = archive.into_inner()
                .map_err(|_| DrawerError::TarFailed)?;
            gz.finish()
                .map_err(|_| DrawerError::TarFailed)?;
        }

        let encrypted = self.encrypt(&buf)?;

        std::fs::write(&self.drawer_file, &encrypted)
            .map_err(|_| DrawerError::WriteFailed(self.drawer_file.clone()))
    }

    fn encrypt(&self, buf: &[u8]) -> Result<Vec<u8>, DrawerError> {
        let (_, identity) = self.identity()?;
        let recipient = age::ssh::Recipient::try_from(identity)
            .map_err(|_| DrawerError::EncryptFailed)?;

        let mut encrypted: Vec<u8> = Vec::new();
        let encryptor = age::Encryptor::with_recipients(
            std::iter::once(&recipient as &dyn age::Recipient)
        ).map_err(|_| DrawerError::EncryptFailed)?;
        let mut writer = encryptor.wrap_output(&mut encrypted)
            .map_err(|_| DrawerError::EncryptFailed)?;
        writer.write_all(buf).map_err(|_| DrawerError::EncryptFailed)?;
        writer.finish().map_err(|_| DrawerError::EncryptFailed)?;
        Ok(encrypted)
    }

    pub fn validate(&self) -> Result<(), DrawerError> {
        if self.operation_type != OperationType::Key {
            if self.drawer_file.extension().and_then(|e| e.to_str()) != Some("drawer") {
                return Err(DrawerError::DrawerFileInvalidExtension(self.drawer_file.clone()));
            }
            self.validate_key()?;
        }
        match self.operation_type {
            OperationType::Open => self.validate_open(),
            OperationType::Close => self.validate_close(),
            OperationType::Key => self.validate_key_gen(),
        }
    }

    fn validate_key_gen(&self) -> Result<(), DrawerError> {
        if self.drawer_file.exists() {
            return Err(DrawerError::KeyOutputFailed(self.drawer_file.clone()));
        }
        let parent = match self.drawer_file.parent() {
            None => Path::new("."),
            Some(p) => {
                if p == Path::new("") {
                    Path::new(".")
                } else {
                    p
                }
            },
        };
        if !parent.exists() {
            return Err(DrawerError::KeyOutputFailed(self.drawer_file.clone()));
        }
        Ok(())
    }

    fn key_path(&self) -> Result<PathBuf, DrawerError> {
        match &self.key {
            EnvVar => {
                let val = std::env::var("DRAWER_KEY").map_err(|_| DrawerError::NoKeyPath)?;
                Ok(PathBuf::from(val))
            }
            Path(p) => Ok(p.clone()),
        }
    }

    fn validate_key(&self) -> Result<(), DrawerError> {
        let key_path = self.key_path()?;
        let file = std::fs::File::open(&key_path)
            .map_err(|_| DrawerError::KeyInvalid(key_path.clone()))?;
        let reader = std::io::BufReader::new(file);
        let identity = age::ssh::Identity::from_buffer(reader, None)
            .map_err(|_| DrawerError::KeyNotValidSsh)?;
        age::ssh::Recipient::try_from(identity)
            .map_err(|_| DrawerError::KeyNotValidSsh)?;
        Ok(())
    }

    fn validate_open(&self) -> Result<(), DrawerError> {
        if self.target_inferred {
            println!("(assuming into {})", self.target_path.display());
        }
        if !self.drawer_file.exists() {
            return Err(DrawerError::DrawerFileNotFound(self.drawer_file.clone()));
        }
        if !self.force && self.target_path.exists() {
            if self.target_path.is_file() {
                return Err(DrawerError::TargetNotEmpty(self.target_path.clone()));
            }
            let mut entries = std::fs::read_dir(&self.target_path)
                .map_err(|_| DrawerError::TargetInvalid(self.target_path.clone()))?;
            if entries.next().is_some() {
                return Err(DrawerError::TargetNotEmpty(self.target_path.clone()));
            }
        }
        Ok(())
    }

    fn perform_key(&self) -> Result<(), DrawerError> {
        let private_key = PrivateKey::random(&mut rand::rngs::OsRng, Algorithm::Ed25519)
            .map_err(|_| DrawerError::KeyGenerateFailed)?;
        let pem = private_key.to_openssh(LineEnding::LF)
            .map_err(|_| DrawerError::KeyGenerateFailed)?;
        std::fs::write(&self.drawer_file, pem.as_bytes())
            .map_err(|_| DrawerError::KeyOutputFailed(self.drawer_file.clone()))
    }

    fn validate_close(&self) -> Result<(), DrawerError> {
        if self.target_inferred {
            println!("(assuming from {})", self.target_path.display());
        }
        let meta = std::fs::metadata(&self.target_path)
            .map_err(|_| DrawerError::TargetInvalid(self.target_path.clone()))?;
        if !meta.is_dir() && !meta.is_file() {
            return Err(DrawerError::TargetInvalid(self.target_path.clone()));
        }
        if meta.is_dir() && std::fs::read_dir(&self.target_path).is_err() {
            return Err(DrawerError::TargetInvalid(self.target_path.clone()));
        }
        let canon_target = self.target_path.canonicalize()
            .unwrap_or_else(|_| self.target_path.clone());
        let drawer_parent = match self.drawer_file.parent() {
            None => Path::new("."),
            Some(p) if p == Path::new("") => {
                Path::new(".")
            },
            Some(p) => p,
        };
        let canon_drawer = drawer_parent.canonicalize()
            .unwrap_or_else(|_| drawer_parent.to_path_buf())
            .join(self.drawer_file.file_name().unwrap_or_default());
        if canon_drawer.starts_with(&canon_target) {
            return Err(DrawerError::DrawerInsideTarget);
        }
        Ok(())
    }
}