summaryrefslogtreecommitdiff
path: root/src/lossless/lsb.rs
blob: 1c9a2136db4611c200a60a2623a00ec00d939a2e (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
use std::{cmp::Ordering, io::{BufWriter, Cursor}};

use image::{DynamicImage, GenericImageView, Pixel};

use crate::{codec::Codec, Error};

/// Least-significant bit (LSB) steganography encodes data in the least-significant bits of colors
/// in an image. This implementation reduces the colors in the carrier (irreversibly) in order to
/// allow a byte of data to fit in each pixel of the image. 3 bits of data are encoded per pixel,
/// and the 9th bit is used to signal the end of data.
#[derive(Debug, Default)]
pub struct LsbCodec;

impl Codec for LsbCodec {
    fn encode(&self, carrier: &[u8], payload: &[u8]) -> Result<Vec<u8>, Error>
    {
        let image_format = image::guess_format(carrier).unwrap();
        let mut image: DynamicImage = image::load_from_memory(carrier).unwrap();
        let payload: &[u8] = payload;

        if image.pixels().count() < payload.len() {
            return Err(Error::DataInvalid("Payload Too Big for Carrier".into()));
        }

        let mut payload_iter = payload.iter();

        match image {
            DynamicImage::ImageRgba8(ref mut image) => {
                for pixel in image.pixels_mut() {
                    if let Some(payload_byte) = payload_iter.next() {
                        encode_pixel(pixel, *payload_byte, false);
                    } else {
                        encode_pixel(pixel, 0, true);
                    }
                }
            },
            DynamicImage::ImageRgb8(ref mut image) => {
                for pixel in image.pixels_mut() {
                    if let Some(payload_byte) = payload_iter.next() {
                        encode_pixel(pixel, *payload_byte, false);
                    } else {
                        encode_pixel(pixel, 0, true);
                    }
                }
            },
            _ => return Err(Error::DataInvalid("Unsupported Image Color Format".into()))
        }

        let mut buf = BufWriter::new(Cursor::new(Vec::<u8>::new()));
        if let Err(e) = image.write_to(&mut buf, image_format) {
            return Err(Error::DependencyError(e.to_string()))
        }
        Ok(buf.into_inner().unwrap().into_inner())
    }

    fn decode(&self, carrier: &[u8]) -> Result<(Vec<u8>, Vec<u8>), Error>
    {
        let image_format = image::guess_format(carrier).unwrap();
        let mut image: DynamicImage = image::load_from_memory(carrier).unwrap();
        let mut payload: Vec<u8> = Vec::new();

        match image {
            DynamicImage::ImageRgba8(ref mut image) => {
                for pixel in image.pixels_mut() {
                    if let Some(payload_byte) = decode_pixel(pixel) {
                        payload.push(payload_byte);
                    } else {
                        break;
                    }
                }
            },
            DynamicImage::ImageRgb8(ref mut image) => {
                for pixel in image.pixels_mut() {
                    if let Some(payload_byte) = decode_pixel(pixel) {
                        payload.push(payload_byte);
                    } else {
                        break;
                    }
                }
            },
            _ => return Err(Error::DataInvalid("Unsupported Image Color Format".into()))
        }
        
        let mut buf = BufWriter::new(Cursor::new(Vec::<u8>::new()));
        if let Err(e) = image.write_to(&mut buf, image_format) {
            return Err(Error::DependencyError(e.to_string()))
        }
        Ok((buf.into_inner().unwrap().into_inner(), payload))
    }
}

fn encode_pixel<P: Pixel<Subpixel = u8>>(pixel: &mut P, payload_byte: u8, end_of_data: bool) {
    let mut bits_remaining: i32 = 8;
    for channel in pixel.channels_mut() {
        *channel &= 0b11111000;
        bits_remaining -= 3;
        if bits_remaining <= -3 {
            break;
        }

        let mask = match bits_remaining.cmp(&0) {
            Ordering::Less => payload_byte << -bits_remaining,
            _ => payload_byte >> bits_remaining,
        } & 0b00000111;

        *channel |= mask;
    }

    // Add end-of-data marker to final bit if necessary
    if end_of_data {
        *pixel.channels_mut().last_mut().unwrap() |= 1;
    }
}

fn decode_pixel<P: Pixel<Subpixel = u8>>(pixel: &mut P) -> Option<u8> {
    
    // Final bit as end-of-data marker
    if pixel.channels().last().unwrap() & 1 == 1 {
        return None;
    }

    let mut bits_remaining: i32 = 8;
    let mut payload_byte: u8 = 0;
    for channel in pixel.channels_mut() {
        bits_remaining -= 3;
        if bits_remaining <= -3 {
            break;
        }

        let channel_bits = *channel & 0b00000111;
        *channel &= 0b11111000;
        let mask = match bits_remaining.cmp(&0) {
            Ordering::Less => channel_bits >> -bits_remaining,
            _ => channel_bits << bits_remaining,
        };
        payload_byte |= mask;
    }
    Some(payload_byte)
}