ptero/cli/
decoder.rs

1use std::{convert::TryInto, error::Error, fs::File, io::Read, sync::mpsc::channel};
2
3use clap::Clap;
4use log::info;
5
6use crate::{
7    context::PivotByRawLineContext,
8    decoder::Decoder,
9    method::complex::{eluv::ELUVMethodBuilder, extended_line::ExtendedLineMethodBuilder},
10};
11
12use super::{encoder::{ELUVCharacterSet, get_character_set_type, validate_pivot_smaller_than_text}, progress::{new_progress_bar, spawn_progress_thread, ProgressStatus}};
13
14/// Decode secret from the stegotext
15#[derive(Clap)]
16pub struct DecodeSubCommand {
17    /// Path to stegotext from which data will be decoded
18    #[clap(short, long)]
19    text: String,
20
21    /// Pivot i.e. line length used to encode with extended line algorithm
22    #[clap(short, long)]
23    pivot: usize,
24
25    /// Use ELUV method for encoding.
26    ///
27    /// This method has 3 variants.
28    #[clap(long, group = "method_args")]
29    eluv: bool,
30
31    /// Override a default set - can only be used with ELUV method!
32    ///
33    /// Provides a different set for the ELUV command to use.
34    /// Please note, that it may change the method's bitrate!
35    #[clap(long, arg_enum, requires = "eluv")]
36    set: Option<ELUVCharacterSet>,
37
38    /// Variant of the method. See concrete method for possible values.
39    ///
40    /// Variant is a permutation of methods that can be used during decoding.
41    #[clap(long, default_value = "1")]
42    variant: u8,
43
44    /// Use Extended Line method for encoding.
45    /// 
46    /// This method has 3 variants.
47    #[clap(long = "eline", group = "method_args")]
48    #[allow(dead_code)]
49    extended_line: bool,
50}
51
52impl DecodeSubCommand {
53    pub fn run(&self) -> Result<Vec<u8>, Box<dyn Error>> {
54        let stego_text_file_input = File::open(&self.text)?;
55
56        self.do_decode(stego_text_file_input)
57    }
58
59    pub(crate) fn get_method(
60        &self,
61    ) -> Result<Box<dyn Decoder<PivotByRawLineContext>>, Box<dyn Error>> {
62        Ok(if self.eluv {
63            Box::new(
64                ELUVMethodBuilder::new()
65                    .character_set(get_character_set_type(&self.set))
66                    .variant(self.variant.try_into()?)
67                    .build(),
68            )
69        } else {
70            Box::new(
71                ExtendedLineMethodBuilder::new()
72                    .variant(self.variant.try_into()?)
73                    .build(),
74            )
75        })
76    }
77
78    pub fn do_decode(&self, mut stego_input: impl Read) -> Result<Vec<u8>, Box<dyn Error>> {
79        let mut stego_text = String::new();
80
81        stego_input.read_to_string(&mut stego_text)?;
82
83        validate_pivot_smaller_than_text(self.pivot, &stego_text)?;
84
85        let decoder = self.get_method()?;
86        info!("Using method variant {}", self.variant);
87        let mut context = PivotByRawLineContext::new(stego_text.as_str(), self.pivot);
88
89        let progress_bar = new_progress_bar(stego_text.len() as u64);
90        let (tx, rx) = channel::<ProgressStatus>();
91        progress_bar.set_message("Decoding cover text...");
92        spawn_progress_thread(progress_bar.clone(), rx);
93
94        let result = decoder.decode(&mut context, Some(&tx));
95
96        tx.send(ProgressStatus::Finished).ok();
97        progress_bar.finish_with_message("Finished decoding");
98
99        result
100    }
101}
102
103#[allow(unused_imports)]
104mod test {
105    use crate::binary::Bit;
106    use std::{error::Error, io::Read};
107
108    use super::DecodeSubCommand;
109
110    #[test]
111    fn decodes_zeroes_if_not_data_encoded_extended_line() -> Result<(), Box<dyn Error>> {
112        let stego_input = "a b";
113
114        let command = DecodeSubCommand {
115            text: "stub".into(),
116            pivot: 3,
117            eluv: false,
118            extended_line: true,
119            set: None,
120            variant: 1,
121        };
122
123        let result = command.do_decode(stego_input.as_bytes());
124        assert_eq!(result.ok(), Some(vec![0]));
125        Ok(())
126    }
127
128    #[test]
129    fn decodes_zeroes_if_not_data_encoded_eluv() -> Result<(), Box<dyn Error>> {
130        let stego_input = "a b";
131
132        let command = DecodeSubCommand {
133            text: "stub".into(),
134            pivot: 3,
135            eluv: true,
136            extended_line: false,
137            set: None,
138            variant: 1,
139        };
140
141        let result = command.do_decode(stego_input.as_bytes());
142        assert_eq!(result.ok(), Some(vec![0]));
143        Ok(())
144    }
145
146    #[test]
147    fn fails_when_pivot_is_too_large() -> Result<(), Box<dyn Error>> {
148        let stego_input = "aaaaa";
149
150        let command = DecodeSubCommand {
151            text: "stub".into(),
152            pivot: 6,
153            eluv: false,
154            extended_line: true,
155            set: None,
156            variant: 1,
157        };
158
159        let result = command.do_decode(stego_input.as_bytes());
160        assert!(result.is_err());
161        Ok(())
162    }
163}
ptero/cli/decoder.rs

ptero/cli/
decoder.rs
