datman/

tree.rs

/*
This file is part of Yama.

Yama is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

Yama is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with Yama.  If not, see <https://www.gnu.org/licenses/>.
*/


use std::collections::BTreeMap;
use std::fmt::Debug;
use std::fs::{read_link, symlink_metadata, DirEntry, Metadata};
use std::io::ErrorKind;
use std::os::unix::fs::MetadataExt;
use std::path::Path;

use anyhow::anyhow;
use indicatif::{ProgressBar, ProgressDrawTarget, ProgressStyle};
use log::warn;
use serde::{Deserialize, Serialize};

pub use yama::definitions::FilesystemOwnership;
pub use yama::definitions::FilesystemPermissions;

#[derive(Debug, Clone, Serialize, Deserialize, Eq, PartialEq)]
pub enum FileTree<NMeta, DMeta, SMeta, Other>
where
    NMeta: Debug + Clone + Eq + PartialEq,
    DMeta: Debug + Clone + Eq + PartialEq,
    SMeta: Debug + Clone + Eq + PartialEq,
    Other: Debug + Clone + Eq + PartialEq,
{
    NormalFile {
        /// modification time in ms
        mtime: u64,
        ownership: FilesystemOwnership,
        permissions: FilesystemPermissions,
        meta: NMeta,
    },
    Directory {
        ownership: FilesystemOwnership,
        permissions: FilesystemPermissions,
        children: BTreeMap<String, FileTree<NMeta, DMeta, SMeta, Other>>,
        meta: DMeta,
    },
    SymbolicLink {
        ownership: FilesystemOwnership,
        target: String,
        meta: SMeta,
    },
    Other(Other),
}

pub type FileTree1<A> = FileTree<A, A, A, ()>;

impl<NMeta, DMeta, SMeta, Other> FileTree<NMeta, DMeta, SMeta, Other>
where
    NMeta: Debug + Clone + Eq + PartialEq,
    DMeta: Debug + Clone + Eq + PartialEq,
    SMeta: Debug + Clone + Eq + PartialEq,
    Other: Debug + Clone + Eq + PartialEq,
{
    pub fn is_dir(&self) -> bool {
        match self {
            FileTree::NormalFile { .. } => false,
            FileTree::Directory { .. } => true,
            FileTree::SymbolicLink { .. } => false,
            FileTree::Other(_) => false,
        }
    }

    pub fn is_symlink(&self) -> bool {
        match self {
            FileTree::NormalFile { .. } => false,
            FileTree::Directory { .. } => false,
            FileTree::SymbolicLink { .. } => true,
            FileTree::Other(_) => false,
        }
    }

    pub fn get_by_path(&self, path: &String) -> Option<&FileTree<NMeta, DMeta, SMeta, Other>> {
        let mut node = self;
        for piece in path.split('/') {
            if piece.is_empty() {
                continue;
            }
            match node {
                FileTree::Directory { children, .. } => match children.get(piece) {
                    None => {
                        return None;
                    }
                    Some(new_node) => {
                        node = new_node;
                    }
                },
                _ => {
                    return None;
                }
            }
        }
        Some(node)
    }

    pub fn replace_meta<Replacement: Clone + Debug + Eq + PartialEq>(
        &self,
        replacement: &Replacement,
    ) -> FileTree<Replacement, Replacement, Replacement, Other> {
        match self {
            FileTree::NormalFile {
                mtime,
                ownership,
                permissions,
                ..
            } => FileTree::NormalFile {
                mtime: *mtime,
                ownership: *ownership,
                permissions: *permissions,
                meta: replacement.clone(),
            },
            FileTree::Directory {
                ownership,
                permissions,
                children,
                ..
            } => {
                let children = children
                    .iter()
                    .map(|(str, ft)| (str.clone(), ft.replace_meta(replacement)))
                    .collect();

                FileTree::Directory {
                    ownership: ownership.clone(),
                    permissions: permissions.clone(),
                    children,
                    meta: replacement.clone(),
                }
            }
            FileTree::SymbolicLink {
                ownership, target, ..
            } => FileTree::SymbolicLink {
                ownership: ownership.clone(),
                target: target.clone(),
                meta: replacement.clone(),
            },
            FileTree::Other(other) => FileTree::Other(other.clone()),
        }
    }

    /// Filters the tree in-place by removing nodes that do not satisfy the predicate.
    /// 'Inclusive' in the sense that if a directory does not satisfy the predicate but one of its
    /// descendants does, then the directory will be included anyway.
    /// (So nodes that satisfy the predicate will never be excluded because of a parent not doing so.)
    ///
    /// Returns true if this node should be included, and false if it should not be.
    pub fn filter_inclusive<F>(&mut self, predicate: &mut F) -> bool
    where
        F: FnMut(&Self) -> bool,
    {
        match self {
            FileTree::Directory { children, .. } => {
                let mut to_remove = Vec::new();
                for (name, child) in children.iter_mut() {
                    if !child.filter_inclusive(predicate) {
                        to_remove.push(name.clone());
                    }
                }
                for name in to_remove {
                    children.remove(&name);
                }
                !children.is_empty() || predicate(&self)
            }
            _ => predicate(&self),
        }
    }
}

impl<X: Debug + Clone + Eq, YAny: Debug + Clone + Eq> FileTree<X, X, X, YAny> {
    pub fn get_metadata(&self) -> Option<&X> {
        match self {
            FileTree::NormalFile { meta, .. } => Some(meta),
            FileTree::Directory { meta, .. } => Some(meta),
            FileTree::SymbolicLink { meta, .. } => Some(meta),
            FileTree::Other(_) => None,
        }
    }

    pub fn set_metadata(&mut self, new_meta: X) {
        match self {
            FileTree::NormalFile { meta, .. } => {
                *meta = new_meta;
            }
            FileTree::Directory { meta, .. } => {
                *meta = new_meta;
            }
            FileTree::SymbolicLink { meta, .. } => {
                *meta = new_meta;
            }
            FileTree::Other(_) => {
                // nop
            }
        }
    }
}

/// Given a file's metadata, returns the mtime in milliseconds.
pub fn mtime_msec(metadata: &Metadata) -> u64 {
    (metadata.mtime() * 1000 + metadata.mtime_nsec() / 1_000_000) as u64
}

/// Scan the filesystem to produce a Tree, using a default progress bar.
pub fn scan(path: &Path) -> anyhow::Result<Option<FileTree<(), (), (), ()>>> {
    let pbar = ProgressBar::with_draw_target(0, ProgressDrawTarget::stdout_with_hz(2));
    pbar.set_style(ProgressStyle::default_spinner().template("{spinner} {pos:7} {msg}"));
    pbar.set_message("dir scan");

    let result = scan_with_progress_bar(path, &pbar);
    pbar.finish_at_current_pos();
    result
}

/// Scan the filesystem to produce a Tree, using the specified progress bar.
pub fn scan_with_progress_bar(
    path: &Path,
    progress_bar: &ProgressBar,
) -> anyhow::Result<Option<FileTree<(), (), (), ()>>> {
    let metadata_res = symlink_metadata(path);
    progress_bar.inc(1);
    if let Err(e) = &metadata_res {
        match e.kind() {
            ErrorKind::NotFound => {
                warn!("vanished: {:?}", path);
                return Ok(None);
            }
            ErrorKind::PermissionDenied => {
                warn!("permission denied: {:?}", path);
                return Ok(None);
            }
            _ => { /* nop */ }
        }
    }
    let metadata = metadata_res?;
    let filetype = metadata.file_type();

    /*let name = path
    .file_name()
    .ok_or(anyhow!("No filename, wat"))?
    .to_str()
    .ok_or(anyhow!("Filename can't be to_str()d"))?
    .to_owned();*/

    let ownership = FilesystemOwnership {
        uid: metadata.uid() as u16,
        gid: metadata.gid() as u16,
    };

    let permissions = FilesystemPermissions {
        mode: metadata.mode(),
    };

    if filetype.is_file() {
        // Leave an unpopulated file node. It's not my responsibility to chunk it right now.
        Ok(Some(FileTree::NormalFile {
            mtime: mtime_msec(&metadata),
            ownership,
            permissions,
            meta: (),
        }))
    } else if filetype.is_dir() {
        let mut children = BTreeMap::new();
        progress_bar.set_message(&format!("{:?}", path));
        let dir_read = path.read_dir();

        if let Err(e) = &dir_read {
            match e.kind() {
                ErrorKind::NotFound => {
                    warn!("vanished/: {:?}", path);
                    return Ok(None);
                }
                ErrorKind::PermissionDenied => {
                    warn!("permission denied/: {:?}", path);
                    return Ok(None);
                }
                _ => { /* nop */ }
            }
        }

        for entry in dir_read? {
            let entry: DirEntry = entry?;
            let scanned = scan_with_progress_bar(&entry.path(), progress_bar)?;
            if let Some(scanned) = scanned {
                children.insert(
                    entry
                        .file_name()
                        .into_string()
                        .expect("OsString not String"),
                    scanned,
                );
            }
        }

        Ok(Some(FileTree::Directory {
            ownership,
            permissions,
            children,
            meta: (),
        }))
    } else if filetype.is_symlink() {
        let target = read_link(path)?
            .to_str()
            .ok_or(anyhow!("target path cannot be to_str()d"))?
            .to_owned();

        Ok(Some(FileTree::SymbolicLink {
            ownership,
            target,
            meta: (),
        }))
    } else {
        Ok(None)
    }
}