From 5602d3fbf466ea2776701c5ff0ef936a6ed80096 Mon Sep 17 00:00:00 2001
From: Chris Pryer <cnpryer@gmail.com>
Date: Sat, 24 Feb 2024 10:07:31 -0500
Subject: [PATCH] Fall back to Twine's keystring in publish command

* Resolve credentials (file, cli)
* Resolve repository (file, cli)
* Simple logic test
* Simple snapshot test
---
 CHANGELOG.md              |   4 +
 rye/src/cli/publish.rs    | 415 +++++++++++++++++++++++++++-----------
 rye/tests/common/mod.rs   |  15 ++
 rye/tests/test_publish.rs |  18 ++
 4 files changed, 338 insertions(+), 114 deletions(-)
 create mode 100644 rye/tests/test_publish.rs

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1ff91dc110..e6e8758385 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -9,6 +9,10 @@ _Unreleased_
 
 <!-- released start -->
 
+- `publish` will fall back to dispatching username and repository url for Twine's keyring support. # 759
+
+- `publish` now includes `--skip-save-credentials`. #759
+
 ## 0.27.0
 
 Released on 2024-02-26
diff --git a/rye/src/cli/publish.rs b/rye/src/cli/publish.rs
index 45bc4756e2..8b65eae653 100644
--- a/rye/src/cli/publish.rs
+++ b/rye/src/cli/publish.rs
@@ -16,14 +16,18 @@ use crate::platform::{get_credentials, write_credentials};
 use crate::pyproject::PyProject;
 use crate::utils::{escape_string, get_venv_python_bin, tui_theme, CommandOutput};
 
+const DEFAULT_REPOSITORY: &str = "pypi";
+const DEFAULT_REPOSITORY_DOMAIN: &str = "upload.pypi.org";
+const DEFAULT_REPOSITORY_URL: &str = "https://upload.pypi.org/legacy/";
+
 /// Publish packages to a package repository.
 #[derive(Parser, Debug)]
 pub struct Args {
     /// The distribution files to upload to the repository (defaults to <workspace-root>/dist/*).
     dist: Option<Vec<PathBuf>>,
     /// The repository to publish to.
-    #[arg(short, long, default_value = "pypi")]
-    repository: String,
+    #[arg(short, long)]
+    repository: Option<String>,
     /// The repository url to publish to.
     #[arg(long)]
     repository_url: Option<Url>,
@@ -42,6 +46,9 @@ pub struct Args {
     /// Path to alternate CA bundle.
     #[arg(long)]
     cert: Option<PathBuf>,
+    /// Skip saving to credentials file.
+    #[arg(long)]
+    skip_save_credentials: bool,
     /// Skip prompts.
     #[arg(short, long)]
     yes: bool,
@@ -68,96 +75,130 @@ pub fn execute(cmd: Args) -> Result<(), Error> {
         None => vec![project.workspace_path().join("dist").join("*")],
     };
 
-    // a. Get token from arguments and offer encryption, then store in credentials file.
-    // b. Get token from ~/.rye/credentials keyed by provided repository and provide decryption option.
-    // c. Otherwise prompt for token and provide encryption option, storing the result in credentials.
-    let repository = &cmd.repository;
-    let mut credentials = get_credentials()?;
-    credentials
-        .entry(repository)
-        .or_insert(Item::Table(Table::new()));
-
-    let repository_url = match cmd.repository_url {
-        Some(url) => url,
-        None => {
-            let default_repository_url = Url::parse("https://upload.pypi.org/legacy/")?;
-            credentials
-                .get(repository)
-                .and_then(|table| table.get("repository-url"))
-                .map(|url| match Url::parse(&escape_string(url.to_string())) {
-                    Ok(url) => url,
-                    Err(_) => default_repository_url.clone(),
-                })
-                .unwrap_or(default_repository_url)
+    // Resolve credentials from CLI and credentials file
+    let mut credentials_file = get_credentials()?;
+    let credentials_table = cmd
+        .repository
+        .as_ref()
+        .and_then(|it| credentials_file.get(it));
+
+    let token = cmd.token.map(Secret::new);
+    let mut credentials =
+        resolve_credentials(credentials_table, cmd.username.as_ref(), token.as_ref());
+
+    // Resolve repo url from CLI and credentials file
+    let mut repository = resolve_repository(credentials_table, cmd.repository, cmd.repository_url)?;
+
+    // Token is from file so we may need to prompt decryption
+    let should_decrypt = token.is_none() && credentials.password.is_some();
+
+    let mut should_encrypt = token.is_some();
+
+    // Fallback prompts
+    if credentials.password.is_none() && !cmd.yes {
+        if repository.name.is_none() || is_default_repository(&repository) {
+            echo!("No access token found, generate one at: https://pypi.org/manage/account/token/");
+        }
+        let new_token = prompt_token()?;
+        if new_token.is_some() && credentials.username.is_none() {
+            credentials.username = Some("__token__".to_string());
         }
+        should_encrypt = new_token.is_some();
+        credentials.password = new_token;
+    }
+
+    let passphrase = if cmd.yes || (!should_decrypt && !should_encrypt) {
+        None
+    } else if should_decrypt {
+        prompt_decrypt_passphrase()?
+    } else {
+        prompt_encrypt_passphrase()?
     };
 
-    // If -r is pypi but the url isn't pypi then bail
-    if repository == "pypi" && repository_url.domain() != Some("upload.pypi.org") {
-        bail!("invalid pypi url {} (use -h for help)", repository_url);
+    if credentials.password.is_none() && !cmd.yes {
+        if credentials.username.is_none() {
+            credentials.username = prompt_username()?;
+        }
+        credentials.password = prompt_password()?;
     }
 
-    let username = match cmd.username {
-        Some(username) => username,
-        None => credentials
-            .get(repository)
-            .and_then(|table| table.get("username"))
-            .map(|username| username.to_string())
-            .map(escape_string)
-            .unwrap_or("__token__".to_string()),
-    };
+    if repository.url.is_none() && !cmd.yes {
+        repository.url = prompt_repository_url()?;
+    }
 
-    let token = if let Some(token) = cmd.token {
-        let secret = Secret::new(token);
-        let maybe_encrypted = maybe_encrypt(&secret, cmd.yes)?;
-        let maybe_encoded = maybe_encode(&secret, &maybe_encrypted);
-        credentials[repository]["token"] = Item::Value(maybe_encoded.expose_secret().into());
-        write_credentials(&credentials)?;
-
-        secret
-    } else if let Some(token) = credentials
-        .get(repository)
-        .and_then(|table| table.get("token"))
-        .map(|token| token.to_string())
-        .map(escape_string)
+    // If no token/password is resolved we can fallback to keyring with username and url.
+    if credentials.password.is_none() && credentials.username.is_none() && repository.url.is_none()
     {
-        let secret = Secret::new(token);
+        bail!(
+            "no configuration was found for repository '{}'",
+            repository.name.unwrap_or_default()
+        );
+    }
 
-        maybe_decrypt(&secret, cmd.yes)?
-    } else {
-        echo!("No access token found, generate one at: https://pypi.org/manage/account/token/");
-        let token = if !cmd.yes {
-            prompt_for_token()?
-        } else {
-            "".to_string()
-        };
-        if token.is_empty() {
-            bail!("an access token is required")
+    if !cmd.skip_save_credentials && repository.name.is_some() {
+        // Can expect repository name if we want to write to the file
+        let repo_name = repository.name.expect("repo name");
+        let table = credentials_file
+            .entry(&repo_name)
+            .or_insert(Item::Table(Table::new()));
+
+        if let Some(it) = credentials.password.as_ref() {
+            let mut final_token = it.expose_secret().clone();
+            if let Some(phrase) = passphrase.as_ref() {
+                if should_encrypt {
+                    final_token = hex::encode(encrypt(it, phrase)?.expose_secret());
+                }
+            }
+            if !final_token.is_empty() {
+                table["token"] = Item::Value(final_token.into());
+            }
         }
-        let secret = Secret::new(token);
-        let maybe_encrypted = maybe_encrypt(&secret, cmd.yes)?;
-        let maybe_encoded = maybe_encode(&secret, &maybe_encrypted);
-        credentials[repository]["token"] = Item::Value(maybe_encoded.expose_secret().into());
 
-        secret
-    };
+        if let Some(usr) = credentials.username.as_ref() {
+            if !usr.is_empty() {
+                table["username"] = Item::Value(usr.clone().into());
+            }
+        }
 
-    credentials[repository]["repository-url"] = Item::Value(repository_url.to_string().into());
-    credentials[repository]["username"] = Item::Value(username.clone().into());
-    write_credentials(&credentials)?;
+        if let Some(url) = repository.url.as_ref() {
+            table["repository-url"] = Item::Value(url.to_string().into());
+        }
+
+        write_credentials(&credentials_file)?;
+    }
 
     let mut publish_cmd = Command::new(get_venv_python_bin(&venv));
+
+    // Build Twine command
     publish_cmd
         .arg("-mtwine")
         .arg("--no-color")
         .arg("upload")
-        .args(files)
-        .arg("--username")
-        .arg(username)
-        .arg("--password")
-        .arg(token.expose_secret())
-        .arg("--repository-url")
-        .arg(repository_url.to_string());
+        .arg("--non-interactive")
+        .args(files);
+
+    // If a username is provided use it, if a password is provided without a username then use __token__ with Twine.
+    if let Some(usr) = credentials.username.or(credentials
+        .password
+        .as_ref()
+        .map(|_| "__token__".to_string()))
+    {
+        publish_cmd.arg("--username").arg(usr);
+    }
+    if let Some(pwd) = credentials.password.as_ref() {
+        publish_cmd.arg("--password");
+
+        if should_decrypt {
+            if let Some(phrase) = passphrase {
+                publish_cmd.arg(decrypt(pwd, &phrase)?.expose_secret());
+            }
+        } else {
+            publish_cmd.arg(pwd.expose_secret());
+        }
+    }
+    if let Some(url) = repository.url.as_ref() {
+        publish_cmd.arg("--repository-url").arg(url.to_string());
+    }
     if cmd.sign {
         publish_cmd.arg("--sign");
     }
@@ -181,30 +222,148 @@ pub fn execute(cmd: Args) -> Result<(), Error> {
     Ok(())
 }
 
-fn prompt_for_token() -> Result<String, Error> {
-    eprint!("Access token: ");
+struct Credentials {
+    username: Option<String>,
+    password: Option<Secret<String>>,
+}
+
+struct Repository {
+    name: Option<String>,
+    url: Option<Url>,
+}
+
+fn is_default_repository(repository: &Repository) -> bool {
+    repository
+        .name
+        .as_ref()
+        .map_or(false, |it| it == DEFAULT_REPOSITORY)
+        || repository
+            .url
+            .as_ref()
+            .map_or(false, |it| it.domain() == Some(DEFAULT_REPOSITORY_DOMAIN))
+}
+
+fn resolve_credentials(
+    credentials_table: Option<&Item>,
+    username: Option<&String>,
+    password: Option<&Secret<String>>,
+) -> Credentials {
+    if username.is_some() && password.is_some() {
+        return Credentials {
+            username: username.cloned(),
+            password: password.cloned(),
+        };
+    }
+
+    let mut credentials = Credentials {
+        username: None,
+        password: None,
+    };
+
+    if password.is_none() {
+        credentials.password = credentials_table.and_then(|it| {
+            it.get("token")
+                .map(Item::to_string)
+                .map(escape_string)
+                .map(Secret::new)
+        });
+    } else {
+        credentials.password = password.cloned();
+    }
+
+    if credentials.username.is_none() {
+        credentials.username = credentials_table
+            .and_then(|it| it.get("username").map(Item::to_string).map(escape_string));
+    }
+
+    credentials
+}
+
+fn resolve_repository(
+    credentials_table: Option<&Item>,
+    name: Option<String>,
+    url: Option<Url>,
+) -> Result<Repository, Error> {
+    let mut repository = Repository { name, url };
+
+    if repository.url.is_some() {
+        return Ok(repository);
+    }
+
+    if let Some(url) = credentials_table.and_then(|it| {
+        it.get("repository-url")
+            .map(Item::to_string)
+            .map(escape_string)
+    }) {
+        repository.url = Some(Url::parse(&url)?);
+    }
+
+    if repository.url.is_none()
+        && repository
+            .name
+            .as_ref()
+            .map_or(false, |it| it == DEFAULT_REPOSITORY)
+    {
+        repository.url = Some(Url::parse(DEFAULT_REPOSITORY_URL)?);
+    }
+
+    Ok(repository)
+}
+
+fn prompt_repository_url() -> Result<Option<Url>, Error> {
+    echo!("Repository URL: ");
+    let url = get_trimmed_user_input().context("failed to read provided url")?;
+
+    if url.is_empty() {
+        Ok(None)
+    } else {
+        Ok(Some(Url::parse(&url)?))
+    }
+}
+
+fn prompt_token() -> Result<Option<Secret<String>>, Error> {
+    echo!("Access token: ");
+    // TODO(cnpryer): Use dialoguer?
     let token = get_trimmed_user_input().context("failed to read provided token")?;
 
-    Ok(token)
+    if token.is_empty() {
+        Ok(None)
+    } else {
+        Ok(Some(Secret::new(token)))
+    }
 }
 
-fn maybe_encrypt(secret: &Secret<String>, yes: bool) -> Result<Secret<Vec<u8>>, Error> {
-    let phrase = if !yes {
-        dialoguer::Password::with_theme(tui_theme())
-            .with_prompt("Encrypt with passphrase (optional)")
-            .allow_empty_password(true)
-            .report(false)
-            .interact()
-            .map(Secret::new)?
+fn prompt_username() -> Result<Option<String>, Error> {
+    echo!("Username: ");
+    let username = get_trimmed_user_input().context("failed to read provided username")?;
+
+    if username.is_empty() {
+        Ok(None)
     } else {
-        Secret::new("".to_string())
-    };
+        Ok(Some(username))
+    }
+}
+
+fn prompt_password() -> Result<Option<Secret<String>>, Error> {
+    let password = dialoguer::Password::with_theme(tui_theme())
+        .with_prompt("Password")
+        .allow_empty_password(true)
+        .report(false)
+        .interact()?;
 
+    if password.is_empty() {
+        Ok(None)
+    } else {
+        Ok(Some(Secret::new(password)))
+    }
+}
+
+fn encrypt(secret: &Secret<String>, phrase: &Secret<String>) -> Result<Secret<Vec<u8>>, Error> {
     let token = if phrase.expose_secret().is_empty() {
         secret.expose_secret().as_bytes().to_vec()
     } else {
         // Do the encryption
-        let encryptor = Encryptor::with_user_passphrase(phrase);
+        let encryptor = Encryptor::with_user_passphrase(phrase.clone());
         let mut encrypted = vec![];
         let mut writer = encryptor.wrap_output(&mut encrypted)?;
         writer.write_all(secret.expose_secret().as_bytes())?;
@@ -216,18 +375,21 @@ fn maybe_encrypt(secret: &Secret<String>, yes: bool) -> Result<Secret<Vec<u8>>,
     Ok(Secret::new(token.to_vec()))
 }
 
-fn maybe_decrypt(secret: &Secret<String>, yes: bool) -> Result<Secret<String>, Error> {
-    let phrase = if !yes {
-        dialoguer::Password::with_theme(tui_theme())
-            .with_prompt("Decrypt with passphrase (optional)")
-            .allow_empty_password(true)
-            .report(false)
-            .interact()
-            .map(Secret::new)?
+fn prompt_encrypt_passphrase() -> Result<Option<Secret<String>>, Error> {
+    let phrase = dialoguer::Password::with_theme(tui_theme())
+        .with_prompt("Encrypt with passphrase (optional)")
+        .allow_empty_password(true)
+        .report(false)
+        .interact()?;
+
+    if phrase.is_empty() {
+        Ok(None)
     } else {
-        Secret::new("".to_string())
-    };
+        Ok(Some(Secret::new(phrase)))
+    }
+}
 
+fn decrypt(secret: &Secret<String>, phrase: &Secret<String>) -> Result<Secret<String>, Error> {
     if phrase.expose_secret().is_empty() {
         return Ok(secret.clone());
     }
@@ -237,7 +399,7 @@ fn maybe_decrypt(secret: &Secret<String>, yes: bool) -> Result<Secret<String>, E
     if let Decryptor::Passphrase(decryptor) = Decryptor::new(bytes.as_slice())? {
         // Do the decryption
         let mut decrypted = vec![];
-        let mut reader = decryptor.decrypt(&phrase, None)?;
+        let mut reader = decryptor.decrypt(phrase, None)?;
         reader.read_to_end(&mut decrypted)?;
 
         let token = String::from_utf8(decrypted).context("failed to parse utf-8")?;
@@ -249,6 +411,20 @@ fn maybe_decrypt(secret: &Secret<String>, yes: bool) -> Result<Secret<String>, E
     bail!("failed to decrypt")
 }
 
+fn prompt_decrypt_passphrase() -> Result<Option<Secret<String>>, Error> {
+    let phrase = dialoguer::Password::with_theme(tui_theme())
+        .with_prompt("Decrypt with passphrase (optional)")
+        .allow_empty_password(true)
+        .report(false)
+        .interact()?;
+
+    if phrase.is_empty() {
+        Ok(None)
+    } else {
+        Ok(Some(Secret::new(phrase)))
+    }
+}
+
 fn get_trimmed_user_input() -> Result<String, Error> {
     std::io::stderr().flush()?;
     let mut input = String::new();
@@ -257,20 +433,6 @@ fn get_trimmed_user_input() -> Result<String, Error> {
     Ok(input.trim().to_string())
 }
 
-/// Helper function to manage potentially encoding secret data.
-///
-/// If the original secret data (bytes) are not the same as the new secret's
-/// then we encode, assuming the new data is encrypted data. Otherwise return
-/// a new secret with the same string.
-fn maybe_encode(original_secret: &Secret<String>, new_secret: &Secret<Vec<u8>>) -> Secret<String> {
-    if original_secret.expose_secret().as_bytes() != new_secret.expose_secret() {
-        let encoded = hex::encode(new_secret.expose_secret());
-        return Secret::new(encoded);
-    }
-
-    original_secret.clone()
-}
-
 fn pad_hex(s: String) -> String {
     if s.len() % 2 == 1 {
         format!("0{}", s)
@@ -278,3 +440,28 @@ fn pad_hex(s: String) -> String {
         s
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    // TODO(cnpryer)
+    #[test]
+    fn test_repository_config_resolution() {
+        // Resolve without credentials file
+        let credentials_table = Item::Table(Table::new());
+
+        let cli_repo = "pypi".to_string();
+        let cli_repo_url = None;
+        let cli_username = None;
+        let cli_token = None;
+
+        let config = resolve_credentials(Some(&credentials_table), cli_username, cli_token);
+        let repository =
+            resolve_repository(Some(&credentials_table), Some(cli_repo), cli_repo_url).unwrap();
+
+        assert!(config.username.is_none());
+        assert!(config.password.is_none());
+        assert_eq!(repository.url.unwrap().to_string(), DEFAULT_REPOSITORY_URL);
+    }
+}
diff --git a/rye/tests/common/mod.rs b/rye/tests/common/mod.rs
index 3ffbb45b9f..e50023dc1d 100644
--- a/rye/tests/common/mod.rs
+++ b/rye/tests/common/mod.rs
@@ -61,6 +61,20 @@ toolchain = "cpython@3.12.2"
         .unwrap();
     }
 
+    // write the credentials file
+    let credentials_file = home.join("credentials");
+    if !credentials_file.is_file() {
+        fs::write(
+            credentials_file,
+            r#"
+[pypi]
+# Update pypi repository url to avoid using it during tests
+repository-url = "don't use"
+"#,
+        )
+        .unwrap();
+    }
+
     // fetch the most important interpreters
     for version in ["cpython@3.8.17", "cpython@3.11.8", "cpython@3.12.2"] {
         if home.join("py").join(version).is_dir() {
@@ -102,6 +116,7 @@ pub fn get_bin() -> PathBuf {
     get_cargo_bin("rye")
 }
 
+#[derive(Debug)]
 pub struct Space {
     #[allow(unused)]
     tempdir: TempDir,
diff --git a/rye/tests/test_publish.rs b/rye/tests/test_publish.rs
new file mode 100644
index 0000000000..e9a678689f
--- /dev/null
+++ b/rye/tests/test_publish.rs
@@ -0,0 +1,18 @@
+use crate::common::{rye_cmd_snapshot, Space};
+
+mod common;
+
+#[test]
+fn test_publish_commands() {
+    let space = Space::new();
+    space.init("my-project");
+
+    rye_cmd_snapshot!(space.rye_cmd().arg("publish").arg("--repository").arg("none").arg("--skip-save-credentials").arg("-y"), @r###"
+    success: false
+    exit_code: 1
+    ----- stdout -----
+
+    ----- stderr -----
+    error: no configuration was found for repository 'none'
+    "###);
+}