// Prevents additional console window on Windows in release, DO NOT REMOVE!! #![cfg_attr(not(debug_assertions), windows_subsystem = "windows")] mod config; use config::{get_config_path, load_config, save_config, Config}; use futures_util::{SinkExt, StreamExt}; use serde_json::Value; use std::{env, process::Command, sync::Arc, time::Duration}; use tauri::{ menu::{Menu, MenuItem}, tray::{MouseButton, MouseButtonState, TrayIconBuilder, TrayIconEvent}, Emitter, Manager, }; use tokio::{ net::{TcpListener, TcpStream}, sync::{oneshot, Mutex}, time::sleep, }; use tokio_tungstenite::accept_async; struct ResponseChannel { sender: Option>, } struct WebSocketState { sender: Option< futures_util::stream::SplitSink< tokio_tungstenite::WebSocketStream, tokio_tungstenite::tungstenite::Message, >, >, response_channel: ResponseChannel, server_abort: Option>, config: Config, } async fn start_websocket_server(app_handle: tauri::AppHandle, port: u16) -> Result<(), String> { let addr = format!("127.0.0.1:{}", port); // First ensure any existing server is stopped { let state = app_handle.state::>>(); let mut state = state.lock().await; if let Some(old_abort) = state.server_abort.take() { let _ = old_abort.send(()); // Give it a moment to shut down sleep(Duration::from_millis(200)).await; } } // Now try to bind to the port let listener = TcpListener::bind(&addr) .await .map_err(|e| format!("Failed to bind to port {}: {}", port, e))?; let (abort_sender, mut abort_receiver) = tokio::sync::oneshot::channel(); // Store the new abort sender { let state = app_handle.state::>>(); let mut state = state.lock().await; state.server_abort = Some(abort_sender); } // Spawn the server task tokio::spawn(async move { println!("Starting WebSocket server on port {}", port); loop { tokio::select! { accept_result = listener.accept() => { match accept_result { Ok((stream, _)) => { let app_handle = app_handle.clone(); tokio::spawn(handle_connection(stream, app_handle)); } Err(e) => { println!("Error accepting connection: {}", e); break; } } } _ = &mut abort_receiver => { println!("WebSocket server shutting down on port {}...", port); break; } } } }); // Wait a moment to ensure the server has started sleep(Duration::from_millis(100)).await; Ok(()) } #[tauri::command] async fn get_config(state: tauri::State<'_, Arc>>) -> Result { let state = state.lock().await; Ok(state.config.clone()) } #[tauri::command] fn get_config_file_path() -> Result { match get_config_path() { Some(path) => Ok(path.to_string_lossy().into_owned()), None => Err("Could not determine config path".to_string()), } } #[tauri::command] async fn update_config( new_config: Config, state: tauri::State<'_, Arc>>, app_handle: tauri::AppHandle, ) -> Result { // Save the new config first save_config(&new_config)?; // Update the state with new config { let mut state = state.lock().await; state.config = new_config.clone(); } // Start the new server (this will also handle stopping the old one) start_websocket_server(app_handle, new_config.port).await?; Ok(new_config) } #[tauri::command] async fn reset_config( state: tauri::State<'_, Arc>>, app_handle: tauri::AppHandle, ) -> Result { let config = Config::default(); save_config(&config)?; { let mut state = state.lock().await; state.config = config.clone(); } start_websocket_server(app_handle, config.port).await?; Ok(config) } #[tauri::command] async fn send_to_extension( message: String, state: tauri::State<'_, Arc>>, ) -> Result<(), String> { let mut state = state.lock().await; if let Some(sender) = &mut state.sender { sender .send(tokio_tungstenite::tungstenite::Message::Text(message)) .await .map_err(|e| format!("Failed to send message: {}", e))?; Ok(()) } else { Err("No active WebSocket connection".to_string()) } } #[tauri::command] async fn receive_frontend_response( response: String, state: tauri::State<'_, Arc>>, ) -> Result<(), String> { let mut state = state.lock().await; if let Some(sender) = state.response_channel.sender.take() { sender .send(response) .map_err(|e| format!("Failed to send response: {:?}", e))?; } Ok(()) } #[tauri::command] fn fetch_video_info(url: String) { #[cfg(target_os = "windows")] { let command = format!("pytubepp \"{}\" -i", &url); Command::new("cmd") .args(["/k", command.as_str()]) .spawn() .unwrap(); } #[cfg(target_os = "linux")] { let command = format!("pytubepp \"{}\" -i", &url); Command::new("gnome-terminal") .args(["--", "bash", "-c", command.as_str()]) .spawn() .unwrap(); } #[cfg(target_os = "macos")] { let command = format!("pytubepp \"{}\" -i", &url); let escaped_command = command.replace("\"", "\\\""); let applescript = format!( "tell application \"Terminal\"\n\ do script \"{}\"\n\ activate\n\ end tell", escaped_command ); Command::new("osascript") .arg("-e") .arg(applescript) .spawn() .unwrap(); } } #[tauri::command] fn install_program(icommand: String) { #[cfg(target_os = "windows")] { let command = format!("{}", &icommand); Command::new("cmd") .args(["/k", command.as_str()]) .spawn() .unwrap(); } #[cfg(target_os = "linux")] { let command = format!("{}", &icommand); Command::new("gnome-terminal") .args(["--", "bash", "-c", command.as_str()]) .spawn() .unwrap(); } #[cfg(target_os = "macos")] { let command = format!("{}", &icommand); let escaped_command = command.replace("\"", "\\\""); let applescript = format!( "tell application \"Terminal\"\n\ do script \"{}\"\n\ activate\n\ end tell", escaped_command ); Command::new("osascript") .arg("-e") .arg(applescript) .spawn() .unwrap(); } } #[tauri::command] fn download_stream(url: String, stream: String, caption: Option) { let caption = caption.unwrap_or("none".to_string()); #[cfg(target_os = "windows")] { let command = format!("pytubepp \"{}\" -s {} -c {}", &url, &stream, &caption); Command::new("cmd") .args(["/k", command.as_str()]) .spawn() .unwrap(); } #[cfg(target_os = "linux")] { let command = format!("pytubepp \"{}\" -s {} -c {}", &url, &stream, &caption); Command::new("gnome-terminal") .args(["--", "bash", "-c", command.as_str()]) .spawn() .unwrap(); } #[cfg(target_os = "macos")] { let command = format!("pytubepp \"{}\" -s {} -c {}", &url, &stream, &caption); let escaped_command = command.replace("\"", "\\\""); let applescript = format!( "tell application \"Terminal\"\n\ do script \"{}\"\n\ activate\n\ end tell", escaped_command ); Command::new("osascript") .arg("-e") .arg(applescript) .spawn() .unwrap(); } } #[cfg_attr(mobile, tauri::mobile_entry_point)] pub async fn run() { let _ = fix_path_env::fix(); let config = load_config(); let port = config.port; let websocket_state = Arc::new(Mutex::new(WebSocketState { sender: None, response_channel: ResponseChannel { sender: None }, server_abort: None, config, })); let args: Vec = env::args().collect(); let start_hidden = args.contains(&"--hidden".to_string()); tauri::Builder::default() .plugin(tauri_plugin_single_instance::init(|app, _args, _cwd| { // Focus the main window when attempting to launch another instance if let Some(window) = app.get_webview_window("main") { let _ = window.show(); let _ = window.set_focus(); } })) .plugin(tauri_plugin_os::init()) .plugin(tauri_plugin_fs::init()) .plugin(tauri_plugin_shell::init()) .manage(websocket_state.clone()) .setup(move |app| { // Create menu items let quit = MenuItem::with_id(app, "quit", "Quit", true, None::<&str>) .map_err(|e| format!("Failed to create quit menu item: {}", e))?; let show = MenuItem::with_id(app, "show", "Show", true, None::<&str>) .map_err(|e| format!("Failed to create show menu item: {}", e))?; // Create the menu let menu = Menu::with_items(app, &[&show, &quit]) .map_err(|e| format!("Failed to create menu: {}", e))?; // Create and store the tray icon let tray = TrayIconBuilder::with_id("main") .icon(app.default_window_icon().unwrap().clone()) .menu(&menu) .tooltip("PytubePP Helper") .on_menu_event(|app, event| match event.id.as_ref() { "show" => { if let Some(window) = app.get_webview_window("main") { let _ = window.show(); let _ = window.set_focus(); } } "quit" => { app.exit(0); } _ => {} }) .on_tray_icon_event(|tray, event| { if let TrayIconEvent::Click { button: MouseButton::Left, button_state: MouseButtonState::Up, .. } = event { let app = tray.app_handle(); if let Some(window) = app.get_webview_window("main") { let _ = window.show(); let _ = window.set_focus(); } } }) .build(app) .map_err(|e| format!("Failed to create tray: {}", e))?; // Store the tray handle in the app state app.manage(tray); let window = app.get_webview_window("main").unwrap(); if start_hidden { window.hide().unwrap(); } // Start the initial WebSocket server let app_handle = app.handle().clone(); tokio::spawn(async move { if let Err(e) = start_websocket_server(app_handle, port).await { println!("Failed to start initial WebSocket server: {}", e); } }); Ok(()) }) .invoke_handler(tauri::generate_handler![ send_to_extension, fetch_video_info, install_program, download_stream, receive_frontend_response, get_config, update_config, reset_config, get_config_file_path ]) .run(tauri::generate_context!()) .expect("error while running tauri application"); } async fn handle_connection(stream: TcpStream, app_handle: tauri::AppHandle) { let ws_stream = accept_async(stream).await.unwrap(); let (ws_sender, mut ws_receiver) = ws_stream.split(); // Store the sender in the shared state { let state = app_handle.state::>>(); let mut state = state.lock().await; state.sender = Some(ws_sender); } println!("New WebSocket connection established"); while let Some(msg) = ws_receiver.next().await { if let Ok(msg) = msg { if let Ok(text) = msg.to_text() { println!("Received message: {}", text); // Parse the JSON message if let Ok(json_value) = serde_json::from_str::(text) { // Create a new channel for this request let (response_sender, response_receiver) = oneshot::channel(); { let state = app_handle.state::>>(); let mut state = state.lock().await; state.response_channel.sender = Some(response_sender); } // Emit an event to the frontend app_handle .emit_to("main", "websocket-message", json_value) .unwrap(); // Wait for the response from the frontend let response = response_receiver .await .unwrap_or_else(|e| format!("Error receiving response: {:?}", e)); // Send the response back through WebSocket let state = app_handle.state::>>(); let mut state = state.lock().await; if let Some(sender) = &mut state.sender { let _ = sender .send(tokio_tungstenite::tungstenite::Message::Text(response)) .await; } } } } } println!("WebSocket connection closed"); // Remove the sender from the shared state when the connection closes let state = app_handle.state::>>(); let mut state = state.lock().await; state.sender = None; }