From 20059ea6fa9e03a3ac6d9f6600888d6e1fb56f8b Mon Sep 17 00:00:00 2001
From: firebadnofire <william@archuser.org>
Date: Fri, 19 Apr 2024 09:08:12 -0400
Subject: [PATCH] commit

---
 src/main.rs | 49 ++++++++++++++++++++++++++++---------------------
 1 file changed, 28 insertions(+), 21 deletions(-)

diff --git a/src/main.rs b/src/main.rs
index 138bd40..0665946 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,8 +1,9 @@
 use std::env;
-use std::net::TcpStream;
 use std::io::{self, Write, BufRead, BufReader};
+use std::net::TcpStream;
 use chrono::Local;
 use std::process;
+use byteorder::{BigEndian, WriteBytesExt};
 
 fn main() {
     let args: Vec<String> = env::args().collect();
@@ -13,8 +14,8 @@ fn main() {
 
     let mut ip = "127.0.0.1"; // Default IP
     let mut port = "7890";     // Default port
-
     let mut i = 1;
+
     while i < args.len() {
         match args[i].as_str() {
             "-p" | "--port" if i + 1 < args.len() => {
@@ -28,32 +29,24 @@ fn main() {
         }
     }
 
-    let mut stream = TcpStream::connect(format!("{}:{}", ip, port)).unwrap_or_else(|_| {
-        eprintln!("Failed to connect to {}:{}", ip, port);
+    let address = format!("{}:{}", ip, port);
+    let mut stream = TcpStream::connect(&address).unwrap_or_else(|_| {
+        eprintln!("Failed to connect to {}", address);
         process::exit(1);
     });
 
-    println!("Connected to {}:{}", ip, port);
+    println!("Connected to {}", address);
     
     let stdin = io::stdin();
     let handle = stdin.lock();
 
-    // Checking if there's piped input by peeking the buffer
+    // Handling piped input by peeking the buffer
     let mut buffer = Vec::new();
     let mut reader = BufReader::new(handle);
     let bytes_available = reader.read_until(b'\n', &mut buffer).expect("Failed to read from input");
 
     if bytes_available == 0 {
-        // No piped input, proceed to interactive mode
-        loop {
-            let mut input = String::new();
-            println!("Enter your message (or 'exit' to quit):");
-            io::stdin().read_line(&mut input).unwrap();
-            if input.trim().eq("exit") {
-                break;
-            }
-            send_message(&mut stream, input.trim());
-        }
+        interactive_mode(&mut stream);
     } else {
         // Handle piped input
         send_message(&mut stream, std::str::from_utf8(&buffer).unwrap().trim());
@@ -64,13 +57,27 @@ fn main() {
     }
 }
 
+fn interactive_mode(stream: &mut TcpStream) {
+    loop {
+        let mut input = String::new();
+        println!("Enter your message (or 'exit' to quit):");
+        io::stdin().read_line(&mut input).unwrap();
+        if input.trim().eq("exit") {
+            break;
+        }
+        send_message(stream, input.trim());
+    }
+}
+
 fn send_message(stream: &mut TcpStream, message: &str) {
     let datetime = Local::now();
     let full_message = format!("{}: {}", datetime.format("%Y-%m-%d %H:%M:%S"), message);
+    let message_len = full_message.as_bytes().len() as u32;
 
-    if let Err(e) = stream.write_all(full_message.as_bytes()) {
-        eprintln!("Failed to send message: {:?}", e);
-        process::exit(1);
-    }
-    println!("Message sent!");
+    println!("Sending message, type: 1, len: {}, msg: {}", message_len, full_message);
+
+    stream.write_u8(1).unwrap(); // Message type 1 for normal messages
+    stream.write_u32::<BigEndian>(message_len).unwrap(); // Message length
+    stream.write_all(full_message.as_bytes()).unwrap(); // Actual message
+    stream.flush().unwrap(); // Ensure the message is sent immediately
 }