use crate::debug::{DebugLog, DebugTransport};
use crate::error::{self, CliError, Operation, Phase, Plane, SdkContext, SecretSet};
use crate::tenant::config::{TenantConfig, js_trim};
use futures_util::StreamExt;
use futures_util::future::BoxFuture;
use scm::airs::management::ManagementClient;
use scm::airs::scan::Scanner;
use scm::auth::{AuthContext, Authentication, OAuthProvider, TokenLease, TokenProvider};
use scm::config::DestinationPolicy;
use scm::transport::{ReqwestTransport, Transport, TransportRequest, TransportResponse};
use scm::{Client, Endpoint, ErrorKind, Limits, OriginPolicy, RequestOptions, Secret, Session};
use std::sync::Arc;
use std::sync::atomic::{AtomicU16, Ordering};
use std::time::Duration;
pub const DEFAULT_TOKEN_ENDPOINT: &str =
"https://auth.apps.paloaltonetworks.com/oauth2/access_token";
pub const DEFAULT_MGMT_ENDPOINT: &str = "https://api.apps.paloaltonetworks.com/aisec";
pub const DEFAULT_SCAN_ENDPOINT: &str = "https://service.api.aisecurity.paloaltonetworks.com";
pub const DEFAULT_RED_TEAM_DATA_ENDPOINT: &str =
"https://api.apps.paloaltonetworks.com/ai-red-teaming/data-plane";
pub const DEFAULT_RED_TEAM_MGMT_ENDPOINT: &str =
"https://api.apps.paloaltonetworks.com/ai-red-teaming/mgmt-plane";
pub const DEFAULT_RED_TEAM_BROKER_ENDPOINT: &str =
"https://api.apps.paloaltonetworks.com/ai-red-teaming/data-plane/network-broker";
pub const SINGLE_REQUEST_BUDGET: Duration = Duration::from_secs(60);
pub const TOKEN_BUDGET: Duration = Duration::from_secs(30);
pub const WALK_BUDGET: Duration = Duration::from_secs(300);
pub const PROBE_BUDGET: Duration = Duration::from_secs(5);
const MANAGEMENT_RETRIES: u32 = 5;
const MAX_SCANNER_CREDENTIAL: usize = 2048;
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct SetupError {
pub key: Option<&'static str>,
pub error: CliError,
}
impl SetupError {
fn config(key: &'static str, message: impl Into<String>) -> Self {
Self {
key: Some(key),
error: CliError::operational(message),
}
}
fn sdk(message: &str) -> Self {
Self {
key: None,
error: CliError::operational(message),
}
}
pub fn message(&self) -> String {
self.error.message()
}
}
impl From<SetupError> for CliError {
fn from(value: SetupError) -> Self {
value.error
}
}
fn invalid_endpoint(key: &'static str) -> SetupError {
SetupError::config(
key,
format!(
"Invalid {key}: endpoints must use https (http is allowed only for loopback hosts) and must not contain credentials, a query or a fragment"
),
)
}
pub fn admit(
key: &'static str,
configured: Option<&str>,
default: &str,
) -> Result<Endpoint, SetupError> {
Endpoint::loopback(configured.unwrap_or(default)).map_err(|_| invalid_endpoint(key))
}
#[derive(Clone, Debug, Default)]
pub struct LastStatus(Arc<AtomicU16>);
impl LastStatus {
pub fn get(&self) -> Option<u16> {
match self.0.load(Ordering::SeqCst) {
0 => None,
status => Some(status),
}
}
}
struct ObservedTransport {
inner: Arc<dyn Transport>,
last: LastStatus,
}
impl Transport for ObservedTransport {
fn send<'a>(
&'a self,
request: TransportRequest,
) -> BoxFuture<'a, scm::Result<TransportResponse>> {
Box::pin(async move {
let response = self.inner.send(request).await?;
self.last
.0
.store(response.status().as_u16(), Ordering::SeqCst);
Ok(response)
})
}
}
fn transport(
endpoints: &[&Endpoint],
last: &LastStatus,
debug: Option<&DebugLog>,
) -> Result<Arc<dyn Transport>, SetupError> {
let failed = || SetupError::sdk("AISEC_SDK_ERROR:The HTTP transport could not be initialised");
let inner: Arc<dyn Transport> = if endpoints
.iter()
.any(|endpoint| endpoint.url().scheme() == "http")
{
Arc::new(DirectTransport::new().ok_or_else(failed)?)
} else {
Arc::new(ReqwestTransport::new().map_err(|_| failed())?)
};
Ok(Arc::new(ObservedTransport {
inner: DebugTransport::wrap(inner, debug),
last: last.clone(),
}))
}
const SDK_USER_AGENT: &str = "PAN-AIRS/0.1.0-rust-sdk";
struct DirectTransport {
client: reqwest::Client,
}
impl DirectTransport {
fn new() -> Option<Self> {
reqwest::Client::builder()
.no_proxy()
.redirect(reqwest::redirect::Policy::none())
.retry(reqwest::retry::never())
.connect_timeout(Duration::from_secs(10))
.pool_max_idle_per_host(32)
.pool_idle_timeout(Duration::from_secs(90))
.http1_max_headers(128)
.http2_max_header_list_size(16 * 1024)
.no_brotli()
.no_deflate()
.no_zstd()
.gzip(true)
.referer(false)
.user_agent(SDK_USER_AGENT)
.build()
.ok()
.map(|client| Self { client })
}
}
impl Transport for DirectTransport {
fn send<'a>(
&'a self,
request: TransportRequest,
) -> BoxFuture<'a, scm::Result<TransportResponse>> {
Box::pin(async move {
let response = self
.client
.request(request.method().clone(), request.url().clone())
.headers(request.headers().clone())
.body(request.body().clone())
.send()
.await
.map_err(|_| scm::Error::new(ErrorKind::Transport))?;
let status = response.status();
let headers = response.headers().clone();
let stream = response
.bytes_stream()
.map(|chunk| chunk.map_err(|_| scm::Error::new(ErrorKind::Transport)));
Ok(TransportResponse::new(status, headers, Box::pin(stream)))
})
}
}
pub struct RecordingProvider {
inner: Arc<OAuthProvider>,
secrets: SecretSet,
debug: Option<DebugLog>,
}
impl RecordingProvider {
pub async fn get_token(&self, options: Option<&RequestOptions>) -> scm::Result<TokenLease> {
let options = match options {
Some(options) => options.clone(),
None => RequestOptions::with_timeout(TOKEN_BUDGET)?,
};
self.token(&options).await
}
}
impl TokenProvider for RecordingProvider {
fn token<'a>(&'a self, options: &'a RequestOptions) -> BoxFuture<'a, scm::Result<TokenLease>> {
Box::pin(async move {
let hold = self.debug.as_ref().map(DebugLog::hold_tokens);
let result = self.inner.token(options).await;
if let Ok(lease) = &result {
self.secrets.add(lease.secret().expose());
}
if let Some(hold) = hold {
hold.release();
}
result
})
}
fn context(&self) -> &AuthContext {
self.inner.context()
}
fn invalidate(&self, generation: u64) -> bool {
self.inner.invalidate(generation)
}
}
pub struct ManagementClients {
pub provider: Arc<RecordingProvider>,
pub management: ManagementClient,
pub last_status: LastStatus,
}
impl ManagementClients {
pub fn build(
config: &TenantConfig,
session: &Session,
secrets: &SecretSet,
debug: Option<&DebugLog>,
) -> Result<Self, SetupError> {
let api = admit(
"mgmtEndpoint",
config.effective_str("mgmtEndpoint"),
DEFAULT_MGMT_ENDPOINT,
)?;
let issuer = admit(
"mgmtTokenEndpoint",
config.effective_str("mgmtTokenEndpoint"),
DEFAULT_TOKEN_ENDPOINT,
)?;
let last_status = LastStatus::default();
let OAuthParts {
provider,
transport,
policy,
} = oauth_parts(
config,
&issuer,
std::slice::from_ref(&api),
session,
secrets,
debug,
&last_status,
)?;
let mut limits = Limits::default();
limits.max_retries = MANAGEMENT_RETRIES;
let client = Client::builder(api)
.transport(transport)
.destination_policy(policy)
.authentication(Authentication::oauth(
provider.clone() as Arc<dyn TokenProvider>,
false,
))
.session(session.clone())
.limits(limits)
.build()
.map_err(|_| {
SetupError::sdk("AISEC_SDK_ERROR:The management client could not be configured")
})?;
let management = ManagementClient::new(client).map_err(|_| {
SetupError::sdk("AISEC_SDK_ERROR:The management client could not be configured")
})?;
Ok(Self {
provider,
management,
last_status,
})
}
pub async fn authenticate(
&self,
options: Option<&RequestOptions>,
budget: Duration,
) -> Result<(), CliError> {
match self.provider.get_token(options).await {
Ok(_) => Ok(()),
Err(failure) => Err(error::from_sdk(
&failure,
&SdkContext {
plane: Plane::Management,
phase: Phase::Token,
operation: Operation::Read,
timeout: budget,
secrets: &self.provider.secrets,
},
)),
}
}
}
struct OAuthParts {
provider: Arc<RecordingProvider>,
transport: Arc<dyn Transport>,
policy: Arc<dyn DestinationPolicy>,
}
fn oauth_parts(
config: &TenantConfig,
issuer: &Endpoint,
apis: &[Endpoint],
session: &Session,
secrets: &SecretSet,
debug: Option<&DebugLog>,
last_status: &LastStatus,
) -> Result<OAuthParts, SetupError> {
let tsg = config.effective_str("mgmtTsgId").unwrap_or_default();
if tsg.is_empty() || tsg.len() > 32 || !tsg.bytes().all(|b| b.is_ascii_digit()) {
return Err(SetupError::config(
"mgmtTsgId",
"Invalid mgmtTsgId: the Rust SDK requires 1-32 ASCII digits",
));
}
let client_id = config.effective_str("mgmtClientId").unwrap_or_default();
let context = AuthContext::new(issuer.clone(), client_id, 0, tsg).map_err(|_| {
SetupError::config(
"mgmtClientId",
"Invalid mgmtClientId: the Rust SDK requires at most 1024 bytes without control characters",
)
})?;
let invalid_secret = || {
SetupError::config(
"mgmtClientSecret",
"Invalid mgmtClientSecret: the Rust SDK requires at most 16 KiB without control characters",
)
};
let secret = Secret::new(config.effective_str("mgmtClientSecret").unwrap_or_default())
.map_err(|_| invalid_secret())?;
let mut endpoints: Vec<&Endpoint> = apis.iter().collect();
endpoints.push(issuer);
let transport = transport(&endpoints, last_status, debug)?;
let mut admitted: Vec<Endpoint> = apis.to_vec();
admitted.push(issuer.clone());
let policy: Arc<dyn DestinationPolicy> = Arc::new(OriginPolicy::new(admitted));
let provider = Arc::new(RecordingProvider {
inner: Arc::new(
OAuthProvider::new(
context,
secret,
transport.clone(),
policy.clone(),
session.clone(),
)
.map_err(|_| invalid_secret())?,
),
secrets: secrets.clone(),
debug: debug.cloned(),
});
Ok(OAuthParts {
provider,
transport,
policy,
})
}
pub struct RedTeamClients {
pub provider: Arc<RecordingProvider>,
pub red_team: scm::airs::red_team::RedTeamClient,
}
impl RedTeamClients {
pub fn build(
config: &TenantConfig,
session: &Session,
secrets: &SecretSet,
debug: Option<&DebugLog>,
) -> Result<Self, SetupError> {
let data = admit(
"redTeamDataEndpoint",
config.effective_str("redTeamDataEndpoint"),
DEFAULT_RED_TEAM_DATA_ENDPOINT,
)?;
let management = admit(
"redTeamMgmtEndpoint",
config.effective_str("redTeamMgmtEndpoint"),
DEFAULT_RED_TEAM_MGMT_ENDPOINT,
)?;
let broker = admit(
"redTeamNetworkBrokerEndpoint",
config.effective_str("redTeamNetworkBrokerEndpoint"),
DEFAULT_RED_TEAM_BROKER_ENDPOINT,
)?;
let issuer = admit(
"mgmtTokenEndpoint",
config.effective_str("mgmtTokenEndpoint"),
DEFAULT_TOKEN_ENDPOINT,
)?;
let OAuthParts {
provider,
transport,
policy,
} = oauth_parts(
config,
&issuer,
&[data.clone(), management.clone(), broker.clone()],
session,
secrets,
debug,
&LastStatus::default(),
)?;
let failed =
|| SetupError::sdk("AISEC_SDK_ERROR:The red team client could not be configured");
let mut limits = Limits::default();
limits.max_retries = MANAGEMENT_RETRIES;
let client = Client::builder(data)
.transport(transport)
.destination_policy(policy)
.authentication(Authentication::oauth(
provider.clone() as Arc<dyn TokenProvider>,
false,
))
.session(session.clone())
.limits(limits)
.build()
.map_err(|_| failed())?;
let red_team = scm::airs::red_team::RedTeamClient::new(client, management, broker)
.map_err(|_| failed())?;
Ok(Self { provider, red_team })
}
pub async fn authenticate(
&self,
options: Option<&RequestOptions>,
budget: Duration,
) -> Result<(), CliError> {
match self.provider.get_token(options).await {
Ok(_) => Ok(()),
Err(failure) => Err(error::from_sdk(
&failure,
&SdkContext {
plane: Plane::Management,
phase: Phase::Token,
operation: Operation::Read,
timeout: budget,
secrets: &self.provider.secrets,
},
)),
}
}
}
fn scanner_credential(
config: &TenantConfig,
key: &'static str,
label: &str,
) -> Result<Option<Secret>, SetupError> {
let Some(value) = config
.effective_str(key)
.map(js_trim)
.filter(|v| !v.is_empty())
else {
return Ok(None);
};
if value.encode_utf16().nth(MAX_SCANNER_CREDENTIAL).is_some() {
return Err(SetupError::config(
key,
format!(
"AISEC_USER_REQUEST_PAYLOAD_ERROR:{label} exceeds max length of {MAX_SCANNER_CREDENTIAL}"
),
));
}
if !value.bytes().all(|byte| byte.is_ascii_graphic()) {
return Err(SetupError::config(
key,
format!("Invalid {key}: the Rust SDK requires printable ASCII without spaces"),
));
}
Secret::new(value)
.map(Some)
.map_err(|_| SetupError::config(key, format!("Invalid {key}")))
}
pub fn build_scanner(
config: &TenantConfig,
session: &Session,
debug: Option<&DebugLog>,
) -> Result<Scanner, SetupError> {
build_observed_scanner(config, session, debug).map(|(scanner, _)| scanner)
}
pub fn build_observed_scanner(
config: &TenantConfig,
session: &Session,
debug: Option<&DebugLog>,
) -> Result<(Scanner, LastStatus), SetupError> {
let endpoint = admit(
"airsApiEndpoint",
config.effective_str("airsApiEndpoint"),
DEFAULT_SCAN_ENDPOINT,
)?;
let api_key = scanner_credential(config, "airsApiKey", "apiKey")?;
let api_token = scanner_credential(config, "airsApiToken", "apiToken")?;
let mut options = scm::configuration::InitOptions::default();
options.api_key = api_key;
options.api_token = api_token;
options.num_retries = config.num_retries();
let last_status = LastStatus::default();
let transport = transport(&[&endpoint], &last_status, debug)?;
options.api_endpoint = Some(endpoint);
let configuration =
scm::configuration::init(options).map_err(|_| invalid_endpoint("airsApiEndpoint"))?;
let failed = || SetupError::sdk("AISEC_SDK_ERROR:The scanner client could not be configured");
let client = configuration
.client_builder()
.map_err(|_| failed())?
.transport(transport)
.session(session.clone())
.build()
.map_err(|_| failed())?;
Ok((Scanner::new(client).map_err(|_| failed())?, last_status))
}
#[cfg(test)]
mod tests {
use super::*;
use crate::tenant::config::{ConfigValue, schema_key};
fn config(text: &str) -> TenantConfig {
let serde_json::Value::Object(entries) = serde_json::from_str(text).unwrap() else {
panic!("object expected");
};
TenantConfig::from_values(entries.into_iter().map(|(key, value)| {
let key = schema_key(&key).unwrap();
let value = match value {
serde_json::Value::String(text) => ConfigValue::Text(text),
serde_json::Value::Number(number) => ConfigValue::Integer(number.as_i64().unwrap()),
other => panic!("unsupported {other}"),
};
(key, value)
}))
}
#[test]
fn endpoint_admission() {
for url in [
"https://example.com/aisec",
"http://127.0.0.1:8080/aisec",
"http://127.9.9.9",
"http://localhost:1",
"http://[::1]:9/x",
] {
assert!(
admit("mgmtEndpoint", Some(url), DEFAULT_MGMT_ENDPOINT).is_ok(),
"{url}"
);
}
for url in [
"http://example.test",
"http://0.0.0.0:1",
"https://user:pw@example.com",
"https://example.com/?q=1",
"https://example.com/#f",
"not a url",
] {
let error = admit("mgmtEndpoint", Some(url), DEFAULT_MGMT_ENDPOINT).unwrap_err();
assert_eq!(error.key, Some("mgmtEndpoint"));
assert!(!error.message().contains(url), "the value is never echoed");
assert!(
error
.message()
.starts_with("Invalid mgmtEndpoint: endpoints must use https")
);
}
assert!(admit("x", None, DEFAULT_SCAN_ENDPOINT).is_ok());
}
#[test]
fn management_configuration_errors_are_local() {
let session = Session::new();
let secrets = SecretSet::default();
let bad_tsg = config(r#"{"mgmtClientId":"c","mgmtClientSecret":"s","mgmtTsgId":"abc"}"#);
let error = ManagementClients::build(&bad_tsg, &session, &secrets, None)
.err()
.unwrap();
assert_eq!(
error.message(),
"Invalid mgmtTsgId: the Rust SDK requires 1-32 ASCII digits"
);
let bad_endpoint = config(
r#"{"mgmtClientId":"c","mgmtClientSecret":"s","mgmtTsgId":"1","mgmtTokenEndpoint":"http://example.test/t"}"#,
);
let error = ManagementClients::build(&bad_endpoint, &session, &secrets, None)
.err()
.unwrap();
assert_eq!(error.key, Some("mgmtTokenEndpoint"));
let control = config(r#"{"mgmtClientId":"c\n","mgmtClientSecret":"s","mgmtTsgId":"1"}"#);
assert_eq!(
ManagementClients::build(&control, &session, &secrets, None)
.err()
.unwrap()
.key,
Some("mgmtClientId")
);
let good = config(
r#"{"mgmtClientId":"c","mgmtClientSecret":"s","mgmtTsgId":"1","mgmtEndpoint":"http://127.0.0.1:1/aisec","mgmtTokenEndpoint":"http://127.0.0.1:1/token"}"#,
);
assert!(ManagementClients::build(&good, &session, &secrets, None).is_ok());
}
#[test]
fn red_team_clients_admit_the_three_planes_or_their_defaults() {
let session = Session::new();
let secrets = SecretSet::default();
let defaults = config(r#"{"mgmtClientId":"c","mgmtClientSecret":"s","mgmtTsgId":"1"}"#);
assert!(RedTeamClients::build(&defaults, &session, &secrets, None).is_ok());
let loopback = config(
r#"{"mgmtClientId":"c","mgmtClientSecret":"s","mgmtTsgId":"1","redTeamDataEndpoint":"http://127.0.0.1:1/d","redTeamMgmtEndpoint":"http://127.0.0.1:1/m","redTeamNetworkBrokerEndpoint":"http://127.0.0.1:1/b","mgmtTokenEndpoint":"http://127.0.0.1:1/t"}"#,
);
assert!(RedTeamClients::build(&loopback, &session, &secrets, None).is_ok());
for key in [
"redTeamDataEndpoint",
"redTeamMgmtEndpoint",
"redTeamNetworkBrokerEndpoint",
] {
let bad = config(&format!(
r#"{{"mgmtClientId":"c","mgmtClientSecret":"s","mgmtTsgId":"1","{key}":"http://example.test/x"}}"#
));
let error = RedTeamClients::build(&bad, &session, &secrets, None)
.err()
.unwrap();
assert_eq!(error.key, Some(key));
}
let bad_tsg = config(r#"{"mgmtClientId":"c","mgmtClientSecret":"s","mgmtTsgId":"x"}"#);
assert_eq!(
RedTeamClients::build(&bad_tsg, &session, &secrets, None)
.err()
.unwrap()
.key,
Some("mgmtTsgId")
);
}
#[test]
fn scanner_credentials_are_trimmed_and_bounded() {
let session = Session::new();
let good = config(r#"{"airsApiKey":" key-1 \n","airsApiEndpoint":"http://127.0.0.1:1"}"#);
assert!(build_scanner(&good, &session, None).is_ok());
let long = format!(r#"{{"airsApiKey":"{}"}}"#, "k".repeat(2049));
assert_eq!(
build_scanner(&config(&long), &session, None)
.err()
.unwrap()
.message(),
"AISEC_USER_REQUEST_PAYLOAD_ERROR:apiKey exceeds max length of 2048"
);
let exact = format!(r#"{{"airsApiToken":"{}"}}"#, "k".repeat(2048));
assert!(build_scanner(&config(&exact), &session, None).is_ok());
let spaced = config(r#"{"airsApiToken":"a b"}"#);
assert_eq!(
build_scanner(&spaced, &session, None).err().unwrap().key,
Some("airsApiToken")
);
let endpoint = config(r#"{"airsApiKey":"k","airsApiEndpoint":"http://example.test"}"#);
assert_eq!(
build_scanner(&endpoint, &session, None).err().unwrap().key,
Some("airsApiEndpoint")
);
}
}