Unified messages when pivot is larger than text

Author: bluebears-dev

Cargo.lock

@@ -4,7 +4,7 @@ edition = "2018"
 name = "ptero-cli"
 description = "A text steganography CLI tool for Social Media"
 keywords = ["steganography", "encoding", "decoding", "text", "cli"]
-version = "0.4.0"
+version = "0.4.1"
 license-file = "LICENSE"
 readme = "README.md"
 include = [

Cargo.toml

@@ -9,11 +9,7 @@ use crate::{
     method::complex::{eluv::ELUVMethod, extended_line::ExtendedLineMethod},
 };
 
-use super::{
-    encoder::determine_pivot_size,
-    progress::{new_progress_bar, spawn_progress_thread, ProgressStatus},
-    writer::Writer,
-};
+use super::{encoder::{determine_pivot_size, validate_pivot_smaller_than_text}, progress::{new_progress_bar, spawn_progress_thread, ProgressStatus}, writer::Writer};
 
 /// Calculate the minimal capacity for the cover text and given pivot
 #[derive(Clap)]
@@ -27,10 +23,14 @@ pub struct GetCapacityCommand {
     pivot: usize,
 
     /// Use ELUV method for encoding.
+    ///
+    /// Does not support different sets or variants, currently.
     #[clap(long, group = "method_args")]
     eluv: bool,
 
     /// Use Extended Line method for encoding.
+    ///
+    /// Does not support different variants, currently.
     #[clap(long = "eline", group = "method_args")]
     #[allow(dead_code)]
     extended_line: bool,
@@ -54,17 +54,12 @@ impl GetCapacityCommand {
         let mut text_fragment_count = 0;
 
         let max_word_length = determine_pivot_size(cover_text.split_whitespace());
-        let text_length = cover_text
-            .split_whitespace()
-            .map(|string| string.chars())
-            .flatten()
-            .count();
+        validate_pivot_smaller_than_text(self.pivot, &cover_text)?;
+
         debug!("Longest word in the cover text is {}", max_word_length);
 
         if max_word_length > self.pivot {
             Writer::warn("This pivot might not guarantee the secret data will be encodable!");
-        } else if self.pivot >= text_length {
-            return Err("Pivot is greater than the cover text length.".into());
         }
 
         let progress_bar = new_progress_bar(cover_text.len() as u64);
@@ -133,4 +128,21 @@ mod test {
         assert_eq!(result.ok(), Some(6 as u32));
         Ok(())
     }
+
+
+    #[test]
+    fn fails_when_pivot_is_too_large() -> Result<(), Box<dyn Error>> {
+        let stego_input = "aaaaa";
+
+        let command = GetCapacityCommand {
+            cover: "stub".into(),
+            pivot: 6,
+            eluv: false,
+            extended_line: true,
+        };
+
+        let result = command.get_cover_text_capacity(stego_input.as_bytes());
+        assert!(result.is_err());
+        Ok(())
+    }
 }

src/cli/capacity.rs

@@ -9,10 +9,7 @@ use crate::{
     method::complex::{eluv::ELUVMethodBuilder, extended_line::ExtendedLineMethodBuilder},
 };
 
-use super::{
-    encoder::{get_character_set_type, ELUVCharacterSet},
-    progress::{new_progress_bar, spawn_progress_thread, ProgressStatus},
-};
+use super::{encoder::{ELUVCharacterSet, get_character_set_type, validate_pivot_smaller_than_text}, progress::{new_progress_bar, spawn_progress_thread, ProgressStatus}};
 
 /// Decode secret from the stegotext
 #[derive(Clap)]
@@ -82,6 +79,9 @@ impl DecodeSubCommand {
         let mut stego_text = String::new();
 
         stego_input.read_to_string(&mut stego_text)?;
+
+        validate_pivot_smaller_than_text(self.pivot, &stego_text)?;
+
         let decoder = self.get_method()?;
         info!("Using method variant {}", self.variant);
         let mut context = PivotByRawLineContext::new(stego_text.as_str(), self.pivot);
@@ -142,4 +142,22 @@ mod test {
         assert_eq!(result.ok(), Some(vec![0]));
         Ok(())
     }
+
+    #[test]
+    fn fails_when_pivot_is_too_large() -> Result<(), Box<dyn Error>> {
+        let stego_input = "aaaaa";
+
+        let command = DecodeSubCommand {
+            text: "stub".into(),
+            pivot: 6,
+            eluv: false,
+            extended_line: true,
+            set: None,
+            variant: 1,
+        };
+
+        let result = command.do_decode(stego_input.as_bytes());
+        assert!(result.is_err());
+        Ok(())
+    }
 }

src/cli/decoder.rs

@@ -1,10 +1,4 @@
-use std::{
-    convert::{TryInto},
-    error::Error,
-    fs::File,
-    io::Read,
-    sync::mpsc::channel,
-};
+use std::{convert::TryInto, error::Error, fs::File, io::Read, sync::mpsc::channel};
 
 use clap::Clap;
 use log::{info, trace};
@@ -91,6 +85,18 @@ pub struct EncodeSubCommand {
     extended_line: bool,
 }
 
+pub fn validate_pivot_smaller_than_text(
+    pivot: usize,
+    cover_text: &str,
+) -> Result<(), Box<dyn Error>> {
+    let text_length = cover_text.len();
+
+    if pivot > text_length {
+        return Err("Pivot is greater than the cover text length.".into());
+    }
+    Ok(())
+}
+
 impl EncodeSubCommand {
     pub fn run(&self) -> Result<Vec<u8>, Box<dyn Error>> {
         let cover_file_input = File::open(&self.cover)?;
@@ -117,6 +123,8 @@ impl EncodeSubCommand {
             determine_pivot_size(cover_text.split_whitespace()),
         )?;
 
+        validate_pivot_smaller_than_text(pivot, &cover_text)?;
+
         let capacity_msg = format!(
             "Required cover text capacity: {}",
             BitIterator::new(&data).count()
@@ -250,6 +258,26 @@ mod test {
         Ok(())
     }
 
+    #[test]
+    fn fails_when_pivot_is_too_large() -> Result<(), Box<dyn Error>> {
+        let cover_input = "aaaaa";
+        let data_input: Vec<u8> = vec![0b11111111];
+
+        let command = EncodeSubCommand {
+            cover: "stub".into(),
+            data: "stub".into(),
+            pivot: Some(6),
+            eluv: false,
+            extended_line: true,
+            set: None,
+            variant: 1,
+        };
+
+        let result = command.do_encode(cover_input.as_bytes(), data_input.as_slice());
+        assert!(result.is_err());
+        Ok(())
+    }
+
     #[test]
     fn get_character_set_type_returns_default_when_none_is_provided() -> Result<(), Box<dyn Error>>
     {