From 7ec28b778d7c87e8988265dac6c74a1079791426 Mon Sep 17 00:00:00 2001
From: Arthur Beck <averse.abfun@gmail.com>
Date: Thu, 20 Mar 2025 15:02:19 +0000
Subject: [PATCH] Added support for passwords

---
 Cargo.toml  |  1 +
 src/main.rs | 44 ++++++++++++++++++++++++++++++++++++++++++--
 2 files changed, 43 insertions(+), 2 deletions(-)

diff --git a/Cargo.toml b/Cargo.toml
index 68e85e3..101290f 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -5,4 +5,5 @@ edition = "2024"
 
 [dependencies]
 clap = { version = "4.5.31", features = ["derive"] }
+sha2 = "0.10.8"
 time = { version = "0.3.37", features = ["formatting", "local-offset"] }
diff --git a/src/main.rs b/src/main.rs
index 0aa1c18..5d7c04a 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,6 +1,7 @@
 #![feature(random)]
 
 use clap::Parser;
+use sha2::{Digest, Sha512};
 use std::io::{ErrorKind, prelude::*};
 use std::net::{TcpListener, TcpStream};
 use std::random;
@@ -13,6 +14,10 @@ struct Args {
     /// Port to bind on.
     #[arg(short, long, default_value_t = 7867)]
     port: u16,
+    /// Password. Will be stored as a SHA512 hashed string.
+    /// Will not require a password if none provided.
+    #[arg(short = 'P', long, default_value = None)]
+    password: Option<String>
 }
 
 #[derive(Clone, PartialEq, Eq)]
@@ -86,7 +91,7 @@ fn display_msg(stream: &mut TcpStream, msg: &Message) -> bool {
     true
 }
 
-fn handle_client(mut stream: TcpStream, channel: Channel) {
+fn handle_client(mut stream: TcpStream, channel: Channel, password: Option<Vec<u8>>) {
     stream
         .write_all(b"\x1b[33mEnter your chosen nickname(up to 30 characters):\x1b[0m\x1b[4m\n")
         .unwrap();
@@ -172,6 +177,34 @@ fn handle_client(mut stream: TcpStream, channel: Channel) {
         }
     }
 
+    if password.is_some() {
+        let password = password.unwrap();
+
+        stream
+            .write_all(b"\x1b[0m\x1b[33mThis server requires a password. Enter the password:\x1b[0m\x1b[4m\n")
+            .unwrap();
+
+        let mut buf = vec![];
+        let mut buf1 = [0u8];
+        'bufreadloop3: while buf1[0] != b'\n' || i == 0 {
+            stream.read_exact(&mut buf1).unwrap();
+            if buf1[0] == b'\n' {
+                let mut sha512 = Sha512::new();
+                sha512.update(buf);
+
+                if password != sha512.finalize().to_vec() {
+                    stream
+                        .write_all(b"\x1b[0m\x1b[31mIncorrect.\x1b[0m\n")
+                        .unwrap();
+                    
+                    return;
+                }
+                break 'bufreadloop3;
+            }
+            buf.push(buf1[0]);
+        }
+    }
+
     {
         let lock = channel.lock().unwrap();
         for msg in &lock.messages {
@@ -476,6 +509,12 @@ fn handle_client(mut stream: TcpStream, channel: Channel) {
 
 fn main() -> std::io::Result<()> {
     let args = Args::parse();
+    let mut password = None;
+    if args.password.is_some() {
+        let mut sha512 = sha2::Sha512::new();
+        sha512.update(args.password.unwrap());
+        password = Some(sha512.finalize().to_vec());
+    }
 
     let listener = TcpListener::bind(format!("0.0.0.0:{}", args.port))?;
     let channel: Channel = Arc::new(Mutex::new(UnwrappedChannel {
@@ -486,8 +525,9 @@ fn main() -> std::io::Result<()> {
 
     for stream in listener.incoming() {
         let channel: Channel = channel.clone();
+        let password = password.clone();
         thread::spawn(|| {
-            handle_client(stream.unwrap(), channel);
+            handle_client(stream.unwrap(), channel, password);
         });
     }