Initial Commit

8 files changed, 247 insertions, 0 deletions

diff --git a/.gitignore b/.gitignore
new file mode 100644
index 0000000..ea8c4bf
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1 @@
+/target
diff --git a/Cargo.toml b/Cargo.toml
new file mode 100644
index 0000000..9a471d6
--- /dev/null
+++ b/Cargo.toml
@@ -0,0 +1,16 @@
+[package]
+name = "occule"
+version = "0.1.0"
+edition = "2021"
+
+[features]
+default = ["jpg", "png", "lossless"]
+jpg = []
+png = []
+lossless = []
+
+[dependencies]
+colored = "2.1.0"
+image = "0.24"
+img-parts = "0.3.0"
+thiserror = "1.0.61"
diff --git a/src/codec.rs b/src/codec.rs
new file mode 100644
index 0000000..03cdf15
--- /dev/null
+++ b/src/codec.rs
@@ -0,0 +1,14 @@
+pub trait Codec {
+    type Carrier;
+    type Payload;
+    type Output;
+    type Error;
+
+    fn encode(
+        &self,
+        carrier: impl Into<Self::Carrier>,
+        payload: impl Into<Self::Payload>,
+    ) -> Result<Self::Output, Self::Error>;
+
+    fn decode(&self, encoded: impl Into<Self::Output>) -> Result<(Self::Carrier, Self::Payload), Self::Error>;
+}
diff --git a/src/jpg/mod.rs b/src/jpg/mod.rs
new file mode 100644
index 0000000..3d6bdc3
--- /dev/null
+++ b/src/jpg/mod.rs
@@ -0,0 +1,2 @@
+mod segment;
+pub use segment::*;
diff --git a/src/jpg/segment.rs b/src/jpg/segment.rs
new file mode 100644
index 0000000..f54d0e2
--- /dev/null
+++ b/src/jpg/segment.rs
@@ -0,0 +1,66 @@
+use std::{mem::size_of, usize};
+
+use img_parts::jpeg::{markers, Jpeg, JpegSegment};
+use thiserror::Error;
+
+use crate::codec::Codec;
+
+#[derive(Debug, PartialEq, Eq)]
+pub struct JpegSegmentCodec {
+    pub start_index: usize,
+}
+
+impl Codec for JpegSegmentCodec {
+    type Carrier = Vec<u8>;
+    type Payload = Vec<u8>;
+    type Output = Self::Carrier;
+    type Error = JpegSegmentError;
+
+    fn encode(&self, carrier: impl Into<Self::Carrier>, payload: impl Into<Self::Payload>) -> Result<Self::Output, Self::Error> {
+        let mut jpeg = match Jpeg::from_bytes(carrier.into().into()) {
+            Ok(image) => image,
+            Err(err) => return Err(JpegSegmentError::ParseFailed { inner: err })
+        };
+        let mut payload_bytes: Self::Carrier = payload.into();
+        let segment_count = ((payload_bytes.len() + size_of::<u64>()) as u64).div_ceil((u16::MAX as usize - size_of::<u16>()) as u64);
+        payload_bytes.splice(0..0, segment_count.to_le_bytes());
+        for (index, payload_chunk) in payload_bytes.chunks(u16::MAX as usize - size_of::<u16>()).enumerate() {
+            let segment = JpegSegment::new_with_contents(markers::COM, payload_chunk.to_vec().into());
+            jpeg.segments_mut().insert(self.start_index + index, segment);
+        }
+        Ok(jpeg.encoder().bytes().to_vec())
+    }
+
+    fn decode(&self, encoded: impl Into<Self::Output>) -> Result<(Self::Carrier, Self::Payload), Self::Error> {
+        let mut jpeg = match Jpeg::from_bytes(encoded.into().into()) {
+            Ok(image) => image,
+            Err(err) => return Err(JpegSegmentError::ParseFailed { inner: err })
+        };
+        let segment = jpeg.segments_mut().remove(self.start_index);
+        let segment_bytes = segment.contents();
+        let segment_count = u64::from_le_bytes(segment_bytes[0..size_of::<u64>()].try_into().unwrap()) as usize;
+        let mut payload_vec: Vec<u8> = Vec::with_capacity((u16::MAX as usize - size_of::<u16>()) * segment_count);
+        payload_vec.extend(segment_bytes[size_of::<u64>()..].to_vec());
+
+        for _ in 0..segment_count-1 {
+            let segment = jpeg.segments_mut().remove(self.start_index);
+            payload_vec.extend(segment.contents());
+        }
+
+        Ok((jpeg.encoder().bytes().to_vec(), payload_vec))
+    }
+}
+
+impl Default for JpegSegmentCodec {
+    fn default() -> Self {
+        Self {
+            start_index: 3,
+        }
+    }
+}
+
+#[derive(Error, Debug)]
+pub enum JpegSegmentError {
+    #[error("Failed to parse JPEG data: {inner:?}")]
+    ParseFailed { inner: img_parts::Error }
+}
diff --git a/src/lib.rs b/src/lib.rs
new file mode 100644
index 0000000..a1ed909
--- /dev/null
+++ b/src/lib.rs
@@ -0,0 +1,7 @@
+pub mod codec;
+
+#[cfg(feature = "jpg")]
+pub mod jpg;
+
+#[cfg(feature = "lossless")]
+pub mod lossless;
diff --git a/src/lossless/champleve.rs b/src/lossless/champleve.rs
new file mode 100644
index 0000000..948408c
--- /dev/null
+++ b/src/lossless/champleve.rs
@@ -0,0 +1,139 @@
+use std::cmp::Ordering;
+
+use image::{ColorType, DynamicImage, GenericImageView, Pixel};
+use thiserror::Error;
+
+use crate::codec::Codec;
+
+#[derive(Debug)]
+pub struct ChampleveCodec;
+
+impl Codec for ChampleveCodec {
+    type Carrier = DynamicImage;
+    type Payload = Vec<u8>;
+    type Output = Self::Carrier;
+    type Error = ChampleveError;
+
+    fn encode(&self, carrier: impl Into<Self::Carrier>, payload: impl Into<Self::Payload>) -> Result<Self::Output, Self::Error> {
+        let mut image: DynamicImage = carrier.into();
+        let payload: Vec<u8> = payload.into();
+
+        if image.pixels().count() < payload.len() {
+            return Err(ChampleveError::PayloadTooBig);
+        }
+
+        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(ChampleveError::UnsupportedFormat { format: image.color() })
+        }
+
+        Ok(image)
+    }
+
+    fn decode(&self, carrier: impl Into<Self::Output>) -> Result<(Self::Carrier, Self::Payload), ChampleveError> {
+        let mut image: DynamicImage = carrier.into();
+        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(ChampleveError::UnsupportedFormat { format: image.color() })
+        }
+        
+        Ok((image, 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)
+}
+
+#[derive(Error, Debug)]
+pub enum ChampleveError {
+    #[error("Payload is too big for the carrier. Choose a smaller payload or an image with greater pixel dimensions.")]
+    PayloadTooBig,
+    #[error("Specified image format ({format:?}) is unsupported.")]
+    UnsupportedFormat {
+        format: ColorType
+    },
+}
diff --git a/src/lossless/mod.rs b/src/lossless/mod.rs
new file mode 100644
index 0000000..a9ae225
--- /dev/null
+++ b/src/lossless/mod.rs
@@ -0,0 +1,2 @@
+mod champleve;
+pub use champleve::*;