use crate::connection::Connection; use crate::file::{ChunkHeader, FileHandle, FileOffer, StdFileHandle}; use crate::handshake::Handshake; use crate::message::{FileOfferPayload, FileRequestPayload, Message}; use crate::ui::prompt_user_for_file_confirmation; use anyhow::{anyhow, Result}; use std::path::PathBuf; use tokio::fs::File; use tokio::net::TcpStream; pub async fn send(file_paths: &Vec, password: &String) -> Result<()> { // Fail early if there are problems generating file handles let handles = FileHandle::get_file_handles(file_paths).await?; // Establish connection to server let socket = TcpStream::connect("127.0.0.1:8080").await?; let (handshake, s1) = Handshake::from_password(password); // Complete handshake, returning key used for encryption let (socket, key) = handshake.negotiate(socket, s1).await?; let mut connection = Connection::new(socket, key); // Offer files, wait for requested file response let requested_chunks = offer_files(&mut connection, &handles).await?; // Upload negotiated files let std_file_handles = FileHandle::to_stds(handles, requested_chunks).await; connection.upload_files(std_file_handles).await?; println!("Done uploading."); // Exit Ok(()) } pub async fn receive(password: &String) -> Result<()> { // Establish connection to server let socket = TcpStream::connect("127.0.0.1:8080").await?; let (handshake, s1) = Handshake::from_password(password); // Complete handshake, returning key used for encryption let (socket, key) = handshake.negotiate(socket, s1).await?; let mut connection = Connection::new(socket, key); // Wait for offered files, respond with desired files let desired_files = request_specific_files(&mut connection).await?; // Create files let std_file_handles = create_or_find_files(desired_files).await?; // Download them connection.download_files(std_file_handles).await?; return Ok(()); } pub async fn offer_files( conn: &mut Connection, file_handles: &Vec, ) -> Result> { // Collect file offer let mut files = vec![]; for handle in file_handles { files.push(handle.to_file_offer()?); } let msg = Message::FileOffer(FileOfferPayload { files }); // Send file offer conn.send_msg(msg).await?; // Wait for reply let reply = conn.await_msg().await?; // Return requested chunks match reply { Message::FileRequest(file_request_payload) => Ok(file_request_payload.chunks), _ => Err(anyhow!("Expecting file request message back")), } } pub async fn request_specific_files(conn: &mut Connection) -> Result> { // Wait for offer message let offer_message = conn.await_msg().await?; let offered_files: Vec = match offer_message { Message::FileOffer(file_offer_payload) => file_offer_payload.files, _ => return Err(anyhow!("Expecting file offer message")), }; // Prompt user for confirmation of files let desired_files = prompt_user_for_file_confirmation(offered_files).await; let file_request_msg = Message::FileRequest(FileRequestPayload { chunks: desired_files .iter() .map(|file| ChunkHeader { id: file.id, start: 0, }) .collect(), }); conn.send_msg(file_request_msg).await?; Ok(desired_files) } pub async fn create_or_find_files(desired_files: Vec) -> Result> { let mut v = Vec::new(); for desired_file in desired_files { let mut filename = desired_file.path; filename.push_str(".part"); let file = match File::open(filename.clone()).await { Ok(file) => { println!( "File {:?} already exists. Attempting to resume download.", filename ); file } Err(_) => File::create(filename).await?, }; let metadata = file.metadata().await?; let std_file_handle = StdFileHandle::new(desired_file.id, file, metadata.len(), desired_file.size).await?; v.push(std_file_handle) } return Ok(v); }