scpcap v0

src/naming.rs

@@ -0,0 +1,258 @@
+use anyhow::{bail, Result};
+use std::path::Path;
+
+/// Either a local filesystem path or a `[user@]host:path` remote target.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum Target {
+    Local(String),
+    Remote { host: String, path: String },
+}
+
+impl Target {
+    pub fn parse(arg: &str) -> Target {
+        // scp-style: split on the first ':' that comes before the first '/'.
+        // A bare Windows-style drive letter isn't a concern (Linux-only).
+        if let Some(colon) = arg.find(':') {
+            let before_colon = &arg[..colon];
+            if !before_colon.is_empty() && !before_colon.contains('/') {
+                return Target::Remote {
+                    host: before_colon.to_string(),
+                    path: arg[colon + 1..].to_string(),
+                };
+            }
+        }
+        Target::Local(arg.to_string())
+    }
+
+    pub fn path(&self) -> &str {
+        match self {
+            Target::Local(p) => p,
+            Target::Remote { path, .. } => path,
+        }
+    }
+}
+
+/// Strip the partial-suffix (e.g. ".partial") first, then a trailing ".zst" if
+/// present. Order matters: "cap.pcap.zst.partial" must be recognized as a
+/// growing, zst-compressed source, not misclassified.
+pub struct StrippedName {
+    pub logical: String,
+    pub was_partial: bool,
+    pub was_zst: bool,
+}
+
+pub fn strip_name(name: &str, partial_ext: &str) -> StrippedName {
+    let suffix = format!(".{partial_ext}");
+    let (after_partial, was_partial) = match name.strip_suffix(&suffix) {
+        Some(stripped) => (stripped, true),
+        None => (name, false),
+    };
+    let (logical, was_zst) = match after_partial.strip_suffix(".zst") {
+        Some(stripped) => (stripped, true),
+        None => (after_partial, false),
+    };
+    StrippedName {
+        logical: logical.to_string(),
+        was_partial,
+        was_zst,
+    }
+}
+
+/// Final (non-partial) path for a source, given its literal argument path.
+pub fn final_source_path(source_path: &str, partial_ext: &str) -> String {
+    let suffix = format!(".{partial_ext}");
+    match source_path.strip_suffix(&suffix) {
+        Some(stripped) => stripped.to_string(),
+        None => source_path.to_string(),
+    }
+}
+
+pub struct DestPaths {
+    pub final_path: String,
+    pub temp_path: String,
+}
+
+/// Resolve the destination's final and temp (in-flight) paths from the literal
+/// DEST argument path, the source's stripped logical name, and whether
+/// compression is actually being applied.
+pub fn resolve_dest_paths(
+    dest_path: &str,
+    source_stripped: &StrippedName,
+    effective_compress: bool,
+    partial_ext: &str,
+) -> DestPaths {
+    // DEST may name a directory (trailing '/') -- in that case reuse the
+    // source's logical basename underneath it. Otherwise DEST names the file
+    // directly.
+    let dest_logical = if dest_path.ends_with('/') {
+        let base = Path::new(&source_stripped.logical)
+            .file_name()
+            .map(|n| n.to_string_lossy().to_string())
+            .unwrap_or_else(|| source_stripped.logical.clone());
+        format!("{dest_path}{base}")
+    } else {
+        dest_path.to_string()
+    };
+
+    let final_path = if effective_compress && !dest_logical.ends_with(".zst") {
+        format!("{dest_logical}.zst")
+    } else {
+        dest_logical
+    };
+
+    let temp_path = format!("{final_path}.{partial_ext}");
+
+    DestPaths {
+        final_path,
+        temp_path,
+    }
+}
+
+/// Parse a human chunk-size string like "256k", "1m", or a bare byte count.
+pub fn parse_size(s: &str) -> Result<usize> {
+    let s = s.trim();
+    if s.is_empty() {
+        bail!("empty size");
+    }
+    let (digits, mult) = match s.chars().last().unwrap() {
+        'k' | 'K' => (&s[..s.len() - 1], 1024),
+        'm' | 'M' => (&s[..s.len() - 1], 1024 * 1024),
+        'g' | 'G' => (&s[..s.len() - 1], 1024 * 1024 * 1024),
+        _ => (s, 1),
+    };
+    let n: usize = digits
+        .trim()
+        .parse()
+        .map_err(|_| anyhow::anyhow!("invalid size: {s:?}"))?;
+    Ok(n * mult)
+}
+
+/// Single-quote a string for safe inclusion as one argument in a remote shell
+/// command line, e.g. for `ssh host sh -c '...'`-style invocations.
+pub fn shell_quote(s: &str) -> String {
+    let mut out = String::with_capacity(s.len() + 2);
+    out.push('\'');
+    for c in s.chars() {
+        if c == '\'' {
+            out.push_str("'\\''");
+        } else {
+            out.push(c);
+        }
+    }
+    out.push('\'');
+    out
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn target_parse_local() {
+        assert_eq!(Target::parse("/a/b/c"), Target::Local("/a/b/c".into()));
+        assert_eq!(Target::parse("relative/path"), Target::Local("relative/path".into()));
+    }
+
+    #[test]
+    fn target_parse_remote() {
+        assert_eq!(
+            Target::parse("host:/a/b"),
+            Target::Remote { host: "host".into(), path: "/a/b".into() }
+        );
+        assert_eq!(
+            Target::parse("user@host:/a/b"),
+            Target::Remote { host: "user@host".into(), path: "/a/b".into() }
+        );
+    }
+
+    #[test]
+    fn target_parse_local_with_colon_after_slash() {
+        // A colon appearing after a '/' isn't a host separator.
+        assert_eq!(Target::parse("/a/b:c"), Target::Local("/a/b:c".into()));
+    }
+
+    #[test]
+    fn strip_order_partial_then_zst() {
+        let s = strip_name("cap.pcap.zst.partial", "partial");
+        assert_eq!(s.logical, "cap.pcap");
+        assert!(s.was_partial);
+        assert!(s.was_zst);
+    }
+
+    #[test]
+    fn strip_plain_partial() {
+        let s = strip_name("cap.pcap.partial", "partial");
+        assert_eq!(s.logical, "cap.pcap");
+        assert!(s.was_partial);
+        assert!(!s.was_zst);
+    }
+
+    #[test]
+    fn strip_already_final() {
+        let s = strip_name("cap.pcap", "partial");
+        assert_eq!(s.logical, "cap.pcap");
+        assert!(!s.was_partial);
+        assert!(!s.was_zst);
+    }
+
+    #[test]
+    fn strip_custom_extension() {
+        let s = strip_name("cap.pcap.zst.inprogress", "inprogress");
+        assert_eq!(s.logical, "cap.pcap");
+        assert!(s.was_partial);
+        assert!(s.was_zst);
+    }
+
+    #[test]
+    fn final_source_path_strips_suffix() {
+        assert_eq!(final_source_path("/x/cap.pcap.partial", "partial"), "/x/cap.pcap");
+        assert_eq!(final_source_path("/x/cap.pcap", "partial"), "/x/cap.pcap");
+    }
+
+    #[test]
+    fn dest_paths_append_zst_when_compressing() {
+        let src = strip_name("cap.pcap.partial", "partial");
+        let d = resolve_dest_paths("/out/cap.pcap", &src, true, "partial");
+        assert_eq!(d.final_path, "/out/cap.pcap.zst");
+        assert_eq!(d.temp_path, "/out/cap.pcap.zst.partial");
+    }
+
+    #[test]
+    fn dest_paths_no_double_zst() {
+        let src = strip_name("cap.pcap.zst.partial", "partial");
+        let d = resolve_dest_paths("/out/cap.pcap.zst", &src, false, "partial");
+        assert_eq!(d.final_path, "/out/cap.pcap.zst");
+        assert_eq!(d.temp_path, "/out/cap.pcap.zst.partial");
+    }
+
+    #[test]
+    fn dest_paths_no_compress() {
+        let src = strip_name("cap.pcap.partial", "partial");
+        let d = resolve_dest_paths("/out/cap.pcap", &src, false, "partial");
+        assert_eq!(d.final_path, "/out/cap.pcap");
+        assert_eq!(d.temp_path, "/out/cap.pcap.partial");
+    }
+
+    #[test]
+    fn dest_paths_directory_target() {
+        let src = strip_name("cap.pcap.partial", "partial");
+        let d = resolve_dest_paths("/out/", &src, false, "partial");
+        assert_eq!(d.final_path, "/out/cap.pcap");
+    }
+
+    #[test]
+    fn size_parsing() {
+        assert_eq!(parse_size("256k").unwrap(), 256 * 1024);
+        assert_eq!(parse_size("1m").unwrap(), 1024 * 1024);
+        assert_eq!(parse_size("2M").unwrap(), 2 * 1024 * 1024);
+        assert_eq!(parse_size("512").unwrap(), 512);
+        assert!(parse_size("").is_err());
+        assert!(parse_size("abc").is_err());
+    }
+
+    #[test]
+    fn quoting() {
+        assert_eq!(shell_quote("/plain/path"), "'/plain/path'");
+        assert_eq!(shell_quote("it's"), "'it'\\''s'");
+    }
+}