diff --git a/crates/goose/src/agents/extension_malware_check.rs b/crates/goose/src/agents/extension_malware_check.rs index 6191122e3..07512f16b 100644 --- a/crates/goose/src/agents/extension_malware_check.rs +++ b/crates/goose/src/agents/extension_malware_check.rs @@ -81,11 +81,18 @@ pub async fn deny_if_malicious( fn parse_first_package_arg(ecosystem: &str, args: &[String]) -> Option<(String, Option)> { let is_flag = |s: &str| s.starts_with('-'); - let token = args - .iter() - .find(|a| !is_flag(a.as_str()))? - .trim() - .to_string(); + let positional = || { + args.iter() + .find(|a| !is_flag(a.as_str())) + .map(String::as_str) + }; + let token = if ecosystem == "PyPI" { + uvx_from_package(args).or_else(positional)? + } else { + positional()? + } + .trim() + .to_string(); if token.is_empty() { return None; } @@ -96,6 +103,108 @@ fn parse_first_package_arg(ecosystem: &str, args: &[String]) -> Option<(String, } } +fn uvx_from_package(args: &[String]) -> Option<&str> { + let mut index = 0; + while let Some(argument) = args.get(index).map(String::as_str) { + if argument == "--" || !argument.starts_with('-') { + break; + } + if let Some(package) = argument + .strip_prefix("--from=") + .filter(|package| !package.is_empty()) + { + return Some(package); + } + if argument == "--from" { + return args.get(index + 1).map(String::as_str); + } + + index += if uvx_option_takes_separate_value(argument) + && args + .get(index + 1) + .is_some_and(|value| !value.starts_with('-')) + { + 2 + } else { + 1 + }; + } + None +} + +fn uvx_option_takes_separate_value(argument: &str) -> bool { + if let Some(cluster) = argument + .strip_prefix('-') + .filter(|cluster| !cluster.starts_with('-')) + { + for (index, option) in cluster.char_indices() { + if matches!(option, 'C' | 'P' | 'b' | 'c' | 'f' | 'i' | 'p' | 'w') { + return index + option.len_utf8() == cluster.len(); + } + } + return false; + } + if argument.contains('=') { + return false; + } + if argument.strip_prefix("--extra-") == Some("index-url") { + return true; + } + + matches!( + argument, + "--allow-insecure-host" + | "--build-constraint" + | "--build-constraints" + | "--cache-dir" + | "--color" + | "--config-file" + | "--config-setting" + | "--config-settings" + | "--config-settings-package" + | "--constraint" + | "--constraints" + | "--default-index" + | "--directory" + | "--env-file" + | "--exclude-newer" + | "--exclude-newer-package" + | "--find-links" + | "--fork-strategy" + | "--generate-shell-completion" + | "--index" + | "--index-strategy" + | "--index-url" + | "--keyring-provider" + | "--link-mode" + | "--no-binary-package" + | "--no-build-isolation-package" + | "--no-build-package" + | "--no-sources-package" + | "--override" + | "--overrides" + | "--prerelease" + | "--prerelease-package" + | "--preview-feature" + | "--preview-features" + | "--project" + | "--python" + | "--python-fetch" + | "--python-platform" + | "--python-preference" + | "--refresh-package" + | "--reinstall-package" + | "--resolution" + | "--torch-backend" + | "--trusted-host" + | "--upgrade-group" + | "--upgrade-package" + | "--with" + | "--with-editable" + | "--with-requirements" + ) +} + fn parse_npm_token(token: &str) -> Option<(String, Option)> { // Handles: // react@18.3.1 @@ -506,6 +615,39 @@ mod tests { assert!(format!("{:?}", result.unwrap_err()).contains("MAL-240")); } + #[tokio::test] + #[serial_test::serial] + async fn cmd_args_pypi_checks_attached_from_selector() { + let server = MockServer::start().await; + Mock::given(method("POST")) + .and(path("/v1/query")) + .and(body_json(json!({ + "package": { + "name": "evil-pkg", + "ecosystem": "PyPI" + } + }))) + .respond_with(ResponseTemplate::new(200).set_body_json(json!({ + "vulns": [ { "id": "MAL-557", "summary": "Malicious package" } ], + "next_page_token": null + }))) + .mount(&server) + .await; + + let _env = TempEnvVar::set("OSV_ENDPOINT", &format!("{}/v1/query", server.uri())); + let result = deny_if_malicious_cmd_args( + "uvx", + &[ + "--from=Evil_Pkg[cli]>=0".to_string(), + "console-script".to_string(), + ], + ) + .await; + + assert!(result.is_err()); + assert!(format!("{:?}", result.unwrap_err()).contains("MAL-557")); + } + #[tokio::test] #[serial_test::serial] async fn cmd_args_pypi_preserves_uvx_exact_version() { @@ -796,4 +938,83 @@ mod tests { Some(("evil-pkg".into(), None)) ); } + + #[test] + fn parse_first_pypi_argument_prefers_uvx_from_selector() { + for args in [ + vec![ + "--from=Evil_Pkg[cli]>=0".to_string(), + "console-script".to_string(), + ], + vec![ + "--from".to_string(), + "Evil_Pkg[cli]>=0".to_string(), + "console-script".to_string(), + ], + ] { + assert_eq!( + super::parse_first_package_arg("PyPI", &args), + Some(("evil-pkg".into(), None)) + ); + } + } + + #[test] + fn parse_first_pypi_argument_ignores_from_after_command() { + for args in [ + vec![ + "malicious-package".to_string(), + "--from".to_string(), + "benign-package".to_string(), + ], + vec![ + "--isolated".to_string(), + "malicious-package".to_string(), + "--from".to_string(), + "benign-package".to_string(), + ], + vec![ + "-wP".to_string(), + "malicious-package".to_string(), + "--from".to_string(), + "benign-package".to_string(), + ], + ] { + assert_eq!( + super::parse_first_package_arg("PyPI", &args), + Some(("malicious-package".into(), None)) + ); + } + } + + #[test] + fn parse_first_pypi_argument_skips_uvx_option_values_before_from() { + for args in [ + vec![ + "--directory".to_string(), + "/tmp".to_string(), + "--from".to_string(), + "Evil_Pkg[cli]>=0".to_string(), + "console-script".to_string(), + ], + vec![ + "-p".to_string(), + "3.12".to_string(), + "--from=Evil_Pkg[cli]>=0".to_string(), + "console-script".to_string(), + ], + vec![ + "-qP".to_string(), + "benign-package".to_string(), + "--from".to_string(), + "Evil_Pkg[cli]>=0".to_string(), + "console-script".to_string(), + ], + ] { + assert_eq!( + super::parse_first_package_arg("PyPI", &args), + Some(("evil-pkg".into(), None)) + ); + } + } }