1  use regex::Regex;
  2  use reqwest::blocking::{Client, Response};
  3  use serde::{Deserialize, Serialize};
  4  use std::{thread, time::Duration};
  5  
  6  use crate::{clear_screen, encryption::decrypt_data, MessageData};
  7  
  8  // Structures for public keys and ciphertext
  9  #[derive(Serialize, Deserialize)]
 10  struct Message {
 11      message: String,
 12      room_id: String,
 13  }
 14  
 15  // Function to create the reqwest blocking client with proxy
 16  pub fn create_client_with_proxy(proxy: &str) -> Client {
 17      // Create the reqwest client with custom transport handling the proxy
 18      let transport = reqwest::blocking::Client::builder()
 19          .danger_accept_invalid_certs(false) // Reject invalid certificates
 20          .proxy(reqwest::Proxy::all(proxy).expect("Invalid proxy address")) // Set proxy
 21          // Route through Tor/I2P proxy
 22          .build()
 23          .unwrap();
 24  
 25      transport
 26  }
 27  
 28  pub fn fetch_kyber_pubkey(password: &str, server_url: &str) -> Option<String> {
 29      // Check if the server URL contains `.i2p`
 30      let proxy = if server_url.contains(".i2p") {
 31          "http://127.0.0.1:4444" // I2P Proxy address
 32      } else {
 33          "socks5h://127.0.0.1:9050" // SOCKS5 Proxy address (Tor)
 34      };
 35      let client = create_client_with_proxy(proxy);
 36      
 37      let url = format!("{}/messages?room_id={}", server_url, password);
 38      let mut retries = 0;
 39      let max_retries = 3;
 40  
 41      loop {
 42          let res: Response = match client.get(&url).timeout(Duration::from_secs(60)).send() {
 43              Ok(response) => response,
 44              Err(_) => {
 45                  retries += 1;
 46                  if retries > max_retries {
 47                      return None; // Return None after 3 failed attempts
 48                  }
 49                  println!("Error while fetching public key. Retrying...");
 50                  thread::sleep(Duration::from_secs(2)); // Wait before retrying
 51                  continue;
 52              }
 53          };
 54  
 55          if res.status().is_success() {
 56              let body = match res.text() {
 57                  Ok(text) => text,
 58                  Err(_) => {
 59                      retries += 1;
 60                      if retries > max_retries {
 61                          return None;
 62                      }
 63                      println!("Error while reading response body. Retrying...");
 64                      thread::sleep(Duration::from_secs(2)); // Wait before retrying
 65                      continue;
 66                  }
 67              };
 68  
 69              if let Some(public_key_start) = body.find("KYBER_PUBLIC_KEY:") {
 70                  let public_key = &body[public_key_start + "KYBER_PUBLIC_KEY:".len()..]; // Remove marker
 71                  if let Some(end_data) = public_key.find("[END DATA]") {
 72                      return Some(public_key[0..end_data].to_string()); // Remove [END DATA] marker
 73                  }
 74              }
 75          }
 76  
 77          retries += 1;
 78          if retries > max_retries {
 79              return None; // Return None after 3 failed attempts
 80          }
 81  
 82          println!("Public key not found. Retrying...");
 83          thread::sleep(Duration::from_secs(2)); // Sleep for 2 seconds before retrying
 84      }
 85  }
 86  
 87  pub fn fetch_dilithium_pubkeys(password: &str, server_url: &str) -> Vec<String> {
 88      // Check if the server URL contains `.i2p`
 89      let proxy = if server_url.contains(".i2p") {
 90          "http://127.0.0.1:4444" // I2P Proxy address
 91      } else {
 92          "socks5h://127.0.0.1:9050" // SOCKS5 Proxy address (Tor)
 93      };
 94      let client = create_client_with_proxy(proxy);
 95  
 96      let url = format!("{}/messages?room_id={}", server_url, password);
 97      let mut retries = 0;
 98      let max_retries = 3;
 99  
100      loop {
101          let res: Response = match client.get(&url).timeout(Duration::from_secs(60)).send() {
102              Ok(response) => response,
103              Err(_) => {
104                  retries += 1;
105                  if retries > max_retries {
106                      eprintln!("Failed to fetch public keys after {} retries.", max_retries);
107                      return Vec::new(); // Return an empty vector on failure
108                  }
109                  println!("Error while fetching public keys. Retrying...");
110                  thread::sleep(Duration::from_secs(2)); // Wait before retrying
111                  continue;
112              }
113          };
114  
115          if res.status().is_success() {
116              let body = match res.text() {
117                  Ok(text) => text,
118                  Err(_) => {
119                      retries += 1;
120                      if retries > max_retries {
121                          eprintln!("Failed to read response body after {} retries.", max_retries);
122                          return Vec::new();
123                      }
124                      println!("Error while reading response body. Retrying...");
125                      thread::sleep(Duration::from_secs(2)); // Wait before retrying
126                      continue;
127                  }
128              };
129  
130              let mut public_keys = Vec::new();
131              for key_data in body.split("DILITHIUM_PUBLIC_KEY:") {
132                  if let Some(end_data) = key_data.find("[END DATA]") {
133                      let key = key_data[0..end_data].trim().to_string();
134                      public_keys.push(key);
135                  }
136              }
137  
138              if !public_keys.is_empty() {
139                  return public_keys; // Return all valid public keys
140              }
141          }
142  
143          retries += 1;
144          if retries > max_retries {
145              eprintln!("Public keys not found after {} retries.", max_retries);
146              return Vec::new(); // Return an empty vector on failure
147          }
148  
149          println!("No valid public keys found in response. Retrying...");
150          thread::sleep(Duration::from_secs(2)); // Sleep for 2 seconds before retrying
151      }
152  }
153  
154  pub fn fetch_eddsa_pubkeys(password: &str, server_url: &str) -> Vec<String> {
155      // Check if the server URL contains `.i2p`
156      let proxy = if server_url.contains(".i2p") {
157          "http://127.0.0.1:4444" // I2P Proxy address
158      } else {
159          "socks5h://127.0.0.1:9050" // SOCKS5 Proxy address (Tor)
160      };
161      let client = create_client_with_proxy(proxy);
162  
163      let url = format!("{}/messages?room_id={}", server_url, password);
164      let mut retries = 0;
165      let max_retries = 3;
166  
167      loop {
168          let res: Response = match client.get(&url).timeout(Duration::from_secs(60)).send() {
169              Ok(response) => response,
170              Err(_) => {
171                  retries += 1;
172                  if retries > max_retries {
173                      eprintln!("Failed to fetch public keys after {} retries.", max_retries);
174                      return Vec::new(); // Return an empty vector on failure
175                  }
176                  println!("Error while fetching public keys. Retrying...");
177                  thread::sleep(Duration::from_secs(2)); // Wait before retrying
178                  continue;
179              }
180          };
181  
182          if res.status().is_success() {
183              let body = match res.text() {
184                  Ok(text) => text,
185                  Err(_) => {
186                      retries += 1;
187                      if retries > max_retries {
188                          eprintln!("Failed to read response body after {} retries.", max_retries);
189                          return Vec::new();
190                      }
191                      println!("Error while reading response body. Retrying...");
192                      thread::sleep(Duration::from_secs(2)); // Wait before retrying
193                      continue;
194                  }
195              };
196  
197              let mut public_keys = Vec::new();
198              for key_data in body.split("EDDSA_PUBLIC_KEY:") {
199                  if let Some(end_data) = key_data.find("[END DATA]") {
200                      let key = key_data[0..end_data].trim().to_string();
201                      public_keys.push(key);
202                  }
203              }
204  
205              if !public_keys.is_empty() {
206                  return public_keys; // Return all valid public keys
207              }
208          }
209  
210          retries += 1;
211          if retries > max_retries {
212              eprintln!("Public keys not found after {} retries.", max_retries);
213              return Vec::new(); // Return an empty vector on failure
214          }
215  
216          println!("No valid public keys found in response. Retrying...");
217          thread::sleep(Duration::from_secs(2)); // Sleep for 2 seconds before retrying
218      }
219  }
220  
221  pub fn fetch_ciphertext(password: &str, server_url: &str) -> String {
222      // Check if the server URL contains `.i2p`
223      let proxy = if server_url.contains(".i2p") {
224          "http://127.0.0.1:4444" // I2P Proxy address
225      } else {
226          "socks5h://127.0.0.1:9050" // SOCKS5 Proxy address (Tor)
227      };
228      let client = create_client_with_proxy(proxy);
229  
230      let url = format!("{}/messages?room_id={}", server_url, password);
231  
232      loop {
233          let res: Response = match client.get(&url).timeout(Duration::from_secs(60)).send() {
234              Ok(response) => response,
235              Err(err) => {
236                  println!("Error while fetching ciphertext: {}. Retrying...", err);
237                  thread::sleep(Duration::from_secs(2)); // Wait before retrying
238                  continue;
239              }
240          };
241  
242          if res.status().is_success() {
243              let body = match res.text() {
244                  Ok(text) => text,
245                  Err(err) => {
246                      println!("Error while reading response body: {}. Retrying...", err);
247                      thread::sleep(Duration::from_secs(2)); // Wait before retrying
248                      continue;
249                  }
250              };
251  
252              if let Some(ciphertext_start) = body.find("KYBER_PUBLIC_KEY:CIPHERTEXT:") {
253                  let ciphertext = &body[ciphertext_start + "KYBER_PUBLIC_KEY:CIPHERTEXT:".len()..]; // Remove marker
254                  if let Some(end_data) = ciphertext.find("[END DATA]") {
255                      return ciphertext[0..end_data].to_string(); // Remove [END DATA] marker
256                  }
257              }
258          }
259  
260          // Wait for 2 seconds before retrying
261          println!("Ciphertext not found. Retrying...");
262          thread::sleep(Duration::from_secs(2)); // Sleep for 2 seconds before retrying
263      }
264  }
265  
266  pub fn send_kyber_pubkey(room_id: &str, public_key: &str, url: &str) {
267      // Check if the server URL contains `.i2p`
268      let proxy = if url.contains(".i2p") {
269          "http://127.0.0.1:4444" // I2P Proxy address
270      } else {
271          "socks5h://127.0.0.1:9050" // SOCKS5 Proxy address (Tor)
272      };
273      let client = create_client_with_proxy(proxy);
274  
275      let full_url = format!("{}/send", url); // Append /send to the URL
276      let message = Message {
277          message: format!("KYBER_PUBLIC_KEY:{}[END DATA]", public_key),
278          room_id: room_id.to_string(),
279      };
280  
281      let res = client.post(&full_url).json(&message).timeout(Duration::from_secs(60)).send(); // Use the full URL
282  
283      match res {
284          Ok(response) if response.status().is_success() => {
285              println!("Kyber1024 public key sent successfully!");
286          }
287          Ok(response) => {
288              println!("Failed to send public key. Status: {}", response.status());
289          }
290          Err(e) => {
291              println!("Failed to send public key. Error: {}", e);
292          }
293      }
294  }
295  
296  pub fn send_dilithium_pubkey(room_id: &str, public_key: &str, url: &str) {
297      // Check if the server URL contains `.i2p`
298      let proxy = if url.contains(".i2p") {
299          "http://127.0.0.1:4444" // I2P Proxy address
300      } else {
301          "socks5h://127.0.0.1:9050" // SOCKS5 Proxy address (Tor)
302      };
303      let client = create_client_with_proxy(proxy);
304  
305      let full_url = format!("{}/send", url); // Append /send to the URL
306      let message = Message {
307          message: format!("DILITHIUM_PUBLIC_KEY:{}[END DATA]", public_key),
308          room_id: room_id.to_string(),
309      };
310  
311      let res = client.post(&full_url).json(&message).timeout(Duration::from_secs(60)).send(); // Use the full URL
312  
313      match res {
314          Ok(response) if response.status().is_success() => {
315              println!("Dilithium5 public key sent successfully!");
316          }
317          Ok(response) => {
318              println!("Failed to send public key. Status: {}", response.status());
319          }
320          Err(e) => {
321              println!("Failed to send public key. Error: {}", e);
322          }
323      }
324  }
325  
326  pub fn send_eddsa_pubkey(room_id: &str, public_key: &str, url: &str) {
327      // Check if the server URL contains `.i2p`
328      let proxy = if url.contains(".i2p") {
329          "http://127.0.0.1:4444" // I2P Proxy address
330      } else {
331          "socks5h://127.0.0.1:9050" // SOCKS5 Proxy address (Tor)
332      };
333      let client = create_client_with_proxy(proxy);
334  
335      let full_url = format!("{}/send", url); // Append /send to the URL
336      let message = Message {
337          message: format!("EDDSA_PUBLIC_KEY:{}[END DATA]", public_key),
338          room_id: room_id.to_string(),
339      };
340  
341      let res: Response = match client.post(&full_url).json(&message).timeout(Duration::from_secs(60)).send() {
342          Ok(response) => response,
343          Err(_) => {
344              println!("Failed to send the public key.");
345              return;
346          }
347      };
348  
349      if res.status().is_success() {
350          println!("EdDSA public key sent successfully!");
351      } else {
352          println!("Failed to send public key.");
353      }
354  }
355  
356  pub fn send_ciphertext(room_id: &str, ciphertext: &str, url: &str) {
357      // Check if the server URL contains `.i2p`
358      let proxy = if url.contains(".i2p") {
359          "http://127.0.0.1:4444" // I2P Proxy address
360      } else {
361          "socks5h://127.0.0.1:9050" // SOCKS5 Proxy address (Tor)
362      };
363      let client = create_client_with_proxy(proxy);
364  
365      let full_url = format!("{}/send", url); // Append /send to the URL
366      let message = Message {
367          message: format!("KYBER_PUBLIC_KEY:CIPHERTEXT:{}[END DATA]", ciphertext),
368          room_id: room_id.to_string(),
369      };
370  
371      let res: Response = client.post(&full_url).json(&message).timeout(Duration::from_secs(60)).send().unwrap(); // Use the full URL
372  
373      if res.status().is_success() {
374          println!("Ciphertext sent successfully!");
375      } else {
376          println!("Failed to send ciphertext");
377      }
378  }
379  
380  pub fn send_encrypted_message(
381      encrypted_message: &str,
382      room_id: &str,
383      server_url: &str,
384  ) -> Result<(), Box<dyn std::error::Error>> {
385      // Check if the server URL contains `.i2p`
386      let proxy = if server_url.contains(".i2p") {
387          "http://127.0.0.1:4444" // I2P Proxy address
388      } else {
389          "socks5h://127.0.0.1:9050" // SOCKS5 Proxy address (Tor)
390      };
391      let client = create_client_with_proxy(proxy);
392  
393      // Format the encrypted message with the BEGIN and END markers
394      let formatted_encrypted_message = format!(
395          "-----BEGIN ENCRYPTED MESSAGE-----{}-----END ENCRYPTED MESSAGE-----",
396          encrypted_message
397      );
398  
399      // Create the message data to send
400      let message_data = MessageData {
401          message: formatted_encrypted_message,
402          room_id: room_id.to_string(),
403      };
404  
405      // Construct the full URL for sending the message
406      let send_url = format!("{}/send", server_url);
407  
408      // Send the message via HTTP POST request
409      let res = client
410          .post(&send_url)
411          .json(&message_data)
412          .timeout(Duration::from_secs(60)) // Set a timeout for the request
413          .send()?;
414  
415      // Check if the request was successful and print the result
416      if res.status().is_success() {
417          println!("Message sent successfully.");
418      } else {
419          eprintln!("Failed to send message: {}", res.status());
420      }
421  
422      Ok(())
423  }
424  
425  pub fn receive_and_fetch_messages(
426      room_id: &str,
427      shared_secret: &str,
428      server_url: &str,
429      gui: bool,
430  ) -> Result<Vec<String>, Box<dyn std::error::Error>> {
431      // SOCKS5 Proxy setup
432      // Check if the server URL contains `.i2p`
433      let proxy = if server_url.contains(".i2p") {
434          "http://127.0.0.1:4444" // I2P Proxy address
435      } else {
436          "socks5h://127.0.0.1:9050" // SOCKS5 Proxy address (Tor)
437      };
438      let client = create_client_with_proxy(proxy);
439  
440      // Build the URL with the provided room password and server URL
441      let url = format!("{}/messages?room_id={}", server_url, room_id);
442  
443      // Send a synchronous GET request to fetch messages
444      let res = client
445          .get(&url)
446          .timeout(std::time::Duration::from_secs(30)) // Set a timeout for the request
447          .send()?;
448  
449      // Declare the vector to store messages outside the response block
450      let mut messages = Vec::new();
451  
452      // Check if the request was successful
453      if res.status().is_success() {
454          clear_screen();
455          // Get the body of the HTML response
456          let body = res.text()?;
457  
458          // Define a regular expression to capture messages between the markers
459          let re = Regex::new(r"-----BEGIN ENCRYPTED MESSAGE-----\s*(.*?)\s*-----END ENCRYPTED MESSAGE-----")
460              .map_err(|e| Box::new(e) as Box<dyn std::error::Error>)?;
461  
462          // Iterate over all matches in the HTML body
463          for cap in re.captures_iter(&body) {
464              if let Some(encrypted_message) = cap.get(1) {
465                  // Step 1: Get the encrypted message without the markers
466                  let cleaned_message = encrypted_message.as_str().trim();
467  
468                  // Step 2: Decrypt the message (ignore the markers, only pass the actual content)
469                  match decrypt_data(cleaned_message, shared_secret) {
470                      Ok(decrypted_message) => {
471                          fn unpad_message(message: &str) -> String {
472                              // Remove everything inside the <padding>...</padding> part and trim any extra spaces from the message
473                              if let Some(start) = message.find("<padding>") {
474                                  if let Some(end) = message.find("</padding>") {
475                                      let (message_before_padding, _) = message.split_at(start); // Part before <padding>
476                                      let (_, message_after_padding) = message.split_at(end + 10); // Skip past </padding> (length 10 including '>')
477                                      return format!("{}{}", message_before_padding, message_after_padding);
478                                  }
479                              }
480                              message.to_string()  // In case there are no <padding>...</padding> markers
481                          }                                  
482                          
483                          let unpadded_message = unpad_message(&decrypted_message);
484                          
485                          // Ignore messages containing `[DUMMY_DATA]:`
486                          if unpadded_message.contains("[DUMMY_DATA]:") {
487                              continue;
488                          }
489  
490                          // If gui is false, skip messages containing `<media>`
491                          if !gui && unpadded_message.contains("<media>") {
492                              continue;
493                          }
494  
495                          // If gui is false, skip messages containing `<pfp>`
496                          if !gui && unpadded_message.contains("<pfp>") {
497                              continue;
498                          }
499  
500                          // If gui is true, do not replace <strong> tags
501                          let final_message = if gui {
502                              unpadded_message.to_string()
503                          } else {
504                              // Step 3: Replace <strong> tags with ANSI escape codes for bold text
505                              let strong_re = Regex::new(r"<strong>(.*?)</strong>").unwrap();
506                              strong_re.replace_all(&unpadded_message, |caps: &regex::Captures| {
507                                  // Replace <strong>...</strong> with ANSI escape codes for bold text
508                                  format!("\x1b[1m{}\x1b[0m", &caps[1])
509                              }).to_string()
510                          };
511  
512                          // Add the messages to the list
513                          messages.push(final_message);
514                      }
515                      Err(_e) => {
516                          // Ignore decryption failure
517                      }
518                  }
519              }
520          }
521      } else {
522          // Provide more detailed error info for failed requests
523          eprintln!("Failed to fetch messages: {} - {}", res.status(), res.text()?);
524      }
525  
526      // Return the collected messages
527      Ok(messages)
528  }