tool: Implement 'split' subcommand to split OpenPGP messages.

4 files changed, 85 insertions, 1 deletions

diff --git a/tool/src/cli.rs b/tool/src/cli.rs
index f965e4e7..eb1f4004 100644
--- a/tool/src/cli.rs
+++ b/tool/src/cli.rs
@@ -97,6 +97,18 @@ pub fn build() -> App<'static, 'static> {
                          .long("hex")
                          .short("x")
                          .help("Print a hexdump")))
+        .subcommand(SubCommand::with_name("split")
+                    .about("Splits a message into OpenPGP packets")
+                    .arg(Arg::with_name("input").value_name("FILE")
+                         .long("input")
+                         .short("i")
+                         .help("Sets the input file to use"))
+                    .arg(Arg::with_name("prefix").value_name("FILE")
+                         .long("prefix")
+                         .short("p")
+                         .help("Sets the prefix to use for output files \
+                                (defaults to the input filename with a dash, \
+                                or 'output')")))
         .subcommand(SubCommand::with_name("keyserver")
                     .about("Interacts with keyservers")
                     .arg(Arg::with_name("server").value_name("URI")
diff --git a/tool/src/commands.rs b/tool/src/commands.rs
index 2ad9ebda..e992904d 100644
--- a/tool/src/commands.rs
+++ b/tool/src/commands.rs
@@ -1,5 +1,7 @@
 use failure::{self, ResultExt};
-use std::io;
+use std::cmp::Ordering;
+use std::fs::File;
+use std::io::{self, Write};
 use rpassword;
 
 extern crate openpgp;
@@ -193,6 +195,49 @@ pub fn dump(input: &mut io::Read, output: &mut io::Write, map: bool)
     Ok(())
 }
 
+pub fn split(input: &mut io::Read, prefix: &str)
+             -> Result<(), failure::Error> {
+    // We (ab)use the mapping feature to create byte-accurate dumps of
+    // nested packets.
+    let mut ppo =
+        openpgp::parse::PacketParserBuilder::from_reader(input)?
+        .map(true).finalize()?;
+
+    // This encodes our position in the tree.
+    let mut pos = vec![0];
+
+    while let Some(pp) = ppo {
+        if let Some(ref map) = pp.map {
+            let filename = format!(
+                "{}{}--{:?}", prefix,
+                pos.iter().map(|n| format!("{}", n))
+                    .collect::<Vec<String>>().join("-"),
+                pp.packet.tag());
+            let mut sink = File::create(filename)
+                .context("Failed to create output file")?;
+
+            // Write all the bytes.
+            for (_, buf) in map.iter() {
+                sink.write_all(buf)?;
+            }
+        }
+
+        let (_, old_depth, ppo_, new_depth) = pp.recurse()?;
+        ppo = ppo_;
+
+        // Update pos.
+        match old_depth.cmp(&new_depth) {
+            Ordering::Less =>
+                pos.push(0),
+            Ordering::Equal =>
+                *pos.last_mut().unwrap() += 1,
+            Ordering::Greater => {
+                pos.pop();
+            },
+        }
+    }
+    Ok(())
+}
 
 struct HexDumper {
     offset: usize,
diff --git a/tool/src/sq-usage.rs b/tool/src/sq-usage.rs
index a48b6ddb..f6e4220d 100644
--- a/tool/src/sq-usage.rs
+++ b/tool/src/sq-usage.rs
@@ -26,6 +26,7 @@
 //!     help         Prints this message or the help of the given subcommand(s)
 //!     keyserver    Interacts with keyservers
 //!     list         Lists key stores and known keys
+//!     split        Splits a message into OpenPGP packets
 //!     store        Interacts with key stores
 //! ```
 //!
@@ -255,6 +256,24 @@
 //!     <PREFIX>    List only stores with the given domain prefix
 //! ```
 //!
+//! ## Subcommand split
+//!
+//! ```text
+//! Splits a message into OpenPGP packets
+//!
+//! USAGE:
+//!     sq split [OPTIONS]
+//!
+//! FLAGS:
+//!     -h, --help       Prints help information
+//!     -V, --version    Prints version information
+//!
+//! OPTIONS:
+//!     -i, --input <FILE>     Sets the input file to use
+//!     -p, --prefix <FILE>    Sets the prefix to use for output files (defaults to the input filename with a dash, or
+//!                            'output')
+//! ```
+//!
 //! ## Subcommand store
 //!
 //! ```text
diff --git a/tool/src/sq.rs b/tool/src/sq.rs
index ea484e70..cf528cc5 100644
--- a/tool/src/sq.rs
+++ b/tool/src/sq.rs
@@ -105,6 +105,14 @@ fn real_main() -> Result<(), failure::Error> {
             let mut input = openpgp::Reader::from_reader(input)?;
             commands::dump(&mut input, &mut output, m.is_present("hex"))?;
         },
+        ("split",  Some(m)) => {
+            let input = open_or_stdin(m.value_of("input"))?;
+            let prefix = m.value_of("prefix").map(|p| p.to_owned()).
+                unwrap_or(
+                    m.value_of("input").unwrap_or("output").to_owned() + "-");
+            let mut input = openpgp::Reader::from_reader(input)?;
+            commands::split(&mut input, &prefix)?;
+        },
         ("keyserver",  Some(m)) => {
             let mut ks = if let Some(uri) = m.value_of("server") {
                 KeyServer::new(&ctx, &uri)