From b40e67b43fe121f546f54328a5a2accaecb6bf29 Mon Sep 17 00:00:00 2001
From: clizia <xanthippe@cock.lu>
Date: Tue, 12 Nov 2024 22:50:18 +0100
Subject: [PATCH] made it usable

---
 src/chunk.rs |  6 ++---
 src/main.rs  | 65 ++++++++++++++++++++++++++++++++++------------------
 src/png.rs   |  8 +++----
 3 files changed, 50 insertions(+), 29 deletions(-)

diff --git a/src/chunk.rs b/src/chunk.rs
index e5f5240..41e0cf1 100644
--- a/src/chunk.rs
+++ b/src/chunk.rs
@@ -4,9 +4,9 @@ use crc::Crc;
 use crate::chunk_type::ChunkType;
 #[derive(Debug)]
 pub struct Chunk {
-    length:     u32,
+    pub length:     u32,
     pub chunk_type: ChunkType,
-    chunk_data: Vec<u8>,
+    pub chunk_data: Vec<u8>,
     pub crc:    u32,
 }
 
@@ -136,7 +136,7 @@ impl TryFrom<&[u8]> for Chunk {
 
 impl Display for Chunk {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
-        write!(f, "chunk length:{}\nchunk type: {}\n, chunk message {}\nCRC u32: {}",
+        write!(f, "chunk length:{}\nchunk type: {}\nchunk message {}\nCRC u32: {}",
             self.length,
             self.chunk_type,
             self.data_as_string().unwrap(),
diff --git a/src/main.rs b/src/main.rs
index 524582d..fa428ed 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -5,53 +5,74 @@ mod commands;
 mod png;
 
 use core::panic;
+use std::io::Read;
 use std::path::PathBuf;
 use std::fs::File;
+use std::str::FromStr;
+use std::io::Write;
 
 use args::Mode;
+use chunk::Chunk;
 use chunk_type::ChunkType;
+use png::Png;
 use clap::Parser;
 use crate::args::Args;
 
 pub type Error = Box<dyn std::error::Error>;
 pub type Result<T> = std::result::Result<T, Error>;
 
-fn get_image(path: &str) -> File {
+fn get_path(path: &str) -> PathBuf {
     let image_path: PathBuf = PathBuf::from(path);
-    let path_display = image_path.display();
-
-    let image = match File::open(&image_path) {
-        Err(why) => panic!("couldn't open {}: {}", path_display, why),
-        Ok(file) => file
-    };
-
-    image
+    image_path
 }
 
-fn encode(image: File, chunk_type: ChunkType, message: &str) {
+fn encode(path: PathBuf, chunk_type: &str, message: &str) -> Result<()> {
+    let mut image = File::open(path)?;
+    let mut image_bytes: Vec<u8> = Vec::new();
+    image.read_to_end(&mut image_bytes)?;
+
+    let mut png = Png::try_from(image_bytes.as_ref())?;
+
+    let chunk = Chunk::new(ChunkType::from_str(chunk_type)?, message.as_bytes().to_vec());
+
+    png.append_chunk(chunk);
+
+    let mut new_png = File::create_new("secret.png")?;
+    new_png.write_all(png.as_bytes().as_ref())?;
+
+    Ok(())
+}
+
+fn decode(path: PathBuf, chunk_type: &str) -> Result<()> {
+    let mut image = File::open(path)?;
+    let mut image_bytes: Vec<u8> = Vec::new();
+    image.read_to_end(&mut image_bytes)?;
+
+    let mut png = Png::try_from(image_bytes.as_ref())?;
+
+    let chunk = png.remove_first_chunk(chunk_type)?;
+
+    println!("MESSAGE: {}", chunk.data_as_string().unwrap());
+
+    Ok(())
+}
+
+fn print_message(path: PathBuf) {
     todo!()
 }
 
-fn decode(image: File, chunk_type: ChunkType) {
-    todo!()
-}
-
-fn print_message(image: File) {
-    todo!()
-}
-
-fn remove_message(image: File, chunk_type: ChunkType) {
+fn remove_message(path: PathBuf, chunk_type: ChunkType) {
     todo!()
 }
 
 fn main() -> Result<()> {
     let args = Args::parse();
 
-    let image = get_image(&args.input_file.unwrap());
+    let path = get_path(&args.input_file.unwrap());
 
     match args.mode {
-        Mode::Encode    => todo!(),
-        Mode::Decode    => todo!(),
+        Mode::Encode    => encode(path, &args.chunk_type.unwrap(), &args.message.unwrap())?,
+        Mode::Decode    => decode(path, &args.chunk_type.unwrap())?,
         Mode::Print     => todo!(),
         Mode::Remove    => todo!(),
         _               => Err("not a valid mode")?
diff --git a/src/png.rs b/src/png.rs
index 4751922..4ded0f9 100644
--- a/src/png.rs
+++ b/src/png.rs
@@ -4,7 +4,7 @@ use crc::Crc;
 use crate::{chunk::{self, Chunk}, chunk_type::{self, ChunkType}};
 
 
-struct Png {
+pub struct Png {
     png_signature: [u8; 8],
     png_chunks: Vec<Chunk>,
 }
@@ -19,11 +19,11 @@ impl Png {
         }
     }
 
-    fn append_chunk(&mut self, chunk: Chunk) {
+    pub fn append_chunk(&mut self, chunk: Chunk) {
         self.png_chunks.push(chunk)
     }
 
-    fn remove_first_chunk(&mut self, chunk_type: &str) -> Result<Chunk, &'static str> {
+    pub fn remove_first_chunk(&mut self, chunk_type: &str) -> Result<Chunk, &'static str> {
         let mut x = 0;
         let mut index = 0;
 
@@ -76,7 +76,7 @@ impl Png {
 
     }
 
-    fn as_bytes(&self) -> Vec<u8> {
+    pub fn as_bytes(&self) -> Vec<u8> {
         let mut final_vec: Vec<u8> = Vec::new();
 
         for i in Self::STANDARD_HEADER {