ptero/
encoder.rs

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/// Possible results of data encoding
8#[derive(Debug, Clone)]
9pub enum EncoderResult {
10    Success,
11    NoDataLeft,
12}
13
14/// Trait that should be implemented by the [Encoders](crate::econder::Encoder).
15/// Gives amount of bits that are encoded per text fragment.
16/// Concrete text fragment (e.g. line) is determined by the [Context](crate::context::Context).
17pub trait Capacity {
18    /// Returns how many bits are encoded per text fragment.
19    fn bitrate(&self) -> usize;
20}
21
22/// Base trait for all data encoders.
23/// The generic type should contain data need by the encoder implementation.
24pub trait Encoder<E>: Capacity
25where
26    E: Context,
27{
28    /// Encodes bits provided by `data` iterator.
29    /// Every Encoder has Context which exposes access to cover text. See [Context] for more info.
30    ///
31    /// # Arguments
32    ///
33    /// * `context` - context of the steganography method, can contain various needed info like pivot etc.
34    /// * `data` - data iterator which return [Bit] with each iteration
35    ///
36    /// # Returns
37    /// It returns whether the encoding was successful. See [EncoderResult] and [EncodingError].
38    ///
39    /// [Context]: crate::context::Context
40    /// [EncoderResult]: EncoderResult
41    /// [EncodingError]: EncodingError
42    /// [Bit]: crate::binary::Bit
43    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        // Append the rest of possible missing cover text
76        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/// Enum for data encoding errors types
93#[derive(Debug, Clone)]
94pub enum EncodingErrorKind {
95    CapacityTooLow,
96    NoWordsLeft,
97}
98
99/// Represents encoding error. Concrete error if differentiated by the [EncodingErrorKind](EncodingErrorKind)
100#[derive(Debug, Clone)]
101pub struct EncodingError {
102    kind: EncodingErrorKind,
103}
104
105impl EncodingError {
106    /// Facade for creating [CapacityTooLow](EncodingErrorKind) error.
107    pub fn capacity_error() -> Self {
108        EncodingError {
109            kind: EncodingErrorKind::CapacityTooLow,
110        }
111    }
112    /// Facade for creating [NoWordsLeft](EncodingErrorKind) error.
113    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 {}