diff --git a/Cargo.lock b/Cargo.lock index 950bdc4c..087427d4 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4490,6 +4490,7 @@ dependencies = [ name = "ts_python" version = "0.2.0" dependencies = [ + "hex", "pyo3", "pyo3-async-runtimes", "tailscale", diff --git a/ts_python/Cargo.toml b/ts_python/Cargo.toml index 6b582b4d..398fc792 100644 --- a/ts_python/Cargo.toml +++ b/ts_python/Cargo.toml @@ -13,6 +13,7 @@ rust-version.workspace = true [dependencies] tailscale.workspace = true +hex.workspace = true pyo3 = { version = "0.28", features = ["bytes", "abi3-py312"] } pyo3-async-runtimes = { version = "0.28", features = ["tokio-runtime", "attributes"] } tracing-subscriber = { version = "0.3", features = ["env-filter"] } diff --git a/ts_python/src/key_state.rs b/ts_python/src/key_state.rs new file mode 100644 index 00000000..67b77c10 --- /dev/null +++ b/ts_python/src/key_state.rs @@ -0,0 +1,95 @@ +use ts::keys::{DiscoPrivateKey, MachinePrivateKey, NetworkLockPrivateKey, NodePrivateKey}; + +/// Tailscale keys. +#[derive(Debug, Clone, PartialEq, Eq)] +#[pyo3::pyclass(frozen, get_all, from_py_object, module = "tailscale")] +pub struct Keystate { + /// Machine key. + pub machine: Vec, + /// Node (device) key. + pub node: Vec, + /// Disco key. + pub disco: Vec, + /// Network lock key. + pub network_lock: Vec, +} + +#[pyo3::pymethods] +impl Keystate { + #[new] + #[pyo3(signature = (machine=None, node=None, disco=None, network_lock=None))] + pub fn new( + machine: Option>, + node: Option>, + disco: Option>, + network_lock: Option>, + ) -> Self { + let mut out = Self { + ..ts::keys::NodeState::default().into() + }; + + if let Some(machine) = machine { + out.machine = machine; + } + + if let Some(node) = node { + out.node = node; + } + + if let Some(disco) = disco { + out.disco = disco; + } + + if let Some(network_lock) = network_lock { + out.network_lock = network_lock; + } + + out + } + + pub fn __repr__(&self) -> String { + match tailscale::keys::NodeState::try_from(self) { + Ok(state) => { + format!( + "tailscale.Keystate(machine={}, node={}, disco={}, network_lock={})", + hex::encode(state.machine_keys.public.to_bytes()), + hex::encode(state.node_keys.public.to_bytes()), + hex::encode(state.disco_keys.public.to_bytes()), + hex::encode(state.network_lock_keys.public.to_bytes()), + ) + } + Err(_) => "tailscale.Keystate()".to_owned(), + } + } +} + +impl From for Keystate { + fn from(value: tailscale::keys::NodeState) -> Self { + Self { + machine: value.machine_keys.private.to_bytes().into(), + node: value.node_keys.private.to_bytes().into(), + disco: value.disco_keys.private.to_bytes().into(), + network_lock: value.network_lock_keys.private.to_bytes().into(), + } + } +} + +impl TryFrom<&Keystate> for tailscale::keys::NodeState { + type Error = (); + + fn try_from(value: &Keystate) -> Result { + fn key(v: &[u8]) -> Result + where + T: From<[u8; 32]>, + { + Ok(<[u8; 32]>::try_from(v).map_err(|_| ())?.into()) + } + + Ok(Self { + machine_keys: key::(&value.machine)?.into(), + node_keys: key::(&value.node)?.into(), + disco_keys: key::(&value.disco)?.into(), + network_lock_keys: key::(&value.network_lock)?.into(), + }) + } +} diff --git a/ts_python/src/lib.rs b/ts_python/src/lib.rs index 33a5c1a3..dcd8e325 100644 --- a/ts_python/src/lib.rs +++ b/ts_python/src/lib.rs @@ -16,10 +16,12 @@ extern crate tailscale as ts; type PyFut<'p> = PyResult>; mod ip_or_str; +mod key_state; mod node_info; mod tcp; mod udp; +use key_state::Keystate; use node_info::NodeInfo; /// Tailscale API. @@ -28,15 +30,23 @@ pub mod _internal { use super::*; #[pymodule_export] use crate::{ - Device, + Device, Keystate, tcp::{TcpListener, TcpStream}, udp::UdpSocket, }; - /// Connect to tailscale using the specified config file and optional auth key. + /// Connect to tailscale using the specified parameters. #[pyfunction] - #[pyo3(signature = (config_path, auth_key=None))] - pub fn connect(py: Python<'_>, config_path: String, auth_key: Option) -> PyFut<'_> { + #[pyo3(signature = (key_file_path=None, /, auth_key=None, *, control_server_url=None, hostname=None, tags=None, keys=None))] + pub fn connect( + py: Python<'_>, + key_file_path: Option, + auth_key: Option, + control_server_url: Option, + hostname: Option, + tags: Option>, + keys: Option, + ) -> PyFut<'_> { static TRACING_ONCE: Once = Once::new(); TRACING_ONCE.call_once(|| { tracing_subscriber::fmt() @@ -49,13 +59,31 @@ pub mod _internal { }); future_into_py(py, async move { - let config = ts::Config { - client_name: Some("ts_python".to_owned()), - ..ts::Config::default_with_key_file(config_path) + let mut config = if let Some(key_file_path) = key_file_path { + ts::Config::default_with_key_file(key_file_path) .await .map_err(py_value_err)? + } else { + ts::Config::default() }; + config.client_name = Some("ts_python".to_owned()); + if let Some(control_server_url) = control_server_url { + config.control_server_url = control_server_url.parse().map_err(py_value_err)?; + } + + if let Some(hostname) = hostname { + config.requested_hostname = Some(hostname); + } + + if let Some(tags) = tags { + config.requested_tags = tags; + } + + if let Some(keys) = &keys { + config.key_state = keys.try_into().map_err(|_| py_value_err("invalid keys"))?; + } + let dev = ts::Device::new(&config, auth_key) .await .map_err(py_value_err)?;