use rmcp::schemars::JsonSchema;
use serde::{Deserialize, Serialize};
use serde_json::Value;
use tailscale_rest::models::key::{CREATE_KEY_TYPES, UPDATE_KEY_TYPES};
use crate::context::ToolContext;
use crate::error::ToolResult;
use crate::tools::common::{Done, one_of, path_segment, report};
crate::tools! {
tailnet_key_list => KeyListParams, key_list,
toolset: TailnetKeys, tier: Read, idempotent: true;
tailnet_key_get => KeyParams, key_get,
toolset: TailnetKeys, tier: Read, idempotent: true;
tailnet_key_create => KeyCreateParams, key_create,
toolset: TailnetKeys, tier: Write;
tailnet_key_update => KeyUpdateParams, key_update,
toolset: TailnetKeys, tier: Write, idempotent: true;
tailnet_key_delete => KeyParams, key_delete,
toolset: TailnetKeys, tier: Destructive, idempotent: true;
}
fn key_path(client: &tailscale_rest::Client, id: &str) -> ToolResult<String> {
let id = path_segment("key_id", id)?;
Ok(client.tailnet_path(None, &format!("/keys/{id}")))
}
fn checked_key_type(key_type: Option<&str>, allowed: &[&str]) -> ToolResult<Option<String>> {
key_type
.map(|key_type| one_of("key_type", key_type, allowed))
.transpose()
}
#[derive(Debug, Deserialize, JsonSchema)]
pub struct KeyListParams {
#[serde(default)]
pub all: Option<bool>,
}
async fn key_list(ctx: &ToolContext, params: KeyListParams) -> ToolResult<Value> {
let client = ctx.tailnet()?;
Ok(client
.get(client.tailnet_path(None, "/keys"))
.query("all", params.all.unwrap_or(true))
.send_as::<Value>()
.await?)
}
#[derive(Debug, Deserialize, JsonSchema)]
pub struct KeyParams {
pub key_id: String,
}
async fn key_get(ctx: &ToolContext, params: KeyParams) -> ToolResult<Value> {
let client = ctx.tailnet()?;
Ok(client
.get(key_path(client, ¶ms.key_id)?)
.send_as::<Value>()
.await?)
}
async fn key_delete(ctx: &ToolContext, params: KeyParams) -> ToolResult<Value> {
let client = ctx.tailnet()?;
client
.delete(key_path(client, ¶ms.key_id)?)
.send()
.await?;
report(Done::new("key revoked").about("key_id", params.key_id))
}
#[derive(Debug, Deserialize, JsonSchema)]
pub struct KeyCreateParams {
#[serde(default)]
pub key_type: Option<String>,
#[serde(default)]
pub description: Option<String>,
#[serde(default)]
pub capabilities: Option<Value>,
#[serde(default)]
pub expiry_seconds: Option<i64>,
#[serde(default)]
pub scopes: Option<Vec<String>>,
#[serde(default)]
pub tags: Option<Vec<String>>,
#[serde(default)]
pub issuer: Option<String>,
#[serde(default)]
pub subject: Option<String>,
#[serde(default)]
pub audience: Option<String>,
#[serde(default)]
pub custom_claim_rules: Option<std::collections::BTreeMap<String, String>>,
}
#[derive(Debug, Serialize)]
struct KeyBody {
#[serde(rename = "keyType", skip_serializing_if = "Option::is_none")]
key_type: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
description: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
capabilities: Option<Value>,
#[serde(rename = "expirySeconds", skip_serializing_if = "Option::is_none")]
expiry_seconds: Option<i64>,
#[serde(skip_serializing_if = "Option::is_none")]
scopes: Option<Vec<String>>,
#[serde(skip_serializing_if = "Option::is_none")]
tags: Option<Vec<String>>,
#[serde(skip_serializing_if = "Option::is_none")]
issuer: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
subject: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
audience: Option<String>,
#[serde(rename = "customClaimRules", skip_serializing_if = "Option::is_none")]
custom_claim_rules: Option<std::collections::BTreeMap<String, String>>,
}
async fn key_create(ctx: &ToolContext, params: KeyCreateParams) -> ToolResult<Value> {
let client = ctx.tailnet()?;
let key_type = checked_key_type(params.key_type.as_deref(), CREATE_KEY_TYPES)?;
let body = KeyBody {
key_type,
description: params.description,
capabilities: params.capabilities,
expiry_seconds: params.expiry_seconds,
scopes: params.scopes,
tags: params.tags,
issuer: params.issuer,
subject: params.subject,
audience: params.audience,
custom_claim_rules: params.custom_claim_rules,
};
Ok(client
.post(client.tailnet_path(None, "/keys"))
.json(&body)
.send_as::<Value>()
.await?)
}
#[derive(Debug, Deserialize, JsonSchema)]
pub struct KeyUpdateParams {
pub key_id: String,
#[serde(default)]
pub key_type: Option<String>,
#[serde(default)]
pub description: Option<String>,
#[serde(default)]
pub scopes: Option<Vec<String>>,
#[serde(default)]
pub tags: Option<Vec<String>>,
#[serde(default)]
pub issuer: Option<String>,
#[serde(default)]
pub subject: Option<String>,
#[serde(default)]
pub audience: Option<String>,
#[serde(default)]
pub custom_claim_rules: Option<std::collections::BTreeMap<String, String>>,
}
async fn key_update(ctx: &ToolContext, params: KeyUpdateParams) -> ToolResult<Value> {
let client = ctx.tailnet()?;
let path = key_path(client, ¶ms.key_id)?;
let body = KeyBody {
key_type: checked_key_type(params.key_type.as_deref(), UPDATE_KEY_TYPES)?,
description: params.description,
capabilities: None,
expiry_seconds: None,
scopes: params.scopes,
tags: params.tags,
issuer: params.issuer,
subject: params.subject,
audience: params.audience,
custom_claim_rules: params.custom_claim_rules,
};
Ok(client.put(path).json(&body).send_as::<Value>().await?)
}
#[cfg(test)]
mod tests {
use super::*;
use serde_json::json;
#[test]
fn each_call_is_held_to_its_own_list_of_key_types() {
assert!(checked_key_type(Some("auth"), CREATE_KEY_TYPES).is_ok());
assert!(checked_key_type(Some("api"), CREATE_KEY_TYPES).is_err());
assert!(checked_key_type(Some("auth"), UPDATE_KEY_TYPES).is_err());
assert!(checked_key_type(Some("client"), UPDATE_KEY_TYPES).is_ok());
assert_eq!(
checked_key_type(None, UPDATE_KEY_TYPES).expect("no type"),
None
);
let error = checked_key_type(Some("auth"), UPDATE_KEY_TYPES).expect_err("not updatable");
let reported = serde_json::to_value(&error).expect("reportable");
let message = reported["message"].as_str().expect("a message");
assert!(
message.contains("client") && message.contains("federated"),
"the refusal quotes the list that applies here: {message}"
);
assert!(
!message.contains("api"),
"and not one that does not: {message}"
);
}
#[test]
fn a_capabilities_object_is_sent_exactly_as_it_was_given() {
let documented = json!({
"devices": {"create": {
"reusable": false, "ephemeral": false,
"preauthorized": false, "tags": ["tag:example"],
}},
});
let body = KeyBody {
key_type: Some("auth".to_owned()),
description: Some("example".to_owned()),
capabilities: Some(documented.clone()),
expiry_seconds: Some(86400),
scopes: None,
tags: None,
issuer: None,
subject: None,
audience: None,
custom_claim_rules: None,
};
let sent = serde_json::to_value(&body).expect("it serialises");
assert_eq!(sent["capabilities"], documented);
assert_eq!(
sent,
json!({
"keyType": "auth",
"description": "example",
"capabilities": documented,
"expirySeconds": 86400,
}),
"and nothing else is sent: an absent field is absent, not null"
);
}
}