ptero/cli/
encoder.rs

1use std::{convert::TryInto, error::Error, fs::File, io::Read, sync::mpsc::channel};
2
3use clap::Clap;
4use log::{info, trace};
5
6use crate::{
7    binary::BitIterator,
8    context::PivotByLineContext,
9    encoder::Encoder,
10    method::{
11        complex::{eluv::ELUVMethodBuilder, extended_line::ExtendedLineMethodBuilder},
12        trailing_unicode::character_sets::CharacterSetType,
13    },
14};
15
16use super::{
17    progress::{new_progress_bar, spawn_progress_thread, ProgressStatus},
18    writer::Writer,
19};
20
21#[derive(Clap, Debug, PartialEq)]
22pub enum ELUVCharacterSet {
23    FourBit,
24    ThreeBit,
25    TwoBit,
26    Full,
27    Twitter,
28}
29
30/// Encode the secret into given cover text
31#[derive(Clap)]
32pub struct EncodeSubCommand {
33    /// Path to cover text used to encoding.
34    ///
35    /// Please note that original whitespace may not be preserved!
36    #[clap(short, long)]
37    cover: String,
38
39    /// Path to secret data file which will be encoded.
40    #[clap(short, long)]
41    data: String,
42
43    /// Pivot i.e. line length used for extended line algorithm.
44    ///
45    /// If omitted, program will determine minimum pivot that can be used.
46    #[clap(short, long)]
47    pivot: Option<usize>,
48
49    /// Use ELUV method for encoding.
50    ///
51    /// The ELUV method is a combination of three smaller encoders.
52    /// Random Whitespace - which puts randomly double whitespace between words,
53    /// Line Extend - which uses pivot to determine the size of the line,
54    /// Trailing Unicode - which puts one of the predefined Unicode invisible chars
55    /// at the end of the line during encoding.
56    ///
57    /// It can encode 7 bits in one pass.
58    /// This method has 3 variant.
59    #[clap(long, group = "method_args")]
60    eluv: bool,
61
62    /// Override a default set - can only be used with ELUV method!
63    ///
64    /// Provides a different set for the ELUV command to use.
65    /// Please note, that it may change the method's bitrate!
66    #[clap(long, arg_enum, requires = "eluv")]
67    set: Option<ELUVCharacterSet>,
68
69    /// Variant of the method. See concrete method for possible values.
70    ///
71    /// Variant is a permutation of methods that can be used during encoding.
72    #[clap(long, default_value = "1")]
73    variant: u8,
74
75    /// Use Extended Line method for encoding.
76    ///
77    /// The Extended Line method is a combination of three smaller encoders.
78    /// Random Whitespace - which puts randomly double whitespace between words,
79    /// Line Extend - which uses pivot to determine the size of the line,
80    /// Trailing Whitespace - which puts whitespace at the end of the line during encoding.
81    ///
82    /// It can encode 3 bits in one pass. Relies purely on ASCII characters.
83    /// This method has 3 variant.
84    #[clap(long = "eline", group = "method_args")]
85    #[allow(dead_code)]
86    extended_line: bool,
87}
88
89pub fn validate_pivot_smaller_than_text(
90    pivot: usize,
91    cover_text: &str,
92) -> Result<(), Box<dyn Error>> {
93    let text_length = cover_text.len();
94
95    if pivot > text_length {
96        return Err("Pivot is greater than the cover text length.".into());
97    }
98    Ok(())
99}
100
101impl EncodeSubCommand {
102    pub fn run(&self) -> Result<Vec<u8>, Box<dyn Error>> {
103        let cover_file_input = File::open(&self.cover)?;
104        let data_file_input = File::open(&self.data)?;
105
106        self.do_encode(cover_file_input, data_file_input)
107    }
108
109    pub(crate) fn do_encode(
110        &self,
111        mut cover_input: impl Read,
112        mut data_input: impl Read,
113    ) -> Result<Vec<u8>, Box<dyn Error>> {
114        let mut cover_text = String::new();
115        let mut data = vec![];
116
117        cover_input.read_to_string(&mut cover_text)?;
118        data_input.read_to_end(&mut data)?;
119
120        trace!("text: {:?}", data);
121
122        let pivot = pick_pivot_from(
123            self.pivot,
124            determine_pivot_size(cover_text.split_whitespace()),
125        )?;
126
127        validate_pivot_smaller_than_text(pivot, &cover_text)?;
128
129        let capacity_msg = format!(
130            "Required cover text capacity: {}",
131            BitIterator::new(&data).count()
132        );
133        Writer::warn(&capacity_msg);
134        info!("Encoding secret data");
135
136        let mut data_iterator = BitIterator::new(&data);
137
138        let progress_bar = new_progress_bar(data_iterator.count() as u64);
139        let (tx, rx) = channel::<ProgressStatus>();
140        progress_bar.set_message("Encoding..");
141        spawn_progress_thread(progress_bar.clone(), rx);
142
143        let method = self.get_method()?;
144        info!("Using method variant {}", self.variant);
145        let mut context = PivotByLineContext::new(&cover_text, pivot);
146        let stego_result = method.encode(&mut context, &mut data_iterator, Some(&tx));
147
148        tx.send(ProgressStatus::Finished).ok();
149        progress_bar.finish_with_message("Finished encoding");
150
151        Ok(stego_result?.as_bytes().into())
152    }
153
154    pub(crate) fn get_method(
155        &self,
156    ) -> Result<Box<dyn Encoder<PivotByLineContext>>, Box<dyn Error>> {
157        Ok(if self.eluv {
158            Box::new(
159                ELUVMethodBuilder::new()
160                    .character_set(get_character_set_type(&self.set))
161                    .variant(self.variant.try_into()?)
162                    .build(),
163            )
164        } else {
165            Box::new(
166                ExtendedLineMethodBuilder::new()
167                    .variant(self.variant.try_into()?)
168                    .build(),
169            )
170        })
171    }
172}
173
174pub fn get_character_set_type(set_option: &Option<ELUVCharacterSet>) -> CharacterSetType {
175    if let Some(char_set) = set_option {
176        match char_set {
177            ELUVCharacterSet::FourBit => CharacterSetType::FourBitUnicodeSet,
178            ELUVCharacterSet::ThreeBit => CharacterSetType::ThreeBitUnicodeSet,
179            ELUVCharacterSet::TwoBit => CharacterSetType::TwoBitUnicodeSet,
180            ELUVCharacterSet::Full => CharacterSetType::FullUnicodeSet,
181            ELUVCharacterSet::Twitter => CharacterSetType::TwitterUnicodeSet,
182        }
183    } else {
184        CharacterSetType::FullUnicodeSet
185    }
186}
187
188pub(crate) fn determine_pivot_size<'a>(words: impl Iterator<Item = &'a str>) -> usize {
189    words
190        .into_iter()
191        .map(|string| string.chars().count() + 1)
192        .max()
193        .unwrap_or(0)
194}
195
196pub(crate) fn pick_pivot_from(
197    user_pivot: Option<usize>,
198    calculated_pivot: usize,
199) -> Result<usize, Box<dyn Error>> {
200    Ok(if let Some(value) = user_pivot {
201        if value < calculated_pivot {
202            return Err("Provided pivot is smaller than the largest word in text! Cannot guarantee encoding will succeed.".into());
203        }
204        Writer::info(&format!("Using user provided pivot: {}", value));
205        value
206    } else {
207        Writer::info(&format!(
208            "Using pivot based on the cover text: {}",
209            calculated_pivot
210        ));
211        calculated_pivot
212    })
213}
214
215#[allow(unused_imports)]
216mod test {
217    use std::{error::Error, io::Read};
218
219    use crate::method::trailing_unicode::character_sets::CharacterSetType;
220
221    use super::{get_character_set_type, ELUVCharacterSet, EncodeSubCommand};
222
223    #[test]
224    fn fails_when_there_is_not_enough_cover_text() -> Result<(), Box<dyn Error>> {
225        let cover_input = "a b c ".repeat(2);
226        let data_input: Vec<u8> = vec![0b11111111];
227
228        let command = EncodeSubCommand {
229            cover: "stub".into(),
230            data: "stub".into(),
231            pivot: Some(3),
232            eluv: false,
233            extended_line: true,
234            set: None,
235            variant: 1,
236        };
237
238        let result = command.do_encode(cover_input.as_bytes(), data_input.as_slice());
239        assert!(result.is_err());
240        Ok(())
241    }
242
243    #[test]
244    fn fails_when_pivot_is_too_small() -> Result<(), Box<dyn Error>> {
245        let cover_input = "aaaaa ".repeat(2);
246        let data_input: Vec<u8> = vec![0b11111111];
247
248        let command = EncodeSubCommand {
249            cover: "stub".into(),
250            data: "stub".into(),
251            pivot: Some(3),
252            eluv: false,
253            extended_line: true,
254            set: None,
255            variant: 1,
256        };
257
258        let result = command.do_encode(cover_input.as_bytes(), data_input.as_slice());
259        assert!(result.is_err());
260        Ok(())
261    }
262
263    #[test]
264    fn fails_when_pivot_is_too_large() -> Result<(), Box<dyn Error>> {
265        let cover_input = "aaaaa";
266        let data_input: Vec<u8> = vec![0b11111111];
267
268        let command = EncodeSubCommand {
269            cover: "stub".into(),
270            data: "stub".into(),
271            pivot: Some(6),
272            eluv: false,
273            extended_line: true,
274            set: None,
275            variant: 1,
276        };
277
278        let result = command.do_encode(cover_input.as_bytes(), data_input.as_slice());
279        assert!(result.is_err());
280        Ok(())
281    }
282
283    #[test]
284    fn get_character_set_type_returns_default_when_none_is_provided() -> Result<(), Box<dyn Error>>
285    {
286        assert_eq!(
287            get_character_set_type(&None),
288            CharacterSetType::FullUnicodeSet
289        );
290        Ok(())
291    }
292
293    #[test]
294    fn get_character_set_type_maps_sets_correctly() -> Result<(), Box<dyn Error>> {
295        assert_eq!(
296            get_character_set_type(&Some(ELUVCharacterSet::Full)),
297            CharacterSetType::FullUnicodeSet
298        );
299        assert_eq!(
300            get_character_set_type(&Some(ELUVCharacterSet::FourBit)),
301            CharacterSetType::FourBitUnicodeSet
302        );
303        assert_eq!(
304            get_character_set_type(&Some(ELUVCharacterSet::ThreeBit)),
305            CharacterSetType::ThreeBitUnicodeSet
306        );
307        assert_eq!(
308            get_character_set_type(&Some(ELUVCharacterSet::TwoBit)),
309            CharacterSetType::TwoBitUnicodeSet
310        );
311        Ok(())
312    }
313}