From 42523c1f519e23a061bd9aaa9e8cc52666bb8a38 Mon Sep 17 00:00:00 2001
From: Lucas <me@lhvy.dev>
Date: Fri, 3 Nov 2023 11:27:57 +1100
Subject: [PATCH] Add basic image bit plane analysis

---
 src/lib.rs  | 35 +++++++++++++++++++++++++++++++++++
 src/main.rs | 21 +++++++++++++++++++--
 2 files changed, 54 insertions(+), 2 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index fef00b6..5dbd986 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -149,6 +149,41 @@ pub fn lsb_raw_decode(carrier: &[u8], n_bits: usize) -> Vec<u8> {
     payload
 }
 
+pub fn extract_bit_planes(image: image::RgbImage) {
+    let mut planes = Vec::new();
+    for i in 0..8 {
+        for j in 0..3 {
+            let mut plane = image.clone();
+            for pixel in plane.pixels_mut() {
+                pixel[j] = (pixel[j] & (1 << i)) * 255;
+                for k in 0..3 {
+                    if k != j {
+                        pixel[k] = 0;
+                    }
+                }
+            }
+            planes.push(plane);
+        }
+    }
+
+    // Save into output folder, labelled R, G, B 0-7
+    // Folder first
+    std::fs::create_dir_all("output").unwrap();
+    // Then save the images
+    for (i, plane) in planes.into_iter().enumerate() {
+        let color = match i % 3 {
+            0 => "R",
+            1 => "G",
+            2 => "B",
+            _ => unreachable!(),
+        };
+        plane
+            .save(format!("output/{}{}.png", color, i / 3))
+            .unwrap();
+    }
+    println!("Saved images to ./output");
+}
+
 pub fn bytes_to_wav(bytes: &[u8]) {
     let spec = hound::WavSpec {
         channels: 1,
diff --git a/src/main.rs b/src/main.rs
index ff67409..a8fae89 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,12 +1,12 @@
 mod cipher;
 
 use inquire::validator::Validation;
-use sneaky::{bytes_to_wav, lsb_decode, lsb_encode, wav_to_bytes};
+use sneaky::{bytes_to_wav, extract_bit_planes, lsb_decode, lsb_encode, wav_to_bytes};
 use std::ops::BitXor;
 use std::path::{Path, PathBuf};
 
 fn main() {
-    let options = vec!["Encode a message", "Decode a message"];
+    let options = vec!["Encode a message", "Decode a message", "Analyse a file"];
     let choice = inquire::Select::new("Select an option", options)
         .prompt()
         .unwrap();
@@ -14,6 +14,7 @@ fn main() {
     match choice {
         "Encode a message" => encode(),
         "Decode a message" => decode(),
+        "Analyse a file" => analyse(),
         _ => panic!("Invalid option"),
     }
 
@@ -103,6 +104,22 @@ fn decode() {
     }
 }
 
+fn analyse() {
+    let options = vec!["Analyse an image"];
+    let choice = inquire::Select::new("Select an option", options)
+        .prompt()
+        .unwrap();
+
+    match choice {
+        "Analyse an image's bit planes" => {
+            let image_path = get_path();
+            let image = image::open(image_path).unwrap().to_rgb8();
+            extract_bit_planes(image);
+        }
+        _ => panic!("Invalid option"),
+    }
+}
+
 fn get_bits() -> u8 {
     // How many bits to use for the message, 1-8
     let n_bits = inquire::Text::new("How many bits to use for LSB?")