1use std::{error::Error, fmt, sync::mpsc::Sender};
2
3use log::{debug, trace};
4
5use crate::{binary::Bit, cli::progress::ProgressStatus, context::Context};
6
7#[derive(Debug, Clone)]
9pub enum EncoderResult {
10 Success,
11 NoDataLeft,
12}
13
14pub trait Capacity {
18 fn bitrate(&self) -> usize;
20}
21
22pub trait Encoder<E>: Capacity
25where
26 E: Context,
27{
28 fn partial_encode(
44 &self,
45 context: &mut E,
46 data: &mut dyn Iterator<Item = Bit>,
47 ) -> Result<EncoderResult, Box<dyn Error>>;
48
49 fn encode(
50 &self,
51 context: &mut E,
52 data: &mut dyn Iterator<Item = Bit>,
53 progress_channel: Option<&Sender<ProgressStatus>>,
54 ) -> Result<String, Box<dyn Error>> {
55 let mut stego_text = String::new();
56
57 let mut no_data_left = false;
58 while !no_data_left {
59 context.load_text()?;
60 trace!("Current line '{}'", context.get_current_text()?);
61 match self.partial_encode(context, data)? {
62 EncoderResult::Success => {
63 if let Some(tx) = progress_channel {
64 tx.send(ProgressStatus::Step(self.bitrate() as u64)).ok();
65 }
66 }
67 EncoderResult::NoDataLeft => {
68 debug!("No data left to encode, stopping");
69 no_data_left = true;
70 }
71 }
72 let line = context.get_current_text()?;
73 stego_text.push_str(&format!("{}\n", &line));
74 }
75 let mut appended_line_count = 0;
77 while let Ok(line) = context.load_text() {
78 appended_line_count += 1;
79 stego_text.push_str(&format!("{}\n", &line));
80 }
81 debug!("Appended the {} of left lines", appended_line_count);
82
83 if !no_data_left {
84 debug!("Capacity exceeded by {} bits", data.count());
85 Err(EncodingError::capacity_error().into())
86 } else {
87 Ok(stego_text)
88 }
89 }
90}
91
92#[derive(Debug, Clone)]
94pub enum EncodingErrorKind {
95 CapacityTooLow,
96 NoWordsLeft,
97}
98
99#[derive(Debug, Clone)]
101pub struct EncodingError {
102 kind: EncodingErrorKind,
103}
104
105impl EncodingError {
106 pub fn capacity_error() -> Self {
108 EncodingError {
109 kind: EncodingErrorKind::CapacityTooLow,
110 }
111 }
112 pub fn no_words_error() -> Self {
114 EncodingError {
115 kind: EncodingErrorKind::NoWordsLeft,
116 }
117 }
118}
119
120#[cfg(not(tarpaulin_include))]
121impl fmt::Display for EncodingError {
122 fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
123 match &self.kind {
124 EncodingErrorKind::CapacityTooLow => write!(f, "Exceeded cover text capacity"),
125 EncodingErrorKind::NoWordsLeft => write!(
126 f,
127 "No extra words found in cover text when tried to encode a bit"
128 ),
129 }
130 }
131}
132
133impl Error for EncodingError {}