|
| 1 | +use std::io::{self, Write}; |
| 2 | +use std::path::PathBuf; |
| 3 | +use std::time::Duration; |
| 4 | + |
| 5 | +use eyre::{bail, Context, Result}; |
| 6 | +use indicatif::{ProgressBar, ProgressStyle}; |
| 7 | +use tokio::fs::File; |
| 8 | +use tokio::io::AsyncWriteExt; |
| 9 | + |
| 10 | +use atuin_client::api_client; |
| 11 | +use atuin_client::encryption::{decode_key, encode_key, load_key, Key}; |
| 12 | +use atuin_client::record::sqlite_store::SqliteStore; |
| 13 | +use atuin_client::record::store::Store; |
| 14 | +use atuin_client::settings::Settings; |
| 15 | + |
| 16 | +/// Timeout for the entire auth flow (10 minutes) |
| 17 | +const AUTH_TIMEOUT: Duration = Duration::from_secs(600); |
| 18 | +/// How often to poll for verification |
| 19 | +const POLL_INTERVAL: Duration = Duration::from_secs(2); |
| 20 | + |
| 21 | +/// Run the Hub authentication flow. |
| 22 | +/// |
| 23 | +/// This is used by both `register` and `login` commands when `hub_sync` is enabled. |
| 24 | +/// The flow is identical for both since the Hub web UI handles registration vs login. |
| 25 | +pub async fn run(settings: &Settings, store: &SqliteStore) -> Result<()> { |
| 26 | + println!("Authenticating with Atuin Hub..."); |
| 27 | + |
| 28 | + // 1. Request a code from the hub |
| 29 | + let code_response = api_client::hub_request_code(&settings.hub_address) |
| 30 | + .await |
| 31 | + .context("Failed to request authentication code from Hub")?; |
| 32 | + |
| 33 | + let code = &code_response.code; |
| 34 | + let auth_url = format!("{}/auth/cli?code={}", settings.hub_address, code); |
| 35 | + |
| 36 | + // 2. Try to open the browser and print the URL regardless |
| 37 | + println!("\nOpening your browser to complete authentication..."); |
| 38 | + println!("If it doesn't open, visit this URL:\n"); |
| 39 | + println!(" {auth_url}\n"); |
| 40 | + |
| 41 | + if let Err(e) = open::that(&auth_url) { |
| 42 | + tracing::debug!("Failed to open browser: {}", e); |
| 43 | + } |
| 44 | + |
| 45 | + // 3. Poll for verification with a spinner |
| 46 | + let spinner = ProgressBar::new_spinner(); |
| 47 | + spinner.set_style( |
| 48 | + ProgressStyle::default_spinner() |
| 49 | + .tick_chars("⠋⠙⠹⠸⠼⠴⠦⠧⠇⠏") |
| 50 | + .template("{spinner:.cyan} {msg}") |
| 51 | + .expect("valid template"), |
| 52 | + ); |
| 53 | + spinner.set_message("Waiting for authorization..."); |
| 54 | + spinner.enable_steady_tick(Duration::from_millis(100)); |
| 55 | + |
| 56 | + let start = std::time::Instant::now(); |
| 57 | + let session = loop { |
| 58 | + if start.elapsed() > AUTH_TIMEOUT { |
| 59 | + spinner.finish_with_message("Timed out"); |
| 60 | + bail!("Authentication timed out. Please try again."); |
| 61 | + } |
| 62 | + |
| 63 | + tokio::time::sleep(POLL_INTERVAL).await; |
| 64 | + |
| 65 | + match api_client::hub_verify_code(&settings.hub_address, code).await { |
| 66 | + Ok(verify_response) => { |
| 67 | + if let Some(session) = verify_response.session { |
| 68 | + spinner.finish_with_message("Authorized!"); |
| 69 | + break session; |
| 70 | + } |
| 71 | + // Still pending, continue polling |
| 72 | + } |
| 73 | + Err(e) => { |
| 74 | + // Log the error but keep polling - could be transient |
| 75 | + tracing::debug!("Verification poll failed: {}", e); |
| 76 | + } |
| 77 | + } |
| 78 | + }; |
| 79 | + |
| 80 | + // 4. Save the session token |
| 81 | + let session_path = settings.session_path.as_str(); |
| 82 | + let mut file = File::create(session_path) |
| 83 | + .await |
| 84 | + .context("Failed to create session file")?; |
| 85 | + file.write_all(session.as_bytes()) |
| 86 | + .await |
| 87 | + .context("Failed to write session file")?; |
| 88 | + |
| 89 | + // 5. Handle encryption key - always prompt |
| 90 | + // Users will always have a key file by this point (created on first atuin usage) |
| 91 | + // But it might not be the right key for this account if they registered elsewhere |
| 92 | + let key_path = PathBuf::from(settings.key_path.as_str()); |
| 93 | + |
| 94 | + println!(); |
| 95 | + println!( |
| 96 | + "If you have already registered on another machine, you will need your encryption key." |
| 97 | + ); |
| 98 | + println!("Run 'atuin key' on your other machine to retrieve it."); |
| 99 | + println!(); |
| 100 | + |
| 101 | + let key_input = read_key_input()?; |
| 102 | + |
| 103 | + if let Some(encoded_key) = key_input { |
| 104 | + // User provided a key - check if we need to re-encrypt |
| 105 | + let current_key: [u8; 32] = load_key(settings)?.into(); |
| 106 | + let new_key: [u8; 32] = decode_key(encoded_key.clone()) |
| 107 | + .context("Could not decode provided key")? |
| 108 | + .into(); |
| 109 | + |
| 110 | + if new_key != current_key { |
| 111 | + println!("\nRe-encrypting local store with new key..."); |
| 112 | + |
| 113 | + store.re_encrypt(¤t_key, &new_key).await?; |
| 114 | + |
| 115 | + println!("Writing new key"); |
| 116 | + let mut file = File::create(&key_path) |
| 117 | + .await |
| 118 | + .context("Failed to create key file")?; |
| 119 | + file.write_all(encoded_key.as_bytes()) |
| 120 | + .await |
| 121 | + .context("Failed to write key file")?; |
| 122 | + } |
| 123 | + } |
| 124 | + // If user pressed Enter, we just use the existing key (nothing to do) |
| 125 | + |
| 126 | + println!("\nAuthentication successful!"); |
| 127 | + println!(); |
| 128 | + println!("IMPORTANT: Please make a note of your key (run 'atuin key') and keep it safe."); |
| 129 | + println!("You will need it to log in on other devices, and we cannot help recover it if you lose it."); |
| 130 | + |
| 131 | + Ok(()) |
| 132 | +} |
| 133 | + |
| 134 | +/// Prompt the user for an encryption key. |
| 135 | +/// Returns `Some(encoded_key)` if they provided one, None if they pressed Enter. |
| 136 | +fn read_key_input() -> Result<Option<String>> { |
| 137 | + print!("Enter your encryption key, or press Enter if you don't already have one: "); |
| 138 | + io::stdout().flush()?; |
| 139 | + |
| 140 | + let mut input = String::new(); |
| 141 | + io::stdin().read_line(&mut input)?; |
| 142 | + let input = input.trim(); |
| 143 | + |
| 144 | + if input.is_empty() { |
| 145 | + return Ok(None); |
| 146 | + } |
| 147 | + |
| 148 | + // The key may be EITHER base64, or a bip mnemonic |
| 149 | + // Try to normalize to base64 |
| 150 | + let encoded = match bip39::Mnemonic::from_phrase(input, bip39::Language::English) { |
| 151 | + Ok(mnemonic) => encode_key(Key::from_slice(mnemonic.entropy()))?, |
| 152 | + Err(err) => { |
| 153 | + match err.downcast_ref::<bip39::ErrorKind>() { |
| 154 | + Some(bip_err) => { |
| 155 | + match bip_err { |
| 156 | + // Not a valid mnemonic word - assume they copied the base64 key |
| 157 | + bip39::ErrorKind::InvalidWord => input.to_string(), |
| 158 | + bip39::ErrorKind::InvalidChecksum => { |
| 159 | + bail!("Key mnemonic was not valid") |
| 160 | + } |
| 161 | + bip39::ErrorKind::InvalidKeysize(_) |
| 162 | + | bip39::ErrorKind::InvalidWordLength(_) |
| 163 | + | bip39::ErrorKind::InvalidEntropyLength(_, _) => { |
| 164 | + bail!("Key was not the correct length") |
| 165 | + } |
| 166 | + } |
| 167 | + } |
| 168 | + _ => { |
| 169 | + // Unknown error - assume they copied the base64 key |
| 170 | + input.to_string() |
| 171 | + } |
| 172 | + } |
| 173 | + } |
| 174 | + }; |
| 175 | + |
| 176 | + // Validate the key |
| 177 | + if decode_key(encoded.clone()).is_err() { |
| 178 | + bail!("The provided key was invalid"); |
| 179 | + } |
| 180 | + |
| 181 | + Ok(Some(encoded)) |
| 182 | +} |
0 commit comments