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#[derive(Clap)]
32pub struct EncodeSubCommand {
33 #[clap(short, long)]
37 cover: String,
38
39 #[clap(short, long)]
41 data: String,
42
43 #[clap(short, long)]
47 pivot: Option<usize>,
48
49 #[clap(long, group = "method_args")]
60 eluv: bool,
61
62 #[clap(long, arg_enum, requires = "eluv")]
67 set: Option<ELUVCharacterSet>,
68
69 #[clap(long, default_value = "1")]
73 variant: u8,
74
75 #[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}